Fixes #154881 and #184844
PR is a bit of a big one so I'll try to summarize below:
Each row is an example tensor (scan dim + dtype) and the kernel it lands on. `
1M` = a long axis, the specialized long-scan kernels need axis `>= 65536` with `<= 1024` scans, everything else takes the generic fallback.
| example shape | kernel | what it does |
|---|---|---|
| `[8, 1M]` dim=-1, float/half/bf16 | `scan_contig_decoupled` | each 1M row is cut into 4096-elem tiles, **one threadgroup per tile** (256 threads, ~256 tiles/row), all running at once; a decoupled look-back hands each tile's carry to the next via small atomics. **"single pass" = every element read once + written once** (many threadgroups, *not* one; data never re-read). |
| `[8, 1M]` dim=-1, int/long | `scan_block_reduce` + `scan_block_carry` | int can't use the float look-back, due to not having a sentinel value like float(NaN) to signal the "word" not being ready. Therefore we need **2 passes over the data**: pass 1 splits the row into blocks (one threadgroup each) and reduces each to a sum; pass 2 re-reads + scans each block, seeded by the running total of preceding block sums. |
| `[1M, 2]` dim=0, float | `scan_vec_decoupled` | the contig kernel, but each "element" is the 2-wide row read as one coalesced `float2`; one threadgroup per tile, look-back per component, 1 read + 1 write. |
| `[1M, 8]` dim=0, float | `scan_strided_col_decoupled` | **one simdgroup (32 lanes) per column** (16 strided reads/lane); 8 simdgroups per threadgroup claim adjacent columns so the strided reads coalesce; look-back per column, 1 read + 1 write. |
| `[1M, 8]` dim=0, int/long | `scan_strided_block_*` | int version of the column scan: **3 passes** (reduce blocks -> scan the block sums -> re-read + scan), with 32 x stride shared-memory tiles. |
| `[1M, 1]` dim=0 | reshape -> `scan_contig_decoupled` | trailing size-1 dim => stride 1 => identical to the contiguous case (row 1). |
| `[8192, 16]` dim=-1 | `scan_tiny_innermost` | rows too short to justify a threadgroup each, so **one threadgroup scans ~128 rows** (packs ~2048 elems into shared memory), serial-scanning each. |
| `[128, 4096]` dim=any; any `>4D`; `logcumsumexp`; complex | `scan_innermost_dim` / `scan_outer_dim` | **one threadgroup per scan line**, whole line scanned in threadgroup memory; the simple baseline, handles any dim/op. |
## tiny - scan_tiny_innermost
| dtype | shape | dim | Metal us | MPSGraph us | speedup |
|---|---|--:|--:|--:|--:|
| bf16 | 262144x16 | -1 | 80.0 | 413.7 | 5.17x |
| bf16 | 262144x32 | -1 | 188.8 | 415.6 | 2.20x |
| bf16 | 262144x8 | -1 | 41.0 | 412.9 | 10.08x |
| fp32 | 262144x16 | -1 | 116.2 | 347.6 | 2.99x |
| fp32 | 262144x32 | -1 | 260.2 | 358.3 | 1.38x |
| fp32 | 262144x8 | -1 | 41.6 | 345.5 | 8.30x |
| i32 | 262144x16 | -1 | 178.9 | 540.9 | 3.02x |
| i32 | 262144x32 | -1 | 466.7 | 755.4 | 1.62x |
| i32 | 262144x8 | -1 | 55.8 | 435.8 | 7.81x |
| i64 | 262144x16 | -1 | 246.4 | 1136.0 | 4.61x |
| i64 | 262144x32 | -1 | 555.5 | 1143.0 | 2.06x |
| i64 | 262144x8 | -1 | 110.0 | 1136.1 | 10.32x |
## contig - decoupled (float) / multiblock 3-pass (int)
| dtype | shape | dim | Metal us | MPSGraph us | speedup |
|---|---|--:|--:|--:|--:|
| bf16 | 512x65536 | -1 | 496.2 | 769.0 | 1.55x |
| bf16 | 64x1048576 | -1 | 988.4 | 1528.1 | 1.55x |
| bf16 | 8x1048576 | -1 | 118.2 | 153.2 | 1.30x |
| fp32 | 512x65536 | -1 | 984.1 | 1530.0 | 1.55x |
| fp32 | 64x1048576 | -1 | 2006.3 | 3075.1 | 1.53x |
| fp32 | 8x1048576 | -1 | 252.7 | 374.9 | 1.48x |
| i32 | 512x65536 | -1 | 2007.4 | 3133.6 | 1.56x |
| i32 | 64x1048576 | -1 | 4000.3 | 6216.6 | 1.55x |
| i32 | 8x1048576 | -1 | 502.0 | 769.2 | 1.53x |
| i64 | 512x65536 | -1 | 2850.1 | 3116.5 | 1.09x |
| i64 | 64x1048576 | -1 | 5688.6 | 6211.4 | 1.09x |
| i64 | 8x1048576 | -1 | 731.9 | 851.9 | 1.16x |
## inner_fb - scan_innermost_dim (generic)
| dtype | shape | dim | Metal us | MPSGraph us | speedup |
|---|---|--:|--:|--:|--:|
| bf16 | 1100x65536 | -1 | 1066.5 | 1632.6 | 1.53x |
| bf16 | 256x8192 | -1 | 24.0 | 35.9 | 1.49x |
| bf16 | 4096x4096 | -1 | 248.8 | 251.4 | 1.01x |
| fp32 | 1100x65536 | -1 | 2091.5 | 3284.4 | 1.57x |
| fp32 | 256x8192 | -1 | 36.0 | 56.5 | 1.57x |
| fp32 | 4096x4096 | -1 | 487.4 | 517.9 | 1.06x |
| i32 | 1100x65536 | -1 | 3306.3 | 6697.6 | 2.03x |
| i32 | 256x8192 | -1 | 62.2 | 145.6 | 2.34x |
| i32 | 4096x4096 | -1 | 768.1 | 1323.8 | 1.72x |
| i64 | 1100x65536 | -1 | 4183.8 | 6638.6 | 1.59x |
| i64 | 256x8192 | -1 | 114.0 | 243.2 | 2.13x |
| i64 | 4096x4096 | -1 | 972.1 | 1009.6 | 1.04x |
## vec - vec_decoupled (float) / strided_multiblock (int)
| dtype | shape | dim | Metal us | MPSGraph us | speedup |
|---|---|--:|--:|--:|--:|
| bf16 | 1048576x2 | 0 | 41.8 | 40.4 | 0.97x |
| fp32 | 1048576x2 | 0 | 50.3 | 66.7 | 1.33x |
| i32 | 1048576x2 | 0 | 90.3 | 150.0 | 1.66x |
| i64 | 1048576x2 | 0 | 152.1 | 176.3 | 1.16x |
## col - strided_col decoupled (float) / strided_multiblock (int)
| dtype | shape | dim | Metal us | MPSGraph us | speedup |
|---|---|--:|--:|--:|--:|
| bf16 | 1048576x32 | 0 | 909.2 | 1809.8 | 1.99x |
| bf16 | 1048576x8 | 0 | 136.0 | 218.4 | 1.61x |
| bf16 | 8192x256 | 0 | 49.9 | 88.7 | 1.78x |
| fp32 | 1048576x32 | 0 | 1295.2 | 3441.8 | 2.66x |
| fp32 | 1048576x8 | 0 | 262.4 | 492.9 | 1.88x |
| fp32 | 8192x256 | 0 | 52.2 | 90.2 | 1.73x |
| i32 | 1048576x32 | 0 | 2360.9 | 5134.7 | 2.17x |
| i32 | 1048576x8 | 0 | 541.3 | 949.8 | 1.75x |
| i32 | 8192x256 | 0 | 141.2 | 173.9 | 1.23x |
| i64 | 1048576x32 | 0 | 3031.4 | 5004.6 | 1.65x |
| i64 | 1048576x8 | 0 | 777.9 | 1124.8 | 1.45x |
| i64 | 8192x256 | 0 | 177.9 | 317.8 | 1.79x |
## outer_fb - scan_outer_dim (generic, axis<1024)
| dtype | shape | dim | Metal us | MPSGraph us | speedup |
|---|---|--:|--:|--:|--:|
| bf16 | 512x16384 | 0 | 127.3 | 186.6 | 1.47x |
| bf16 | 512x256 | 0 | 15.2 | 12.6 | 0.83x |
| fp32 | 512x16384 | 0 | 256.5 | 303.0 | 1.18x |
| fp32 | 512x256 | 0 | 15.0 | 12.3 | 0.82x |
| i32 | 512x16384 | 0 | 404.2 | 733.1 | 1.81x |
| i32 | 512x256 | 0 | 14.1 | 15.6 | 1.11x |
| i64 | 512x16384 | 0 | 504.4 | 739.7 | 1.47x |
| i64 | 512x256 | 0 | 14.5 | 27.6 | 1.90x |
## col|outer - strided_col (float) / scan_outer_dim (int, cols>=64)
| dtype | shape | dim | Metal us | MPSGraph us | speedup |
|---|---|--:|--:|--:|--:|
| bf16 | 1024x4096 | 0 | 86.1 | 106.6 | 1.24x |
| bf16 | 2048x8192 | 0 | 386.0 | 414.5 | 1.07x |
| fp32 | 1024x4096 | 0 | 138.0 | 134.9 | 0.98x |
| fp32 | 2048x8192 | 0 | 641.0 | 810.9 | 1.27x |
| i32 | 1024x4096 | 0 | 195.6 | 339.4 | 1.74x |
| i32 | 2048x8192 | 0 | 830.3 | 1710.4 | 2.06x |
| i64 | 1024x4096 | 0 | 253.4 | 323.1 | 1.27x |
| i64 | 2048x8192 | 0 | 1020.7 | 1741.3 | 1.71x |
## reshape - trailing size-1 dims -> contiguous scan
| dtype | shape | dim | Metal us | MPSGraph us | speedup |
|---|---|--:|--:|--:|--:|
| bf16 | 8x65536x1 | 1 | 10.5 | 12.8 | 1.22x |
| fp32 | 8x65536x1 | 1 | 13.9 | 16.3 | 1.17x |
| i32 | 8x65536x1 | 1 | 22.4 | 28.8 | 1.29x |
| i64 | 8x65536x1 | 1 | 31.1 | 37.1 | 1.19x |
## mid3d - middle-dim outer scan
| dtype | shape | dim | Metal us | MPSGraph us | speedup |
|---|---|--:|--:|--:|--:|
| bf16 | 16x70000x5 | 1 | 79.5 | 158.1 | 1.99x |
| bf16 | 64x4096x8 | 1 | 34.0 | 49.7 | 1.46x |
| fp32 | 16x70000x5 | 1 | 173.6 | 381.4 | 2.20x |
| fp32 | 64x4096x8 | 1 | 43.5 | 69.4 | 1.59x |
| i32 | 16x70000x5 | 1 | 504.3 | 857.8 | 1.70x |
| i32 | 64x4096x8 | 1 | 105.7 | 155.2 | 1.47x |
| i64 | 16x70000x5 | 1 | 533.6 | 1069.0 | 2.00x |
| i64 | 64x4096x8 | 1 | 180.6 | 308.4 | 1.71x |
## nc_in - non-contiguous (transposed) input
| dtype | shape | dim | Metal us | MPSGraph us | speedup |
|---|---|--:|--:|--:|--:|
| bf16 | 4096x4096 | -1 | 303.8 | 391.9 | 1.29x |
| fp32 | 4096x4096 | -1 | 540.8 | 673.8 | 1.25x |
| i32 | 4096x4096 | -1 | 840.2 | 1496.5 | 1.78x |
| i64 | 4096x4096 | -1 | 1019.0 | 1338.5 | 1.31x |
## model_seq - [batch, seq]
| dtype | shape | dim | Metal us | MPSGraph us | speedup |
|---|---|--:|--:|--:|--:|
| bf16 | 32x2048 | -1 | 3.1 | 11.7 | 3.79x |
| fp32 | 32x2048 | -1 | 2.8 | 11.8 | 4.17x |
| i32 | 32x2048 | -1 | 3.7 | 14.9 | 4.08x |
| i64 | 32x2048 | -1 | 3.7 | 23.7 | 6.37x |
## model_bhs - [batch, heads, seq]
| dtype | shape | dim | Metal us | MPSGraph us | speedup |
|---|---|--:|--:|--:|--:|
| bf16 | 16x16x4096 | -1 | 14.3 | 21.6 | 1.51x |
| fp32 | 16x16x4096 | -1 | 18.9 | 33.4 | 1.77x |
| i32 | 16x16x4096 | -1 | 28.6 | 58.7 | 2.05x |
| i64 | 16x16x4096 | -1 | 37.0 | 133.2 | 3.60x |
## model_mamba - [d_inner, seq] outer
| dtype | shape | dim | Metal us | MPSGraph us | speedup |
|---|---|--:|--:|--:|--:|
| bf16 | 8192x1024 | 0 | 188.2 | 340.8 | 1.81x |
| fp32 | 8192x1024 | 0 | 305.2 | 497.1 | 1.63x |
| i32 | 8192x1024 | 0 | 655.1 | 905.2 | 1.38x |
| i64 | 8192x1024 | 0 | 833.7 | 1174.8 | 1.41x |
## sweep_1Mx - [1M, n] dim0 inner-stride sweep
| dtype | shape | dim | Metal us | MPSGraph us | speedup |
|---|---|--:|--:|--:|--:|
| bf16 | 1048576x2 | 0 | 41.6 | 40.4 | 0.97x |
| bf16 | 1048576x3 | 0 | 67.6 | 64.3 | 0.95x |
| bf16 | 1048576x4 | 0 | 80.3 | 95.9 | 1.19x |
| bf16 | 1048576x8 | 0 | 136.5 | 216.3 | 1.58x |
| bf16 | 1048576x16 | 0 | 297.5 | 613.7 | 2.06x |
| bf16 | 1048576x24 | 0 | 502.2 | 961.4 | 1.91x |
| bf16 | 1048576x32 | 0 | 904.1 | 1544.7 | 1.71x |
| bf16 | 1048576x48 | 0 | 1239.9 | 2239.2 | 1.81x |
| bf16 | 1048576x64 | 0 | 1847.0 | 4507.1 | 2.44x |
| fp32 | 1048576x2 | 0 | 50.3 | 66.8 | 1.33x |
| fp32 | 1048576x3 | 0 | 90.7 | 94.2 | 1.04x |
| fp32 | 1048576x4 | 0 | 138.0 | 168.1 | 1.22x |
| fp32 | 1048576x8 | 0 | 263.3 | 491.8 | 1.87x |
| fp32 | 1048576x16 | 0 | 643.8 | 1133.0 | 1.76x |
| fp32 | 1048576x24 | 0 | 963.0 | 1813.9 | 1.88x |
| fp32 | 1048576x32 | 0 | 1297.6 | 3365.6 | 2.59x |
| fp32 | 1048576x48 | 0 | 1889.4 | 4937.8 | 2.61x |
| fp32 | 1048576x64 | 0 | 2588.0 | 7967.3 | 3.08x |
| i32 | 1048576x2 | 0 | 93.2 | 145.6 | 1.56x |
| i32 | 1048576x3 | 0 | 167.5 | 232.7 | 1.39x |
| i32 | 1048576x4 | 0 | 225.0 | 356.8 | 1.59x |
| i32 | 1048576x8 | 0 | 541.2 | 933.5 | 1.72x |
| i32 | 1048576x16 | 0 | 1162.5 | 2373.7 | 2.04x |
| i32 | 1048576x24 | 0 | 1969.4 | 3806.6 | 1.93x |
| i32 | 1048576x32 | 0 | 2367.8 | 5066.6 | 2.14x |
| i32 | 1048576x48 | 0 | 4325.3 | 7599.3 | 1.76x |
| i32 | 1048576x64 | 0 | 4637.8 | 11633.8 | 2.51x |
| i64 | 1048576x2 | 0 | 152.2 | 178.1 | 1.17x |
| i64 | 1048576x3 | 0 | 276.1 | 443.2 | 1.61x |
| i64 | 1048576x4 | 0 | 375.2 | 549.3 | 1.46x |
| i64 | 1048576x8 | 0 | 777.5 | 1143.2 | 1.47x |
| i64 | 1048576x16 | 0 | 1480.6 | 2349.6 | 1.59x |
| i64 | 1048576x24 | 0 | 2360.8 | 3733.4 | 1.58x |
| i64 | 1048576x32 | 0 | 3029.8 | 4991.4 | 1.65x |
| i64 | 1048576x48 | 0 | 4791.1 | 7845.6 | 1.64x |
| i64 | 1048576x64 | 0 | 6117.7 | 10750.9 | 1.76x |
## sweep_4x256k - [4, 256k, n] dim1 inner-stride sweep
| dtype | shape | dim | Metal us | MPSGraph us | speedup |
|---|---|--:|--:|--:|--:|
| bf16 | 4x262144x2 | 1 | 40.1 | 37.3 | 0.93x |
| bf16 | 4x262144x4 | 1 | 69.7 | 88.5 | 1.27x |
| bf16 | 4x262144x8 | 1 | 125.9 | 219.2 | 1.74x |
| bf16 | 4x262144x16 | 1 | 288.5 | 618.1 | 2.14x |
| bf16 | 4x262144x32 | 1 | 821.8 | 1627.7 | 1.98x |
| fp32 | 4x262144x2 | 1 | 50.3 | 64.1 | 1.27x |
| fp32 | 4x262144x4 | 1 | 121.7 | 170.8 | 1.40x |
| fp32 | 4x262144x8 | 1 | 257.7 | 486.7 | 1.89x |
| fp32 | 4x262144x16 | 1 | 572.3 | 1171.6 | 2.05x |
| fp32 | 4x262144x32 | 1 | 1239.0 | 3542.3 | 2.86x |
| i32 | 4x262144x2 | 1 | 91.7 | 148.0 | 1.61x |
| i32 | 4x262144x4 | 1 | 224.7 | 363.0 | 1.62x |
| i32 | 4x262144x8 | 1 | 543.2 | 936.7 | 1.72x |
| i32 | 4x262144x16 | 1 | 1192.4 | 2314.3 | 1.94x |
| i32 | 4x262144x32 | 1 | 2413.2 | 5172.1 | 2.14x |
| i64 | 4x262144x2 | 1 | 152.2 | 163.8 | 1.08x |
| i64 | 4x262144x4 | 1 | 372.6 | 524.1 | 1.41x |
| i64 | 4x262144x8 | 1 | 774.6 | 1132.0 | 1.46x |
| i64 | 4x262144x16 | 1 | 1482.6 | 2396.9 | 1.62x |
| i64 | 4x262144x32 | 1 | 3047.8 | 5210.7 | 1.71x |
## sweep_large - [axis, n] dim0 large-stride sweep
| dtype | shape | dim | Metal us | MPSGraph us | speedup |
|---|---|--:|--:|--:|--:|
| bf16 | 262144x64 | 0 | 424.3 | 704.8 | 1.66x |
| bf16 | 262144x128 | 0 | 852.3 | 1472.3 | 1.73x |
| bf16 | 262144x256 | 0 | 1736.2 | 3562.1 | 2.05x |
| bf16 | 131072x512 | 0 | 1602.9 | 3035.3 | 1.89x |
| bf16 | 131072x1024 | 0 | 3149.5 | 7540.6 | 2.39x |
| fp32 | 262144x64 | 0 | 627.6 | 1141.2 | 1.82x |
| fp32 | 262144x128 | 0 | 1265.6 | 3779.8 | 2.99x |
| fp32 | 262144x256 | 0 | 2476.1 | 7887.5 | 3.19x |
| fp32 | 131072x512 | 0 | 2450.2 | 7993.3 | 3.26x |
| fp32 | 131072x1024 | 0 | 5103.4 | 16029.0 | 3.14x |
| i32 | 262144x64 | 0 | 1244.2 | 2363.8 | 1.90x |
| i32 | 262144x128 | 0 | 2360.0 | 5240.1 | 2.22x |
| i32 | 262144x256 | 0 | 4567.1 | 11069.9 | 2.42x |
| i32 | 131072x512 | 0 | 4612.6 | 10710.7 | 2.32x |
| i32 | 131072x1024 | 0 | 9130.5 | 22414.9 | 2.45x |
| i64 | 262144x64 | 0 | 1597.2 | 2493.7 | 1.56x |
| i64 | 262144x128 | 0 | 3048.3 | 5219.8 | 1.71x |
| i64 | 262144x256 | 0 | 6004.4 | 10937.7 | 1.82x |
| i64 | 131072x512 | 0 | 6041.5 | 10788.0 | 1.79x |
| i64 | 131072x1024 | 0 | 12158.0 | 22534.8 | 1.85x |
Pull Request resolved: https://github.com/pytorch/pytorch/pull/185609
Approved by: https://github.com/malfet
