diff --git a/BENCHMARK_H200.md b/BENCHMARK_H200.md index 5c82b66..a2243cd 100644 --- a/BENCHMARK_H200.md +++ b/BENCHMARK_H200.md @@ -1,56 +1,56 @@ # Benchmark Results — Hopper (SM90) -> Auto-generated by `benchmarks/generate_benchmark_hopper_md.py` on 2026-05-19. +> Auto-generated by `benchmarks/generate_benchmark_hopper_md.py` on 2026-08-10. -> **GPU:** NVIDIA H200 | **CUDA:** 12.9 | **PyTorch:** 2.9.1+cu129 +> **GPU:** NVIDIA H200 | **CUDA:** 13.0.48 | **PyTorch:** 2.8.0a0+34c6371d24.nv25.8 > FLA baseline: [flash-linear-attention v0.5.0](https://github.com/fla-org/flash-linear-attention/releases/tag/v0.5.0) -## KDA Fused Forward (Kimi Delta Attention) +## KDA Prefill (Kimi Delta Attention) -Fully-fused KDA forward prefill kernel (sm90). +SM90 K1+K2 two-kernel prefill (sm90). ### Fixed-Length (H=64, D=128, bf16) -| B | T | FLA Triton (ms) | cuLA Fused (ms) | Speedup | +| B | T | FLA Triton (ms) | cuLA (ms) | Speedup | |---|---|-----------------|-----------------|---------| -| 1 | 512 | 0.556 | 0.224 | **2.48x** | -| 1 | 1024 | 0.581 | 0.248 | **2.34x** | -| 1 | 4096 | 0.936 | 0.896 | **1.04x** | -| 1 | 8192 | 1.810 | 1.754 | **1.03x** | -| 1 | 16384 | 3.576 | 3.492 | **1.02x** | -| 2 | 512 | 0.567 | 0.226 | **2.51x** | -| 2 | 1024 | 0.585 | 0.315 | **1.86x** | -| 2 | 4096 | 1.815 | 1.170 | **1.55x** | -| 2 | 8192 | 3.576 | 2.283 | **1.57x** | -| 2 | 16384 | 7.115 | 4.408 | **1.61x** | +| 1 | 512 | 0.658 | 0.137 | **4.79x** | +| 1 | 1024 | 0.666 | 0.159 | **4.18x** | +| 1 | 4096 | 1.038 | 0.567 | **1.83x** | +| 1 | 8192 | 2.004 | 1.106 | **1.81x** | +| 1 | 16384 | 3.963 | 2.208 | **1.79x** | +| 2 | 512 | 0.658 | 0.136 | **4.85x** | +| 2 | 1024 | 0.688 | 0.198 | **3.47x** | +| 2 | 4096 | 1.998 | 0.717 | **2.79x** | +| 2 | 8192 | 3.935 | 1.429 | **2.75x** | +| 2 | 16384 | 7.835 | 2.860 | **2.74x** | ### Variable-Length (H=64, D=128, bf16) -| Config | FLA Triton (ms) | cuLA Fused (ms) | Speedup | +| Config | FLA Triton (ms) | cuLA (ms) | Speedup | |--------|-----------------|-----------------|---------| -| uniform 10seqs T=4096 [409..415] avg=409 | 1.019 | 0.707 | **1.44x** | -| random 10seqs T=4096 [24..1201] avg=409 | 1.013 | 0.669 | **1.51x** | -| skewed 10seqs T=4096 [227..2053] avg=409 | 1.010 | 0.681 | **1.48x** | -| uniform 20seqs T=4096 [204..220] avg=204 | 1.098 | 0.932 | **1.18x** | -| random 20seqs T=4096 [5..787] avg=204 | 1.074 | 0.748 | **1.44x** | -| skewed 20seqs T=4096 [107..2063] avg=204 | 1.048 | 0.732 | **1.43x** | -| uniform 10seqs T=8192 [819..821] avg=819 | 1.851 | 1.174 | **1.58x** | -| random 10seqs T=8192 [48..2401] avg=819 | 1.890 | 1.217 | **1.55x** | -| skewed 10seqs T=8192 [455..4097] avg=819 | 1.905 | 1.225 | **1.55x** | -| uniform 20seqs T=8192 [409..421] avg=409 | 1.960 | 1.406 | **1.39x** | -| random 20seqs T=8192 [9..1574] avg=409 | 1.953 | 1.290 | **1.51x** | -| skewed 20seqs T=8192 [215..4107] avg=409 | 1.957 | 1.300 | **1.51x** | -| uniform 10seqs T=16384 [1638..1642] avg=1638 | 3.642 | 2.162 | **1.68x** | -| random 10seqs T=16384 [95..4802] avg=1638 | 3.609 | 2.279 | **1.58x** | -| skewed 10seqs T=16384 [910..8194] avg=1638 | 3.625 | 2.354 | **1.54x** | -| uniform 20seqs T=16384 [819..823] avg=819 | 3.644 | 2.320 | **1.57x** | -| random 20seqs T=16384 [19..3147] avg=819 | 3.681 | 2.293 | **1.61x** | -| skewed 20seqs T=16384 [431..8195] avg=819 | 3.634 | 2.371 | **1.53x** | - -Summary (28 configs): **avg=1.58x**, min=1.02x, max=2.51x. +| uniform 10seqs T=4096 [409..415] avg=409 | 1.020 | 0.382 | **2.67x** | +| random 10seqs T=4096 [24..1201] avg=409 | 1.013 | 0.451 | **2.25x** | +| skewed 10seqs T=4096 [227..2053] avg=409 | 1.010 | 0.569 | **1.78x** | +| uniform 20seqs T=4096 [204..220] avg=204 | 1.087 | 0.374 | **2.91x** | +| random 20seqs T=4096 [5..787] avg=204 | 1.064 | 0.442 | **2.41x** | +| skewed 20seqs T=4096 [107..2063] avg=204 | 1.040 | 0.594 | **1.75x** | +| uniform 10seqs T=8192 [819..821] avg=819 | 1.844 | 0.733 | **2.52x** | +| random 10seqs T=8192 [48..2401] avg=819 | 1.884 | 0.852 | **2.21x** | +| skewed 10seqs T=8192 [455..4097] avg=819 | 1.903 | 1.093 | **1.74x** | +| uniform 20seqs T=8192 [409..421] avg=409 | 1.944 | 0.687 | **2.83x** | +| random 20seqs T=8192 [9..1574] avg=409 | 1.939 | 0.831 | **2.33x** | +| skewed 20seqs T=8192 [215..4107] avg=409 | 1.956 | 1.123 | **1.74x** | +| uniform 10seqs T=16384 [1638..1642] avg=1638 | 3.603 | 1.450 | **2.48x** | +| random 10seqs T=16384 [95..4802] avg=1638 | 3.617 | 1.693 | **2.14x** | +| skewed 10seqs T=16384 [910..8194] avg=1638 | 3.644 | 2.167 | **1.68x** | +| uniform 20seqs T=16384 [819..823] avg=819 | 3.635 | 1.359 | **2.68x** | +| random 20seqs T=16384 [19..3147] avg=819 | 3.675 | 1.600 | **2.30x** | +| skewed 20seqs T=16384 [431..8195] avg=819 | 3.641 | 2.205 | **1.65x** | + +Summary (28 configs): **avg=2.54x**, min=1.65x, max=4.85x. To reproduce: diff --git a/README.md b/README.md index c5e145f..9f71dc9 100644 --- a/README.md +++ b/README.md @@ -127,13 +127,13 @@ See [BENCHMARK_GB200_CUDA_130.md](BENCHMARK_GB200_CUDA_130.md) tested with CUDA **Hopper (SM90)** -See [BENCHMARK_H200.md](BENCHMARK_H200.md) tested with CUDA 12.9 for detailed results. +See [BENCHMARK_H200.md](BENCHMARK_H200.md) tested with CUDA 13.0 for detailed results. **Highlights:** - **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). -- **KDA Fused Forward (Hopper):** **avg 1.58x** speedup across fixed-length and variable-length sequences. +- **KDA Prefill (Hopper):** **avg 2.54x** speedup across fixed-length and variable-length sequences. To regenerate benchmarks: diff --git a/cula/ops/kda/sm90/_common.py b/cula/ops/kda/sm90/_common.py index 8044a4c..a24beb6 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 900b333..da40f66 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 8b2fdd1..34aa1e4 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 12ccbe3..1cb60cc 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,