Fixes
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:
Community-Analysen & Experten-Meinungen 0
Verwandte Story-Cluster & Quellen (Vektor-KI)
Ähnliche Beiträge
Auch interessante Nachrichten viable/strict/1781045526: [MPS] Metal cumsum cumprod kernels (#185609)
Thematisch verwandte Begriffe: viablestrict1781045526, Metal, cumsum, cumprod · 6 Treffer
kaspersky kernel vuln windows 11
[webapps] C-MOR 6.0104 - Cross-Site Scripting (XSS)
[webapps] Bludit CMS - Stored XSS
Videos werden geladen ...
Beiträge werden geladen ...
Videos werden geladen ...
Beiträge werden geladen ...
Videos werden geladen ...
Beiträge werden geladen ...
Videos werden geladen ...
Beiträge werden geladen ...
Videos werden geladen ...
SOCIAL SHARE CARD GENERATOR