diff --git a/BENCHMARK_H200.md b/BENCHMARK_H200.md index ab9676ce..bae08fc0 100644 --- a/BENCHMARK_H200.md +++ b/BENCHMARK_H200.md @@ -1,6 +1,6 @@ # Benchmark Results — Hopper (SM90) -> Measured on 2026-08-09. +> Measured on 2026-08-12. > **GPU:** NVIDIA H200 141GB, 132 SMs | **CUDA:** 12.9 | **PyTorch:** 2.9.1+cu129 | **Triton:** 3.5.1 @@ -19,42 +19,43 @@ iterations, and 100 measured iterations aggregated with the IQR mean. | B | T | FLA Triton (ms) | cuLA FlashKDA (ms) | Speedup | |---|---:|---:|---:|---:| -| 1 | 512 | 1.0013 | 0.1284 | **7.80x** | -| 1 | 1024 | 0.9747 | 0.1832 | **5.32x** | -| 1 | 4096 | 0.9994 | 0.6537 | **1.53x** | -| 1 | 8192 | 1.7774 | 1.2831 | **1.39x** | -| 1 | 16384 | 3.5038 | 2.5421 | **1.38x** | -| 2 | 512 | 0.9855 | 0.1386 | **7.11x** | -| 2 | 1024 | 1.0283 | 0.2447 | **4.20x** | -| 2 | 4096 | 1.7690 | 0.9010 | **1.96x** | -| 2 | 8192 | 3.4866 | 1.7810 | **1.96x** | -| 2 | 16384 | 6.9413 | 3.5363 | **1.96x** | +| 1 | 512 | 0.9405 | 0.1244 | **7.56x** | +| 1 | 1024 | 0.9372 | 0.1579 | **5.94x** | +| 1 | 4096 | 0.9344 | 0.5701 | **1.64x** | +| 1 | 8192 | 1.7701 | 1.1150 | **1.59x** | +| 1 | 16384 | 3.4846 | 2.2266 | **1.56x** | +| 2 | 512 | 0.9274 | 0.1262 | **7.35x** | +| 2 | 1024 | 0.9300 | 0.1971 | **4.72x** | +| 2 | 4096 | 1.7628 | 0.7223 | **2.44x** | +| 2 | 8192 | 3.4622 | 1.4306 | **2.42x** | +| 2 | 16384 | 6.8872 | 2.8817 | **2.39x** | ### Variable-Length | Config | FLA Triton (ms) | cuLA FlashKDA (ms) | Speedup | |---|---:|---:|---:| -| uniform 10 seqs, T=4096, 409–415 tokens | 1.0295 | 0.4748 | **2.17x** | -| random 10 seqs, T=4096, 24–1201 tokens | 1.0378 | 0.5544 | **1.87x** | -| skewed 10 seqs, T=4096, 227–2053 tokens | 1.0126 | 0.6700 | **1.51x** | -| uniform 20 seqs, T=4096, 204–220 tokens | 1.0510 | 0.4639 | **2.27x** | -| random 20 seqs, T=4096, 5–787 tokens | 1.0250 | 0.5513 | **1.86x** | -| skewed 20 seqs, T=4096, 107–2063 tokens | 1.0048 | 0.6767 | **1.48x** | -| uniform 10 seqs, T=8192, 819–821 tokens | 1.7923 | 0.9051 | **1.98x** | -| random 10 seqs, T=8192, 48–2401 tokens | 1.8277 | 1.0525 | **1.74x** | -| skewed 10 seqs, T=8192, 455–4097 tokens | 1.8457 | 1.2728 | **1.45x** | -| uniform 20 seqs, T=8192, 409–421 tokens | 1.8855 | 0.8628 | **2.19x** | -| random 20 seqs, T=8192, 9–1574 tokens | 1.8785 | 1.0189 | **1.84x** | -| skewed 20 seqs, T=8192, 215–4107 tokens | 1.8875 | 1.2913 | **1.46x** | -| uniform 10 seqs, T=16384, 1638–1642 tokens | 3.5091 | 1.7505 | **2.00x** | -| random 10 seqs, T=16384, 95–4802 tokens | 3.5175 | 2.0473 | **1.72x** | -| skewed 10 seqs, T=16384, 910–8194 tokens | 3.5350 | 2.4892 | **1.42x** | -| uniform 20 seqs, T=16384, 819–823 tokens | 3.5335 | 1.6558 | **2.13x** | -| random 20 seqs, T=16384, 19–3147 tokens | 3.5694 | 1.9728 | **1.81x** | -| skewed 20 seqs, T=16384, 431–8195 tokens | 3.5403 | 2.4949 | **1.42x** | - -Across all 28 fixed-length and variable-length configs, FlashKDA averages -**2.39x** over FLA (minimum **1.38x**, maximum **7.80x**). +| uniform 10 seqs, T=4096, 409–415 tokens | 0.9819 | 0.3802 | **2.58x** | +| random 10 seqs, T=4096, 24–1201 tokens | 0.9749 | 0.4574 | **2.13x** | +| skewed 10 seqs, T=4096, 227–2053 tokens | 0.9730 | 0.5753 | **1.69x** | +| uniform 20 seqs, T=4096, 204–220 tokens | 1.0478 | 0.3768 | **2.78x** | +| random 20 seqs, T=4096, 5–787 tokens | 1.0242 | 0.4426 | **2.31x** | +| skewed 20 seqs, T=4096, 107–2063 tokens | 0.9984 | 0.5940 | **1.68x** | +| uniform 10 seqs, T=8192, 819–821 tokens | 1.7868 | 0.7329 | **2.44x** | +| random 10 seqs, T=8192, 48–2401 tokens | 1.8201 | 0.8595 | **2.12x** | +| skewed 10 seqs, T=8192, 455–4097 tokens | 1.8388 | 1.0881 | **1.69x** | +| uniform 20 seqs, T=8192, 409–421 tokens | 1.8781 | 0.6879 | **2.73x** | +| random 20 seqs, T=8192, 9–1574 tokens | 1.8714 | 0.8146 | **2.30x** | +| skewed 20 seqs, T=8192, 215–4107 tokens | 1.8815 | 1.1248 | **1.67x** | +| uniform 10 seqs, T=16384, 1638–1642 tokens | 3.5032 | 1.4460 | **2.42x** | +| random 10 seqs, T=16384, 95–4802 tokens | 3.5105 | 1.6967 | **2.07x** | +| skewed 10 seqs, T=16384, 910–8194 tokens | 3.5259 | 2.1618 | **1.63x** | +| uniform 20 seqs, T=16384, 819–823 tokens | 3.5272 | 1.3552 | **2.60x** | +| random 20 seqs, T=16384, 19–3147 tokens | 3.5600 | 1.5926 | **2.24x** | +| skewed 20 seqs, T=16384, 431–8195 tokens | 3.5296 | 2.2084 | **1.60x** | + +Across all 28 fixed-length and variable-length configs, the arithmetic mean of +the per-row speedups is **2.72x** over FLA (geometric mean **2.43x**; +minimum **1.56x**, maximum **7.56x**). To reproduce: @@ -72,41 +73,41 @@ H200's SM array more effectively. Each row uses BF16, `D=128`, | Config | H | CP engaged | CP off (ms) | CP auto (ms) | Speedup | |---|---:|:---:|---:|---:|---:| -| T=1023 | 4 | no | 0.1322 | 0.1447 | 0.91x | -| T=1025 | 4 | no | 0.1364 | 0.1388 | 0.98x | -| T=4K | 4 | yes | 0.4374 | 0.2163 | **2.02x** | -| T=8K | 4 | yes | 0.8636 | 0.2820 | **3.06x** | -| T=16K | 4 | yes | 1.7099 | 0.4078 | **4.19x** | -| T=32K | 4 | yes | 3.3888 | 0.5973 | **5.67x** | -| T=64K | 4 | yes | 6.7518 | 0.9658 | **6.99x** | -| T=64K+1 | 4 | yes | 7.3033 | 1.0126 | **7.21x** | -| 2x16K | 4 | yes | 1.7631 | 0.5436 | **3.24x** | -| 32K+4K | 4 | yes | 3.4000 | 0.6682 | **5.09x** | -| 32K+1K | 4 | yes | 3.3923 | 0.6046 | **5.61x** | -| 32K+1023+1025 | 4 | yes | 3.6605 | 0.6821 | **5.37x** | -| 64K+1K | 4 | yes | 6.7713 | 1.0644 | **6.36x** | -| 64K+2x1K | 4 | yes | 6.7412 | 1.0803 | **6.24x** | -| 64K+5x1K | 4 | yes | 6.7648 | 1.1064 | **6.11x** | -| 64K+1+1023+1025 | 4 | yes | 7.3010 | 1.1200 | **6.52x** | -| T=1023 | 8 | no | 0.1365 | 0.1395 | 0.98x | -| T=1025 | 8 | no | 0.1376 | 0.1391 | 0.99x | -| T=4K | 8 | yes | 0.4557 | 0.2356 | **1.93x** | -| T=8K | 8 | yes | 0.8941 | 0.3177 | **2.81x** | -| T=16K | 8 | yes | 1.7600 | 0.5132 | **3.43x** | -| T=32K | 8 | yes | 3.4958 | 0.8774 | **3.98x** | -| T=64K | 8 | yes | 6.9656 | 1.6239 | **4.29x** | -| T=64K+1 | 8 | yes | 7.5212 | 1.7221 | **4.37x** | -| 2x16K | 8 | yes | 1.8764 | 0.8328 | **2.25x** | -| 32K+4K | 8 | yes | 3.5256 | 1.0458 | **3.37x** | -| 32K+1K | 8 | yes | 3.5080 | 0.9871 | **3.55x** | -| 32K+1023+1025 | 8 | yes | 3.7879 | 1.0731 | **3.53x** | -| 64K+1K | 8 | yes | 6.9812 | 1.8345 | **3.81x** | -| 64K+2x1K | 8 | yes | 6.9297 | 1.8368 | **3.77x** | -| 64K+5x1K | 8 | yes | 6.9738 | 1.9318 | **3.61x** | -| 64K+1+1023+1025 | 8 | yes | 7.4816 | 1.9172 | **3.90x** | +| T=1023 | 4 | no | 0.1308 | 0.1456 | 0.90x | +| T=1025 | 4 | no | 0.1312 | 0.1362 | 0.96x | +| T=4K | 4 | yes | 0.4070 | 0.2105 | **1.93x** | +| T=8K | 4 | yes | 0.8141 | 0.2659 | **3.06x** | +| T=16K | 4 | yes | 1.6149 | 0.3792 | **4.26x** | +| T=32K | 4 | yes | 3.2087 | 0.5455 | **5.88x** | +| T=64K | 4 | yes | 6.3907 | 0.8634 | **7.40x** | +| T=64K+1 | 4 | yes | 7.3199 | 0.9349 | **7.83x** | +| 2x16K | 4 | yes | 1.6487 | 0.4960 | **3.32x** | +| 32K+4K | 4 | yes | 3.2272 | 0.6156 | **5.24x** | +| 32K+1K | 4 | yes | 3.2135 | 0.5434 | **5.91x** | +| 32K+1023+1025 | 4 | yes | 3.6700 | 0.6341 | **5.79x** | +| 64K+1K | 4 | yes | 6.4081 | 0.9721 | **6.59x** | +| 64K+2x1K | 4 | yes | 6.4054 | 0.9865 | **6.49x** | +| 64K+5x1K | 4 | yes | 6.4184 | 1.0059 | **6.38x** | +| 64K+1+1023+1025 | 4 | yes | 7.3167 | 1.0531 | **6.95x** | +| T=1023 | 8 | no | 0.1326 | 0.1361 | 0.97x | +| T=1025 | 8 | no | 0.1335 | 0.1362 | 0.98x | +| T=4K | 8 | yes | 0.4246 | 0.2174 | **1.95x** | +| T=8K | 8 | yes | 0.8376 | 0.2854 | **2.94x** | +| T=16K | 8 | yes | 1.6496 | 0.4520 | **3.65x** | +| T=32K | 8 | yes | 3.2839 | 0.7789 | **4.22x** | +| T=64K | 8 | yes | 6.5424 | 1.4323 | **4.57x** | +| T=64K+1 | 8 | yes | 7.4745 | 1.5786 | **4.73x** | +| 2x16K | 8 | yes | 1.7236 | 0.7318 | **2.36x** | +| 32K+4K | 8 | yes | 3.3146 | 0.9499 | **3.49x** | +| 32K+1K | 8 | yes | 3.2958 | 0.8917 | **3.70x** | +| 32K+1023+1025 | 8 | yes | 3.7542 | 0.9980 | **3.76x** | +| 64K+1K | 8 | yes | 6.5675 | 1.6512 | **3.98x** | +| 64K+2x1K | 8 | yes | 6.5733 | 1.6640 | **3.95x** | +| 64K+5x1K | 8 | yes | 6.5782 | 1.7464 | **3.77x** | +| 64K+1+1023+1025 | 8 | yes | 7.4786 | 1.7998 | **4.16x** | Intracard CP engages for 28 of the 32 configs. On the engaged subset it -delivers a **4.11x geometric-mean** speedup, with a range of **1.93–7.21x**. +delivers a **4.29x geometric-mean** speedup, with a range of **1.93–7.83x**. The non-CHUNK-aligned `T=64K+1` cases retain essentially the same benefit as `T=64K`; the four 1023/1025-token rows stay on the serial path and expose only the auto-planner overhead. diff --git a/README.md b/README.md index b154782a..91b25a08 100644 --- a/README.md +++ b/README.md @@ -162,8 +162,8 @@ See [BENCHMARK_H200.md](BENCHMARK_H200.md) for CuTe DSL FlashKDA results on an H - **KDA Modular Forward (Blackwell):** **avg 1.33x** speedup on fixed-length, **avg 1.35x** on variable-length (18 configs, uniform/skewed/random). - **Lightning Attention Prefill (Blackwell):** up to **2.08x** speedup (B=2). - **Lightning Attention Varlen (Blackwell):** **avg 1.47x** speedup across 126 configs (uniform/skewed/random). -- **FlashKDA Prefill (Hopper):** **avg 2.39x** speedup over FLA across 28 fixed-length and variable-length configs, up to **7.80x**. -- **FlashKDA Intracard CP (Hopper):** **4.11x geo-mean** speedup over serial FlashKDA on 28 CP-engaged long-sequence configs, up to **7.21x**. +- **FlashKDA Prefill (Hopper):** **avg 2.72x** speedup over FLA across 28 fixed-length and variable-length configs, up to **7.56x**. +- **FlashKDA Intracard CP (Hopper):** **4.29x geo-mean** speedup over serial FlashKDA on 28 CP-engaged long-sequence configs, up to **7.83x**. To reproduce the benchmark suites directly: diff --git a/cula/ops/kda/sm90/_common.py b/cula/ops/kda/sm90/_common.py index 8044a4c6..a24beb60 100644 --- a/cula/ops/kda/sm90/_common.py +++ b/cula/ops/kda/sm90/_common.py @@ -3,13 +3,55 @@ """Shared low-level helpers for the SM90 FlashKDA kernels.""" +import re +from importlib.metadata import PackageNotFoundError +from importlib.metadata import version as package_version + import cutlass +import cutlass.cute as cute import torch from cutlass import Int32 from cutlass._mlir.dialects import llvm as _llvm from cutlass.cutlass_dsl import T as _T +def _parse_cutedsl_version(raw_version: str) -> tuple[int, int, int]: + """Return the numeric CuTeDSL version from a release or dev version string.""" + match = re.match(r"^(\d+)\.(\d+)(?:\.(\d+))?", raw_version) + if match is None: + raise RuntimeError(f"Unable to parse the installed CuTeDSL version: {raw_version!r}") + return tuple(int(component or 0) for component in match.groups()) + + +def _installed_cutedsl_version() -> tuple[int, int, int]: + """Read the CuTeDSL version at runtime without relying on a private symbol.""" + raw_version = getattr(cutlass, "__version__", None) + if raw_version is None: + try: + raw_version = package_version("nvidia-cutlass-dsl") + except PackageNotFoundError as exc: + raise RuntimeError("nvidia-cutlass-dsl is required by the SM90 FlashKDA backend") from exc + return _parse_cutedsl_version(raw_version) + + +# CuTeDSL 4.6.0 added the missing elect_one inside cute.copy for async bulk +# atoms. Older releases need an explicit elect_one, while nesting one around +# cute.copy is incorrect in 4.6+. Keep this decision as a compile-time +# constant after detecting the installed runtime version, so the same source +# supports both API behaviours. +_CUTEDSL_VERSION = _installed_cutedsl_version() +_CUTE_COPY_AUTO_ELECTS_BULK = _CUTEDSL_VERSION >= (4, 6, 0) + + +def copy_async_bulk(atom, src, dst, **kwargs) -> None: + """Issue a CuTeDSL async bulk copy across supported elect_one APIs.""" + if cutlass.const_expr(_CUTE_COPY_AUTO_ELECTS_BULK): + cute.copy(atom, src, dst, **kwargs) + else: + with cute.arch.elect_one(): + cute.copy(atom, src, dst, **kwargs) + + def _stream_key(device: torch.device) -> tuple[str, int]: return str(device), int(torch.cuda.current_stream(device).cuda_stream) diff --git a/cula/ops/kda/sm90/cp/pre_scan.py b/cula/ops/kda/sm90/cp/pre_scan.py index 900b3335..da40f663 100644 --- a/cula/ops/kda/sm90/cp/pre_scan.py +++ b/cula/ops/kda/sm90/cp/pre_scan.py @@ -30,7 +30,7 @@ from cutlass.cute.nvgpu.warpgroup import SmemLayoutAtomKind, make_smem_layout_atom from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream -from cula.ops.kda.sm90._common import movm_t_b16 +from cula.ops.kda.sm90._common import copy_async_bulk, movm_t_b16 from cula.ops.kda.sm90.k2 import ( CHUNK, D, @@ -48,16 +48,14 @@ def pre_scan_kernel( tma_atom_v: cute.CopyAtom, tma_tensor_v: cute.Tensor, - tma_atom_kd: cute.CopyAtom, - tma_tensor_kd: cute.Tensor, - tma_atom_kr: cute.CopyAtom, - tma_tensor_kr: cute.Tensor, tma_atom_inv: cute.CopyAtom, tma_tensor_inv: cute.Tensor, tma_atom_gt: cute.CopyAtom, tma_tensor_gt: cute.Tensor, tma_atom_beta: cute.CopyAtom, tma_tensor_beta: cute.Tensor, + ws_kd: cute.Tensor, + ws_kr: cute.Tensor, H: cutlass.Constexpr[int], total_tiles: cutlass.Int32, T_total: cutlass.Int32, @@ -118,22 +116,6 @@ def pre_scan_kernel( cute.group_modes(sV, 0, 2), cute.group_modes(gSrc_v, 0, 2), ) - gSrc_kd = cute.local_tile(tma_tensor_kd, (CHUNK, D), (None, None, None)) - tKDs, tKDg = cpasync.tma_partition( - tma_atom_kd, - 0, - cute.make_layout(1), - cute.group_modes(sKd, 0, 2), - cute.group_modes(gSrc_kd, 0, 2), - ) - gSrc_kr = cute.local_tile(tma_tensor_kr, (CHUNK, D), (None, None, None)) - tKRs, tKRg = cpasync.tma_partition( - tma_atom_kr, - 0, - cute.make_layout(1), - cute.group_modes(sKr, 0, 2), - cute.group_modes(gSrc_kr, 0, 2), - ) gSrc_inv = cute.local_tile(tma_tensor_inv, (CHUNK, CHUNK), (None, None, None)) tIs, tIg = cpasync.tma_partition( tma_atom_inv, @@ -158,6 +140,19 @@ def pre_scan_kernel( cute.group_modes(sBeta, 0, 2), cute.group_modes(gSrc_beta, 0, 2), ) + raw_copy_atom = cute.make_copy_atom(cpasync.CopyBulkG2SOp(), cutlass.BFloat16) + raw_stage_layout = cute.make_layout( + (CHUNK * D, STAGES), + stride=(1, CHUNK * D), + ) + raw_gmem_layout = cute.make_layout( + (CHUNK * D, total_tiles * H), + stride=(1, CHUNK * D), + ) + sKD_raw = cute.make_tensor(sKd.iterator, raw_stage_layout) + sKR_raw = cute.make_tensor(sKr.iterator, raw_stage_layout) + gKD_raw = cute.make_tensor(ws_kd.iterator, raw_gmem_layout) + gKR_raw = cute.make_tensor(ws_kr.iterator, raw_gmem_layout) # sState=0, sM=I if tidx < D: @@ -273,8 +268,8 @@ def pre_scan_kernel( cute.copy(tma_atom_v, tVg_seq[(None, t, 0, head_idx)], tVs_seq[(None, s_dyn_l)], tma_bar_ptr=bar_l) else: cute.copy(tma_atom_v, tVg[(None, tg_l, 0, head_idx)], tVs[(None, s_dyn_l)], tma_bar_ptr=bar_l) - cute.copy(tma_atom_kd, tKDg[(None, 0, 0, wt_l)], tKDs[(None, s_dyn_l)], tma_bar_ptr=bar_l) - cute.copy(tma_atom_kr, tKRg[(None, 0, 0, wt_l)], tKRs[(None, s_dyn_l)], tma_bar_ptr=bar_l) + copy_async_bulk(raw_copy_atom, gKD_raw[(None, wt_l)], sKD_raw[(None, s_dyn_l)], mbar_ptr=bar_l) + copy_async_bulk(raw_copy_atom, gKR_raw[(None, wt_l)], sKR_raw[(None, s_dyn_l)], mbar_ptr=bar_l) cute.copy(tma_atom_inv, tIg[(None, 0, 0, wt_l)], tIs[(None, s_dyn_l)], tma_bar_ptr=bar_l) cute.copy(tma_atom_gt, tGTg[(None, 0, 0, wt_l)], tGTs[(None, s_dyn_l)], tma_bar_ptr=bar_l) cute.copy(tma_atom_beta, tBg[(None, 0, 0, wt_l)], tBs[(None, s_dyn_l)], tma_bar_ptr=bar_l) @@ -487,13 +482,6 @@ def make_thd_atom(t, op): ) return cpasync.make_tiled_tma_atom(op, view, kinter_smem, (CHUNK, D)) - def make_ws_qkd_atom(t): - view = cute.make_tensor( - t.iterator, - cute.make_layout((CHUNK, D, total_tiles * H), stride=(D, 1, CHUNK * D)), - ) - return cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileG2SOp(), view, kinter_smem, (CHUNK, D)) - def make_ws_cc_atom(t): view = cute.make_tensor( t.iterator, @@ -502,8 +490,6 @@ def make_ws_cc_atom(t): return cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileG2SOp(), view, cc_smem, (CHUNK, CHUNK)) tma_atom_v, tma_tensor_v = make_thd_atom(v, cpasync.CopyBulkTensorTileG2SOp()) - tma_atom_kd, tma_tensor_kd = make_ws_qkd_atom(ws_kd) - tma_atom_kr, tma_tensor_kr = make_ws_qkd_atom(ws_kr) tma_atom_inv, tma_tensor_inv = make_ws_cc_atom(ws_inv) gt_smem = cute.make_layout((D, 1), stride=(1, D)) @@ -551,16 +537,14 @@ def make_beta_atom(t): pre_scan_kernel( tma_atom_v, tma_tensor_v, - tma_atom_kd, - tma_tensor_kd, - tma_atom_kr, - tma_tensor_kr, tma_atom_inv, tma_tensor_inv, tma_atom_gt, tma_tensor_gt, tma_atom_beta, tma_tensor_beta, + ws_kd, + ws_kr, H, total_tiles, T_total, diff --git a/cula/ops/kda/sm90/k1.py b/cula/ops/kda/sm90/k1.py index 8b2fdd17..34aa1e43 100644 --- a/cula/ops/kda/sm90/k1.py +++ b/cula/ops/kda/sm90/k1.py @@ -30,7 +30,7 @@ from cutlass.cute.nvgpu.warpgroup import SmemLayoutAtomKind, make_smem_layout_atom from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream -from cula.ops.kda.sm90._common import _stream_key, add_f16x2_u32, movm_t_b16 +from cula.ops.kda.sm90._common import _stream_key, add_f16x2_u32, copy_async_bulk, movm_t_b16 CHUNK: int = 16 D: int = 128 @@ -45,16 +45,13 @@ def k1_kernel( tma_tensor_k: cute.Tensor, tma_atom_g: cute.CopyAtom, tma_tensor_g: cute.Tensor, - tma_atom_ws_qd: cute.CopyAtom, - tma_tensor_ws_qd: cute.Tensor, - tma_atom_ws_kd: cute.CopyAtom, - tma_tensor_ws_kd: cute.Tensor, - tma_atom_ws_kr: cute.CopyAtom, - tma_tensor_ws_kr: cute.Tensor, tma_atom_ws_inv: cute.CopyAtom, tma_tensor_ws_inv: cute.Tensor, tma_atom_ws_mqk: cute.CopyAtom, tma_tensor_ws_mqk: cute.Tensor, + ws_qd: cute.Tensor, + ws_kd: cute.Tensor, + ws_kr: cute.Tensor, a_log: cute.Tensor, dt_bias: cute.Tensor, beta: cute.Tensor, @@ -148,30 +145,6 @@ def k1_kernel( cute.group_modes(gSrc_g, 0, 2), ) - gDst_qd = cute.local_tile(tma_tensor_ws_qd, (CHUNK, D), (None, None, None)) - tQDws_s, tQDws_g = cpasync.tma_partition( - tma_atom_ws_qd, - 0, - cute.make_layout(1), - cute.group_modes(s_q_decayed, 0, 2), - cute.group_modes(gDst_qd, 0, 2), - ) - gDst_kd = cute.local_tile(tma_tensor_ws_kd, (CHUNK, D), (None, None, None)) - tKDws_s, tKDws_g = cpasync.tma_partition( - tma_atom_ws_kd, - 0, - cute.make_layout(1), - cute.group_modes(s_k_decayed, 0, 2), - cute.group_modes(gDst_kd, 0, 2), - ) - gDst_kr = cute.local_tile(tma_tensor_ws_kr, (CHUNK, D), (None, None, None)) - tKRws_s, tKRws_g = cpasync.tma_partition( - tma_atom_ws_kr, - 0, - cute.make_layout(1), - cute.group_modes(s_k_restored, 0, 2), - cute.group_modes(gDst_kr, 0, 2), - ) gDst_inv = cute.local_tile(tma_tensor_ws_inv, (CHUNK, CHUNK), (None, None, None)) tINVws_s, tINVws_g = cpasync.tma_partition( tma_atom_ws_inv, @@ -189,6 +162,18 @@ def k1_kernel( cute.group_modes(gDst_mqk, 0, 2), ) ws_slot = head_idx * total_tiles + tile_idx + raw_copy_atom = cute.make_copy_atom(cpasync.CopyBulkS2GOp(), cutlass.BFloat16) + raw_smem_layout = cute.make_layout((CHUNK * D,), stride=(1,)) + raw_gmem_layout = cute.make_layout( + (CHUNK * D, total_tiles * H), + stride=(1, CHUNK * D), + ) + sQDws_raw = cute.make_tensor(s_q_decayed.iterator, raw_smem_layout) + sKDws_raw = cute.make_tensor(s_k_decayed.iterator, raw_smem_layout) + sKRws_raw = cute.make_tensor(s_k_restored.iterator, raw_smem_layout) + gQDws_raw = cute.make_tensor(ws_qd.iterator, raw_gmem_layout) + gKDws_raw = cute.make_tensor(ws_kd.iterator, raw_gmem_layout) + gKRws_raw = cute.make_tensor(ws_kr.iterator, raw_gmem_layout) if warp_idx == 0: with cute.arch.elect_one(): @@ -533,12 +518,16 @@ def k1_kernel( cute.arch.fence_view_async_shared() cute.arch.barrier() - # TMA bulk store all 5 workspace tensors (one elect_one, one thread). + # Preserve the physical K_INTER byte image for qd/kd/kr; inv/mqk remain + # layout-aware tensor TMA stores. This raw-workspace transport idea comes + # from Flash-Flash-KDA: https://github.com/Itssshikhar/Flash-Flash-KDA if warp_idx == 0: + # copy_async_bulk supplies elect_one only for pre-4.6 CuTeDSL. CuTeDSL + # 4.6+ emits it inside cute.copy, where an outer elect_one is invalid. + copy_async_bulk(raw_copy_atom, sQDws_raw, gQDws_raw[(None, ws_slot)]) + copy_async_bulk(raw_copy_atom, sKDws_raw, gKDws_raw[(None, ws_slot)]) + copy_async_bulk(raw_copy_atom, sKRws_raw, gKRws_raw[(None, ws_slot)]) with cute.arch.elect_one(): - cute.copy(tma_atom_ws_qd, tQDws_s[(None,)], tQDws_g[(None, 0, 0, ws_slot)]) - cute.copy(tma_atom_ws_kd, tKDws_s[(None,)], tKDws_g[(None, 0, 0, ws_slot)]) - cute.copy(tma_atom_ws_kr, tKRws_s[(None,)], tKRws_g[(None, 0, 0, ws_slot)]) cute.copy(tma_atom_ws_inv, tINVws_s[(None,)], tINVws_g[(None, 0, 0, ws_slot)]) cute.copy(tma_atom_ws_mqk, tMQKws_s[(None,)], tMQKws_g[(None, 0, 0, ws_slot)]) cute.arch.cp_async_bulk_commit_group() @@ -571,9 +560,6 @@ def run_k1( stream: cuda_drv.CUstream, ): smem_layout_qk = cute.make_layout((CHUNK, D), stride=(D, 1)) - # K_INTER swizzled layout — must match kernel SMEM layout for TMA stores. - kinter_atom = make_smem_layout_atom(SmemLayoutAtomKind.K_INTER, cutlass.BFloat16) - smem_layout_qk_kinter = cute.tile_to_shape(kinter_atom, (CHUNK, D), order=(0, 1)) def make_atom(t): view = cute.make_tensor( @@ -587,18 +573,6 @@ def make_atom(t): (CHUNK, D), ) - def make_ws_store_atom(t): - view = cute.make_tensor( - t.iterator, - cute.make_layout((CHUNK, D, total_tiles * H), stride=(D, 1, CHUNK * D)), - ) - return cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileS2GOp(), - view, - smem_layout_qk_kinter, - (CHUNK, D), - ) - # (CHUNK, CHUNK) bf16 plain layout for ws_inv / ws_mqk TMA bulk store. smem_layout_cc = cute.make_layout((CHUNK, CHUNK), stride=(CHUNK, 1)) @@ -620,9 +594,6 @@ def make_ws_cc_store_atom(t): tma_atom_q, tma_tensor_q = make_atom(q) tma_atom_k, tma_tensor_k = make_atom(k) tma_atom_g, tma_tensor_g = make_atom(g) - tma_atom_ws_qd, tma_tensor_ws_qd = make_ws_store_atom(ws_qd) - tma_atom_ws_kd, tma_tensor_ws_kd = make_ws_store_atom(ws_kd) - tma_atom_ws_kr, tma_tensor_ws_kr = make_ws_store_atom(ws_kr) tma_atom_ws_inv, tma_tensor_ws_inv = make_ws_cc_store_atom(ws_inv) tma_atom_ws_mqk, tma_tensor_ws_mqk = make_ws_cc_store_atom(ws_mqk) @@ -635,16 +606,13 @@ def make_ws_cc_store_atom(t): tma_tensor_k, tma_atom_g, tma_tensor_g, - tma_atom_ws_qd, - tma_tensor_ws_qd, - tma_atom_ws_kd, - tma_tensor_ws_kd, - tma_atom_ws_kr, - tma_tensor_ws_kr, tma_atom_ws_inv, tma_tensor_ws_inv, tma_atom_ws_mqk, tma_tensor_ws_mqk, + ws_qd, + ws_kd, + ws_kr, a_log, dt_bias, beta, diff --git a/cula/ops/kda/sm90/k2.py b/cula/ops/kda/sm90/k2.py index 12ccbe3c..1cb60cc5 100644 --- a/cula/ops/kda/sm90/k2.py +++ b/cula/ops/kda/sm90/k2.py @@ -44,7 +44,7 @@ def _make_state_smem_layout(): return cute.tile_to_shape(atom, (D, D), (0, 1)) -from cula.ops.kda.sm90._common import _stream_key, movm_t_b16 # noqa: E402 +from cula.ops.kda.sm90._common import _stream_key, copy_async_bulk, movm_t_b16 # noqa: E402 def _make_out_kinter_one_stage(): @@ -63,16 +63,13 @@ def _make_out_kinter_one_stage(): def k2_kernel( tma_atom_v: cute.CopyAtom, tma_tensor_v: cute.Tensor, - tma_atom_kd: cute.CopyAtom, - tma_tensor_kd: cute.Tensor, - tma_atom_qd: cute.CopyAtom, - tma_tensor_qd: cute.Tensor, - tma_atom_kr: cute.CopyAtom, - tma_tensor_kr: cute.Tensor, tma_atom_inv: cute.CopyAtom, tma_tensor_inv: cute.Tensor, tma_atom_mqk: cute.CopyAtom, tma_tensor_mqk: cute.Tensor, + ws_qd: cute.Tensor, + ws_kd: cute.Tensor, + ws_kr: cute.Tensor, tma_atom_out: cute.CopyAtom, tma_tensor_out: cute.Tensor, out_gmem: cute.Tensor, @@ -156,30 +153,6 @@ def k2_kernel( cute.group_modes(sV, 0, 2), cute.group_modes(gSrc_v, 0, 2), ) - gSrc_kd = cute.local_tile(tma_tensor_kd, (CHUNK, D), (None, None, None)) - tKDs, tKDg = cpasync.tma_partition( - tma_atom_kd, - 0, - cute.make_layout(1), - cute.group_modes(sKd, 0, 2), - cute.group_modes(gSrc_kd, 0, 2), - ) - gSrc_qd = cute.local_tile(tma_tensor_qd, (CHUNK, D), (None, None, None)) - tQDs, tQDg = cpasync.tma_partition( - tma_atom_qd, - 0, - cute.make_layout(1), - cute.group_modes(sQd, 0, 2), - cute.group_modes(gSrc_qd, 0, 2), - ) - gSrc_kr = cute.local_tile(tma_tensor_kr, (CHUNK, D), (None, None, None)) - tKRs, tKRg = cpasync.tma_partition( - tma_atom_kr, - 0, - cute.make_layout(1), - cute.group_modes(sKr, 0, 2), - cute.group_modes(gSrc_kr, 0, 2), - ) gSrc_inv = cute.local_tile(tma_tensor_inv, (CHUNK, CHUNK), (None, None, None)) tIs, tIg = cpasync.tma_partition( tma_atom_inv, @@ -220,6 +193,21 @@ def k2_kernel( cute.group_modes(sBeta, 0, 2), cute.group_modes(gSrc_beta, 0, 2), ) + raw_copy_atom = cute.make_copy_atom(cpasync.CopyBulkG2SOp(), cutlass.BFloat16) + raw_stage_layout = cute.make_layout( + (CHUNK * D, STAGES), + stride=(1, CHUNK * D), + ) + raw_gmem_layout = cute.make_layout( + (CHUNK * D, total_tiles * H), + stride=(1, CHUNK * D), + ) + sKD_raw = cute.make_tensor(sKd.iterator, raw_stage_layout) + sQD_raw = cute.make_tensor(sQd.iterator, raw_stage_layout) + sKR_raw = cute.make_tensor(sKr.iterator, raw_stage_layout) + gKD_raw = cute.make_tensor(ws_kd.iterator, raw_gmem_layout) + gQD_raw = cute.make_tensor(ws_qd.iterator, raw_gmem_layout) + gKR_raw = cute.make_tensor(ws_kr.iterator, raw_gmem_layout) # Load initial_state -> sState[K_in, D_out] if has_initial_state: @@ -360,9 +348,11 @@ def k2_kernel( cute.copy(tma_atom_v, tVg_seq[(None, t, 0, head_idx)], tVs_seq[(None, s_dyn_l)], tma_bar_ptr=bar_l) else: cute.copy(tma_atom_v, tVg[(None, tg_l, 0, head_idx)], tVs[(None, s_dyn_l)], tma_bar_ptr=bar_l) - cute.copy(tma_atom_kd, tKDg[(None, 0, 0, wt_l)], tKDs[(None, s_dyn_l)], tma_bar_ptr=bar_l) - cute.copy(tma_atom_qd, tQDg[(None, 0, 0, wt_l)], tQDs[(None, s_dyn_l)], tma_bar_ptr=bar_l) - cute.copy(tma_atom_kr, tKRg[(None, 0, 0, wt_l)], tKRs[(None, s_dyn_l)], tma_bar_ptr=bar_l) + # Restore the byte-identical K_INTER images produced by K1. The + # raw-workspace transport idea is credited there to Flash-Flash-KDA. + copy_async_bulk(raw_copy_atom, gKD_raw[(None, wt_l)], sKD_raw[(None, s_dyn_l)], mbar_ptr=bar_l) + copy_async_bulk(raw_copy_atom, gQD_raw[(None, wt_l)], sQD_raw[(None, s_dyn_l)], mbar_ptr=bar_l) + copy_async_bulk(raw_copy_atom, gKR_raw[(None, wt_l)], sKR_raw[(None, s_dyn_l)], mbar_ptr=bar_l) cute.copy(tma_atom_inv, tIg[(None, 0, 0, wt_l)], tIs[(None, s_dyn_l)], tma_bar_ptr=bar_l) cute.copy(tma_atom_mqk, tMg[(None, 0, 0, wt_l)], tMs[(None, s_dyn_l)], tma_bar_ptr=bar_l) cute.copy(tma_atom_gt, tGTg[(None, 0, 0, wt_l)], tGTs[(None, s_dyn_l)], tma_bar_ptr=bar_l) @@ -633,13 +623,6 @@ def make_thd_atom(t, op, t_total: cutlass.Int32): ) return cpasync.make_tiled_tma_atom(op, view, kinter_smem, (CHUNK, D)) - def make_ws_qkd_atom(t): - view = cute.make_tensor( - t.iterator, - cute.make_layout((CHUNK, D, total_tiles * H), stride=(D, 1, CHUNK * D)), - ) - return cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileG2SOp(), view, kinter_smem, (CHUNK, D)) - def make_ws_cc_atom(t): view = cute.make_tensor( t.iterator, @@ -649,9 +632,6 @@ def make_ws_cc_atom(t): tma_atom_v, tma_tensor_v = make_thd_atom(v, cpasync.CopyBulkTensorTileG2SOp(), V_T_total) tma_atom_out, tma_tensor_out = make_thd_atom(out, cpasync.CopyBulkTensorTileS2GOp(), O_T_total) - tma_atom_kd, tma_tensor_kd = make_ws_qkd_atom(ws_kd) - tma_atom_qd, tma_tensor_qd = make_ws_qkd_atom(ws_qd) - tma_atom_kr, tma_tensor_kr = make_ws_qkd_atom(ws_kr) tma_atom_inv, tma_tensor_inv = make_ws_cc_atom(ws_inv) tma_atom_mqk, tma_tensor_mqk = make_ws_cc_atom(ws_mqk) @@ -701,16 +681,13 @@ def make_beta_atom(t): k2_kernel( tma_atom_v, tma_tensor_v, - tma_atom_kd, - tma_tensor_kd, - tma_atom_qd, - tma_tensor_qd, - tma_atom_kr, - tma_tensor_kr, tma_atom_inv, tma_tensor_inv, tma_atom_mqk, tma_tensor_mqk, + ws_qd, + ws_kd, + ws_kr, tma_atom_out, tma_tensor_out, out,