- read the 8 Q elements each lane needs straight from global into the mma A-operand layout, pre-scaled, instead of staging through shared sQ - removes the sQ smem area (20KB->16KB) and the serialized per-warp prologue with its WARPS __syncthreads barriers - result vs torch SDPA: prefill 0.70-0.82x -> 0.85-1.00x (matches torch at seq=128; 2048 1.251->1.159ms), correctness unchanged across head dims |
||
|---|---|---|
| .. | ||
| kernels | ||
| tests | ||
| __init__.py | ||
| build.py | ||