diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 048e59ddb763..0a783ac01d23 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1469,7 +1469,7 @@ legacy-files: &legacy_files |
# list; the hook's own `files:` pattern only gates *when* the hook triggers.
# Global exclude: vendored code + trtllm-gen FMHA artifacts (cubin pointers, export headers, cuda_ptx)
-exclude: '(^cpp/tensorrt_llm/common/sha256/|^triton_kernels/|^tensorrt_llm/_torch/attention/backends/prims_ts/|^tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/|trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo\.h$|trtllmGenKernels/gemm/trtllmGen_gemm_export/KernelMetaInfo\.h$|\.cubin\.tar\.zst$)'
+exclude: '(^tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_chunked_bt16(_small_bh)?\.py$|^cpp/tensorrt_llm/common/sha256/|^triton_kernels/|^tensorrt_llm/_torch/attention/backends/prims_ts/|^tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/|trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo\.h$|trtllmGenKernels/gemm/trtllmGen_gemm_export/KernelMetaInfo\.h$|\.cubin\.tar\.zst$)'
default_install_hook_types: [pre-commit, commit-msg]
repos:
diff --git a/3rdparty/vendor_patches/flashinfer-kda-bt16.patch b/3rdparty/vendor_patches/flashinfer-kda-bt16.patch
new file mode 100644
index 000000000000..587961c19b07
--- /dev/null
+++ b/3rdparty/vendor_patches/flashinfer-kda-bt16.patch
@@ -0,0 +1,207 @@
+diff --git a/kda_chunked_bt16_small_bh.py b/kda_chunked_bt16_small_bh.py
+index e8bc9feb5ece3dbfe3156c90bb27e68fa5a4521a..7478dc054a5a9afb9de5d416588a4b074971f38c 100644
+--- a/kda_chunked_bt16_small_bh.py
++++ b/kda_chunked_bt16_small_bh.py
+@@ -1,3 +1,4 @@
++# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ # Copyright (c) 2026 by FlashInfer team.
+ #
+ # Licensed under the Apache License, Version 2.0 (the "License");
+@@ -12,6 +13,16 @@
+ # See the License for the specific language governing permissions and
+ # limitations under the License.
+
++"""FlashInfer Small-BH KDA, with indexed recurrent-state pool access.
++
++Vendored from flashinfer-ai/flashinfer commit
++b8e81a7dae6c2aa5c9aa19c06e81e9a729271493,
++flashinfer/kda_kernels/kda_chunked_bt16_small_bh.py (Apache-2.0).
++Local changes add state-slot indirection and padded pool-row strides, reject
++unwarmed graph specializations, and use the active device SM count for scheduling.
++The K1/K2 arithmetic and layouts are retained from upstream.
++"""
++
+ import functools
+ from typing import Optional
+
+@@ -24,7 +35,6 @@ from cutlass._mlir.dialects import llvm
+ from cutlass.cute.core import _pack_shape
+ from cutlass.cute.nvgpu import cpasync, warp
+ from cutlass.cutlass_dsl import dsl_user_op
+-from flashinfer.cute_dsl.utils import get_max_active_clusters
+
+
+ RCP_LN2 = 1.4426950408889634
+@@ -1803,6 +1813,7 @@ class ChunkKdaFwdK2:
+ mS_in: Optional[cute.Tensor] = None, # (B, H, D, D)
+ mS_out: Optional[cute.Tensor] = None, # (B, H, D, D)
+ mCuSeqlens: Optional[cute.Tensor] = None, # (B + 1,)
++ mStateIndices: Optional[cute.Tensor] = None, # (B,)
+ stream=None,
+ ):
+ if cutlass.const_expr(mCuSeqlens is not None):
+@@ -2076,6 +2087,7 @@ class ChunkKdaFwdK2:
+ mS_in,
+ mS_out,
+ mCuSeqlens,
++ mStateIndices,
+ sQK_layout_staged,
+ sQKt_layout_staged,
+ sV_layout_staged,
+@@ -2123,6 +2135,7 @@ class ChunkKdaFwdK2:
+ mS_in: Optional[cute.Tensor],
+ mS_out: Optional[cute.Tensor],
+ mCuSeqlens: Optional[cute.Tensor],
++ mStateIndices: Optional[cute.Tensor],
+ sQK_layout_staged: cute.ComposedLayout,
+ sQKt_layout_staged: cute.ComposedLayout,
+ sV_layout_staged: cute.ComposedLayout,
+@@ -2148,6 +2161,9 @@ class ChunkKdaFwdK2:
+ warpgroup_idx = cute.arch.make_warp_uniform(tidx // self.warpgroup_threads)
+
+ v_idx, h_idx, b_idx = cute.arch.block_idx()
++ state_idx = b_idx
++ if cutlass.const_expr(mStateIndices is not None):
++ state_idx = mStateIndices[b_idx]
+
+ seq_len = cute.size(tma_tensor_Q, mode=[0])
+ seq_offset = 0
+@@ -2475,11 +2491,11 @@ class ChunkKdaFwdK2:
+ if cutlass.const_expr(mS_in is not None):
+ if cutlass.const_expr(self.transpose_S):
+ gS_in = cute.local_tile(
+- mS_in[None, None, h_idx, b_idx], (self.BV, self.D), (v_idx, 0)
++ mS_in[None, None, h_idx, state_idx], (self.BV, self.D), (v_idx, 0)
+ )
+ else:
+ gS_in = cute.local_tile(
+- mS_in[None, None, h_idx, b_idx], (self.D, self.BV), (0, v_idx)
++ mS_in[None, None, h_idx, state_idx], (self.D, self.BV), (0, v_idx)
+ )
+ gS_in = cute.make_tensor(
+ gS_in.iterator, cute.select(gS_in.layout, mode=[1, 0])
+@@ -2763,11 +2779,11 @@ class ChunkKdaFwdK2:
+ if cutlass.const_expr(mS_out is not None):
+ if cutlass.const_expr(self.transpose_S):
+ gS_out = cute.local_tile(
+- mS_out[None, None, h_idx, b_idx], (self.BV, self.D), (v_idx, 0)
++ mS_out[None, None, h_idx, state_idx], (self.BV, self.D), (v_idx, 0)
+ )
+ else:
+ gS_out = cute.local_tile(
+- mS_out[None, None, h_idx, b_idx], (self.D, self.BV), (0, v_idx)
++ mS_out[None, None, h_idx, state_idx], (self.D, self.BV), (0, v_idx)
+ )
+ gS_out = cute.make_tensor(
+ gS_out.iterator, cute.select(gS_out.layout, mode=[1, 0])
+@@ -2808,12 +2824,12 @@ def _cute_dtype(dtype: torch.dtype):
+ raise ValueError(f"unsupported dtype: {dtype}")
+
+
+-def _fake(dtype, shape):
++def _fake(dtype, shape, assumed_align=16):
+ return cute.runtime.make_fake_compact_tensor(
+ dtype,
+ shape,
+ stride_order=tuple(reversed(range(len(shape)))),
+- assumed_align=16,
++ assumed_align=assumed_align,
+ )
+
+
+@@ -2838,6 +2854,8 @@ def _compile_k1(
+ is_varlen: bool,
+ max_active_clusters: int,
+ ):
++ if torch.cuda.is_current_stream_capturing():
++ raise RuntimeError("Warm Small-BH KDA before CUDA graph capture")
+ cu_dtype = _cute_dtype(dtype)
+ b, t, tp, nt = (cute.sym_int() for _ in range(4))
+ feature: tuple[object, ...]
+@@ -2904,7 +2922,11 @@ def _compile_k2(
+ has_state_output: bool,
+ transpose_state: bool,
+ is_varlen: bool,
++ state_row_stride: int | None,
++ has_state_indices: bool,
+ ):
++ if torch.cuda.is_current_stream_capturing():
++ raise RuntimeError("Warm Small-BH KDA before CUDA graph capture")
+ cu_dtype = _cute_dtype(dtype)
+ b, bs, t, tp, h, nt = (cute.sym_int() for _ in range(6))
+ feature: tuple[object, ...]
+@@ -2925,6 +2947,14 @@ def _compile_k2(
+ cu_seqlens = None
+ state = (b, h, d, d)
+
++ def fake_state():
++ if state_row_stride is None:
++ return _fake(_cute_dtype(state_dtype), state)
++ return cute.runtime.make_fake_tensor(
++ _cute_dtype(state_dtype), state,
++ stride=(state_row_stride, d * d, d, 1), assumed_align=16,
++ )
++
+ q, k, kr = [_fake(cu_dtype, feature) for _ in range(3)]
+ kernel = ChunkKdaFwdK2(D=d, BT=_BT, BV=bv, transpose_S=transpose_state)
+ stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+@@ -2938,9 +2968,10 @@ def _compile_k2(
+ _fake(cu_dtype, matrix),
+ _fake(cutlass.Float32, gate),
+ _fake(cu_dtype, output),
+- _fake(_cute_dtype(state_dtype), state) if has_state_input else None,
+- _fake(_cute_dtype(state_dtype), state) if has_state_output else None,
++ fake_state() if has_state_input else None,
++ fake_state() if has_state_output else None,
+ cu_seqlens,
++ _fake(cutlass.Int32, (cute.sym_int(),), assumed_align=4) if has_state_indices else None,
+ stream,
+ options="--enable-tvm-ffi",
+ )
+@@ -3040,6 +3071,7 @@ def chunk_kda_fwd_k2(
+ output: torch.Tensor | None = None,
+ state_input: torch.Tensor | None = None,
+ state_output: torch.Tensor | None = None,
++ state_indices: torch.Tensor | None = None,
+ max_active_clusters: int = 152,
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ is_varlen = cu_seqlens is not None
+@@ -3072,6 +3104,8 @@ def chunk_kda_fwd_k2(
+ has_output,
+ transpose_state,
+ is_varlen,
++ state_input.stride(0) if state_indices is not None else None,
++ state_indices is not None,
+ )
+ compiled(
+ q,
+@@ -3085,6 +3119,7 @@ def chunk_kda_fwd_k2(
+ state_input,
+ state_output,
+ cu_seqlens,
++ state_indices,
+ )
+ return output, state_output
+
+@@ -3108,11 +3143,12 @@ def chunk_kda_fwd(
+ output: torch.Tensor | None = None,
+ state_input: torch.Tensor | None = None,
+ state_output: torch.Tensor | None = None,
++ state_indices: torch.Tensor | None = None,
+ workspace_tensors: tuple[torch.Tensor, ...] | None = None,
+ max_active_clusters: int = 0,
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ if max_active_clusters <= 0:
+- max_active_clusters = get_max_active_clusters(1)
++ max_active_clusters = torch.cuda.get_device_properties(q.device).multi_processor_count
+ qd, kd, kr, mqk, mkk, gk = chunk_kda_fwd_k1(
+ q,
+ k,
+@@ -3142,6 +3178,7 @@ def chunk_kda_fwd(
+ output=output,
+ state_input=state_input,
+ state_output=state_output,
++ state_indices=state_indices,
+ max_active_clusters=max_active_clusters,
+ )
+
diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml
index 56ebefa3a6a0..ea34d65b41be 100644
--- a/3rdparty/vendor_sources.lock.yaml
+++ b/3rdparty/vendor_sources.lock.yaml
@@ -1,5 +1,17 @@
schema_version: 1
vendors:
+ flashinfer-kda-bt16:
+ url: https://github.com/flashinfer-ai/flashinfer.git
+ branch: main
+ commit: b8e81a7dae6c2aa5c9aa19c06e81e9a729271493
+ source: flashinfer/kda_kernels
+ destination: tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda
+ include:
+ - kda_chunked_bt16.py
+ - kda_chunked_bt16_small_bh.py
+ patch: 3rdparty/vendor_patches/flashinfer-kda-bt16.patch
+ patch_digest: sha256:bc9b96a69873e0ddfcba66ab899c666cf52766976f24d850048d2a0bf8d967ff
+ digest: sha256-tree-v1:d2d382335ccb68b55078523e63835c28b699607227afb81308cb7bb1d85845ac
flashinfer-prims-ts:
url: https://github.com/yuxianq/flashinfer.git
branch: trtllm-prims-ts-dev-20261012
diff --git a/pyproject.toml b/pyproject.toml
index dcfb1369b991..88de9143fb5c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -40,6 +40,9 @@ fix = true
force-exclude = true
exclude = [
"**3rdparty/**",
+ # Preserve vendored FlashInfer KDA sources as recorded in the vendor lock.
+ "tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_chunked_bt16.py",
+ "tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_chunked_bt16_small_bh.py",
"tensorrt_llm/_torch/attention/backends/prims_ts/**",
"tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn/**",
"triton_kernels/**",
diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/akk_inverse.py b/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/akk_inverse.py
deleted file mode 100644
index a8a96aa43bac..000000000000
--- a/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/akk_inverse.py
+++ /dev/null
@@ -1,1075 +0,0 @@
-# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-Akk 64×64 Lower Triangular Block Inversion — mixed FP32/BF16.
-
-BF16 input/output and packed BF16 shared storage. The four 16×16 diagonal
-blocks use FP32 forward substitution to avoid unstable BF16 Neumann
-cancellation; the cross-block products retain BF16 MMA with FP32 accumulators.
-
-Architecture:
- - 4 warps (128 threads) per CTA, each CTA processes one 64×64 Akk matrix
- - sAkk [64,36] fp32 — each FP32 slot holds packed bf16x2 (stride 36)
- Total SMEM: 64*36*4 = 9216 bytes (vs 64*72*4 = 18432 in FP32 version)
- - sTemp [16, 24, 2] fp32: inter-warp FP32 accumulator communication
- - cp.async: BF16 global → packed bf16x2 in FP32 SMEM (raw bytes align)
- - FP32 forward substitution for diagonal blocks
- - BF16 MMA m16n8k16 with FP32 accumulators for cross-block products
- - movmatrix.sync.aligned.m8n8.trans.b16 for A→B layout conversion
- - C→A chain: cvt.rn.bf16x2.f32 to pack FP32 accum → bf16x2 A-operand
-
-Block layout in sAkk (4×4 sub-blocks of 16×16, packed bf16x2):
- Upper tri = INPUT, Diagonal = in-place inversion, Lower tri = OUTPUT
-
-Stages:
- 0. cp.async load bf16 64×64 → sAkk (packed bf16x2)
- 1. Invert 4 diagonal blocks via FP32 forward substitution (2 warps)
- 2. Warps 0-2: Ai10, Ai21, Ai32 via chain MMA (C→A pack)
- 3. Warps 0+2 → Ai20, warps 1+3 → Ai31 (parallel pairs, sTemp)
- 4. Warps 0+1+2 → Ai30 (sTemp aggregation)
- 5. All warps: store bf16x2 SMEM → bf16 global
-
-Inputs: A_in [B, T, H, BT] bf16
-Outputs: A_out [B, T, H, BT] bf16
-"""
-
-import cuda.bindings.driver as cuda
-import cutlass
-import cutlass.cute as cute
-from cutlass._mlir import ir
-from cutlass._mlir.dialects import llvm
-from cutlass.cute.nvgpu import cpasync
-from cutlass.cutlass_dsl import T, dsl_user_op
-
-# ===========================================================================
-# Constants
-# ===========================================================================
-BS = 64
-SB = 16
-THREADS = 128
-TEMP_PAD = 8
-TEMP_COLS = SB + TEMP_PAD # 24
-NUM_TEMPS = 2
-AKK_PAD = 4
-AKK_STRIDE = BS // 2 + AKK_PAD # 36 (in FP32 units = 72 bf16 elements)
-
-
-# ===========================================================================
-# BF16 MMA m16n8k16 with FP32 accumulator
-# ===========================================================================
-@dsl_user_op
-def mma_bf16_m16n8k16(
- a0,
- a1,
- a2,
- a3,
- b0,
- b1,
- c0,
- c1,
- c2,
- c3,
- *,
- loc=None,
- ip=None,
-):
- """mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32
- A: 4×i32 (bf16x2 pairs), B: 2×i32, C/D: 4×f32 (FP32 accum)."""
- a0b = llvm.bitcast(T.i32(), a0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- a1b = llvm.bitcast(T.i32(), a1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- a2b = llvm.bitcast(T.i32(), a2.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- a3b = llvm.bitcast(T.i32(), a3.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- b0b = llvm.bitcast(T.i32(), b0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- b1b = llvm.bitcast(T.i32(), b1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- result = llvm.inline_asm(
- ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"),
- [
- a0b,
- a1b,
- a2b,
- a3b,
- b0b,
- b1b,
- c0.ir_value(loc=loc, ip=ip),
- c1.ir_value(loc=loc, ip=ip),
- c2.ir_value(loc=loc, ip=ip),
- c3.ir_value(loc=loc, ip=ip),
- ],
- """{
- mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32
- {$0, $1, $2, $3},
- {$4, $5, $6, $7},
- {$8, $9},
- {$10, $11, $12, $13};
- }""",
- "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
- has_side_effects=True,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
- d0 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip))
- d1 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip))
- d2 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip))
- d3 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip))
- return d0, d1, d2, d3
-
-
-# ===========================================================================
-# movmatrix.sync.aligned.m8n8.trans.b16 — hardware A→B transpose (warp-level)
-# Works on any 16-bit format including BF16.
-# ===========================================================================
-@dsl_user_op
-def _movmatrix_trans(src, *, loc=None, ip=None):
- src_b = llvm.bitcast(T.i32(), src.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- result = llvm.inline_asm(
- T.i32(),
- [src_b],
- "movmatrix.sync.aligned.m8n8.trans.b16 $0, $1;",
- "=r,r",
- has_side_effects=True,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
- return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip))
-
-
-# ===========================================================================
-# BF16x2 pack/unpack helpers
-# ===========================================================================
-@dsl_user_op
-def _pack_bf16x2(lo_f32, hi_f32, *, loc=None, ip=None):
- """Pack two FP32 values into one i32 holding two BF16 values.
- cvt.rn.bf16x2.f32 converts and packs in one instruction."""
- result = llvm.inline_asm(
- T.i32(),
- [lo_f32.ir_value(loc=loc, ip=ip), hi_f32.ir_value(loc=loc, ip=ip)],
- "cvt.rn.bf16x2.f32 $0, $2, $1;",
- "=r,f,f",
- has_side_effects=False,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
- return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip))
-
-
-@dsl_user_op
-def _mask_packed_ltri(packed, row, pair, *, loc=None, ip=None):
- """Apply lower-triangular mask to packed bf16x2.
- Zero out elements where row < col. col0=2*pair, col1=2*pair+1."""
- p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- result = llvm.inline_asm(
- T.i32(),
- [p, row.ir_value(loc=loc, ip=ip), pair.ir_value(loc=loc, ip=ip)],
- """{
- .reg .b32 %c0, %c1, %mlo, %mhi, %mask;
- .reg .pred %p0, %p1;
- shl.b32 %c0, $3, 1;
- add.u32 %c1, %c0, 1;
- setp.ge.s32 %p0, $2, %c0;
- setp.ge.s32 %p1, $2, %c1;
- selp.b32 %mlo, 0xFFFF, 0, %p0;
- selp.b32 %mhi, 0xFFFF0000, 0, %p1;
- or.b32 %mask, %mlo, %mhi;
- and.b32 $0, $1, %mask;
- }""",
- "=r,r,r,r",
- has_side_effects=False,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
- return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip))
-
-
-@dsl_user_op
-def _unpack_bf16x2_lo(packed, *, loc=None, ip=None):
- """Unpack lower BF16 from packed i32 to FP32.
- BF16 is upper 16 bits of FP32, so shift left by 16."""
- p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- result = llvm.inline_asm(
- T.i32(),
- [p],
- """{
- shl.b32 $0, $1, 16;
- }""",
- "=r,r",
- has_side_effects=False,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
- return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip))
-
-
-@dsl_user_op
-def _unpack_bf16x2_hi(packed, *, loc=None, ip=None):
- """Unpack upper BF16 from packed i32 to FP32.
- Upper 16 bits are already in the right position for FP32."""
- p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- result = llvm.inline_asm(
- T.i32(),
- [p],
- """{
- and.b32 $0, $1, 0xFFFF0000;
- }""",
- "=r,r",
- has_side_effects=False,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
- return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip))
-
-
-# ===========================================================================
-# Neumann diagonal 16×16 inversion using BF16 MMA m16n8k16 + FP32 accum
-#
-# (I+L)⁻¹ = (I-L)(I+L²)(I+L⁴)(I+L⁸)
-#
-# All intermediate values kept in FP32. Pack to bf16x2 only at MMA boundaries.
-# sAkk is packed bf16x2: sAkk[row, pair] = {bf16[row, 2*pair], bf16[row, 2*pair+1]}
-# ===========================================================================
-@dsl_user_op
-def _invert_diag_neumann(sAkk: cute.Tensor, block_idx, lane_id, *, loc=None, ip=None):
- r_off = block_idx * 16
- c_off = block_idx * 8 # packed bf16x2: 16 cols → 8 pairs
- gid = lane_id // 4
- tid = lane_id % 4
-
- # --- Load A from packed bf16x2 SMEM → unpack to 8 FP32 values ---
- packed0 = cutlass.Float32(sAkk[r_off + gid, c_off + tid])
- packed1 = cutlass.Float32(sAkk[r_off + gid + 8, c_off + tid])
- packed2 = cutlass.Float32(sAkk[r_off + gid, c_off + 4 + tid])
- packed3 = cutlass.Float32(sAkk[r_off + gid + 8, c_off + 4 + tid])
-
- A_f0 = _unpack_bf16x2_lo(packed0) # A[gid, 2*tid]
- A_f1 = _unpack_bf16x2_hi(packed0) # A[gid, 2*tid+1]
- A_f2 = _unpack_bf16x2_lo(packed1) # A[gid+8, 2*tid]
- A_f3 = _unpack_bf16x2_hi(packed1) # A[gid+8, 2*tid+1]
- A_f4 = _unpack_bf16x2_lo(packed2) # A[gid, 8+2*tid]
- A_f5 = _unpack_bf16x2_hi(packed2) # A[gid, 8+2*tid+1]
- A_f6 = _unpack_bf16x2_lo(packed3) # A[gid+8, 8+2*tid]
- A_f7 = _unpack_bf16x2_hi(packed3) # A[gid+8, 8+2*tid+1]
-
- # Build identity (FP32)
- _one = cutlass.Float32(1.0)
- _zero = cutlass.Float32(0.0)
- I_f0 = _one * cutlass.Float32(gid == 2 * tid) + _zero * cutlass.Float32(gid != 2 * tid)
- I_f1 = _one * cutlass.Float32(gid == 2 * tid + 1) + _zero * cutlass.Float32(gid != 2 * tid + 1)
- I_f2 = _one * cutlass.Float32(gid + 8 == 2 * tid) + _zero * cutlass.Float32(gid + 8 != 2 * tid)
- I_f3 = _one * cutlass.Float32(gid + 8 == 2 * tid + 1) + _zero * cutlass.Float32(
- gid + 8 != 2 * tid + 1
- )
- I_f4 = _one * cutlass.Float32(gid == 8 + 2 * tid) + _zero * cutlass.Float32(gid != 8 + 2 * tid)
- I_f5 = _one * cutlass.Float32(gid == 8 + 2 * tid + 1) + _zero * cutlass.Float32(
- gid != 8 + 2 * tid + 1
- )
- I_f6 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid) + _zero * cutlass.Float32(
- gid + 8 != 8 + 2 * tid
- )
- I_f7 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid + 1) + _zero * cutlass.Float32(
- gid + 8 != 8 + 2 * tid + 1
- )
-
- # L = A - I (FP32)
- L_f0 = A_f0 - I_f0
- L_f1 = A_f1 - I_f1
- L_f2 = A_f2 - I_f2
- L_f3 = A_f3 - I_f3
- L_f4 = A_f4 - I_f4
- L_f5 = A_f5 - I_f5
- L_f6 = A_f6 - I_f6
- L_f7 = A_f7 - I_f7
-
- # INV = I - L = 2I - A (FP32)
- INV_f0 = I_f0 - L_f0
- INV_f1 = I_f1 - L_f1
- INV_f2 = I_f2 - L_f2
- INV_f3 = I_f3 - L_f3
- INV_f4 = I_f4 - L_f4
- INV_f5 = I_f5 - L_f5
- INV_f6 = I_f6 - L_f6
- INV_f7 = I_f7 - L_f7
-
- _zf = cutlass.Float32(0.0)
-
- # Pack L and INV → bf16x2 for MMA
- L_a0 = _pack_bf16x2(L_f0, L_f1)
- L_a1 = _pack_bf16x2(L_f2, L_f3)
- L_a2 = _pack_bf16x2(L_f4, L_f5)
- L_a3 = _pack_bf16x2(L_f6, L_f7)
-
- INV_a0 = _pack_bf16x2(INV_f0, INV_f1)
- INV_a1 = _pack_bf16x2(INV_f2, INV_f3)
- INV_a2 = _pack_bf16x2(INV_f4, INV_f5)
- INV_a3 = _pack_bf16x2(INV_f6, INV_f7)
-
- # === Iteration 1: L² = L × L, then INV = INV + INV × L² ===
- L_b0 = _movmatrix_trans(L_a0)
- L_b1 = _movmatrix_trans(L_a1)
- L_b2 = _movmatrix_trans(L_a2)
- L_b3 = _movmatrix_trans(L_a3)
-
- Lp_c0, Lp_c1, Lp_c2, Lp_c3 = mma_bf16_m16n8k16(
- L_a0, L_a1, L_a2, L_a3, L_b0, L_b1, _zf, _zf, _zf, _zf
- )
- Lp_c4, Lp_c5, Lp_c6, Lp_c7 = mma_bf16_m16n8k16(
- L_a0, L_a1, L_a2, L_a3, L_b2, L_b3, _zf, _zf, _zf, _zf
- )
-
- Lp_a0 = _pack_bf16x2(Lp_c0, Lp_c1)
- Lp_a1 = _pack_bf16x2(Lp_c2, Lp_c3)
- Lp_a2 = _pack_bf16x2(Lp_c4, Lp_c5)
- Lp_a3 = _pack_bf16x2(Lp_c6, Lp_c7)
-
- Lp_b0 = _movmatrix_trans(Lp_a0)
- Lp_b1 = _movmatrix_trans(Lp_a1)
- Lp_b2 = _movmatrix_trans(Lp_a2)
- Lp_b3 = _movmatrix_trans(Lp_a3)
-
- # mm = INV × L² (FP32 accum output)
- mm_c0, mm_c1, mm_c2, mm_c3 = mma_bf16_m16n8k16(
- INV_a0, INV_a1, INV_a2, INV_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf
- )
- mm_c4, mm_c5, mm_c6, mm_c7 = mma_bf16_m16n8k16(
- INV_a0, INV_a1, INV_a2, INV_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf
- )
-
- # INV += mm (FP32)
- INV_f0 = INV_f0 + mm_c0
- INV_f1 = INV_f1 + mm_c1
- INV_f2 = INV_f2 + mm_c2
- INV_f3 = INV_f3 + mm_c3
- INV_f4 = INV_f4 + mm_c4
- INV_f5 = INV_f5 + mm_c5
- INV_f6 = INV_f6 + mm_c6
- INV_f7 = INV_f7 + mm_c7
-
- INV_a0 = _pack_bf16x2(INV_f0, INV_f1)
- INV_a1 = _pack_bf16x2(INV_f2, INV_f3)
- INV_a2 = _pack_bf16x2(INV_f4, INV_f5)
- INV_a3 = _pack_bf16x2(INV_f6, INV_f7)
-
- # === Iteration 2: L⁴ = L² × L², then INV = INV + INV × L⁴ ===
- L4_c0, L4_c1, L4_c2, L4_c3 = mma_bf16_m16n8k16(
- Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf
- )
- L4_c4, L4_c5, L4_c6, L4_c7 = mma_bf16_m16n8k16(
- Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf
- )
-
- L4_a0 = _pack_bf16x2(L4_c0, L4_c1)
- L4_a1 = _pack_bf16x2(L4_c2, L4_c3)
- L4_a2 = _pack_bf16x2(L4_c4, L4_c5)
- L4_a3 = _pack_bf16x2(L4_c6, L4_c7)
-
- L4_b0 = _movmatrix_trans(L4_a0)
- L4_b1 = _movmatrix_trans(L4_a1)
- L4_b2 = _movmatrix_trans(L4_a2)
- L4_b3 = _movmatrix_trans(L4_a3)
-
- mm_c0, mm_c1, mm_c2, mm_c3 = mma_bf16_m16n8k16(
- INV_a0, INV_a1, INV_a2, INV_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf
- )
- mm_c4, mm_c5, mm_c6, mm_c7 = mma_bf16_m16n8k16(
- INV_a0, INV_a1, INV_a2, INV_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf
- )
-
- INV_f0 = INV_f0 + mm_c0
- INV_f1 = INV_f1 + mm_c1
- INV_f2 = INV_f2 + mm_c2
- INV_f3 = INV_f3 + mm_c3
- INV_f4 = INV_f4 + mm_c4
- INV_f5 = INV_f5 + mm_c5
- INV_f6 = INV_f6 + mm_c6
- INV_f7 = INV_f7 + mm_c7
-
- INV_a0 = _pack_bf16x2(INV_f0, INV_f1)
- INV_a1 = _pack_bf16x2(INV_f2, INV_f3)
- INV_a2 = _pack_bf16x2(INV_f4, INV_f5)
- INV_a3 = _pack_bf16x2(INV_f6, INV_f7)
-
- # === Iteration 3: L⁸ = L⁴ × L⁴, then INV = INV + INV × L⁸ ===
- L8_c0, L8_c1, L8_c2, L8_c3 = mma_bf16_m16n8k16(
- L4_a0, L4_a1, L4_a2, L4_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf
- )
- L8_c4, L8_c5, L8_c6, L8_c7 = mma_bf16_m16n8k16(
- L4_a0, L4_a1, L4_a2, L4_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf
- )
-
- L8_a0 = _pack_bf16x2(L8_c0, L8_c1)
- L8_a1 = _pack_bf16x2(L8_c2, L8_c3)
- L8_a2 = _pack_bf16x2(L8_c4, L8_c5)
- L8_a3 = _pack_bf16x2(L8_c6, L8_c7)
-
- L8_b0 = _movmatrix_trans(L8_a0)
- L8_b1 = _movmatrix_trans(L8_a1)
- L8_b2 = _movmatrix_trans(L8_a2)
- L8_b3 = _movmatrix_trans(L8_a3)
-
- mm_c0, mm_c1, mm_c2, mm_c3 = mma_bf16_m16n8k16(
- INV_a0, INV_a1, INV_a2, INV_a3, L8_b0, L8_b1, _zf, _zf, _zf, _zf
- )
- mm_c4, mm_c5, mm_c6, mm_c7 = mma_bf16_m16n8k16(
- INV_a0, INV_a1, INV_a2, INV_a3, L8_b2, L8_b3, _zf, _zf, _zf, _zf
- )
-
- INV_f0 = INV_f0 + mm_c0
- INV_f1 = INV_f1 + mm_c1
- INV_f2 = INV_f2 + mm_c2
- INV_f3 = INV_f3 + mm_c3
- INV_f4 = INV_f4 + mm_c4
- INV_f5 = INV_f5 + mm_c5
- INV_f6 = INV_f6 + mm_c6
- INV_f7 = INV_f7 + mm_c7
-
- # --- Store INV to packed bf16x2 SMEM ---
- sAkk[r_off + gid, c_off + tid] = _pack_bf16x2(INV_f0, INV_f1)
- sAkk[r_off + gid + 8, c_off + tid] = _pack_bf16x2(INV_f2, INV_f3)
- sAkk[r_off + gid, c_off + 4 + tid] = _pack_bf16x2(INV_f4, INV_f5)
- sAkk[r_off + gid + 8, c_off + 4 + tid] = _pack_bf16x2(INV_f6, INV_f7)
-
-
-@dsl_user_op
-def _load_packed_bf16(sAkk: cute.Tensor, row, packed_col_base, col, *, loc=None, ip=None):
- """Load one logical BF16 value from packed-bf16x2 shared storage."""
- valid_i = cutlass.Int32(col >= 0)
- safe_col = col * valid_i
- packed = cutlass.Float32(sAkk[row, packed_col_base + safe_col // 2])
- lo = _unpack_bf16x2_lo(packed)
- hi = _unpack_bf16x2_hi(packed)
- use_hi = cutlass.Float32(safe_col % 2)
- value = lo * (cutlass.Float32(1.0) - use_hi) + hi * use_hi
- return value * cutlass.Float32(valid_i)
-
-
-@dsl_user_op
-def _invert_diag_forward_fp32(
- sAkk: cute.Tensor,
- block_idx,
- lane_id,
- *,
- loc=None,
- ip=None,
-):
- """FP32 forward substitution with one final BF16 pack per inverse row."""
- row = lane_id % 16
- halfwarp_base = (lane_id // 16) * 16
- r_off = block_idx * 16
- c_off = block_idx * 8
- inv = cute.make_rmem_tensor(cute.make_layout((16,), stride=(1,)), cutlass.Float32)
- inv[0] = cutlass.Float32(1.0)
- for d in range(1, 16):
- inv[d] = cutlass.Float32(0.0)
-
- for d in range(1, 16):
- col_d = row - d
- valid = cutlass.Float32(col_d >= 0)
- a_val = _load_packed_bf16(sAkk, r_off + row, c_off, col_d)
- acc = cutlass.Float32(0.0)
- for j in range(1, d):
- a_re = _load_packed_bf16(sAkk, r_off + row, c_off, row - (d - j))
- inv_prev = cute.arch.shuffle_sync(inv[j], halfwarp_base + row - d + j)
- acc = acc + a_re * inv_prev
- inv[d] = (-a_val - acc) * valid
-
- for pair in range(8):
- col0 = pair * 2
- col1 = col0 + 1
- d0 = row - col0
- d1 = row - col1
- v0 = cutlass.Float32(row == col0)
- v1 = cutlass.Float32(row == col1)
- for d in range(1, 16):
- v0 = v0 + inv[d] * cutlass.Float32(d0 == d)
- v1 = v1 + inv[d] * cutlass.Float32(d1 == d)
- sAkk[r_off + row, c_off + pair] = _pack_bf16x2(v0, v1)
-
-
-# ===========================================================================
-# 16×16 matmul: load A & B from packed bf16x2 sAkk, BF16 MMA, return FP32 C
-# ===========================================================================
-@dsl_user_op
-def _matmul_AB(sAkk: cute.Tensor, br_A, bc_A, br_B, bc_B, lane_id, *, loc=None, ip=None):
- gid = lane_id // 4
- tid = lane_id % 4
- _zf = cutlass.Float32(0.0)
- rA = br_A * 16
- cA = bc_A * 8 # packed: 16 cols → 8 pairs
- rB = br_B * 16
- cB = bc_B * 8
-
- # Load A-operand (packed bf16x2, direct from SMEM)
- a0 = cutlass.Float32(sAkk[rA + gid, cA + tid])
- a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + tid])
- a2 = cutlass.Float32(sAkk[rA + gid, cA + 4 + tid])
- a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 4 + tid])
-
- # Load B-operand: load in A-layout then movmatrix → B-layout
- bA0 = cutlass.Float32(sAkk[rB + gid, cB + tid])
- bA1 = cutlass.Float32(sAkk[rB + gid + 8, cB + tid])
- bA2 = cutlass.Float32(sAkk[rB + gid, cB + 4 + tid])
- bA3 = cutlass.Float32(sAkk[rB + gid + 8, cB + 4 + tid])
- b0 = _movmatrix_trans(bA0)
- b1 = _movmatrix_trans(bA1)
- b2 = _movmatrix_trans(bA2)
- b3 = _movmatrix_trans(bA3)
-
- # 16×16 = 2 × m16n8k16
- cn0_0, cn0_1, cn0_2, cn0_3 = mma_bf16_m16n8k16(a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf)
- cn1_0, cn1_1, cn1_2, cn1_3 = mma_bf16_m16n8k16(a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf)
-
- # Return 8 FP32 C-registers (C-layout of m16n8k16 FP32 accum)
- return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3
-
-
-# ===========================================================================
-# Chain MMA: pre-loaded A (from C→A pack), load B from sAkk (Stage 2)
-# A-operand already packed as bf16x2 from previous C result.
-# ===========================================================================
-@dsl_user_op
-def _chain_mma_B(sAkk: cute.Tensor, br_B, bc_B, a0, a1, a2, a3, lane_id, *, loc=None, ip=None):
- gid = lane_id // 4
- tid = lane_id % 4
- _zf = cutlass.Float32(0.0)
- rB = br_B * 16
- cB = bc_B * 8
-
- bA0 = cutlass.Float32(sAkk[rB + gid, cB + tid])
- bA1 = cutlass.Float32(sAkk[rB + gid + 8, cB + tid])
- bA2 = cutlass.Float32(sAkk[rB + gid, cB + 4 + tid])
- bA3 = cutlass.Float32(sAkk[rB + gid + 8, cB + 4 + tid])
- b0 = _movmatrix_trans(bA0)
- b1 = _movmatrix_trans(bA1)
- b2 = _movmatrix_trans(bA2)
- b3 = _movmatrix_trans(bA3)
-
- cn0_0, cn0_1, cn0_2, cn0_3 = mma_bf16_m16n8k16(a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf)
- cn1_0, cn1_1, cn1_2, cn1_3 = mma_bf16_m16n8k16(a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf)
-
- return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3
-
-
-# ===========================================================================
-# Chain MMA: load A from sAkk, pre-loaded B (from C→B shuffle) (Stages 3-4)
-# B-operand from shuffle is still FP32 TF32-layout → need to convert.
-# Actually in BF16 version we use movmatrix, so B comes from pack+movmatrix.
-# This function takes B already in B-layout (packed bf16x2 after movmatrix).
-# ===========================================================================
-@dsl_user_op
-def _chain_mma_A(sAkk: cute.Tensor, br_A, bc_A, b0, b1, b2, b3, lane_id, *, loc=None, ip=None):
- gid = lane_id // 4
- tid = lane_id % 4
- _zf = cutlass.Float32(0.0)
- rA = br_A * 16
- cA = bc_A * 8
-
- a0 = cutlass.Float32(sAkk[rA + gid, cA + tid])
- a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + tid])
- a2 = cutlass.Float32(sAkk[rA + gid, cA + 4 + tid])
- a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 4 + tid])
-
- cn0_0, cn0_1, cn0_2, cn0_3 = mma_bf16_m16n8k16(a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf)
- cn1_0, cn1_1, cn1_2, cn1_3 = mma_bf16_m16n8k16(a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf)
-
- return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3
-
-
-# ===========================================================================
-# Store negated C result (16×16, FP32 accum) to packed bf16x2 sAkk
-# C-layout for m16n8k16 FP32 accum:
-# cn0_0 = C[gid, 0..7 left half col0], cn0_1 = C[gid+8, left half col0]
-# cn0_2 = C[gid, left half col1], cn0_3 = C[gid+8, left half col1]
-# cn1_* = right half
-# Packing: negate FP32 then pack pairs → bf16x2
-# ===========================================================================
-@dsl_user_op
-def _store_neg_C(
- sAkk: cute.Tensor, br, bc, c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None
-):
- gid = lane_id // 4
- tid = lane_id % 4
- r = br * 16
- c = bc * 8 # packed
-
- # Pack negated FP32 pairs → bf16x2, then store
- sAkk[r + gid, c + tid] = _pack_bf16x2(-c0, -c1)
- sAkk[r + gid + 8, c + tid] = _pack_bf16x2(-c2, -c3)
- sAkk[r + gid, c + 4 + tid] = _pack_bf16x2(-c4, -c5)
- sAkk[r + gid + 8, c + 4 + tid] = _pack_bf16x2(-c6, -c7)
-
-
-# ===========================================================================
-# Pack FP32 C-accum → bf16x2 A-operand for C→A chain
-# C-layout (FP32 accum, m16n8k16): 8 floats → 4 bf16x2 A-regs
-# c0,c1 → a0 (rows gid/gid+8, left-half k0..7)
-# c2,c3 → a1
-# c4,c5 → a2 (right-half k8..15)
-# c6,c7 → a3
-# ===========================================================================
-@dsl_user_op
-def _pack_C_to_A(c0, c1, c2, c3, c4, c5, c6, c7, *, loc=None, ip=None):
- a0 = _pack_bf16x2(c0, c1)
- a1 = _pack_bf16x2(c2, c3)
- a2 = _pack_bf16x2(c4, c5)
- a3 = _pack_bf16x2(c6, c7)
- return a0, a1, a2, a3
-
-
-# ===========================================================================
-# Convert FP32 C-accum → bf16x2 B-operand via pack + movmatrix
-# ===========================================================================
-@dsl_user_op
-def _pack_C_to_B(c0, c1, c2, c3, c4, c5, c6, c7, *, loc=None, ip=None):
- a0 = _pack_bf16x2(c0, c1)
- a1 = _pack_bf16x2(c2, c3)
- a2 = _pack_bf16x2(c4, c5)
- a3 = _pack_bf16x2(c6, c7)
- b0 = _movmatrix_trans(a0)
- b1 = _movmatrix_trans(a1)
- b2 = _movmatrix_trans(a2)
- b3 = _movmatrix_trans(a3)
- return b0, b1, b2, b3
-
-
-# ===========================================================================
-# sTemp helpers (FP32, non-swizzled, for inter-warp accumulator exchange)
-# ===========================================================================
-@dsl_user_op
-def _store_C_temp(
- sT: cute.Tensor,
- buf,
- c0,
- c1,
- c2,
- c3,
- c4,
- c5,
- c6,
- c7,
- lane_id,
- *,
- loc=None,
- ip=None,
-):
- gid = lane_id // 4
- tid = lane_id % 4
- sT[gid, 2 * tid, buf] = c0
- sT[gid, 2 * tid + 1, buf] = c1
- sT[gid + 8, 2 * tid, buf] = c2
- sT[gid + 8, 2 * tid + 1, buf] = c3
- sT[gid, 8 + 2 * tid, buf] = c4
- sT[gid, 8 + 2 * tid + 1, buf] = c5
- sT[gid + 8, 8 + 2 * tid, buf] = c6
- sT[gid + 8, 8 + 2 * tid + 1, buf] = c7
-
-
-@dsl_user_op
-def _load_C_temp(sT: cute.Tensor, buf, lane_id, *, loc=None, ip=None):
- gid = lane_id // 4
- tid = lane_id % 4
- c0 = cutlass.Float32(sT[gid, 2 * tid, buf])
- c1 = cutlass.Float32(sT[gid, 2 * tid + 1, buf])
- c2 = cutlass.Float32(sT[gid + 8, 2 * tid, buf])
- c3 = cutlass.Float32(sT[gid + 8, 2 * tid + 1, buf])
- c4 = cutlass.Float32(sT[gid, 8 + 2 * tid, buf])
- c5 = cutlass.Float32(sT[gid, 8 + 2 * tid + 1, buf])
- c6 = cutlass.Float32(sT[gid + 8, 8 + 2 * tid, buf])
- c7 = cutlass.Float32(sT[gid + 8, 8 + 2 * tid + 1, buf])
- return c0, c1, c2, c3, c4, c5, c6, c7
-
-
-# ===========================================================================
-# Main kernel
-# ===========================================================================
-@cute.kernel
-def akk_inv_kernel(
- g2s_copy: cute.TiledCopy,
- gA_tensor: cute.Tensor,
- mOut: cute.Tensor,
- mBeta: cute.Tensor,
- akk_smem_layout: cute.Layout,
- temp_layout: cute.Layout,
- NT: cutlass.Int32,
- H: int,
- mCuSeqlens: cute.Tensor,
- mChunkIndices: cute.Tensor,
- IS_VARLEN: cutlass.Constexpr[int],
-):
- tidx, _, _ = cute.arch.thread_idx()
- warp_idx = cute.arch.warp_idx()
- warp_idx = cute.arch.make_warp_uniform(warp_idx)
- lane_id = tidx % 32
- h_idx, nt_idx, b_idx = cute.arch.block_idx()
-
- # ===== SMEM allocation =====
- smem = cutlass.utils.SmemAllocator()
- sAkk = smem.allocate_tensor(cutlass.Float32, akk_smem_layout, 128)
- sTemp = smem.allocate_tensor(cutlass.Float32, temp_layout, 128)
- sBeta = smem.allocate_tensor(cutlass.Float32, cute.make_layout(BS, stride=1), 128)
-
- # ===== Stage 0: cp.async load bf16 global → packed bf16x2 sAkk =====
- ld_bnt = nt_idx
- if IS_VARLEN:
- ld_seq_id = cutlass.Int32(mChunkIndices[nt_idx, 0])
- ld_local = cutlass.Int32(mChunkIndices[nt_idx, 1])
- ld_bos = cutlass.Int32(mCuSeqlens[ld_seq_id])
- ld_bnt = ld_bos + ld_local * BS
- gA_batch = gA_tensor[(None, None, h_idx, ld_bnt, b_idx)]
-
- thr_g2s = g2s_copy.get_slice(tidx)
- thr_gSrc = thr_g2s.partition_S(gA_batch)
- thr_sDst = thr_g2s.partition_D(sAkk)
- cute.copy(g2s_copy, thr_gSrc, thr_sDst)
- cute.arch.cp_async_commit_group()
- cute.arch.cp_async_wait_group(0)
- cute.arch.barrier()
-
- # Zero out-of-bounds rows for varlen partial chunks.
- if IS_VARLEN:
- _z_seq = cutlass.Int32(mChunkIndices[nt_idx, 0])
- _z_local = cutlass.Int32(mChunkIndices[nt_idx, 1])
- _z_bos = cutlass.Int32(mCuSeqlens[_z_seq])
- _z_eos = cutlass.Int32(mCuSeqlens[_z_seq + 1])
- _z_cs = _z_bos + _z_local * BS
- _z_vr = _z_eos - _z_cs
- _zr_start = warp_idx * SB
- for ri in cutlass.range_constexpr(SB):
- row = _zr_start + ri
- if row >= _z_vr:
- # Zero packed bf16x2 slots (each slot = 2 bf16 elements)
- c0 = lane_id
- if c0 < AKK_STRIDE:
- sAkk[row, c0] = cutlass.Float32(0.0)
- cute.arch.barrier()
-
- # ===== Stage 0b: Load 64 per-token betas into sBeta (fp32) =====
- _b_chunk_start = nt_idx * BS
- _b_eos = cutlass.Int32(_b_chunk_start + BS)
- if IS_VARLEN:
- _b_seq_id = cutlass.Int32(mChunkIndices[nt_idx, 0])
- _b_local = cutlass.Int32(mChunkIndices[nt_idx, 1])
- _b_bos = cutlass.Int32(mCuSeqlens[_b_seq_id])
- _b_eos = cutlass.Int32(mCuSeqlens[_b_seq_id + 1])
- _b_chunk_start = _b_bos + _b_local * BS
-
- if warp_idx == 0:
- _bcol_lo = lane_id * 2
- _bcol_hi = _bcol_lo + 1
- if _bcol_lo < BS:
- _bt_lo = _b_chunk_start + _bcol_lo
- _bt_hi = _b_chunk_start + _bcol_hi
- if IS_VARLEN:
- if _bt_lo < _b_eos:
- sBeta[_bcol_lo] = cutlass.Float32(mBeta[b_idx, _bt_lo, h_idx])
- else:
- sBeta[_bcol_lo] = cutlass.Float32(0.0)
- if _bt_hi < _b_eos:
- sBeta[_bcol_hi] = cutlass.Float32(mBeta[b_idx, _bt_hi, h_idx])
- else:
- sBeta[_bcol_hi] = cutlass.Float32(0.0)
- else:
- sBeta[_bcol_lo] = cutlass.Float32(mBeta[b_idx, _bt_lo, h_idx])
- sBeta[_bcol_hi] = cutlass.Float32(mBeta[b_idx, _bt_hi, h_idx])
- cute.arch.barrier()
-
- # ===== Stage 1: FP32 forward solve, two diagonal blocks per warp =====
- if warp_idx == 0:
- _invert_diag_forward_fp32(sAkk, lane_id // 16, lane_id)
- if warp_idx == 1:
- _invert_diag_forward_fp32(sAkk, 2 + lane_id // 16, lane_id)
-
- cute.arch.barrier()
-
- # ===== Stage 2: First batch — Ai10, Ai21, Ai32 =====
- # C→A chain: pack FP32 C-accum → bf16x2 A-operand via _pack_C_to_A
- if warp_idx == 0:
- t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 1, 1, 0, 1, lane_id)
- a0, a1, a2, a3 = _pack_C_to_A(t0, t1, t2, t3, t4, t5, t6, t7)
- r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_B(sAkk, 0, 0, a0, a1, a2, a3, lane_id)
- _store_neg_C(sAkk, 1, 0, r0, r1, r2, r3, r4, r5, r6, r7, lane_id)
-
- if warp_idx == 1:
- t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 2, 2, 1, 2, lane_id)
- a0, a1, a2, a3 = _pack_C_to_A(t0, t1, t2, t3, t4, t5, t6, t7)
- r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_B(sAkk, 1, 1, a0, a1, a2, a3, lane_id)
- _store_neg_C(sAkk, 2, 1, r0, r1, r2, r3, r4, r5, r6, r7, lane_id)
-
- if warp_idx == 2:
- t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 3, 3, 2, 3, lane_id)
- a0, a1, a2, a3 = _pack_C_to_A(t0, t1, t2, t3, t4, t5, t6, t7)
- r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_B(sAkk, 2, 2, a0, a1, a2, a3, lane_id)
- _store_neg_C(sAkk, 3, 2, r0, r1, r2, r3, r4, r5, r6, r7, lane_id)
-
- cute.arch.barrier()
-
- # ===== Stage 3: Second batch — Ai20, Ai31 (warp pairs via sTemp) =====
- _z = cutlass.Float32(0.0)
- t0 = _z
- t1 = _z
- t2 = _z
- t3 = _z
- t4 = _z
- t5 = _z
- t6 = _z
- t7 = _z
-
- # --- Ai20 = -Ai22 @ (Akk20 @ Ai00 + Akk21 @ Ai10) ---
- if warp_idx == 0:
- t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 0, 2, 0, 0, lane_id)
-
- if warp_idx == 2:
- s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk, 1, 2, 1, 0, lane_id)
- _store_C_temp(sTemp, 0, s0, s1, s2, s3, s4, s5, s6, s7, lane_id)
-
- # --- Ai31 = -Ai33 @ (Akk31 @ Ai11 + Akk32 @ Ai21) ---
- if warp_idx == 1:
- t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 1, 3, 1, 1, lane_id)
-
- if warp_idx == 3:
- s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk, 2, 3, 2, 1, lane_id)
- _store_C_temp(sTemp, 1, s0, s1, s2, s3, s4, s5, s6, s7, lane_id)
-
- cute.arch.barrier()
-
- # Warp 0: accumulate T1+T2, pack→B, multiply by Ai22
- if warp_idx == 0:
- e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 0, lane_id)
- t0 = t0 + e0
- t1 = t1 + e1
- t2 = t2 + e2
- t3 = t3 + e3
- t4 = t4 + e4
- t5 = t5 + e5
- t6 = t6 + e6
- t7 = t7 + e7
- b0, b1, b2, b3 = _pack_C_to_B(t0, t1, t2, t3, t4, t5, t6, t7)
- r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_A(sAkk, 2, 2, b0, b1, b2, b3, lane_id)
- _store_neg_C(sAkk, 2, 0, r0, r1, r2, r3, r4, r5, r6, r7, lane_id)
-
- # Warp 1: accumulate T1'+T2', pack→B, multiply by Ai33
- if warp_idx == 1:
- e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 1, lane_id)
- t0 = t0 + e0
- t1 = t1 + e1
- t2 = t2 + e2
- t3 = t3 + e3
- t4 = t4 + e4
- t5 = t5 + e5
- t6 = t6 + e6
- t7 = t7 + e7
- b0, b1, b2, b3 = _pack_C_to_B(t0, t1, t2, t3, t4, t5, t6, t7)
- r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_A(sAkk, 3, 3, b0, b1, b2, b3, lane_id)
- _store_neg_C(sAkk, 3, 1, r0, r1, r2, r3, r4, r5, r6, r7, lane_id)
-
- cute.arch.barrier()
-
- # ===== Stage 4: Third batch — Ai30 =====
- t0 = _z
- t1 = _z
- t2 = _z
- t3 = _z
- t4 = _z
- t5 = _z
- t6 = _z
- t7 = _z
-
- if warp_idx == 0:
- t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk, 0, 3, 0, 0, lane_id)
-
- if warp_idx == 1:
- s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk, 1, 3, 1, 0, lane_id)
- _store_C_temp(sTemp, 0, s0, s1, s2, s3, s4, s5, s6, s7, lane_id)
-
- if warp_idx == 2:
- s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk, 2, 3, 2, 0, lane_id)
- _store_C_temp(sTemp, 1, s0, s1, s2, s3, s4, s5, s6, s7, lane_id)
-
- cute.arch.barrier()
-
- # Warp 0: accumulate all three, pack→B, multiply by Ai33
- if warp_idx == 0:
- e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 0, lane_id)
- t0 = t0 + e0
- t1 = t1 + e1
- t2 = t2 + e2
- t3 = t3 + e3
- t4 = t4 + e4
- t5 = t5 + e5
- t6 = t6 + e6
- t7 = t7 + e7
- e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 1, lane_id)
- t0 = t0 + e0
- t1 = t1 + e1
- t2 = t2 + e2
- t3 = t3 + e3
- t4 = t4 + e4
- t5 = t5 + e5
- t6 = t6 + e6
- t7 = t7 + e7
- b0, b1, b2, b3 = _pack_C_to_B(t0, t1, t2, t3, t4, t5, t6, t7)
- r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_A(sAkk, 3, 3, b0, b1, b2, b3, lane_id)
- _store_neg_C(sAkk, 3, 0, r0, r1, r2, r3, r4, r5, r6, r7, lane_id)
-
- cute.arch.barrier()
-
- # ===== Stage 5: Store packed bf16x2 sAkk → global (b32 with direct bit-mask) =====
- vl_chunk_start = nt_idx * BS
- vl_eos = cutlass.Int32(vl_chunk_start + BS)
- if IS_VARLEN:
- vl_seq_id = cutlass.Int32(mChunkIndices[nt_idx, 0])
- vl_local = cutlass.Int32(mChunkIndices[nt_idx, 1])
- vl_bos = cutlass.Int32(mCuSeqlens[vl_seq_id])
- vl_eos = cutlass.Int32(mCuSeqlens[vl_seq_id + 1])
- vl_chunk_start = vl_bos + vl_local * BS
-
- row_start = warp_idx * SB
-
- # PDL: hint downstream K4 to pre-launch so its setup work overlaps with our
- # gmem writes below. K4 has griddepcontrol_wait after its setup but before
- # reading any gmem (TMA loads / gk_last_exp), so ordering is correct.
- cute.arch.griddepcontrol_launch_dependents()
-
- for ri in cutlass.range_constexpr(SB):
- row = row_start + ri
- pair = lane_id
- if pair < BS // 2:
- packed = cutlass.Float32(sAkk[row, pair])
- masked = _mask_packed_ltri(packed, cutlass.Int32(row), cutlass.Int32(pair))
-
- _beta_lo = cutlass.Float32(sBeta[2 * pair])
- _beta_hi = cutlass.Float32(sBeta[2 * pair + 1])
- _lo_f32 = _unpack_bf16x2_lo(masked) * _beta_lo
- _hi_f32 = _unpack_bf16x2_hi(masked) * _beta_hi
- final = _pack_bf16x2(_lo_f32, _hi_f32)
-
- t_row = vl_chunk_start + row
- if IS_VARLEN:
- if t_row < vl_eos:
- mOut[b_idx, t_row, h_idx, pair] = final
- else:
- mOut[b_idx, t_row, h_idx, pair] = final
-
-
-# ===========================================================================
-# Host JIT function
-# ===========================================================================
-@cute.jit
-def akk_inv_host(
- A_in: cute.Tensor,
- A_out: cute.Tensor,
- Beta_in: cute.Tensor,
- # B / NT / T_VAL are runtime scalars (shape-independent compile):
- # baking them recompiled the kernel per batch shape — at eval traffic
- # (unique total T per prefill batch) ~100 s of JIT per batch. Only H
- # and IS_VARLEN remain compile-time specializers.
- B: cutlass.Int32,
- NT: cutlass.Int32,
- H: cutlass.Constexpr[int],
- mCuSeqlens: cute.Tensor,
- mChunkIndices: cute.Tensor,
- IS_VARLEN: cutlass.Constexpr[int],
- T_VAL: cutlass.Int32,
- # Launch stream — runtime argument; launching on the DSL default stream
- # races with the executor's non-blocking execution stream.
- stream: cuda.CUstream,
-):
- # BF16 input: view as FP32 (packed bf16x2). Each FP32 element = 2 BF16.
- # Original BF16 shape: [B, T, H, BS] with stride per-element in BF16.
- # As FP32 packed bf16x2: shape [BS, BS//2, H, dim3, B]
- # The row stride in BF16 units is H*BS. In FP32 (bf16x2) units = H*BS/2.
- # The col stride in BF16 is 1. In bf16x2 pairs: pair stride = 1 (adjacent pairs in memory).
- _dim3_size = IS_VARLEN * T_VAL + (1 - IS_VARLEN) * (T_VAL // BS)
- _dim3_stride_bf16 = IS_VARLEN * (H * BS) + (1 - IS_VARLEN) * (BS * H * BS)
- # All strides in FP32 units (each FP32 = 2 BF16):
- _row_stride = H * BS // 2 # row stride (BF16 row stride = H*BS, /2 for bf16x2)
- _col_stride = 1 # adjacent bf16x2 pairs
- _h_stride = BS // 2 # head stride (BF16 = BS, /2 for bf16x2)
- _dim3_stride = _dim3_stride_bf16 // 2
- # T_VAL is a runtime Int32, so this stride is dynamic. Mathematically it is
- # T_VAL * (H * BS // 2) = T_VAL * 32 * H, i.e. always a multiple of 4 fp32
- # elements (16 B). The IR verifier can't derive that on its own, and without
- # it the b_idx slice offset breaks the 128-bit alignment proof required by
- # the cp.async G2S atom below (ICE: "src ptr alignment (32 bits) does not
- # meet requirement (128 bits)"). Assert divisibility explicitly.
- _batch_stride = cute.assume(T_VAL * (H * BS // 2), divby=H * BS // 2)
-
- view_layout = cute.make_layout(
- (BS, BS // 2, H, _dim3_size, B),
- stride=(_row_stride, _col_stride, _h_stride, _dim3_stride, _batch_stride),
- )
- gA_view = cute.make_tensor(A_in.iterator, view_layout)
-
- # sAkk: packed bf16x2, logical shape (64, 32) with stride 36 for bank-conflict-free access
- akk_smem_2d = cute.make_layout((BS, BS // 2), stride=(AKK_STRIDE, 1))
-
- # cp.async G→S copy: 128-bit (4×fp32 = 8 bf16) vectorised
- copy_atom = cute.make_copy_atom(
- cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),
- cutlass.Float32,
- num_bits_per_copy=128,
- )
- # Source layout: BS rows × BS//2 pairs = 64×32
- # 128 threads, each copies 4 FP32 (128 bits) per iteration
- # Need 64*32/128/4 = 4 iterations → but (16,8) × (1,4) = 128 threads × 4 vals = 512 per iter
- # 64×32 = 2048 elements / 512 = 4 iterations. But SMEM is 64×36, need to handle padding.
- # Actually: the copy maps source shape to SMEM shape.
- # Source is (BS, BS//2) = (64, 32), SMEM is (BS, AKK_STRIDE) = (64, 36).
- # The copy handles the actual data region (64, 32); padding cols 32-35 stay zero.
- g2s_copy = cute.make_tiled_copy_tv(
- copy_atom,
- thr_layout=cute.make_layout((16, 8), stride=(8, 1)),
- val_layout=cute.make_layout((1, 4)),
- )
-
- # sTemp layout (non-swizzled)
- temp_layout = cute.make_layout(
- (SB, TEMP_COLS, NUM_TEMPS), stride=(TEMP_COLS, 1, SB * TEMP_COLS)
- )
-
- smem_bytes = BS * AKK_STRIDE * 4 + SB * TEMP_COLS * NUM_TEMPS * 4 + BS * 4 + 256
-
- out_layout = cute.make_layout(
- (B, T_VAL, H, BS // 2), stride=(T_VAL * H * BS // 2, H * BS // 2, BS // 2, 1)
- )
- out_view = cute.make_tensor(A_out.iterator, out_layout)
-
- akk_inv_kernel(
- g2s_copy,
- gA_view,
- out_view,
- Beta_in,
- akk_smem_2d,
- temp_layout,
- NT,
- H,
- mCuSeqlens,
- mChunkIndices,
- IS_VARLEN,
- ).launch(
- grid=(H, NT, B),
- block=(THREADS, 1, 1),
- smem=smem_bytes,
- stream=stream,
- )
diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/fused_k123.py b/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/fused_k123.py
deleted file mode 100644
index 99f0cd3f535d..000000000000
--- a/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/fused_k123.py
+++ /dev/null
@@ -1,2353 +0,0 @@
-# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-Persistent Fused K1+K2+K3 Kernel for KDA.
-
-Fuses gate activation + cumsum + scaling (K1), intra sub-chunk Aqk/Akk (K2),
-and inter sub-chunk solve + merged inverse (K3) into a single persistent kernel.
-
-Grid: (NUM_SMS, 1, 1) — 148 persistent blocks, each loops over work units
- Total work units = (NT/4) * H * B, distributed round-robin across SMs
- Block i processes work units i, i+NUM_SMS, i+2*NUM_SMS, ...
-Block: 992 threads (31 warps), warp-specialized with setmaxnreg:
- Warps 0-15: K1 compute (8×2, vec2) – 4 WGs, 56 regs
- Warps 16-25: K2 MMA compute – 72 regs
- Warp 26: Dedicated TMA producer (TMA_WARP_ID) – 72 regs
- Warps 27-30: Store/Inversion warps – 24 regs
-
-Pipeline (single for_generate, warp groups separated by if-blocks):
- per work unit:
- Warps 0-15: wait(tma), K1 cumsum, arrive(k1_done),
- Pass 2b Q/K reads and scaling, arrive(stage_reuse)
- Warps 16-25: wait(k1_done)+wait(store_done), MMA, arrive(mma_done+stage_reuse)
- Warp 26: wait(stage_reuse), TMA Q/K/G for each chunk, signal(tma)
- Warps 27-30: wait(mma_done), store sAqk/sAkk→GMEM, arrive(store_done)
- All warp-group invariants are computed inside each group's if-block (not hoisted)
- to eliminate cross-group register pressure — same budget as the _all version.
- Mbarrier phases self-reset after 4 iterations (2 stages × 2 phases).
-
-Mbarriers:
- tma_mbars[2]: count=1, TMA_WARP_ID (26) lane 0 → K1+MMA wait for TMA data
- stage_reuse_mbars[2]: count=832, K1(16)+MMA(10 warps) → TMA waits before reuse
- k1_done_mbars[2]: count=512, K1(16 warps) → MMA waits for g_cumsum ready
- mma_done_mbars[2]: count=320, MMA(10 warps) → Store waits for sAqk/sAkk ready
- store_done_mbars[2]: count=128, Store(4 warps) → MMA waits for sAqk/sAkk stage free
-
-SMEM: ~215KB (q+k+g × [64,128] bf16 × 2 stages + g_cumsum [64,136] fp32 × 2 stages
- + sPartialLast [8,132] fp32 + sAqk [16,168,2] bf16
- + sAkk [64,72,2] fp32 block-transposed upper-tri layout)
-
-Inputs:
- g [B,T,H,K] bf16 raw gate
- k [B,T,H,K] bf16 normalized by caller
- q [B,T,H,K] bf16 normalized by caller
- A_log [H] fp32 per-head log decay
- beta [B,T,H] bf16/fp32, activated in-kernel when requested
- scale fp32 1/sqrt(K)
-
-Outputs (g_cumsum stays in SMEM, not written to GMEM):
- k_scaled [B,T,H,K] bf16
- q_scaled [B,T,H,K] bf16
- kg [B,T,H,K] bf16
- gk_last_exp[B,NT,H,K] fp32
- A_qk [B,T,H,BT] bf16 full merged (diagonal + off-diagonal)
- A_kk [B,T,H,BT] fp32 block-transposed upper-tri (input to akk_inv)
-"""
-
-import cuda.bindings.driver as cuda
-import cutlass
-import cutlass.cute as cute
-from cutlass import for_generate, yield_out
-from cutlass._mlir import ir
-from cutlass._mlir.dialects import llvm
-from cutlass.cute.nvgpu import cpasync, tcgen05
-from cutlass.cutlass_dsl import T, dsl_user_op
-
-BT = 64
-BC = 16
-K_DIM = 128
-K_PAD = 8
-K_STRIDE = K_DIM + K_PAD # 136, padded row stride to avoid bank conflicts
-CHUNKS_PER_BLOCK = 4
-NUM_SMS = 148 # Persistent kernel: one resident block per SM
-
-NUM_K1_TMA_WARPS = 16 # Warps 0-15: K1 compute (4 warpgroups, 8×2) -- TMA offloaded
-NUM_MMA_WARPS = 11 # Warps 16-26: MMA (10 active + 1 TMA producer, dropped idle warp 27)
-NUM_MMA_ACTIVE = 10 # mma_warp 0..9: actual MMA work
-TMA_WARP_ID = NUM_K1_TMA_WARPS + NUM_MMA_ACTIVE # warp 26 = dedicated TMA producer
-NUM_STORE_WARPS = 4 # Warps 27-30: Store/Inversion (1 warpgroup)
-NUM_WARPS = NUM_K1_TMA_WARPS + NUM_MMA_WARPS + NUM_STORE_WARPS # 31
-THREADS = NUM_WARPS * 32 # 992
-
-NUM_SUB_CHUNKS = BT // BC # 4
-NUM_TILES = NUM_SUB_CHUNKS * (NUM_SUB_CHUNKS + 1) // 2 # 10 lower-tri tiles
-MMA_K_TILE = 16
-NUM_MMA_K_TILES = K_DIM // MMA_K_TILE # 8 (bf16 m16n8k16)
-AQK_TILE_PAD = 8
-AQK_TILE_STRIDE = BT + AQK_TILE_PAD # 72 — sAqk now 64x72 row-major (same shape as sAkk)
-
-AKK_PAD = 8
-AKK_STRIDE = BT + AKK_PAD # 72
-
-# For fused akk_inv: sAkk_pkd viewed as fp32 packed bf16x2, stride = BT/2 + small pad
-AKK_PKD_PAD = 4
-AKK_PKD_STRIDE = BT // 2 + AKK_PKD_PAD # 36 fp32 units = 72 bf16 elements per row
-AKK_TEMP_PAD = 8
-AKK_TEMP_COLS = 16 + AKK_TEMP_PAD # 24
-AKK_TEMP_BUFS = 2
-
-K1_ROW_GROUPS = 8
-K1_COL_GROUPS = 2
-ROWS_PER_K1_WARP = BT // K1_ROW_GROUPS # 8
-K1_COLS_PER_WARP = K_DIM // K1_COL_GROUPS # 64
-ROWS_PER_STORE_WARP = BT // NUM_STORE_WARPS # 16
-
-VEC = K1_COLS_PER_WARP // 32 # 2
-K_VEC = K_DIM // VEC # 64
-NUM_STAGES = 2
-PARTIAL_COLS = K_DIM + 4 # 132
-PARTIAL_COLS_PER_WARP = K_DIM // NUM_K1_TMA_WARPS # 8
-
-_TILE_IQ = [0, 1, 1, 2, 2, 2, 3, 3, 3, 3]
-_TILE_IK = [0, 0, 1, 0, 1, 2, 0, 1, 2, 3]
-
-LOG2E = 1.4426950408889634
-LN2 = 0.6931471805599453
-RCP_LN2 = LOG2E
-
-
-@dsl_user_op
-def k1_internal_barrier(*, loc=None, ip=None):
- """Named barrier for K1+TMA warps (0-15, 512 threads). barrier_id=2."""
- llvm.inline_asm(
- T.i32(),
- [],
- "membar.cta; bar.sync 2, 512; mov.u32 $0, 0;",
- "=r",
- has_side_effects=True,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
-
-
-@dsl_user_op
-def pack_bf16x2_f32(hi_f32, lo_f32, *, loc=None, ip=None):
- """Pack two fp32 values into a bf16x2 u32 register.
- Returns u32 with bf16(hi) in bits [31:16] and bf16(lo) in bits [15:0].
- PTX: cvt.rn.bf16x2.f32 d, a, b -> d[31:16]=bf16(a), d[15:0]=bf16(b)
- """
- result = llvm.inline_asm(
- T.i32(),
- [hi_f32.ir_value(loc=loc, ip=ip), lo_f32.ir_value(loc=loc, ip=ip)],
- "cvt.rn.bf16x2.f32 $0, $1, $2;",
- "=r,f,f",
- has_side_effects=False,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
- return cutlass.Int32(result)
-
-
-@dsl_user_op
-def mma_bf16_m16n8k16(
- a0,
- a1,
- a2,
- a3, # 4 u32 (bf16x2 packed) A operands
- b0,
- b1, # 2 u32 (bf16x2 packed) B operands
- c0,
- c1,
- c2,
- c3, # 4 fp32 accumulators
- *,
- loc=None,
- ip=None,
-):
- """bf16 MMA with fp32 accumulator, shape m16n8k16.
- D_fp32 = A_bf16 * B_bf16 + C_fp32
- """
- # a/b already i32 (from pack_bf16x2_f32) -> no bitcast needed
- result = llvm.inline_asm(
- ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"),
- [
- a0.ir_value(loc=loc, ip=ip),
- a1.ir_value(loc=loc, ip=ip),
- a2.ir_value(loc=loc, ip=ip),
- a3.ir_value(loc=loc, ip=ip),
- b0.ir_value(loc=loc, ip=ip),
- b1.ir_value(loc=loc, ip=ip),
- c0.ir_value(loc=loc, ip=ip),
- c1.ir_value(loc=loc, ip=ip),
- c2.ir_value(loc=loc, ip=ip),
- c3.ir_value(loc=loc, ip=ip),
- ],
- """{
- mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32
- {$0, $1, $2, $3},
- {$4, $5, $6, $7},
- {$8, $9},
- {$10, $11, $12, $13};
- }""",
- "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
- has_side_effects=True,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
-
- d0 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip))
- d1 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip))
- d2 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip))
- d3 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip))
- return d0, d1, d2, d3
-
-
-SHFL_W8_CLAMP = 0x1800
-
-
-@dsl_user_op
-def fast_rcp(x, *, loc=None, ip=None):
- """Hardware fast reciprocal: rcp.approx.ftz.f32 (~2 cycles vs ~20 for div)."""
- result = llvm.inline_asm(
- T.f32(),
- [x.ir_value(loc=loc, ip=ip)],
- "rcp.approx.ftz.f32 $0, $1;",
- "=f,f",
- has_side_effects=False,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
- return cutlass.Float32(result)
-
-
-# ============================================================
-# Akk inverse helpers (fused from Akk_inverse_lower_triangle_bf16.py)
-# ============================================================
-@dsl_user_op
-def store_internal_barrier(*, loc=None, ip=None):
- """Named barrier for Store warps (4 warps, 128 threads). barrier_id=3."""
- llvm.inline_asm(
- T.i32(),
- [],
- "membar.cta; bar.sync 3, 128; mov.u32 $0, 0;",
- "=r",
- has_side_effects=True,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
-
-
-@dsl_user_op
-def _ak_pack_bf16x2(lo_f32, hi_f32, *, loc=None, ip=None):
- """cvt.rn.bf16x2.f32 -- pack two fp32 into bf16x2 (as fp32 bitcast view)."""
- result = llvm.inline_asm(
- T.i32(),
- [lo_f32.ir_value(loc=loc, ip=ip), hi_f32.ir_value(loc=loc, ip=ip)],
- "cvt.rn.bf16x2.f32 $0, $2, $1;",
- "=r,f,f",
- has_side_effects=False,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
- return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip))
-
-
-@dsl_user_op
-def _ak_movmatrix_trans(src, *, loc=None, ip=None):
- src_b = llvm.bitcast(T.i32(), src.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- result = llvm.inline_asm(
- T.i32(),
- [src_b],
- "movmatrix.sync.aligned.m8n8.trans.b16 $0, $1;",
- "=r,r",
- has_side_effects=True,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
- return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip))
-
-
-@dsl_user_op
-def _ak_mask_packed_ltri(packed, row, pair, *, loc=None, ip=None):
- """Mask packed bf16x2 to lower-tri: zero elements where row < col."""
- p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- result = llvm.inline_asm(
- T.i32(),
- [p, row.ir_value(loc=loc, ip=ip), pair.ir_value(loc=loc, ip=ip)],
- """{
- .reg .b32 %c0, %c1, %mlo, %mhi, %mask;
- .reg .pred %p0, %p1;
- shl.b32 %c0, $3, 1;
- add.u32 %c1, %c0, 1;
- setp.ge.s32 %p0, $2, %c0;
- setp.ge.s32 %p1, $2, %c1;
- selp.b32 %mlo, 0xFFFF, 0, %p0;
- selp.b32 %mhi, 0xFFFF0000, 0, %p1;
- or.b32 %mask, %mlo, %mhi;
- and.b32 $0, $1, %mask;
- }""",
- "=r,r,r,r",
- has_side_effects=False,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
- return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip))
-
-
-@dsl_user_op
-def _ak_unpack_bf16x2_lo(packed, *, loc=None, ip=None):
- p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- result = llvm.inline_asm(
- T.i32(),
- [p],
- "shl.b32 $0, $1, 16;",
- "=r,r",
- has_side_effects=False,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
- return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip))
-
-
-@dsl_user_op
-def _ak_unpack_bf16x2_hi(packed, *, loc=None, ip=None):
- p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- result = llvm.inline_asm(
- T.i32(),
- [p],
- "and.b32 $0, $1, 0xFFFF0000;",
- "=r,r",
- has_side_effects=False,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
- return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip))
-
-
-@dsl_user_op
-def _ak_mma(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, *, loc=None, ip=None):
- """BF16 MMA m16n8k16, args are Float32 (bitcast-viewed as i32 for bf16x2)."""
- a0b = llvm.bitcast(T.i32(), a0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- a1b = llvm.bitcast(T.i32(), a1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- a2b = llvm.bitcast(T.i32(), a2.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- a3b = llvm.bitcast(T.i32(), a3.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- b0b = llvm.bitcast(T.i32(), b0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- b1b = llvm.bitcast(T.i32(), b1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- result = llvm.inline_asm(
- ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"),
- [
- a0b,
- a1b,
- a2b,
- a3b,
- b0b,
- b1b,
- c0.ir_value(loc=loc, ip=ip),
- c1.ir_value(loc=loc, ip=ip),
- c2.ir_value(loc=loc, ip=ip),
- c3.ir_value(loc=loc, ip=ip),
- ],
- """{
- mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32
- {$0, $1, $2, $3},
- {$4, $5, $6, $7},
- {$8, $9},
- {$10, $11, $12, $13};
- }""",
- "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
- has_side_effects=True,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
- d0 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip))
- d1 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip))
- d2 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip))
- d3 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip))
- return d0, d1, d2, d3
-
-
-@dsl_user_op
-def _ak_invert_diag_neumann(sAkk, block_idx, lane_id, *, loc=None, ip=None):
- """Invert 16x16 diag block via Neumann series (I+L)^-1 = (I-L)(I+L^2)(I+L^4)(I+L^8).
- Reads/writes sAkk as fp32 packed bf16x2 (stride 36). Each block handled by 1 warp."""
- r_off = block_idx * 16
- c_off = block_idx * 8 # 16 cols = 8 pairs
- gid = lane_id // 4
- tid = lane_id % 4
-
- packed0 = cutlass.Float32(sAkk[r_off + gid, c_off + tid])
- packed1 = cutlass.Float32(sAkk[r_off + gid + 8, c_off + tid])
- packed2 = cutlass.Float32(sAkk[r_off + gid, c_off + 4 + tid])
- packed3 = cutlass.Float32(sAkk[r_off + gid + 8, c_off + 4 + tid])
-
- A_f0 = _ak_unpack_bf16x2_lo(packed0)
- A_f1 = _ak_unpack_bf16x2_hi(packed0)
- A_f2 = _ak_unpack_bf16x2_lo(packed1)
- A_f3 = _ak_unpack_bf16x2_hi(packed1)
- A_f4 = _ak_unpack_bf16x2_lo(packed2)
- A_f5 = _ak_unpack_bf16x2_hi(packed2)
- A_f6 = _ak_unpack_bf16x2_lo(packed3)
- A_f7 = _ak_unpack_bf16x2_hi(packed3)
-
- _one = cutlass.Float32(1.0)
- _zero = cutlass.Float32(0.0)
- I_f0 = _one * cutlass.Float32(gid == 2 * tid) + _zero * cutlass.Float32(gid != 2 * tid)
- I_f1 = _one * cutlass.Float32(gid == 2 * tid + 1) + _zero * cutlass.Float32(gid != 2 * tid + 1)
- I_f2 = _one * cutlass.Float32(gid + 8 == 2 * tid) + _zero * cutlass.Float32(gid + 8 != 2 * tid)
- I_f3 = _one * cutlass.Float32(gid + 8 == 2 * tid + 1) + _zero * cutlass.Float32(
- gid + 8 != 2 * tid + 1
- )
- I_f4 = _one * cutlass.Float32(gid == 8 + 2 * tid) + _zero * cutlass.Float32(gid != 8 + 2 * tid)
- I_f5 = _one * cutlass.Float32(gid == 8 + 2 * tid + 1) + _zero * cutlass.Float32(
- gid != 8 + 2 * tid + 1
- )
- I_f6 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid) + _zero * cutlass.Float32(
- gid + 8 != 8 + 2 * tid
- )
- I_f7 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid + 1) + _zero * cutlass.Float32(
- gid + 8 != 8 + 2 * tid + 1
- )
-
- L_f0 = A_f0 - I_f0
- L_f1 = A_f1 - I_f1
- L_f2 = A_f2 - I_f2
- L_f3 = A_f3 - I_f3
- L_f4 = A_f4 - I_f4
- L_f5 = A_f5 - I_f5
- L_f6 = A_f6 - I_f6
- L_f7 = A_f7 - I_f7
- INV_f0 = I_f0 - L_f0
- INV_f1 = I_f1 - L_f1
- INV_f2 = I_f2 - L_f2
- INV_f3 = I_f3 - L_f3
- INV_f4 = I_f4 - L_f4
- INV_f5 = I_f5 - L_f5
- INV_f6 = I_f6 - L_f6
- INV_f7 = I_f7 - L_f7
-
- _zf = cutlass.Float32(0.0)
- L_a0 = _ak_pack_bf16x2(L_f0, L_f1)
- L_a1 = _ak_pack_bf16x2(L_f2, L_f3)
- L_a2 = _ak_pack_bf16x2(L_f4, L_f5)
- L_a3 = _ak_pack_bf16x2(L_f6, L_f7)
- INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1)
- INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3)
- INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5)
- INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7)
-
- # Iter 1: L^2, INV += INV*L^2
- L_b0 = _ak_movmatrix_trans(L_a0)
- L_b1 = _ak_movmatrix_trans(L_a1)
- L_b2 = _ak_movmatrix_trans(L_a2)
- L_b3 = _ak_movmatrix_trans(L_a3)
- Lp_c0, Lp_c1, Lp_c2, Lp_c3 = _ak_mma(L_a0, L_a1, L_a2, L_a3, L_b0, L_b1, _zf, _zf, _zf, _zf)
- Lp_c4, Lp_c5, Lp_c6, Lp_c7 = _ak_mma(L_a0, L_a1, L_a2, L_a3, L_b2, L_b3, _zf, _zf, _zf, _zf)
- Lp_a0 = _ak_pack_bf16x2(Lp_c0, Lp_c1)
- Lp_a1 = _ak_pack_bf16x2(Lp_c2, Lp_c3)
- Lp_a2 = _ak_pack_bf16x2(Lp_c4, Lp_c5)
- Lp_a3 = _ak_pack_bf16x2(Lp_c6, Lp_c7)
- Lp_b0 = _ak_movmatrix_trans(Lp_a0)
- Lp_b1 = _ak_movmatrix_trans(Lp_a1)
- Lp_b2 = _ak_movmatrix_trans(Lp_a2)
- Lp_b3 = _ak_movmatrix_trans(Lp_a3)
- mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma(
- INV_a0, INV_a1, INV_a2, INV_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf
- )
- mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma(
- INV_a0, INV_a1, INV_a2, INV_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf
- )
- INV_f0 = INV_f0 + mm_c0
- INV_f1 = INV_f1 + mm_c1
- INV_f2 = INV_f2 + mm_c2
- INV_f3 = INV_f3 + mm_c3
- INV_f4 = INV_f4 + mm_c4
- INV_f5 = INV_f5 + mm_c5
- INV_f6 = INV_f6 + mm_c6
- INV_f7 = INV_f7 + mm_c7
- INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1)
- INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3)
- INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5)
- INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7)
-
- # Iter 2: L^4, INV += INV*L^4
- L4_c0, L4_c1, L4_c2, L4_c3 = _ak_mma(
- Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf
- )
- L4_c4, L4_c5, L4_c6, L4_c7 = _ak_mma(
- Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf
- )
- L4_a0 = _ak_pack_bf16x2(L4_c0, L4_c1)
- L4_a1 = _ak_pack_bf16x2(L4_c2, L4_c3)
- L4_a2 = _ak_pack_bf16x2(L4_c4, L4_c5)
- L4_a3 = _ak_pack_bf16x2(L4_c6, L4_c7)
- L4_b0 = _ak_movmatrix_trans(L4_a0)
- L4_b1 = _ak_movmatrix_trans(L4_a1)
- L4_b2 = _ak_movmatrix_trans(L4_a2)
- L4_b3 = _ak_movmatrix_trans(L4_a3)
- mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma(
- INV_a0, INV_a1, INV_a2, INV_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf
- )
- mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma(
- INV_a0, INV_a1, INV_a2, INV_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf
- )
- INV_f0 = INV_f0 + mm_c0
- INV_f1 = INV_f1 + mm_c1
- INV_f2 = INV_f2 + mm_c2
- INV_f3 = INV_f3 + mm_c3
- INV_f4 = INV_f4 + mm_c4
- INV_f5 = INV_f5 + mm_c5
- INV_f6 = INV_f6 + mm_c6
- INV_f7 = INV_f7 + mm_c7
- INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1)
- INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3)
- INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5)
- INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7)
-
- # Iter 3: L^8, INV += INV*L^8
- L8_c0, L8_c1, L8_c2, L8_c3 = _ak_mma(
- L4_a0, L4_a1, L4_a2, L4_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf
- )
- L8_c4, L8_c5, L8_c6, L8_c7 = _ak_mma(
- L4_a0, L4_a1, L4_a2, L4_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf
- )
- L8_a0 = _ak_pack_bf16x2(L8_c0, L8_c1)
- L8_a1 = _ak_pack_bf16x2(L8_c2, L8_c3)
- L8_a2 = _ak_pack_bf16x2(L8_c4, L8_c5)
- L8_a3 = _ak_pack_bf16x2(L8_c6, L8_c7)
- L8_b0 = _ak_movmatrix_trans(L8_a0)
- L8_b1 = _ak_movmatrix_trans(L8_a1)
- L8_b2 = _ak_movmatrix_trans(L8_a2)
- L8_b3 = _ak_movmatrix_trans(L8_a3)
- mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma(
- INV_a0, INV_a1, INV_a2, INV_a3, L8_b0, L8_b1, _zf, _zf, _zf, _zf
- )
- mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma(
- INV_a0, INV_a1, INV_a2, INV_a3, L8_b2, L8_b3, _zf, _zf, _zf, _zf
- )
- INV_f0 = INV_f0 + mm_c0
- INV_f1 = INV_f1 + mm_c1
- INV_f2 = INV_f2 + mm_c2
- INV_f3 = INV_f3 + mm_c3
- INV_f4 = INV_f4 + mm_c4
- INV_f5 = INV_f5 + mm_c5
- INV_f6 = INV_f6 + mm_c6
- INV_f7 = INV_f7 + mm_c7
-
- sAkk[r_off + gid, c_off + tid] = _ak_pack_bf16x2(INV_f0, INV_f1)
- sAkk[r_off + gid + 8, c_off + tid] = _ak_pack_bf16x2(INV_f2, INV_f3)
- sAkk[r_off + gid, c_off + 4 + tid] = _ak_pack_bf16x2(INV_f4, INV_f5)
- sAkk[r_off + gid + 8, c_off + 4 + tid] = _ak_pack_bf16x2(INV_f6, INV_f7)
-
-
-@dsl_user_op
-def _ak_invert_diag_neumann_inreg(
- sAkk,
- block_idx,
- lane_id,
- raw_f0,
- raw_f1,
- raw_f2,
- raw_f3,
- raw_f4,
- raw_f5,
- raw_f6,
- raw_f7,
- *,
- loc=None,
- ip=None,
-):
- """In-register variant: invert 16x16 diag block from K2 acc fragment values
- (no SMEM read). Applies I+L mask, runs Neumann iterations, writes INV to sAkk.
-
- raw_f0..raw_f7 are K2 fp32 acc*beta values in C-fragment layout for m16n8k16:
- raw_f0,raw_f1 = (gid, 2*tid), (gid, 2*tid+1)
- raw_f2,raw_f3 = (gid+8, 2*tid), (gid+8, 2*tid+1)
- raw_f4,raw_f5 = (gid, 2*tid+8), (gid, 2*tid+9)
- raw_f6,raw_f7 = (gid+8, 2*tid+8), (gid+8, 2*tid+9)
- """
- r_off = block_idx * 16
- c_off = block_idx * 8
- gid = lane_id // 4
- tid = lane_id % 4
-
- _one = cutlass.Float32(1.0)
- _zero = cutlass.Float32(0.0)
-
- # Identity pattern (matches the SMEM-read variant's I_f computation)
- I_f0 = _one * cutlass.Float32(gid == 2 * tid) + _zero * cutlass.Float32(gid != 2 * tid)
- I_f1 = _one * cutlass.Float32(gid == 2 * tid + 1) + _zero * cutlass.Float32(gid != 2 * tid + 1)
- I_f2 = _one * cutlass.Float32(gid + 8 == 2 * tid) + _zero * cutlass.Float32(
- gid + 8 != 2 * tid
- ) # always 0
- I_f3 = _one * cutlass.Float32(gid + 8 == 2 * tid + 1) + _zero * cutlass.Float32(
- gid + 8 != 2 * tid + 1
- ) # always 0
- I_f4 = _one * cutlass.Float32(gid == 8 + 2 * tid) + _zero * cutlass.Float32(
- gid != 8 + 2 * tid
- ) # always 0
- I_f5 = _one * cutlass.Float32(gid == 8 + 2 * tid + 1) + _zero * cutlass.Float32(
- gid != 8 + 2 * tid + 1
- ) # always 0
- I_f6 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid) + _zero * cutlass.Float32(
- gid + 8 != 8 + 2 * tid
- )
- I_f7 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid + 1) + _zero * cutlass.Float32(
- gid + 8 != 8 + 2 * tid + 1
- )
-
- # Apply I+L mask to raw K2 values:
- # diag (row==col): replace with 1
- # strict upper (rowcol within block): keep raw value (= L value)
- m_gt_a = cutlass.Float32(gid > 2 * tid) # for cols 2*tid (with row=gid)
- m_gt_b = cutlass.Float32(gid > 2 * tid + 1) # for cols 2*tid+1
- # (gid, 2*tid) and (gid, 2*tid+1): mask conditional
- A_f0 = m_gt_a * raw_f0 + I_f0
- A_f1 = m_gt_b * raw_f1 + I_f1
- # (gid+8, 2*tid) and (gid+8, 2*tid+1): always strict lower (gid+8 > 2*tid+1 always)
- A_f2 = raw_f2
- A_f3 = raw_f3
- # (gid, 2*tid+8) and (gid, 2*tid+9): always strict upper (gid <= 7 < 8 <= 2*tid+8)
- A_f4 = _zero
- A_f5 = _zero
- # (gid+8, 2*tid+8): row-col offset = gid+8 - (2*tid+8) = gid - 2*tid -> same as A_f0 mask
- # (gid+8, 2*tid+9): similar -> same as A_f1 mask
- A_f6 = m_gt_a * raw_f6 + I_f6
- A_f7 = m_gt_b * raw_f7 + I_f7
-
- # L = A - I, INV = I - L (first two Neumann terms)
- L_f0 = A_f0 - I_f0
- L_f1 = A_f1 - I_f1
- L_f2 = A_f2 - I_f2
- L_f3 = A_f3 - I_f3
- L_f4 = A_f4 - I_f4
- L_f5 = A_f5 - I_f5
- L_f6 = A_f6 - I_f6
- L_f7 = A_f7 - I_f7
- INV_f0 = I_f0 - L_f0
- INV_f1 = I_f1 - L_f1
- INV_f2 = I_f2 - L_f2
- INV_f3 = I_f3 - L_f3
- INV_f4 = I_f4 - L_f4
- INV_f5 = I_f5 - L_f5
- INV_f6 = I_f6 - L_f6
- INV_f7 = I_f7 - L_f7
-
- _zf = cutlass.Float32(0.0)
- L_a0 = _ak_pack_bf16x2(L_f0, L_f1)
- L_a1 = _ak_pack_bf16x2(L_f2, L_f3)
- L_a2 = _ak_pack_bf16x2(L_f4, L_f5)
- L_a3 = _ak_pack_bf16x2(L_f6, L_f7)
- INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1)
- INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3)
- INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5)
- INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7)
-
- # Iter 1: L^2, INV += INV*L^2
- L_b0 = _ak_movmatrix_trans(L_a0)
- L_b1 = _ak_movmatrix_trans(L_a1)
- L_b2 = _ak_movmatrix_trans(L_a2)
- L_b3 = _ak_movmatrix_trans(L_a3)
- Lp_c0, Lp_c1, Lp_c2, Lp_c3 = _ak_mma(L_a0, L_a1, L_a2, L_a3, L_b0, L_b1, _zf, _zf, _zf, _zf)
- Lp_c4, Lp_c5, Lp_c6, Lp_c7 = _ak_mma(L_a0, L_a1, L_a2, L_a3, L_b2, L_b3, _zf, _zf, _zf, _zf)
- Lp_a0 = _ak_pack_bf16x2(Lp_c0, Lp_c1)
- Lp_a1 = _ak_pack_bf16x2(Lp_c2, Lp_c3)
- Lp_a2 = _ak_pack_bf16x2(Lp_c4, Lp_c5)
- Lp_a3 = _ak_pack_bf16x2(Lp_c6, Lp_c7)
- Lp_b0 = _ak_movmatrix_trans(Lp_a0)
- Lp_b1 = _ak_movmatrix_trans(Lp_a1)
- Lp_b2 = _ak_movmatrix_trans(Lp_a2)
- Lp_b3 = _ak_movmatrix_trans(Lp_a3)
- mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma(
- INV_a0, INV_a1, INV_a2, INV_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf
- )
- mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma(
- INV_a0, INV_a1, INV_a2, INV_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf
- )
- INV_f0 = INV_f0 + mm_c0
- INV_f1 = INV_f1 + mm_c1
- INV_f2 = INV_f2 + mm_c2
- INV_f3 = INV_f3 + mm_c3
- INV_f4 = INV_f4 + mm_c4
- INV_f5 = INV_f5 + mm_c5
- INV_f6 = INV_f6 + mm_c6
- INV_f7 = INV_f7 + mm_c7
- INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1)
- INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3)
- INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5)
- INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7)
-
- # Iter 2: L^4, INV += INV*L^4
- L4_c0, L4_c1, L4_c2, L4_c3 = _ak_mma(
- Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf
- )
- L4_c4, L4_c5, L4_c6, L4_c7 = _ak_mma(
- Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf
- )
- L4_a0 = _ak_pack_bf16x2(L4_c0, L4_c1)
- L4_a1 = _ak_pack_bf16x2(L4_c2, L4_c3)
- L4_a2 = _ak_pack_bf16x2(L4_c4, L4_c5)
- L4_a3 = _ak_pack_bf16x2(L4_c6, L4_c7)
- L4_b0 = _ak_movmatrix_trans(L4_a0)
- L4_b1 = _ak_movmatrix_trans(L4_a1)
- L4_b2 = _ak_movmatrix_trans(L4_a2)
- L4_b3 = _ak_movmatrix_trans(L4_a3)
- mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma(
- INV_a0, INV_a1, INV_a2, INV_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf
- )
- mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma(
- INV_a0, INV_a1, INV_a2, INV_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf
- )
- INV_f0 = INV_f0 + mm_c0
- INV_f1 = INV_f1 + mm_c1
- INV_f2 = INV_f2 + mm_c2
- INV_f3 = INV_f3 + mm_c3
- INV_f4 = INV_f4 + mm_c4
- INV_f5 = INV_f5 + mm_c5
- INV_f6 = INV_f6 + mm_c6
- INV_f7 = INV_f7 + mm_c7
- INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1)
- INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3)
- INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5)
- INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7)
-
- # Iter 3: L^8, INV += INV*L^8
- L8_c0, L8_c1, L8_c2, L8_c3 = _ak_mma(
- L4_a0, L4_a1, L4_a2, L4_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf
- )
- L8_c4, L8_c5, L8_c6, L8_c7 = _ak_mma(
- L4_a0, L4_a1, L4_a2, L4_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf
- )
- L8_a0 = _ak_pack_bf16x2(L8_c0, L8_c1)
- L8_a1 = _ak_pack_bf16x2(L8_c2, L8_c3)
- L8_a2 = _ak_pack_bf16x2(L8_c4, L8_c5)
- L8_a3 = _ak_pack_bf16x2(L8_c6, L8_c7)
- L8_b0 = _ak_movmatrix_trans(L8_a0)
- L8_b1 = _ak_movmatrix_trans(L8_a1)
- L8_b2 = _ak_movmatrix_trans(L8_a2)
- L8_b3 = _ak_movmatrix_trans(L8_a3)
- mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma(
- INV_a0, INV_a1, INV_a2, INV_a3, L8_b0, L8_b1, _zf, _zf, _zf, _zf
- )
- mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma(
- INV_a0, INV_a1, INV_a2, INV_a3, L8_b2, L8_b3, _zf, _zf, _zf, _zf
- )
- INV_f0 = INV_f0 + mm_c0
- INV_f1 = INV_f1 + mm_c1
- INV_f2 = INV_f2 + mm_c2
- INV_f3 = INV_f3 + mm_c3
- INV_f4 = INV_f4 + mm_c4
- INV_f5 = INV_f5 + mm_c5
- INV_f6 = INV_f6 + mm_c6
- INV_f7 = INV_f7 + mm_c7
-
- sAkk[r_off + gid, c_off + tid] = _ak_pack_bf16x2(INV_f0, INV_f1)
- sAkk[r_off + gid + 8, c_off + tid] = _ak_pack_bf16x2(INV_f2, INV_f3)
- sAkk[r_off + gid, c_off + 4 + tid] = _ak_pack_bf16x2(INV_f4, INV_f5)
- sAkk[r_off + gid + 8, c_off + 4 + tid] = _ak_pack_bf16x2(INV_f6, INV_f7)
-
-
-@dsl_user_op
-def _ak_matmul_AB(sAkk, br_A, bc_A, br_B, bc_B, lane_id, *, loc=None, ip=None):
- gid = lane_id // 4
- tid = lane_id % 4
- _zf = cutlass.Float32(0.0)
- rA = br_A * 16
- cA = bc_A * 8
- rB = br_B * 16
- cB = bc_B * 8
- a0 = cutlass.Float32(sAkk[rA + gid, cA + tid])
- a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + tid])
- a2 = cutlass.Float32(sAkk[rA + gid, cA + 4 + tid])
- a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 4 + tid])
- bA0 = cutlass.Float32(sAkk[rB + gid, cB + tid])
- bA1 = cutlass.Float32(sAkk[rB + gid + 8, cB + tid])
- bA2 = cutlass.Float32(sAkk[rB + gid, cB + 4 + tid])
- bA3 = cutlass.Float32(sAkk[rB + gid + 8, cB + 4 + tid])
- b0 = _ak_movmatrix_trans(bA0)
- b1 = _ak_movmatrix_trans(bA1)
- b2 = _ak_movmatrix_trans(bA2)
- b3 = _ak_movmatrix_trans(bA3)
- cn0_0, cn0_1, cn0_2, cn0_3 = _ak_mma(a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf)
- cn1_0, cn1_1, cn1_2, cn1_3 = _ak_mma(a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf)
- return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3
-
-
-@dsl_user_op
-def _ak_chain_mma_B(sAkk, br_B, bc_B, a0, a1, a2, a3, lane_id, *, loc=None, ip=None):
- gid = lane_id // 4
- tid = lane_id % 4
- _zf = cutlass.Float32(0.0)
- rB = br_B * 16
- cB = bc_B * 8
- bA0 = cutlass.Float32(sAkk[rB + gid, cB + tid])
- bA1 = cutlass.Float32(sAkk[rB + gid + 8, cB + tid])
- bA2 = cutlass.Float32(sAkk[rB + gid, cB + 4 + tid])
- bA3 = cutlass.Float32(sAkk[rB + gid + 8, cB + 4 + tid])
- b0 = _ak_movmatrix_trans(bA0)
- b1 = _ak_movmatrix_trans(bA1)
- b2 = _ak_movmatrix_trans(bA2)
- b3 = _ak_movmatrix_trans(bA3)
- cn0_0, cn0_1, cn0_2, cn0_3 = _ak_mma(a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf)
- cn1_0, cn1_1, cn1_2, cn1_3 = _ak_mma(a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf)
- return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3
-
-
-@dsl_user_op
-def _ak_chain_mma_A(sAkk, br_A, bc_A, b0, b1, b2, b3, lane_id, *, loc=None, ip=None):
- gid = lane_id // 4
- tid = lane_id % 4
- _zf = cutlass.Float32(0.0)
- rA = br_A * 16
- cA = bc_A * 8
- a0 = cutlass.Float32(sAkk[rA + gid, cA + tid])
- a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + tid])
- a2 = cutlass.Float32(sAkk[rA + gid, cA + 4 + tid])
- a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 4 + tid])
- cn0_0, cn0_1, cn0_2, cn0_3 = _ak_mma(a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf)
- cn1_0, cn1_1, cn1_2, cn1_3 = _ak_mma(a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf)
- return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3
-
-
-@dsl_user_op
-def _ak_store_neg_C(sAkk, br, bc, c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None):
- gid = lane_id // 4
- tid = lane_id % 4
- r = br * 16
- c = bc * 8
- sAkk[r + gid, c + tid] = _ak_pack_bf16x2(-c0, -c1)
- sAkk[r + gid + 8, c + tid] = _ak_pack_bf16x2(-c2, -c3)
- sAkk[r + gid, c + 4 + tid] = _ak_pack_bf16x2(-c4, -c5)
- sAkk[r + gid + 8, c + 4 + tid] = _ak_pack_bf16x2(-c6, -c7)
-
-
-@dsl_user_op
-def _ak_pack_C_to_A(c0, c1, c2, c3, c4, c5, c6, c7, *, loc=None, ip=None):
- a0 = _ak_pack_bf16x2(c0, c1)
- a1 = _ak_pack_bf16x2(c2, c3)
- a2 = _ak_pack_bf16x2(c4, c5)
- a3 = _ak_pack_bf16x2(c6, c7)
- return a0, a1, a2, a3
-
-
-@dsl_user_op
-def _ak_pack_C_to_B(c0, c1, c2, c3, c4, c5, c6, c7, *, loc=None, ip=None):
- a0 = _ak_pack_bf16x2(c0, c1)
- a1 = _ak_pack_bf16x2(c2, c3)
- a2 = _ak_pack_bf16x2(c4, c5)
- a3 = _ak_pack_bf16x2(c6, c7)
- b0 = _ak_movmatrix_trans(a0)
- b1 = _ak_movmatrix_trans(a1)
- b2 = _ak_movmatrix_trans(a2)
- b3 = _ak_movmatrix_trans(a3)
- return b0, b1, b2, b3
-
-
-@dsl_user_op
-def _ak_store_C_temp(sT, buf, c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None):
- gid = lane_id // 4
- tid = lane_id % 4
- sT[gid, 2 * tid, buf] = c0
- sT[gid, 2 * tid + 1, buf] = c1
- sT[gid + 8, 2 * tid, buf] = c2
- sT[gid + 8, 2 * tid + 1, buf] = c3
- sT[gid, 8 + 2 * tid, buf] = c4
- sT[gid, 8 + 2 * tid + 1, buf] = c5
- sT[gid + 8, 8 + 2 * tid, buf] = c6
- sT[gid + 8, 8 + 2 * tid + 1, buf] = c7
-
-
-@dsl_user_op
-def _ak_load_C_temp(sT, buf, lane_id, *, loc=None, ip=None):
- gid = lane_id // 4
- tid = lane_id % 4
- c0 = cutlass.Float32(sT[gid, 2 * tid, buf])
- c1 = cutlass.Float32(sT[gid, 2 * tid + 1, buf])
- c2 = cutlass.Float32(sT[gid + 8, 2 * tid, buf])
- c3 = cutlass.Float32(sT[gid + 8, 2 * tid + 1, buf])
- c4 = cutlass.Float32(sT[gid, 8 + 2 * tid, buf])
- c5 = cutlass.Float32(sT[gid, 8 + 2 * tid + 1, buf])
- c6 = cutlass.Float32(sT[gid + 8, 8 + 2 * tid, buf])
- c7 = cutlass.Float32(sT[gid + 8, 8 + 2 * tid + 1, buf])
- return c0, c1, c2, c3, c4, c5, c6, c7
-
-
-@dsl_user_op
-def opaque_zero_from_work_id(*, loc=None, ip=None):
- """
- Return 0 via opaque side-effectful asm (no inputs needed).
-
- Because has_side_effects=True, MLIR LICM treats this as having memory
- effects and will NOT hoist it outside the for_generate loop. Any value
- computed from the result (_oz) therefore appears loop-variant to LICM,
- preventing get_slice() and scalar layout-invariant computations from being
- hoisted to the kernel prologue. This keeps prologue register pressure < 64
- and eliminates the 440-byte stack frame that caused ~300-cycle L2 LDL
- penalties per iteration.
- """
- result = llvm.inline_asm(
- T.i32(),
- [],
- "mov.b32 $0, 0;", # output = 0 (opaque to compiler constant-folding)
- "=r",
- has_side_effects=True,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
- return cutlass.Int32(result)
-
-
-@cute.kernel
-def fused_kernel123(
- tma_atom_Q: cute.CopyAtom,
- tma_tensor_Q: cute.Tensor,
- tma_atom_K: cute.CopyAtom,
- tma_tensor_K: cute.Tensor,
- tma_atom_G: cute.CopyAtom,
- tma_tensor_G: cute.Tensor,
- mA_log: cute.Tensor,
- mBeta: cute.Tensor,
- mBetaActivated: cute.Tensor,
- scale: cutlass.Float32,
- mKscaled: cute.Tensor,
- mKg: cute.Tensor,
- mQscaled: cute.Tensor,
- mGkLast: cute.Tensor,
- mAqk: cute.Tensor, # 4D (B, T, H, BT) — used by per-tile pure path
- mAkk: cute.Tensor, # 4D (B, T, H, BT) — used by per-tile pure path
- mAqk_v2: cute.Tensor, # 5D (B, T, H, BT/2, 2) — used by vec autovec non-pure path
- mAkk_v2: cute.Tensor, # 5D (B, T, H, BT/2, 2) — used by vec autovec non-pure path
- tiled_copy_qk_k1,
- tiled_mma_k2,
- tiled_copy_mma_A,
- tiled_copy_mma_B,
- tiled_copy_Gcum_norm,
- tiled_copy_Gcum_gate,
- qk_smem_layout,
- g_smem_layout,
- g_cumsum_layout,
- num_chunks: cutlass.Int32, # runtime — shape-independent compile
- num_heads: int, # baked (single variant in practice)
- batch_size: cutlass.Int32, # runtime — shape-independent compile
- mCuSeqlens: cute.Tensor,
- mChunkIndices: cute.Tensor,
- IS_VARLEN: cutlass.Constexpr[int],
- mDtBias: cute.Tensor,
- lower_bound: cutlass.Float32,
- HAS_BIAS: cutlass.Constexpr[int],
- USE_SAFE_GATE: cutlass.Constexpr[int],
- valid_tokens: cutlass.Int32,
- USE_BETA_SIGMOID: cutlass.Constexpr[int],
- VARLEN_PURE: cutlass.Constexpr[int] = 0,
-):
- block_id, _, _ = cute.arch.block_idx()
- tidx = cute.arch.thread_idx()[0]
- warp_idx = cute.arch.warp_idx()
- warp_idx = cute.arch.make_warp_uniform(warp_idx)
- lane_id = tidx % 32
-
- total_cgs_per_head = cutlass.Int32(0)
- cgs_per_head = cutlass.Int32(0)
- total_cgs = cutlass.Int32(0)
- if IS_VARLEN:
- total_cgs_per_head = (num_chunks + CHUNKS_PER_BLOCK - 1) // CHUNKS_PER_BLOCK
- total_cgs = total_cgs_per_head * num_heads
- else:
- cgs_per_head = num_chunks // CHUNKS_PER_BLOCK
- total_cgs = cgs_per_head * num_heads * batch_size
-
- # =====================================================================
- # SMEM allocation
- # =====================================================================
- smem = cutlass.utils.SmemAllocator()
- sQ = smem.allocate_tensor(
- cutlass.BFloat16, qk_smem_layout.outer, 128, swizzle=qk_smem_layout.inner
- )
- sK = smem.allocate_tensor(
- cutlass.BFloat16, qk_smem_layout.outer, 128, swizzle=qk_smem_layout.inner
- )
- sG = smem.allocate_tensor(cutlass.BFloat16, g_smem_layout, 128)
- sGcum = smem.allocate_tensor(cutlass.Float32, g_cumsum_layout, 128)
- beta_layout = cute.make_layout((BT, NUM_STAGES), stride=(1, BT))
- sBeta = smem.allocate_tensor(cutlass.BFloat16, beta_layout, 128)
- partial_last_layout = cute.make_layout((K1_ROW_GROUPS, PARTIAL_COLS), stride=(PARTIAL_COLS, 1))
- sPartialLast = smem.allocate_tensor(cutlass.Float32, partial_last_layout, 128)
-
- # sAqk: 64x72 row-major (BT rows, BT+pad cols), same shape as sAkk.
- # Each sub-tile (i_q, i_k) sits at SMEM rows [i_q*BC, (i_q+1)*BC) cols
- # [i_k*BC, (i_k+1)*BC) — directly mirrors the 64x64 attention matrix.
- aqk_tile_layout = cute.make_layout(
- (BT, AQK_TILE_STRIDE, NUM_STAGES), stride=(AQK_TILE_STRIDE, 1, BT * AQK_TILE_STRIDE)
- )
- sAqk = smem.allocate_tensor(cutlass.BFloat16, aqk_tile_layout, 128)
-
- akk_tile_layout = cute.make_layout(
- (BT, AKK_STRIDE, NUM_STAGES), stride=(AKK_STRIDE, 1, BT * AKK_STRIDE)
- )
- sAkk = smem.allocate_tensor(cutlass.BFloat16, akk_tile_layout, 128)
- # sAkk_pkd / sTemp removed: akk_inv runs as a separate kernel call (chained back-to-back).
-
- # =====================================================================
- # Mbarrier allocation & init
- # =====================================================================
- tma_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES)
- stage_reuse_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES)
- k1_done_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES)
- mma_done_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES)
- store_done_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES)
-
- bytes_per_stage = BT * K_DIM * 2 * 3
-
- if tidx == 0:
- for s in range(NUM_STAGES):
- cute.arch.mbarrier_init(tma_mbars + s, 1)
- # Both K1 Pass 2b and MMA read Q/K before TMA can reuse a stage.
- # The dedicated TMA warp does not arrive on either barrier.
- cute.arch.mbarrier_init(
- stage_reuse_mbars + s, (NUM_MMA_WARPS - 1 + NUM_K1_TMA_WARPS) * 32
- )
- cute.arch.mbarrier_init(k1_done_mbars + s, NUM_K1_TMA_WARPS * 32)
- cute.arch.mbarrier_init(mma_done_mbars + s, (NUM_MMA_WARPS - 1) * 32)
- cute.arch.mbarrier_init(store_done_mbars + s, NUM_STORE_WARPS * 32)
- cute.arch.mbarrier_init_fence()
- cute.arch.barrier()
-
- # =====================================================================
- # SMEM init: zero out sAqk and sAkk valid 64x64 region (cols 64..71 are
- # padding for SMEM bank-conflict avoidance, never read or written by MMA
- # or store warps). Required for downstream row-major store optimizations
- # — positions outside MMA-written sub-tiles stay at 0.
- #
- # Cooperative pattern (31 warps × 32 lanes = 992 threads):
- # - Rows strided by warp count: warp w owns rows w, w+31, w+62 (< BT)
- # - Each lane owns 2 contiguous bf16 cols (lane*2, lane*2+1)
- # - Adjacent (lane*2, lane*2+1) bf16 pairs are 4-byte aligned →
- # ptxas should fuse into STS.32 wide stores.
- # =====================================================================
- _warp_id_in_cta = tidx >> 5 # tidx // 32, range 0..30
- _lane_id_warp = tidx & 31 # tidx % 32, range 0..31
- _col_lo = _lane_id_warp * 2 # this lane owns cols [_col_lo, _col_lo+1]
- _col_hi = _col_lo + 1
- for _s in cutlass.range_constexpr(NUM_STAGES):
- for _ri in cutlass.range_constexpr((BT + NUM_WARPS - 1) // NUM_WARPS):
- _row = _warp_id_in_cta + _ri * NUM_WARPS
- if _row < BT:
- sAqk[_row, _col_lo, _s] = cutlass.BFloat16(0.0)
- sAqk[_row, _col_hi, _s] = cutlass.BFloat16(0.0)
- sAkk[_row, _col_lo, _s] = cutlass.BFloat16(0.0)
- sAkk[_row, _col_hi, _s] = cutlass.BFloat16(0.0)
- cute.arch.barrier()
-
- # =====================================================================
- # Pre-arrive before the first TMA load and MMA store
- # stage_reuse_mbars: pre-arrive all 16 K1 and 10 MMA reader warps
- # store_done_mbars: MMA waits before Store arrives → pre-arrive first 4 MMA warps
- # =====================================================================
- if (
- warp_idx >= NUM_K1_TMA_WARPS
- and warp_idx < NUM_K1_TMA_WARPS + NUM_MMA_WARPS
- and warp_idx != TMA_WARP_ID
- ):
- mma_warp_tmp = warp_idx - NUM_K1_TMA_WARPS
- for s in range(NUM_STAGES):
- cute.arch.mbarrier_arrive(stage_reuse_mbars + s)
- if mma_warp_tmp < NUM_STORE_WARPS:
- cute.arch.mbarrier_arrive(store_done_mbars + s)
-
- if warp_idx < NUM_K1_TMA_WARPS:
- for s in range(NUM_STAGES):
- cute.arch.mbarrier_arrive(stage_reuse_mbars + s)
-
- # =================================================================
- # Persistent outer loop. Single for_generate at top level (required).
- # Opaque asm barrier on work_id prevents MLIR LICM from hoisting
- # get_slice() and scalar layout invariants to the kernel prologue,
- # keeping register pressure < 64 and eliminating prologue spill.
- # =================================================================
- for work_id in for_generate(block_id, total_cgs, NUM_SMS):
- i_cg = cutlass.Int32(0)
- i_h = cutlass.Int32(0)
- i_b = cutlass.Int32(0)
- chunk_base = cutlass.Int32(0)
- if IS_VARLEN:
- i_cg = work_id % total_cgs_per_head
- i_h = work_id // total_cgs_per_head
- i_b = cutlass.Int32(0)
- chunk_base = i_cg * CHUNKS_PER_BLOCK
- else:
- i_cg = work_id % cgs_per_head
- i_h = (work_id // cgs_per_head) % num_heads
- i_b = work_id // (cgs_per_head * num_heads)
- chunk_base = i_cg * CHUNKS_PER_BLOCK
-
- # Anti-LICM barrier: _oz is always 0 but appears to depend on work_id.
- # Because this asm has side_effects=True, it stays inside the loop.
- # Any value computed from _oz/_lane/_warp is also loop-variant from
- # LICM's perspective → get_slice() and scalar invariants stay in-loop.
- _oz = opaque_zero_from_work_id()
- _lane = lane_id + _oz
- _warp = warp_idx + _oz
-
- # =============================================================
- # Warps 0-15: Fused TMA + K1
- # =============================================================
- if warp_idx < NUM_K1_TMA_WARPS:
- # Warp-layout invariants (scope-local → no cross-group register spill)
- k1_warp = _warp
- warp_row_group = k1_warp % K1_ROW_GROUPS
- warp_col_group = k1_warp // K1_ROW_GROUPS
- k1_row_start = warp_row_group * ROWS_PER_K1_WARP
- col_base = warp_col_group * K1_COLS_PER_WARP + _lane * VEC
- col_vec_idx = warp_col_group * (K1_COLS_PER_WARP // VEC) + _lane
- cumsum_scale = cutlass.Float32(RCP_LN2)
- thr_copy_k1 = tiled_copy_qk_k1.get_slice(_lane)
-
- rAcc = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32)
- rPrefix = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32)
- rGkLast = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32)
- rKsOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16)
- rQsOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16)
- rKgOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16)
- rGkOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32)
-
- # exp_A depends on i_h (changes per work unit)
- exp_A = cute.exp(mA_log[i_h], fastmath=True)
-
- # Load dt_bias per (head, col) — broadcast across all rows
- rBias = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32)
- if HAS_BIAS:
- for vi in cutlass.range_constexpr(VEC):
- rBias[vi] = mDtBias[i_h, col_base + vi]
- else:
- for vi in cutlass.range_constexpr(VEC):
- rBias[vi] = cutlass.Float32(0.0)
-
- # 3D TMA head slices (fixed for this work unit's head)
- gQ_head = tma_tensor_Q[(None, None, i_h)]
- gK_head = tma_tensor_K[(None, None, i_h)]
- gG_head = tma_tensor_G[(None, None, i_h)]
-
- for chunk_iter in cutlass.range_constexpr(CHUNKS_PER_BLOCK):
- cur_stage = chunk_iter % NUM_STAGES
- cur_phase = chunk_iter // NUM_STAGES % 2
- chunk_idx = chunk_base + chunk_iter
- chunk_start = cutlass.Int32(0)
- ci_eos = cutlass.Int32(0)
- if IS_VARLEN:
- if chunk_idx < num_chunks:
- _sid = cutlass.Int32(mChunkIndices[chunk_idx, 0])
- chunk_start = (
- cutlass.Int32(mCuSeqlens[_sid])
- + cutlass.Int32(mChunkIndices[chunk_idx, 1]) * BT
- )
- ci_eos = cutlass.Int32(mCuSeqlens[_sid + 1])
- else:
- chunk_start = chunk_idx * BT
-
- cute.arch.mbarrier_wait(tma_mbars + cur_stage, cur_phase)
-
- csG = sG[(None, None, cur_stage)]
- csGcum = sGcum[(None, None, cur_stage)]
- csQ = sQ[(None, None, cur_stage)]
- csK = sK[(None, None, cur_stage)]
- rGact = cute.make_rmem_tensor(
- cute.make_layout((ROWS_PER_K1_WARP, VEC)), cutlass.Float32
- )
- for vi in cutlass.range_constexpr(VEC):
- rAcc[vi] = cutlass.Float32(0.0)
-
- # Activate beta once per token/head and retain the bf16
- # result in this stage for K2. The activated GMEM scratch is
- # consumed by akk_inv, beta's second pipeline consumer.
- if k1_warp < 2:
- beta_row = k1_warp * 32 + lane_id
- beta_t = chunk_start + beta_row
- beta_value = cutlass.Float32(0.0)
- if IS_VARLEN:
- if chunk_idx < num_chunks and beta_t < ci_eos:
- beta_value = mBeta[i_b, beta_t, i_h].to(cutlass.Float32)
- if USE_BETA_SIGMOID:
- beta_value = fast_rcp(
- cutlass.Float32(1.0)
- + cute.exp2(
- -beta_value * LOG2E,
- fastmath=True,
- )
- )
- mBetaActivated[i_b, beta_t, i_h] = beta_value.to(cutlass.BFloat16)
- else:
- if beta_t < valid_tokens:
- beta_value = mBeta[i_b, beta_t, i_h].to(cutlass.Float32)
- if USE_BETA_SIGMOID:
- beta_value = fast_rcp(
- cutlass.Float32(1.0)
- + cute.exp2(
- -beta_value * LOG2E,
- fastmath=True,
- )
- )
- if USE_BETA_SIGMOID:
- # Eqlen padding is allocated storage, so publish
- # explicit zeros for its invalid tail. Applying
- # sigmoid to the padded zero logits would produce
- # 0.5 and corrupt the recurrent state.
- mBetaActivated[i_b, beta_t, i_h] = beta_value.to(cutlass.BFloat16)
- sBeta[beta_row, cur_stage] = beta_value.to(cutlass.BFloat16)
-
- for ri in cutlass.range_constexpr(ROWS_PER_K1_WARP):
- row = k1_row_start + ri
- for vi in cutlass.range_constexpr(VEC):
- c = col_base + vi
- g_val = csG[row, c].to(cutlass.Float32)
- if HAS_BIAS:
- g_val = g_val + rBias[vi]
- g_activated = cutlass.Float32(0.0)
- if USE_SAFE_GATE:
- sigmoid_g = fast_rcp(
- cutlass.Float32(1.0)
- + cute.exp2(-exp_A * g_val * LOG2E, fastmath=True)
- )
- g_activated = lower_bound * sigmoid_g
- else:
- softplus_g = (
- cute.log2(
- cutlass.Float32(1.0) + cute.exp2(g_val * LOG2E, fastmath=True),
- fastmath=True,
- )
- * LN2
- )
- g_activated = -exp_A * softplus_g
- # Varlen: zero gate for out-of-bounds rows so cumsum
- # stays flat beyond the last valid position.
- # VARLEN_PURE=1 elides this at compile time — caller
- # guarantees all seq lengths are multiples of BT so no
- # chunk has OOB rows.
- if IS_VARLEN and not VARLEN_PURE:
- if chunk_start + row >= ci_eos:
- g_activated = cutlass.Float32(0.0)
- rGact[ri, vi] = g_activated
- rAcc[vi] = rAcc[vi] + g_activated
-
- for vi in cutlass.range_constexpr(VEC):
- sPartialLast[warp_row_group, col_base + vi] = rAcc[vi]
-
- k1_internal_barrier()
-
- prefix_col_start = k1_warp * PARTIAL_COLS_PER_WARP
- row_in_prefix = lane_id % K1_ROW_GROUPS
- col_in_group = lane_id // K1_ROW_GROUPS
-
- for j in cutlass.range_constexpr(PARTIAL_COLS_PER_WARP // 4):
- col = prefix_col_start + j * 4 + col_in_group
- val = cutlass.Float32(sPartialLast[row_in_prefix, col])
- tmp = cute.arch.shuffle_sync_up(val, 1, mask=-1, mask_and_clamp=SHFL_W8_CLAMP)
- if row_in_prefix >= 1:
- val = val + tmp
- tmp = cute.arch.shuffle_sync_up(val, 2, mask=-1, mask_and_clamp=SHFL_W8_CLAMP)
- if row_in_prefix >= 2:
- val = val + tmp
- tmp = cute.arch.shuffle_sync_up(val, 4, mask=-1, mask_and_clamp=SHFL_W8_CLAMP)
- if row_in_prefix >= 4:
- val = val + tmp
- sPartialLast[row_in_prefix, col] = val
-
- k1_internal_barrier()
-
- for vi in cutlass.range_constexpr(VEC):
- rGkLast[vi] = sPartialLast[K1_ROW_GROUPS - 1, col_base + vi]
-
- for vi in cutlass.range_constexpr(VEC):
- rPrefix[vi] = cutlass.Float32(0.0)
- if warp_row_group > 0:
- for vi in cutlass.range_constexpr(VEC):
- rPrefix[vi] = sPartialLast[warp_row_group - 1, col_base + vi]
-
- # The next chunk's partial sums overwrite sPartialLast, so every
- # K1 warp must have read its prefix and the chunk total first.
- k1_internal_barrier()
-
- # ---- Pass 2a: ONLY cumsum + write csGcum (critical path, minimal work) ----
- for vi in cutlass.range_constexpr(VEC):
- rAcc[vi] = rPrefix[vi]
-
- for ri in cutlass.range_constexpr(ROWS_PER_K1_WARP):
- row = k1_row_start + ri
- for vi in cutlass.range_constexpr(VEC):
- rAcc[vi] = rAcc[vi] + rGact[ri, vi]
- csGcum[row, col_base + vi] = rAcc[vi] * cumsum_scale
-
- # Signal MMA early: csGcum is ready
- cute.arch.mbarrier_arrive(k1_done_mbars + cur_stage)
-
- # ---- Pass 2b: recompute + write GMEM (overlaps with MMA, off critical path) ----
- for vi in cutlass.range_constexpr(VEC):
- rAcc[vi] = rPrefix[vi]
-
- for ri in cutlass.range_constexpr(ROWS_PER_K1_WARP):
- row = k1_row_start + ri
- t = chunk_start + row
-
- sK_tile = cute.local_tile(
- csK, tiler=(1, K1_COLS_PER_WARP), coord=(row, warp_col_group)
- )
- tCsK = thr_copy_k1.partition_S(sK_tile)
- tCrK = cute.make_fragment_like(tCsK)
- cute.copy(tiled_copy_qk_k1, tCsK, thr_copy_k1.retile(tCrK))
-
- sQ_tile = cute.local_tile(
- csQ, tiler=(1, K1_COLS_PER_WARP), coord=(row, warp_col_group)
- )
- tCsQ = thr_copy_k1.partition_S(sQ_tile)
- tCrQ = cute.make_fragment_like(tCsQ)
- cute.copy(tiled_copy_qk_k1, tCsQ, thr_copy_k1.retile(tCrQ))
-
- for vi in cutlass.range_constexpr(VEC):
- rAcc[vi] = rAcc[vi] + rGact[ri, vi]
- cs = rAcc[vi] * cumsum_scale
-
- k_val = tCrK[vi].to(cutlass.Float32)
- q_val = tCrQ[vi].to(cutlass.Float32)
-
- exp2_cs = cute.exp2(cs, fastmath=True)
- gk_last_cs = rGkLast[vi] * cumsum_scale
- exp2_kg = cute.exp2(gk_last_cs - cs, fastmath=True)
-
- rKsOut[vi] = (k_val * exp2_cs).to(cutlass.BFloat16)
- rQsOut[vi] = (q_val * exp2_cs * scale).to(cutlass.BFloat16)
- rKgOut[vi] = (k_val * exp2_kg).to(cutlass.BFloat16)
-
- if IS_VARLEN and not VARLEN_PURE:
- if t < ci_eos:
- cute.autovec_copy(rKsOut, mKscaled[i_b, t, i_h, col_vec_idx, None])
- cute.autovec_copy(rQsOut, mQscaled[i_b, t, i_h, col_vec_idx, None])
- cute.autovec_copy(rKgOut, mKg[i_b, t, i_h, col_vec_idx, None])
- elif IS_VARLEN and VARLEN_PURE:
- if chunk_idx < num_chunks:
- cute.autovec_copy(rKsOut, mKscaled[i_b, t, i_h, col_vec_idx, None])
- cute.autovec_copy(rQsOut, mQscaled[i_b, t, i_h, col_vec_idx, None])
- cute.autovec_copy(rKgOut, mKg[i_b, t, i_h, col_vec_idx, None])
- else:
- cute.autovec_copy(rKsOut, mKscaled[i_b, t, i_h, col_vec_idx, None])
- cute.autovec_copy(rQsOut, mQscaled[i_b, t, i_h, col_vec_idx, None])
- cute.autovec_copy(rKgOut, mKg[i_b, t, i_h, col_vec_idx, None])
-
- if warp_row_group == 0:
- for vi in cutlass.range_constexpr(VEC):
- rGkOut[vi] = cute.exp2(rGkLast[vi] * cumsum_scale, fastmath=True)
- if IS_VARLEN:
- if ci_eos > cutlass.Int32(0):
- cute.autovec_copy(
- rGkOut, mGkLast[i_b, chunk_idx, i_h, col_vec_idx, None]
- )
- else:
- cute.autovec_copy(rGkOut, mGkLast[i_b, chunk_idx, i_h, col_vec_idx, None])
-
- # K1 still reads Q/K in Pass 2b after signaling k1_done.
- # TMA may reuse this stage only once those reads finish too.
- cute.arch.mbarrier_arrive(stage_reuse_mbars + cur_stage)
-
- # =============================================================
- # Warp 26 (TMA_WARP_ID): dedicated TMA producer.
- # Waits stage_reuse (K1 and MMA readers), issues TMA for Q/K/G,
- # signals tma_mbar. Decouples MMA -> TMA dependency from K1 compute.
- # =============================================================
- if warp_idx == TMA_WARP_ID:
- gQ_head = tma_tensor_Q[(None, None, i_h)]
- gK_head = tma_tensor_K[(None, None, i_h)]
- gG_head = tma_tensor_G[(None, None, i_h)]
-
- # Prefetch chunk 0 -> stage 0 (stage_reuse[0] pre-arrived)
- pf_cs = cutlass.Int32(0)
- if IS_VARLEN:
- pf_seq_id_0 = cutlass.Int32(mChunkIndices[chunk_base, 0])
- pf_local_0 = cutlass.Int32(mChunkIndices[chunk_base, 1])
- pf_bos_0 = cutlass.Int32(mCuSeqlens[pf_seq_id_0])
- pf_cs = pf_bos_0 + pf_local_0 * BT
- else:
- pf_cs = i_b * num_chunks * BT + chunk_base * BT
- cute.arch.mbarrier_wait(stage_reuse_mbars, 0)
- if lane_id == 0:
- cute.arch.mbarrier_expect_tx(tma_mbars, bytes_per_stage)
- sQ_pf = sQ[(None, None, 0)]
- gQ_pf = cute.local_tile(cute.domain_offset((pf_cs, 0), gQ_head), (BT, K_DIM), (0, 0))
- ts_pf, tg_pf = cpasync.tma_partition(
- tma_atom_Q,
- 0,
- cute.make_layout(1),
- cute.group_modes(sQ_pf, 0, 2),
- cute.group_modes(gQ_pf, 0, 2),
- )
- cute.copy(tma_atom_Q, tg_pf, ts_pf, tma_bar_ptr=tma_mbars)
- sK_pf = sK[(None, None, 0)]
- gK_pf = cute.local_tile(cute.domain_offset((pf_cs, 0), gK_head), (BT, K_DIM), (0, 0))
- ts_pf, tg_pf = cpasync.tma_partition(
- tma_atom_K,
- 0,
- cute.make_layout(1),
- cute.group_modes(sK_pf, 0, 2),
- cute.group_modes(gK_pf, 0, 2),
- )
- cute.copy(tma_atom_K, tg_pf, ts_pf, tma_bar_ptr=tma_mbars)
- sG_pf = sG[(None, None, 0)]
- gG_pf = cute.local_tile(cute.domain_offset((pf_cs, 0), gG_head), (BT, K_DIM), (0, 0))
- ts_pf, tg_pf = cpasync.tma_partition(
- tma_atom_G,
- 0,
- cute.make_layout(1),
- cute.group_modes(sG_pf, 0, 2),
- cute.group_modes(gG_pf, 0, 2),
- )
- cute.copy(tma_atom_G, tg_pf, ts_pf, tma_bar_ptr=tma_mbars)
- if lane_id == 0:
- cute.arch.mbarrier_arrive(tma_mbars)
-
- # Issue TMAs for chunks 1..CHUNKS_PER_BLOCK-1
- for next_i in cutlass.range_constexpr(1, CHUNKS_PER_BLOCK):
- next_stage = next_i % NUM_STAGES
- next_phase = next_i // NUM_STAGES % 2
- next_cs = cutlass.Int32(0)
- if IS_VARLEN:
- next_chunk_idx = chunk_base + next_i
- if next_chunk_idx < num_chunks:
- _nsid = cutlass.Int32(mChunkIndices[next_chunk_idx, 0])
- next_cs = (
- cutlass.Int32(mCuSeqlens[_nsid])
- + cutlass.Int32(mChunkIndices[next_chunk_idx, 1]) * BT
- )
- else:
- next_cs = i_b * num_chunks * BT + (chunk_base + next_i) * BT
- cute.arch.mbarrier_wait(stage_reuse_mbars + next_stage, next_phase)
- if lane_id == 0:
- cute.arch.mbarrier_expect_tx(tma_mbars + next_stage, bytes_per_stage)
- sQ_ns = sQ[(None, None, next_stage)]
- gQ_ns = cute.local_tile(
- cute.domain_offset((next_cs, 0), gQ_head), (BT, K_DIM), (0, 0)
- )
- ts_ns, tg_ns = cpasync.tma_partition(
- tma_atom_Q,
- 0,
- cute.make_layout(1),
- cute.group_modes(sQ_ns, 0, 2),
- cute.group_modes(gQ_ns, 0, 2),
- )
- cute.copy(tma_atom_Q, tg_ns, ts_ns, tma_bar_ptr=tma_mbars + next_stage)
- sK_ns = sK[(None, None, next_stage)]
- gK_ns = cute.local_tile(
- cute.domain_offset((next_cs, 0), gK_head), (BT, K_DIM), (0, 0)
- )
- ts_ns, tg_ns = cpasync.tma_partition(
- tma_atom_K,
- 0,
- cute.make_layout(1),
- cute.group_modes(sK_ns, 0, 2),
- cute.group_modes(gK_ns, 0, 2),
- )
- cute.copy(tma_atom_K, tg_ns, ts_ns, tma_bar_ptr=tma_mbars + next_stage)
- sG_ns = sG[(None, None, next_stage)]
- gG_ns = cute.local_tile(
- cute.domain_offset((next_cs, 0), gG_head), (BT, K_DIM), (0, 0)
- )
- ts_ns, tg_ns = cpasync.tma_partition(
- tma_atom_G,
- 0,
- cute.make_layout(1),
- cute.group_modes(sG_ns, 0, 2),
- cute.group_modes(gG_ns, 0, 2),
- )
- cute.copy(tma_atom_G, tg_ns, ts_ns, tma_bar_ptr=tma_mbars + next_stage)
- if lane_id == 0:
- cute.arch.mbarrier_arrive(tma_mbars + next_stage)
-
- # =============================================================
- # Warps 16-26 (excluding TMA_WARP_ID=26): K2 MMA Compute
- # =============================================================
- if (
- warp_idx >= NUM_K1_TMA_WARPS
- and warp_idx < NUM_K1_TMA_WARPS + NUM_MMA_WARPS
- and warp_idx != TMA_WARP_ID
- ):
- # Warp-layout invariants (scope-local → no cross-group register spill)
- _tid_in_group = _lane % 4
- _group_id = _lane // 4
- mma_warp = _warp - NUM_K1_TMA_WARPS
- my_i_q = cutlass.Int32(0)
- my_i_k = cutlass.Int32(0)
- if mma_warp < 1:
- my_i_q = cutlass.Int32(0)
- my_i_k = mma_warp
- elif mma_warp < 3:
- my_i_q = cutlass.Int32(1)
- my_i_k = mma_warp - 1
- elif mma_warp < 6:
- my_i_q = cutlass.Int32(2)
- my_i_k = mma_warp - 3
- elif mma_warp < NUM_MMA_ACTIVE:
- my_i_q = cutlass.Int32(3)
- my_i_k = mma_warp - 6
- q_row_base = my_i_q * BC
- k_row_base = my_i_k * BC
- akk_row_base = k_row_base
- akk_col_base = q_row_base
- norm_row = q_row_base
- if my_i_q == my_i_k:
- norm_row = q_row_base + cutlass.Int32(BC // 2)
- row0 = _group_id
- row1 = _group_id + 8
- col0 = _tid_in_group * 2
- col1 = _tid_in_group * 2 + 1
- col2 = 8 + _tid_in_group * 2
- col3 = 8 + _tid_in_group * 2 + 1
- thr_mma = tiled_mma_k2.get_slice(_lane)
- thr_copy_A = tiled_copy_mma_A.get_slice(_lane)
- thr_copy_B = tiled_copy_mma_B.get_slice(_lane)
- thr_copy_Gn = tiled_copy_Gcum_norm.get_slice(_tid_in_group)
- thr_copy_Ggate = tiled_copy_Gcum_gate.get_slice(_lane)
-
- for chunk_iter in cutlass.range_constexpr(CHUNKS_PER_BLOCK):
- s = chunk_iter % NUM_STAGES
- phase = chunk_iter // NUM_STAGES % 2
- chunk_idx = chunk_base + chunk_iter
- chunk_start = cutlass.Int32(0)
- if IS_VARLEN:
- if chunk_idx < num_chunks:
- _sid = cutlass.Int32(mChunkIndices[chunk_idx, 0])
- chunk_start = (
- cutlass.Int32(mCuSeqlens[_sid])
- + cutlass.Int32(mChunkIndices[chunk_idx, 1]) * BT
- )
- else:
- chunk_start = chunk_idx * BT
-
- cute.arch.mbarrier_wait(k1_done_mbars + s, phase)
- cute.arch.mbarrier_wait(store_done_mbars + s, phase)
-
- if mma_warp < NUM_MMA_ACTIVE:
- csQ = sQ[(None, None, s)]
- csK = sK[(None, None, s)]
- csGcum = sGcum[(None, None, s)]
- csAqk = sAqk[(None, None, s)]
- csAkk = sAkk[(None, None, s)]
-
- _z = cutlass.Float32(0.0)
-
- # K1 stages activated beta with invalid rows set to zero.
- # Reading it from SMEM also avoids repeating sigmoid for
- # every lower-triangular K2 tile that shares a query row.
- beta_row0 = sBeta[q_row_base + row0, s].to(cutlass.Float32)
- beta_row1 = sBeta[q_row_base + row1, s].to(cutlass.Float32)
-
- acc_aqk_n0_0, acc_aqk_n0_1, acc_aqk_n0_2, acc_aqk_n0_3 = _z, _z, _z, _z
- acc_aqk_n1_0, acc_aqk_n1_1, acc_aqk_n1_2, acc_aqk_n1_3 = _z, _z, _z, _z
- acc_akk_n0_0, acc_akk_n0_1, acc_akk_n0_2, acc_akk_n0_3 = _z, _z, _z, _z
- acc_akk_n1_0, acc_akk_n1_1, acc_akk_n1_2, acc_akk_n1_3 = _z, _z, _z, _z
-
- # bf16 m16n8k16 MMA: each k_block covers k=16.
- for k_block in cutlass.range_constexpr(NUM_MMA_K_TILES):
- # ---- Load Q/Kq bf16 fragments (16x16, 8 bf16/thread) ----
- sQ_tile = cute.local_tile(csQ, tiler=(16, 16), coord=(my_i_q, k_block))
- tCrQ = tiled_mma_k2.make_fragment_A(thr_mma.partition_A(sQ_tile))
- cute.copy(
- tiled_copy_mma_A,
- thr_copy_A.partition_S(sQ_tile),
- thr_copy_A.retile(tCrQ),
- )
-
- sKq_tile = cute.local_tile(csK, tiler=(16, 16), coord=(my_i_q, k_block))
- tCrKq = tiled_mma_k2.make_fragment_A(thr_mma.partition_A(sKq_tile))
- cute.copy(
- tiled_copy_mma_A,
- thr_copy_A.partition_S(sKq_tile),
- thr_copy_A.retile(tCrKq),
- )
-
- # ---- Issue K n0/n1 LDSMs early for better ILP ----
- sK_tile_n0 = cute.local_tile(
- csK, tiler=(8, 16), coord=(my_i_k * 2, k_block)
- )
- tCrK_n0 = tiled_mma_k2.make_fragment_B(thr_mma.partition_B(sK_tile_n0))
- cute.copy(
- tiled_copy_mma_B,
- thr_copy_B.partition_S(sK_tile_n0),
- thr_copy_B.retile(tCrK_n0),
- )
-
- sK_tile_n1 = cute.local_tile(
- csK, tiler=(8, 16), coord=(my_i_k * 2 + 1, k_block)
- )
- tCrK_n1 = tiled_mma_k2.make_fragment_B(thr_mma.partition_B(sK_tile_n1))
- cute.copy(
- tiled_copy_mma_B,
- thr_copy_B.partition_S(sK_tile_n1),
- thr_copy_B.retile(tCrK_n1),
- )
-
- # ---- Gate norm (2x k=8 covers k=16) ----
- sGn_a = cute.local_tile(csGcum, tiler=(1, 8), coord=(norm_row, k_block * 2))
- tCsGn_a = thr_copy_Gn.partition_S(sGn_a)
- tCrGn_a = cute.make_fragment_like(tCsGn_a, cutlass.Float32)
- cute.copy(tiled_copy_Gcum_norm, tCsGn_a, thr_copy_Gn.retile(tCrGn_a))
- gn_a0 = tCrGn_a[0]
- gn_a1 = tCrGn_a[1]
-
- sGn_b = cute.local_tile(
- csGcum, tiler=(1, 8), coord=(norm_row, k_block * 2 + 1)
- )
- tCsGn_b = thr_copy_Gn.partition_S(sGn_b)
- tCrGn_b = cute.make_fragment_like(tCsGn_b, cutlass.Float32)
- cute.copy(tiled_copy_Gcum_norm, tCsGn_b, thr_copy_Gn.retile(tCrGn_b))
- gn_b0 = tCrGn_b[0]
- gn_b1 = tCrGn_b[1]
-
- # ---- Gate Q (2x (16,8) partition_C covers m=16,k=16) ----
- sGq_a = cute.local_tile(csGcum, tiler=(16, 8), coord=(my_i_q, k_block * 2))
- tCrGq_a = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGq_a))
- cute.copy(
- tiled_copy_Gcum_gate,
- thr_copy_Ggate.partition_S(sGq_a),
- thr_copy_Ggate.retile(tCrGq_a),
- )
-
- sGq_b = cute.local_tile(
- csGcum, tiler=(16, 8), coord=(my_i_q, k_block * 2 + 1)
- )
- tCrGq_b = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGq_b))
- cute.copy(
- tiled_copy_Gcum_gate,
- thr_copy_Ggate.partition_S(sGq_b),
- thr_copy_Ggate.retile(tCrGq_b),
- )
-
- # 8 Q gate values per thread (matching A bf16 m16n8k16 layout):
- # first half k=0..7 (tCrGq_a): a0=(r0,c0) a1=(r0,c0+1) a2=(r0+8,c0) a3=(r0+8,c0+1)
- # second half k=8..15 (tCrGq_b): a4=(r0,c0+8) a5=(r0,c0+9) a6=(r0+8,c0+8) a7=(r0+8,c0+9)
- gate_q_0 = cute.exp2(tCrGq_a[0] - gn_a0, fastmath=True)
- gate_q_1 = cute.exp2(tCrGq_a[1] - gn_a1, fastmath=True)
- gate_q_2 = cute.exp2(tCrGq_a[2] - gn_a0, fastmath=True)
- gate_q_3 = cute.exp2(tCrGq_a[3] - gn_a1, fastmath=True)
- gate_q_4 = cute.exp2(tCrGq_b[0] - gn_b0, fastmath=True)
- gate_q_5 = cute.exp2(tCrGq_b[1] - gn_b1, fastmath=True)
- gate_q_6 = cute.exp2(tCrGq_b[2] - gn_b0, fastmath=True)
- gate_q_7 = cute.exp2(tCrGq_b[3] - gn_b1, fastmath=True)
-
- # qa fp32 = Q*gate (8 per thread). tCrQ indexing assumption:
- # [0..3] first k-chunk (k=0..7), [4..7] second k-chunk (k=8..15).
- qa0 = tCrQ[0].to(cutlass.Float32) * gate_q_0
- qa1 = tCrQ[1].to(cutlass.Float32) * gate_q_1
- qa2 = tCrQ[2].to(cutlass.Float32) * gate_q_2
- qa3 = tCrQ[3].to(cutlass.Float32) * gate_q_3
- qa4 = tCrQ[4].to(cutlass.Float32) * gate_q_4
- qa5 = tCrQ[5].to(cutlass.Float32) * gate_q_5
- qa6 = tCrQ[6].to(cutlass.Float32) * gate_q_6
- qa7 = tCrQ[7].to(cutlass.Float32) * gate_q_7
-
- # Pack fp32 qa -> 4 u32 bf16x2 for MMA A (k-adjacent pairs per u32)
- qa_u32_0 = pack_bf16x2_f32(qa1, qa0) # reg0 = [qa0 lo | qa1 hi]
- qa_u32_1 = pack_bf16x2_f32(qa3, qa2)
- qa_u32_2 = pack_bf16x2_f32(qa5, qa4)
- qa_u32_3 = pack_bf16x2_f32(qa7, qa6)
-
- ka0 = tCrKq[0].to(cutlass.Float32) * gate_q_0
- ka1 = tCrKq[1].to(cutlass.Float32) * gate_q_1
- ka2 = tCrKq[2].to(cutlass.Float32) * gate_q_2
- ka3 = tCrKq[3].to(cutlass.Float32) * gate_q_3
- ka4 = tCrKq[4].to(cutlass.Float32) * gate_q_4
- ka5 = tCrKq[5].to(cutlass.Float32) * gate_q_5
- ka6 = tCrKq[6].to(cutlass.Float32) * gate_q_6
- ka7 = tCrKq[7].to(cutlass.Float32) * gate_q_7
-
- ka_u32_0 = pack_bf16x2_f32(ka1, ka0)
- ka_u32_1 = pack_bf16x2_f32(ka3, ka2)
- ka_u32_2 = pack_bf16x2_f32(ka5, ka4)
- ka_u32_3 = pack_bf16x2_f32(ka7, ka6)
-
- # ---- Gate K (2x (16,8) partition_C covers m=16,k=16) ----
- sGk_a = cute.local_tile(csGcum, tiler=(16, 8), coord=(my_i_k, k_block * 2))
- tCrGk_a = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGk_a))
- cute.copy(
- tiled_copy_Gcum_gate,
- thr_copy_Ggate.partition_S(sGk_a),
- thr_copy_Ggate.retile(tCrGk_a),
- )
-
- sGk_b = cute.local_tile(
- csGcum, tiler=(16, 8), coord=(my_i_k, k_block * 2 + 1)
- )
- tCrGk_b = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGk_b))
- cute.copy(
- tiled_copy_Gcum_gate,
- thr_copy_Ggate.partition_S(sGk_b),
- thr_copy_Ggate.retile(tCrGk_b),
- )
-
- # n0 uses rows 0..7 of (16,*) tile (tCrGk_*[0,1])
- # n1 uses rows 8..15 of (16,*) tile (tCrGk_*[2,3])
- gk_n0_0 = cute.exp2(gn_a0 - tCrGk_a[0], fastmath=True)
- gk_n0_1 = cute.exp2(gn_a1 - tCrGk_a[1], fastmath=True)
- gk_n0_2 = cute.exp2(gn_b0 - tCrGk_b[0], fastmath=True)
- gk_n0_3 = cute.exp2(gn_b1 - tCrGk_b[1], fastmath=True)
-
- gk_n1_0 = cute.exp2(gn_a0 - tCrGk_a[2], fastmath=True)
- gk_n1_1 = cute.exp2(gn_a1 - tCrGk_a[3], fastmath=True)
- gk_n1_2 = cute.exp2(gn_b0 - tCrGk_b[2], fastmath=True)
- gk_n1_3 = cute.exp2(gn_b1 - tCrGk_b[3], fastmath=True)
-
- # tCrK_n0 bf16 fragment: 4 elems/thread at (n_row, c0), (n_row, c0+1),
- # (n_row, c0+8), (n_row, c0+9) — k-adjacent pairs
- k_n0_b0 = tCrK_n0[0].to(cutlass.Float32) * gk_n0_0
- k_n0_b1 = tCrK_n0[1].to(cutlass.Float32) * gk_n0_1
- k_n0_b2 = tCrK_n0[2].to(cutlass.Float32) * gk_n0_2
- k_n0_b3 = tCrK_n0[3].to(cutlass.Float32) * gk_n0_3
-
- k_n1_b0 = tCrK_n1[0].to(cutlass.Float32) * gk_n1_0
- k_n1_b1 = tCrK_n1[1].to(cutlass.Float32) * gk_n1_1
- k_n1_b2 = tCrK_n1[2].to(cutlass.Float32) * gk_n1_2
- k_n1_b3 = tCrK_n1[3].to(cutlass.Float32) * gk_n1_3
-
- # Pack fp32 k -> 2 u32 bf16x2 for MMA B
- k_n0_u32_0 = pack_bf16x2_f32(k_n0_b1, k_n0_b0)
- k_n0_u32_1 = pack_bf16x2_f32(k_n0_b3, k_n0_b2)
- k_n1_u32_0 = pack_bf16x2_f32(k_n1_b1, k_n1_b0)
- k_n1_u32_1 = pack_bf16x2_f32(k_n1_b3, k_n1_b2)
-
- # ---- 4 bf16 MMA calls ----
- acc_aqk_n0_0, acc_aqk_n0_1, acc_aqk_n0_2, acc_aqk_n0_3 = mma_bf16_m16n8k16(
- qa_u32_0,
- qa_u32_1,
- qa_u32_2,
- qa_u32_3,
- k_n0_u32_0,
- k_n0_u32_1,
- acc_aqk_n0_0,
- acc_aqk_n0_1,
- acc_aqk_n0_2,
- acc_aqk_n0_3,
- )
- acc_aqk_n1_0, acc_aqk_n1_1, acc_aqk_n1_2, acc_aqk_n1_3 = mma_bf16_m16n8k16(
- qa_u32_0,
- qa_u32_1,
- qa_u32_2,
- qa_u32_3,
- k_n1_u32_0,
- k_n1_u32_1,
- acc_aqk_n1_0,
- acc_aqk_n1_1,
- acc_aqk_n1_2,
- acc_aqk_n1_3,
- )
- acc_akk_n0_0, acc_akk_n0_1, acc_akk_n0_2, acc_akk_n0_3 = mma_bf16_m16n8k16(
- ka_u32_0,
- ka_u32_1,
- ka_u32_2,
- ka_u32_3,
- k_n0_u32_0,
- k_n0_u32_1,
- acc_akk_n0_0,
- acc_akk_n0_1,
- acc_akk_n0_2,
- acc_akk_n0_3,
- )
- acc_akk_n1_0, acc_akk_n1_1, acc_akk_n1_2, acc_akk_n1_3 = mma_bf16_m16n8k16(
- ka_u32_0,
- ka_u32_1,
- ka_u32_2,
- ka_u32_3,
- k_n1_u32_0,
- k_n1_u32_1,
- acc_akk_n1_0,
- acc_akk_n1_1,
- acc_akk_n1_2,
- acc_akk_n1_3,
- )
-
- # sQ/sK/sG reads done, signal TMA before SMEM writes
- cute.arch.mbarrier_arrive(stage_reuse_mbars + s)
-
- # Dual-path MMA write (constexpr-gated):
- # - non-pure: apply causal + diag=1 inline so SMEM matches
- # final GMEM layout (pairs with row-major vec autovec
- # store warp that does no causal).
- # - pure: write all 16x16 unconditionally (pairs with
- # per-tile store warp that applies causal + diag in
- # store; this is the original baseline behavior — no
- # extra MMA-write cost).
- _z16 = cutlass.BFloat16(0.0)
- _one16 = cutlass.BFloat16(1.0)
- if IS_VARLEN and not VARLEN_PURE:
- if my_i_q == my_i_k:
- # Diagonal sub-tile: causal-mask sAqk and write
- # diag=1 / strict-lower=MMA*beta / strict-upper=0
- # to sAkk so SMEM is in final form.
- _v_q00 = (acc_aqk_n0_0 * scale).to(cutlass.BFloat16)
- if row0 < col0:
- _v_q00 = _z16
- csAqk[q_row_base + row0, k_row_base + col0] = _v_q00
- _v_q01 = (acc_aqk_n0_1 * scale).to(cutlass.BFloat16)
- if row0 < col1:
- _v_q01 = _z16
- csAqk[q_row_base + row0, k_row_base + col1] = _v_q01
- _v_q02 = (acc_aqk_n0_2 * scale).to(cutlass.BFloat16)
- if row1 < col0:
- _v_q02 = _z16
- csAqk[q_row_base + row1, k_row_base + col0] = _v_q02
- _v_q03 = (acc_aqk_n0_3 * scale).to(cutlass.BFloat16)
- if row1 < col1:
- _v_q03 = _z16
- csAqk[q_row_base + row1, k_row_base + col1] = _v_q03
- _v_q04 = (acc_aqk_n1_0 * scale).to(cutlass.BFloat16)
- if row0 < col2:
- _v_q04 = _z16
- csAqk[q_row_base + row0, k_row_base + col2] = _v_q04
- _v_q05 = (acc_aqk_n1_1 * scale).to(cutlass.BFloat16)
- if row0 < col3:
- _v_q05 = _z16
- csAqk[q_row_base + row0, k_row_base + col3] = _v_q05
- _v_q06 = (acc_aqk_n1_2 * scale).to(cutlass.BFloat16)
- if row1 < col2:
- _v_q06 = _z16
- csAqk[q_row_base + row1, k_row_base + col2] = _v_q06
- _v_q07 = (acc_aqk_n1_3 * scale).to(cutlass.BFloat16)
- if row1 < col3:
- _v_q07 = _z16
- csAqk[q_row_base + row1, k_row_base + col3] = _v_q07
-
- _v_k00 = (acc_akk_n0_0 * beta_row0).to(cutlass.BFloat16)
- if row0 == col0:
- _v_k00 = _one16
- if row0 < col0:
- _v_k00 = _z16
- csAkk[akk_row_base + row0, akk_col_base + col0] = _v_k00
- _v_k01 = (acc_akk_n0_1 * beta_row0).to(cutlass.BFloat16)
- if row0 == col1:
- _v_k01 = _one16
- if row0 < col1:
- _v_k01 = _z16
- csAkk[akk_row_base + row0, akk_col_base + col1] = _v_k01
- _v_k02 = (acc_akk_n0_2 * beta_row1).to(cutlass.BFloat16)
- if row1 == col0:
- _v_k02 = _one16
- if row1 < col0:
- _v_k02 = _z16
- csAkk[akk_row_base + row1, akk_col_base + col0] = _v_k02
- _v_k03 = (acc_akk_n0_3 * beta_row1).to(cutlass.BFloat16)
- if row1 == col1:
- _v_k03 = _one16
- if row1 < col1:
- _v_k03 = _z16
- csAkk[akk_row_base + row1, akk_col_base + col1] = _v_k03
- _v_k04 = (acc_akk_n1_0 * beta_row0).to(cutlass.BFloat16)
- if row0 == col2:
- _v_k04 = _one16
- if row0 < col2:
- _v_k04 = _z16
- csAkk[akk_row_base + row0, akk_col_base + col2] = _v_k04
- _v_k05 = (acc_akk_n1_1 * beta_row0).to(cutlass.BFloat16)
- if row0 == col3:
- _v_k05 = _one16
- if row0 < col3:
- _v_k05 = _z16
- csAkk[akk_row_base + row0, akk_col_base + col3] = _v_k05
- _v_k06 = (acc_akk_n1_2 * beta_row1).to(cutlass.BFloat16)
- if row1 == col2:
- _v_k06 = _one16
- if row1 < col2:
- _v_k06 = _z16
- csAkk[akk_row_base + row1, akk_col_base + col2] = _v_k06
- _v_k07 = (acc_akk_n1_3 * beta_row1).to(cutlass.BFloat16)
- if row1 == col3:
- _v_k07 = _one16
- if row1 < col3:
- _v_k07 = _z16
- csAkk[akk_row_base + row1, akk_col_base + col3] = _v_k07
- else:
- # Non-diag (i_q > i_k): write all 16x16 unchanged.
- csAqk[q_row_base + row0, k_row_base + col0] = (acc_aqk_n0_0 * scale).to(
- cutlass.BFloat16
- )
- csAqk[q_row_base + row0, k_row_base + col1] = (acc_aqk_n0_1 * scale).to(
- cutlass.BFloat16
- )
- csAqk[q_row_base + row1, k_row_base + col0] = (acc_aqk_n0_2 * scale).to(
- cutlass.BFloat16
- )
- csAqk[q_row_base + row1, k_row_base + col1] = (acc_aqk_n0_3 * scale).to(
- cutlass.BFloat16
- )
- csAqk[q_row_base + row0, k_row_base + col2] = (acc_aqk_n1_0 * scale).to(
- cutlass.BFloat16
- )
- csAqk[q_row_base + row0, k_row_base + col3] = (acc_aqk_n1_1 * scale).to(
- cutlass.BFloat16
- )
- csAqk[q_row_base + row1, k_row_base + col2] = (acc_aqk_n1_2 * scale).to(
- cutlass.BFloat16
- )
- csAqk[q_row_base + row1, k_row_base + col3] = (acc_aqk_n1_3 * scale).to(
- cutlass.BFloat16
- )
- csAkk[akk_row_base + row0, akk_col_base + col0] = (
- acc_akk_n0_0 * beta_row0
- ).to(cutlass.BFloat16)
- csAkk[akk_row_base + row0, akk_col_base + col1] = (
- acc_akk_n0_1 * beta_row0
- ).to(cutlass.BFloat16)
- csAkk[akk_row_base + row1, akk_col_base + col0] = (
- acc_akk_n0_2 * beta_row1
- ).to(cutlass.BFloat16)
- csAkk[akk_row_base + row1, akk_col_base + col1] = (
- acc_akk_n0_3 * beta_row1
- ).to(cutlass.BFloat16)
- csAkk[akk_row_base + row0, akk_col_base + col2] = (
- acc_akk_n1_0 * beta_row0
- ).to(cutlass.BFloat16)
- csAkk[akk_row_base + row0, akk_col_base + col3] = (
- acc_akk_n1_1 * beta_row0
- ).to(cutlass.BFloat16)
- csAkk[akk_row_base + row1, akk_col_base + col2] = (
- acc_akk_n1_2 * beta_row1
- ).to(cutlass.BFloat16)
- csAkk[akk_row_base + row1, akk_col_base + col3] = (
- acc_akk_n1_3 * beta_row1
- ).to(cutlass.BFloat16)
- else:
- # PURE: write all 16x16 unconditionally — store warp
- # applies causal+diag in its per-tile loop. This is
- # the original baseline behavior (no extra MMA cost).
- csAqk[q_row_base + row0, k_row_base + col0] = (acc_aqk_n0_0 * scale).to(
- cutlass.BFloat16
- )
- csAqk[q_row_base + row0, k_row_base + col1] = (acc_aqk_n0_1 * scale).to(
- cutlass.BFloat16
- )
- csAqk[q_row_base + row1, k_row_base + col0] = (acc_aqk_n0_2 * scale).to(
- cutlass.BFloat16
- )
- csAqk[q_row_base + row1, k_row_base + col1] = (acc_aqk_n0_3 * scale).to(
- cutlass.BFloat16
- )
- csAqk[q_row_base + row0, k_row_base + col2] = (acc_aqk_n1_0 * scale).to(
- cutlass.BFloat16
- )
- csAqk[q_row_base + row0, k_row_base + col3] = (acc_aqk_n1_1 * scale).to(
- cutlass.BFloat16
- )
- csAqk[q_row_base + row1, k_row_base + col2] = (acc_aqk_n1_2 * scale).to(
- cutlass.BFloat16
- )
- csAqk[q_row_base + row1, k_row_base + col3] = (acc_aqk_n1_3 * scale).to(
- cutlass.BFloat16
- )
- csAkk[akk_row_base + row0, akk_col_base + col0] = (
- acc_akk_n0_0 * beta_row0
- ).to(cutlass.BFloat16)
- csAkk[akk_row_base + row0, akk_col_base + col1] = (
- acc_akk_n0_1 * beta_row0
- ).to(cutlass.BFloat16)
- csAkk[akk_row_base + row1, akk_col_base + col0] = (
- acc_akk_n0_2 * beta_row1
- ).to(cutlass.BFloat16)
- csAkk[akk_row_base + row1, akk_col_base + col1] = (
- acc_akk_n0_3 * beta_row1
- ).to(cutlass.BFloat16)
- csAkk[akk_row_base + row0, akk_col_base + col2] = (
- acc_akk_n1_0 * beta_row0
- ).to(cutlass.BFloat16)
- csAkk[akk_row_base + row0, akk_col_base + col3] = (
- acc_akk_n1_1 * beta_row0
- ).to(cutlass.BFloat16)
- csAkk[akk_row_base + row1, akk_col_base + col2] = (
- acc_akk_n1_2 * beta_row1
- ).to(cutlass.BFloat16)
- csAkk[akk_row_base + row1, akk_col_base + col3] = (
- acc_akk_n1_3 * beta_row1
- ).to(cutlass.BFloat16)
- else:
- cute.arch.mbarrier_arrive(stage_reuse_mbars + s)
-
- cute.arch.mbarrier_arrive(mma_done_mbars + s)
-
- # =============================================================
- # Warps 27-30: Store/Inversion warps
- # =============================================================
- if warp_idx >= NUM_K1_TMA_WARPS + NUM_MMA_WARPS:
- store_warp = warp_idx - (NUM_K1_TMA_WARPS + NUM_MMA_WARPS)
- for chunk_iter in cutlass.range_constexpr(CHUNKS_PER_BLOCK):
- s = chunk_iter % NUM_STAGES
- phase = chunk_iter // NUM_STAGES % 2
- chunk_idx = chunk_base + chunk_iter
- chunk_start = cutlass.Int32(0)
- st_eos = cutlass.Int32(0)
- if IS_VARLEN:
- if chunk_idx < num_chunks:
- _sid = cutlass.Int32(mChunkIndices[chunk_idx, 0])
- chunk_start = (
- cutlass.Int32(mCuSeqlens[_sid])
- + cutlass.Int32(mChunkIndices[chunk_idx, 1]) * BT
- )
- st_eos = cutlass.Int32(mCuSeqlens[_sid + 1])
- else:
- chunk_start = chunk_idx * BT
-
- cute.arch.mbarrier_wait(mma_done_mbars + s, phase)
-
- csAqk = sAqk[(None, None, s)]
- csAkk = sAkk[(None, None, s)]
-
- # Dual-path row-major store. SMEM is in final GMEM layout
- # already (causal mask + diag=1 applied at MMA write; upper-tri
- # SMEM is zero from CTA-startup init). Both paths use vec2
- # autovec_copy → STG.E.32 (4-byte coalesced).
- #
- # mAqk_v2 / mAkk_v2 shape (B, T, H, BT/2, 2): last dim is vec2.
- #
- # NON-PURE: full 64-col row-major + row mask (chunk may
- # overflow seq end). 32 lanes/warp × 1 vec2 = full row.
- # PURE: reduced cols. Warp s writes (s+1)*16 cols/row (no row
- # mask, all rows in seq). Saves ~37% GMEM bandwidth vs full.
- col_lo = lane_id * 2
- col_hi = col_lo + 1
- col_vec_idx = lane_id
- rAqkOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16)
- rAkkOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16)
-
- if IS_VARLEN and not VARLEN_PURE:
- # NON-PURE: row-major full-row vec autovec + row mask. SMEM
- # upper-tri is zero (MMA-masked at write), so writing all 64
- # cols is correct. STG.E.32 (32 lanes × 4 bytes per row).
- row_base_warp = store_warp * (BT // NUM_STORE_WARPS)
- for ri in cutlass.range_constexpr(BT // NUM_STORE_WARPS):
- local_row = row_base_warp + ri
- abs_row = chunk_start + local_row
- rAqkOut[0] = csAqk[local_row, col_lo]
- rAqkOut[1] = csAqk[local_row, col_hi]
- rAkkOut[0] = csAkk[local_row, col_lo]
- rAkkOut[1] = csAkk[local_row, col_hi]
- if abs_row < st_eos:
- cute.autovec_copy(
- rAqkOut, mAqk_v2[i_b, abs_row, i_h, col_vec_idx, None]
- )
- cute.autovec_copy(
- rAkkOut, mAkk_v2[i_b, abs_row, i_h, col_vec_idx, None]
- )
- else:
- # PURE/eqlen: reduced lower-triangular store. Eqlen folds
- # tile_valid to True; pure varlen checks the scheduler's
- # potentially padded tail slot once per store warp.
- tile_valid = True if not IS_VARLEN else chunk_idx < num_chunks
- if tile_valid:
- for tile_idx in cutlass.range_constexpr(NUM_TILES):
- i_q = _TILE_IQ[tile_idx]
- i_k = _TILE_IK[tile_idx]
- is_diag = _TILE_IQ[tile_idx] == _TILE_IK[tile_idx]
- gmem_aqk_row_base = chunk_start + i_q * BC
- gmem_aqk_col_base = i_k * BC
- gmem_akk_row_base = chunk_start + i_k * BC
- gmem_akk_col_base = i_q * BC
- smem_aqk_row_base = i_q * BC
- smem_aqk_col_base = i_k * BC
- smem_akk_row_base = i_k * BC
- smem_akk_col_base = i_q * BC
- for ri in cutlass.range_constexpr(BC // NUM_STORE_WARPS):
- local_row = store_warp * (BC // NUM_STORE_WARPS) + ri
- if lane_id < BC:
- local_col = lane_id
- aqk_val = csAqk[
- smem_aqk_row_base + local_row,
- smem_aqk_col_base + local_col,
- ]
- akk_val = csAkk[
- smem_akk_row_base + local_row,
- smem_akk_col_base + local_col,
- ]
- if is_diag and local_row < local_col:
- aqk_val = cutlass.BFloat16(0.0)
- if is_diag and local_row < local_col:
- akk_val = cutlass.BFloat16(0.0)
- if is_diag and local_row == local_col:
- akk_val = cutlass.BFloat16(1.0)
- mAqk[
- i_b,
- gmem_aqk_row_base + local_row,
- i_h,
- gmem_aqk_col_base + local_col,
- ] = aqk_val
- mAkk[
- i_b,
- gmem_akk_row_base + local_row,
- i_h,
- gmem_akk_col_base + local_col,
- ] = akk_val
-
- cute.arch.mbarrier_arrive(store_done_mbars + s)
-
- yield_out()
-
-
-# =========================================================================
-# Host function
-# =========================================================================
-def make_host_function(
- B,
- NT,
- H,
- is_varlen=False,
- T_padded=None,
- has_bias=False,
- use_safe_gate=False,
- use_beta_sigmoid=False,
- varlen_pure=False,
-):
- """
- `varlen_pure=True` asserts that all seq lengths in the batch are multiples
- of BT (= 64). Under that assumption every chunk has 64 valid rows so the
- four mask sites (K1 row mask, K1 store mask, MMA accumulator zero-fill,
- Store row mask) are guaranteed to never fire and are dead-code eliminated
- at compile time. Caller's data layout is unchanged — this is a hint only.
- """
- _H = H
- _IS_VARLEN = 1 if is_varlen else 0
- _HAS_BIAS = 1 if has_bias else 0
- _USE_SAFE_GATE = 1 if use_safe_gate else 0
- _USE_BETA_SIGMOID = 1 if use_beta_sigmoid else 0
- _VARLEN_PURE = 1 if (is_varlen and varlen_pure) else 0
- if is_varlen:
- assert B == 1, "Varlen requires B=1"
- assert T_padded is not None, "T_padded required for varlen"
-
- # B / NT / T are RUNTIME arguments of host_fn (rt_b / rt_nt /
- # rt_t_total below), not closure constants: baking them keyed every
- # compiled kernel to the batch shape, and eval traffic (a distinct
- # total token count per prefill batch) recompiled ~100 s per batch —
- # observed as GSM8K "hangs" (0/1319 responses). Only genuine constexpr
- # specializers (H, is_varlen, bias/gate/pure variants) stay baked.
- s_row = _H * K_DIM
- s_col = 1
- s_h = K_DIM
-
- @cute.jit
- def host_fn(
- mQ,
- mK,
- mG,
- mA_log,
- mBeta,
- mBetaActivated,
- scale,
- mKscaled,
- mKg,
- mQscaled,
- mGkLast,
- mAqk,
- mAkk,
- mCuSeqlens,
- mChunkIndices,
- mDtBias,
- lower_bound_val,
- rt_nt: cutlass.Int32,
- rt_b: cutlass.Int32,
- rt_t_total: cutlass.Int32,
- valid_tokens: cutlass.Int32,
- # Launch stream — a runtime argument (the executor runs the model on
- # a dedicated non-blocking torch stream; launching on the DSL default
- # stream races with the caller's stream). See _launch_fused_k123_inv.
- stream: cuda.CUstream,
- ):
- # 3D TMA view: (B*T, K_DIM, H). domain_offset in kernel shifts to
- # arbitrary chunk_start — no BT alignment required for varlen.
- # Dynamic T extent: the TMA tensormap is materialized per call from
- # the runtime shape, so one compiled host_fn serves every batch.
- view_layout_3d = cute.make_layout(
- (rt_t_total, K_DIM, _H),
- stride=(s_row, s_col, s_h),
- )
- mQ_view = cute.make_tensor(mQ.iterator, view_layout_3d)
- mK_view = cute.make_tensor(mK.iterator, view_layout_3d)
- mG_view = cute.make_tensor(mG.iterator, view_layout_3d)
-
- smem_atom_qk = tcgen05.make_smem_layout_atom(
- tcgen05.SmemLayoutAtomKind.K_SW128, cutlass.BFloat16
- )
- qk_smem_2d = cute.tile_to_shape(smem_atom_qk, (BT, K_DIM), order=(0, 1))
- qk_smem_3d = cute.tile_to_shape(smem_atom_qk, (BT, K_DIM, NUM_STAGES), order=(0, 1, 2))
-
- g_smem_2d = cute.make_layout((BT, K_DIM), stride=(K_DIM, 1))
- g_smem_3d = cute.make_layout((BT, K_DIM, NUM_STAGES), stride=(K_DIM, 1, BT * K_DIM))
-
- tma_op = cpasync.CopyBulkTensorTileG2SOp(cpasync.CtaGroup.ONE)
- tma_atom_Q, tma_tensor_Q = cpasync.make_tiled_tma_atom(
- tma_op, mQ_view, qk_smem_2d, cute.product_each(qk_smem_2d.shape), num_multicast=1
- )
- tma_atom_K, tma_tensor_K = cpasync.make_tiled_tma_atom(
- tma_op, mK_view, qk_smem_2d, cute.product_each(qk_smem_2d.shape), num_multicast=1
- )
- tma_atom_G, tma_tensor_G = cpasync.make_tiled_tma_atom(
- tma_op, mG_view, g_smem_2d, cute.product_each(g_smem_2d.shape), num_multicast=1
- )
-
- g_cumsum_layout = cute.make_layout(
- (BT, K_DIM, NUM_STAGES), stride=(K_STRIDE, 1, BT * K_STRIDE)
- )
-
- copy_atom_qk_k1 = cute.make_copy_atom(
- cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=32
- )
- tiled_copy_qk_k1 = cute.make_tiled_copy_tv(
- copy_atom_qk_k1,
- thr_layout=cute.make_layout((1, 32)),
- val_layout=cute.make_layout((1, 2)),
- )
-
- # The v2 views are indexed [i_b, per-batch row, ...] by the kernel,
- # so their T extent / batch stride must be PER-BATCH T, not the flat
- # B*T used by the 3D TMA view above. Using rt_t_total here sent every
- # i_b > 0 store rt_b× too far (OOB writes + parity failure for the
- # eqlen B=2 path; varlen is B=1 so it never noticed).
- rt_t_pb = rt_t_total // rt_b
- out_v2_layout = cute.make_layout(
- (rt_b, rt_t_pb, _H, K_VEC, VEC),
- stride=(rt_t_pb * _H * K_DIM, _H * K_DIM, K_DIM, VEC, 1),
- )
- mKscaled_v2 = cute.make_tensor(mKscaled.iterator, out_v2_layout)
- mQscaled_v2 = cute.make_tensor(mQscaled.iterator, out_v2_layout)
- mKg_v2 = cute.make_tensor(mKg.iterator, out_v2_layout)
-
- # vec2 views for mAqk / mAkk (BT dimension instead of K_DIM).
- # Shape: (B, T, H, BT/2, 2). Each VEC=2 slot = 2 contiguous bf16 = 1
- # fp32-aligned 4-byte unit. Used by store warp's autovec_copy →
- # STG.E.32 with 32 lanes coalesced to 1 cache line per row.
- BT_VEC = BT // VEC # 32
- akk_v2_layout = cute.make_layout(
- (rt_b, rt_t_pb, _H, BT_VEC, VEC),
- stride=(rt_t_pb * _H * BT, _H * BT, BT, VEC, 1),
- )
- mAqk_v2 = cute.make_tensor(mAqk.iterator, akk_v2_layout)
- mAkk_v2 = cute.make_tensor(mAkk.iterator, akk_v2_layout)
-
- gklast_v2_layout = cute.make_layout(
- (rt_b, rt_nt, _H, K_VEC, VEC),
- stride=(rt_nt * _H * K_DIM, _H * K_DIM, K_DIM, VEC, 1),
- )
- mGkLast_v2 = cute.make_tensor(mGkLast.iterator, gklast_v2_layout)
-
- mma_op = cute.nvgpu.warp.MmaF16BF16Op(cutlass.BFloat16, cutlass.Float32, (16, 8, 16))
- tiled_mma_k2 = cute.make_tiled_mma(
- mma_op, cute.make_layout((1, 1, 1)), permutation_mnk=(16, 8, 16)
- )
-
- tiled_copy_mma_A = cute.make_tiled_copy_A(
- cute.make_copy_atom(cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 4), cutlass.BFloat16),
- tiled_mma_k2,
- )
- tiled_copy_mma_B = cute.make_tiled_copy_B(
- cute.make_copy_atom(cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 2), cutlass.BFloat16),
- tiled_mma_k2,
- )
-
- copy_atom_Gcum = cute.make_copy_atom(
- cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=64
- )
- tiled_copy_Gcum_norm = cute.make_tiled_copy_tv(
- copy_atom_Gcum, thr_layout=cute.make_layout((1, 4)), val_layout=cute.make_layout((1, 2))
- )
- tiled_copy_Gcum_gate = cute.make_tiled_copy_C(copy_atom_Gcum, tiled_mma_k2)
-
- smem_size = (
- BT * K_DIM * 2 * 2 * NUM_STAGES
- + BT * K_DIM * 2 * NUM_STAGES
- + BT * K_STRIDE * 4 * NUM_STAGES
- + BT * NUM_STAGES * 2 # sBeta bf16
- + K1_ROW_GROUPS * PARTIAL_COLS * 4
- + BT * AQK_TILE_STRIDE * 2 * NUM_STAGES # sAqk bf16 (64x72 row-major)
- + BT * AKK_STRIDE * 2 * NUM_STAGES # sAkk bf16
- + 512
- )
-
- # Constant grid: the persistent for_generate loop bounds work by the
- # runtime total_cgs, so blocks past the work count exit immediately.
- # (A shape-dependent min() here would re-bake the grid per NT.)
- _grid_x = NUM_SMS
-
- fused_kernel123(
- tma_atom_Q,
- tma_tensor_Q,
- tma_atom_K,
- tma_tensor_K,
- tma_atom_G,
- tma_tensor_G,
- mA_log,
- mBeta,
- mBetaActivated,
- scale,
- mKscaled_v2,
- mKg_v2,
- mQscaled_v2,
- mGkLast_v2,
- mAqk,
- mAkk,
- mAqk_v2,
- mAkk_v2,
- tiled_copy_qk_k1,
- tiled_mma_k2,
- tiled_copy_mma_A,
- tiled_copy_mma_B,
- tiled_copy_Gcum_norm,
- tiled_copy_Gcum_gate,
- qk_smem_3d,
- g_smem_3d,
- g_cumsum_layout,
- rt_nt,
- _H,
- rt_b,
- mCuSeqlens,
- mChunkIndices,
- _IS_VARLEN,
- mDtBias,
- lower_bound_val,
- _HAS_BIAS,
- _USE_SAFE_GATE,
- valid_tokens,
- _USE_BETA_SIGMOID,
- _VARLEN_PURE,
- ).launch(
- grid=(_grid_x, 1, 1),
- block=(THREADS, 1, 1),
- smem=smem_size,
- stream=stream,
- )
-
- return host_fn
diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/fused_k1234.py b/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/fused_k1234.py
deleted file mode 100644
index 32dbdb2297c7..000000000000
--- a/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/fused_k1234.py
+++ /dev/null
@@ -1,2278 +0,0 @@
-# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Fused K1+K2+K3+K4 single kernel for KDA forward pass.
-
-Architecture:
- 640 threads (20 warps = 5 warpgroups), grid = (B*H, 1, 1) for eqlen.
- Per-chunk two-phase execution:
- Phase A (K123): K1/K2/K3 warps — gate activation, intra-attention, inversion
- Phase B (K4): WG0 + WG1 — 6 MMAs, readout, state update
-
- K123 writes intermediates to Zone B SMEM + sGkLast. K4 reads from SMEM.
- Only O and S (final outputs) touch GMEM. SU and gk_last are SMEM-only.
- SMEM aliased between phases via flat pool (~225KB, within B200's 228KB).
-
-Warp assignment (5 warpgroups, 128 threads each):
- WG0 (W0-3): K4 MMA(W0) + TMA(W2) + idle(W1,W3) — idle during K123
- WG1 (W4-7): state readout+decay (K123) AND W/NV/O readout (K4)
- WG2 (W8-11): K1 — TMA+gate/cumsum/KG/KS/QS + K3 inversion+store (merged)
- WG3 (W12-15): K2 — intra-attention MMA (first 4 warps)
- WG4 (W16-19): K2 — intra-attention MMA (second 4 warps)
-
- Per-warpgroup register reallocation via setmaxnreg (CUDA 13.1 cu13 libs):
- WG0: dec(56) donates 5120 regs (MMA+TMA+idle)
- WG1: inc(104) claims 1024 (state readout+decay + K4 readout)
- WG2: default 96 (K1+K3 stays at base)
- WG3+4 (K2): inc(112) claims 4096 (heavy intra-attention, 8 warps total)
- Donate 5120 == Claim 5120 (exact balance — INC requires matching DEC).
-
- WG1 sync (mbarriers, asymmetric): st_ready_mbar(WG1→MMA: sST ready),
- state_decayed_mbar(WG1→MMA: TMEM ready), mma6_done_mbar(MMA→WG1),
- gk_last_ready_mbar(K1→WG1: gkLast ready)
-"""
-
-import cuda.bindings.driver as cuda
-import cutlass
-import cutlass.cute as cute
-import cutlass.pipeline as pipeline
-import cutlass.utils as utils
-import cutlass.utils.blackwell_helpers as sm100_utils
-from cutlass._mlir import ir
-from cutlass._mlir.dialects import llvm
-from cutlass.cute.nvgpu import cpasync, tcgen05
-from cutlass.cute.nvgpu.tcgen05 import Field, OperandMajorMode, OperandSource
-from cutlass.cutlass_dsl import T, dsl_user_op
-from cutlass.utils.gemm.sm100 import transform_partitioned_tensor_layout
-
-# =============================================================================
-# Constants
-# =============================================================================
-mma_dtype = cutlass.BFloat16
-acc_dtype = cutlass.Float32
-out_dtype = cutlass.BFloat16
-
-BT = 64 # chunk_size
-BC = 16 # sub-chunk size
-K_DIM = 128 # head dim
-K_PAD = 8
-K_STRIDE = K_DIM + K_PAD # 136
-
-# K4 tile dimensions
-M4 = 64 # chunk_size
-N4 = 128 # head dim (K or V)
-K4_K = 64 # inner dim for MMA1/2/5
-K4_K3 = 128 # inner dim for MMA3/4 (state)
-M6 = 128 # MMA6 M
-N6 = 128 # MMA6 N
-K6 = 64 # MMA6 K
-
-# Debug flags (set to 0 for normal operation)
-# DEBUG_K4_LEVEL: 0=skip K4, 1=TMA only, 2=TMA+MMA, 3=full K4
-DEBUG_K4_LEVEL = 3
-# When True, skip MMA5 (AQC@NV) to isolate: O = OI = QS @ S only.
-# If O becomes correct, bug is in sAQC. If still wrong, bug is in sQS.
-DEBUG_SKIP_MMA5 = False
-
-# Thread/warp counts
-THREADS = 640
-WARP_SIZE = 32
-WG_SIZE = 128
-NUM_WARPS = THREADS // WARP_SIZE # 28
-
-# K123 warp assignment
-NUM_K1_WARPS = 4 # Warps 8-11 (WG2)
-NUM_MMA_WARPS = 8 # Warps 12-19
-NUM_MMA_ACTIVE = 8
-NUM_STORE_WARPS = 4 # Warps 20-23
-K1_FIRST_WARP = 8
-K2_FIRST_WARP = 12
-K3_FIRST_WARP = 20
-
-# K4 warp assignment
-K4_MMA_WARP = 0
-K4_TMA_WARP = 2
-K4_READOUT_WG = 1 # warpgroup index (warps 4-7)
-
-# K123 sub-parameters
-K1_ROW_GROUPS = 4
-K1_COL_GROUPS = 1
-ROWS_PER_K1_WARP = BT // K1_ROW_GROUPS # 16
-K1_COLS_PER_WARP = K_DIM // K1_COL_GROUPS # 128
-ROWS_PER_STORE_WARP = BT // NUM_STORE_WARPS # 16
-VEC = K1_COLS_PER_WARP // 32 # 4
-K_VEC = K_DIM // VEC # 32
-
-
-NUM_SUB_CHUNKS = BT // BC # 4
-NUM_TILES = NUM_SUB_CHUNKS * (NUM_SUB_CHUNKS + 1) // 2 # 10
-MMA_K_TILE = 8
-NUM_MMA_K_TILES = K_DIM // MMA_K_TILE # 16
-AQK_TILE_COLS = NUM_TILES * BC # 160
-AQK_TILE_PAD = 8
-AQK_TILE_STRIDE = AQK_TILE_COLS + AQK_TILE_PAD # 168
-AKK_PAD = 8
-AKK_STRIDE = BT + AKK_PAD # 72
-
-TEMP_PAD = 8
-TEMP_COLS = BC + TEMP_PAD # 24
-NUM_TEMPS = 2
-
-_TILE_IQ = [0, 1, 1, 2, 2, 2, 3, 3, 3, 3]
-_TILE_IK = [0, 0, 1, 0, 1, 2, 0, 1, 2, 3]
-
-LOG2E = 1.4426950408889634
-LN2 = 0.6931471805599453
-RCP_LN2 = LOG2E
-SHFL_W4_CLAMP = 0x1C00
-
-# SMEM layout byte offsets (informational — allocator handles actual placement)
-# Persistent region (0-112KB): sAB(8) + sAQC(8) + sKG(16) + sQS(16) + sKS(16) + sST(32) + sV_ext(16)
-# sQS/sKS: time-shared with K123 TMA Q/K, overwritten by K4-format after K2
-# sST: persistent — readout WG writes during K123
-# Zone A (112-192KB): sW(16) + sNV(16) + sO(16) + scratch(31KB)
-# K123 scratch aliases: sG(16) + sGcum(34) + sAqk(5) + sAkk(18) + sTemp(3)
-# Total peak: ~192KB
-
-
-# =============================================================================
-# dsl_user_op helpers (from fuse_kernel123_no_persistent.py)
-# =============================================================================
-
-
-@dsl_user_op
-def k1_internal_barrier(*, loc=None, ip=None):
- """Named barrier for K1 warps (0-3, 128 threads). barrier_id=2."""
- llvm.inline_asm(
- T.i32(),
- [],
- "membar.cta; bar.sync 2, 128; mov.u32 $0, 0;",
- "=r",
- has_side_effects=True,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
-
-
-@dsl_user_op
-def mma_tf32_m16n8k8(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, *, loc=None, ip=None):
- """TF32 MMA: D = A * B + C, shape m16n8k8"""
- a0_bits = llvm.bitcast(T.i32(), a0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- a1_bits = llvm.bitcast(T.i32(), a1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- a2_bits = llvm.bitcast(T.i32(), a2.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- a3_bits = llvm.bitcast(T.i32(), a3.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- b0_bits = llvm.bitcast(T.i32(), b0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- b1_bits = llvm.bitcast(T.i32(), b1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- result = llvm.inline_asm(
- ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"),
- [
- a0_bits,
- a1_bits,
- a2_bits,
- a3_bits,
- b0_bits,
- b1_bits,
- c0.ir_value(loc=loc, ip=ip),
- c1.ir_value(loc=loc, ip=ip),
- c2.ir_value(loc=loc, ip=ip),
- c3.ir_value(loc=loc, ip=ip),
- ],
- """{
- mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32
- {$0, $1, $2, $3},
- {$4, $5, $6, $7},
- {$8, $9},
- {$10, $11, $12, $13};
- }""",
- "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
- has_side_effects=True,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
- d0 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip))
- d1 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip))
- d2 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip))
- d3 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip))
- return d0, d1, d2, d3
-
-
-@dsl_user_op
-def read_clock(*, loc=None, ip=None):
- """Read globaltimer (ns). Returns i64 as two i32 words packed into fp32 pair."""
- result = llvm.inline_asm(
- T.i64(),
- [],
- "mov.u64 $0, %globaltimer;",
- "=l",
- has_side_effects=True,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
- return result
-
-
-@dsl_user_op
-def fast_rcp(x, *, loc=None, ip=None):
- result = llvm.inline_asm(
- T.f32(),
- [x.ir_value(loc=loc, ip=ip)],
- "rcp.approx.ftz.f32 $0, $1;",
- "=f,f",
- has_side_effects=False,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
- return cutlass.Float32(result)
-
-
-@dsl_user_op
-def inv_internal_barrier(*, loc=None, ip=None):
- llvm.inline_asm(
- T.i32(),
- [],
- "bar.sync 3, 128; mov.u32 $0, 0;",
- "=r",
- has_side_effects=True,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
-
-
-@dsl_user_op
-def _invert_diag(sAkk: cute.Tensor, block_rc, lane_id, *, loc=None, ip=None):
- my_row = lane_id % 16
- halfwarp_base = (lane_id // 16) * 16
- r_off = block_rc * 16
- c_off = block_rc * 16
- rInv = cute.make_rmem_tensor(cute.make_layout((16,), stride=(1,)), cutlass.Float32)
- rInv[0] = cutlass.Float32(1.0)
- for x in range(1, 16):
- rInv[x] = cutlass.Float32(0.0)
- for d in range(1, 16):
- col_d = my_row - d
- valid = cutlass.Float32(col_d >= 0)
- a_val = cutlass.Float32(sAkk[r_off + my_row, c_off + col_d]) * valid
- acc = cutlass.Float32(0.0)
- for j in range(1, d):
- a_re = cutlass.Float32(sAkk[r_off + my_row, c_off + my_row - (d - j)])
- inv_shfl = cute.arch.shuffle_sync(rInv[j], halfwarp_base + my_row - d + j)
- acc = acc + a_re * inv_shfl
- rInv[d] = (-a_val - acc) * valid
- rInv[0] = cutlass.Float32(1.0)
- sAkk[r_off + my_row, c_off + my_row] = rInv[0]
- for d in range(1, 16):
- sAkk[r_off + my_row, c_off + (my_row + 16 - d) % 16] = rInv[d] * cutlass.Float32(
- my_row >= d
- )
-
-
-@dsl_user_op
-def _matmul_AB(sAkk: cute.Tensor, br_A, bc_A, br_B, bc_B, lane_id, *, loc=None, ip=None):
- gid = lane_id // 4
- tid = lane_id % 4
- _z = cutlass.Float32(0.0)
- rA = br_A * 16
- cA = bc_A * 16
- rB = br_B * 16
- cB = bc_B * 16
- a0 = cutlass.Float32(sAkk[rA + gid, cA + 2 * tid])
- a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + 2 * tid])
- a2 = cutlass.Float32(sAkk[rA + gid, cA + 2 * tid + 1])
- a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 2 * tid + 1])
- b0n0 = cutlass.Float32(sAkk[rB + 2 * tid, cB + gid])
- b1n0 = cutlass.Float32(sAkk[rB + 2 * tid + 1, cB + gid])
- b0n1 = cutlass.Float32(sAkk[rB + 2 * tid, cB + 8 + gid])
- b1n1 = cutlass.Float32(sAkk[rB + 2 * tid + 1, cB + 8 + gid])
- cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0n0, b1n0, _z, _z, _z, _z)
- cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0n1, b1n1, _z, _z, _z, _z)
- a0 = cutlass.Float32(sAkk[rA + gid, cA + 8 + 2 * tid])
- a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + 8 + 2 * tid])
- a2 = cutlass.Float32(sAkk[rA + gid, cA + 8 + 2 * tid + 1])
- a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 8 + 2 * tid + 1])
- b0n0 = cutlass.Float32(sAkk[rB + 8 + 2 * tid, cB + gid])
- b1n0 = cutlass.Float32(sAkk[rB + 8 + 2 * tid + 1, cB + gid])
- b0n1 = cutlass.Float32(sAkk[rB + 8 + 2 * tid, cB + 8 + gid])
- b1n1 = cutlass.Float32(sAkk[rB + 8 + 2 * tid + 1, cB + 8 + gid])
- cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8(
- a0, a1, a2, a3, b0n0, b1n0, cn0_0, cn0_1, cn0_2, cn0_3
- )
- cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8(
- a0, a1, a2, a3, b0n1, b1n1, cn1_0, cn1_1, cn1_2, cn1_3
- )
- return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3
-
-
-@dsl_user_op
-def _chain_mma_B(
- sAkk: cute.Tensor,
- br_B,
- bc_B,
- a0k0,
- a1k0,
- a2k0,
- a3k0,
- a0k1,
- a1k1,
- a2k1,
- a3k1,
- lane_id,
- *,
- loc=None,
- ip=None,
-):
- gid = lane_id // 4
- tid = lane_id % 4
- _z = cutlass.Float32(0.0)
- rB = br_B * 16
- cB = bc_B * 16
- b0n0 = cutlass.Float32(sAkk[rB + 2 * tid, cB + gid])
- b1n0 = cutlass.Float32(sAkk[rB + 2 * tid + 1, cB + gid])
- b0n1 = cutlass.Float32(sAkk[rB + 2 * tid, cB + 8 + gid])
- b1n1 = cutlass.Float32(sAkk[rB + 2 * tid + 1, cB + 8 + gid])
- cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8(
- a0k0, a1k0, a2k0, a3k0, b0n0, b1n0, _z, _z, _z, _z
- )
- cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8(
- a0k0, a1k0, a2k0, a3k0, b0n1, b1n1, _z, _z, _z, _z
- )
- b0n0 = cutlass.Float32(sAkk[rB + 8 + 2 * tid, cB + gid])
- b1n0 = cutlass.Float32(sAkk[rB + 8 + 2 * tid + 1, cB + gid])
- b0n1 = cutlass.Float32(sAkk[rB + 8 + 2 * tid, cB + 8 + gid])
- b1n1 = cutlass.Float32(sAkk[rB + 8 + 2 * tid + 1, cB + 8 + gid])
- cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8(
- a0k1, a1k1, a2k1, a3k1, b0n0, b1n0, cn0_0, cn0_1, cn0_2, cn0_3
- )
- cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8(
- a0k1, a1k1, a2k1, a3k1, b0n1, b1n1, cn1_0, cn1_1, cn1_2, cn1_3
- )
- return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3
-
-
-@dsl_user_op
-def _chain_mma_A(
- sAkk: cute.Tensor,
- br_A,
- bc_A,
- b0_k0n0,
- b1_k0n0,
- b0_k0n1,
- b1_k0n1,
- b0_k1n0,
- b1_k1n0,
- b0_k1n1,
- b1_k1n1,
- lane_id,
- *,
- loc=None,
- ip=None,
-):
- gid = lane_id // 4
- tid = lane_id % 4
- _z = cutlass.Float32(0.0)
- rA = br_A * 16
- cA = bc_A * 16
- a0 = cutlass.Float32(sAkk[rA + gid, cA + 2 * tid])
- a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + 2 * tid])
- a2 = cutlass.Float32(sAkk[rA + gid, cA + 2 * tid + 1])
- a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 2 * tid + 1])
- cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0_k0n0, b1_k0n0, _z, _z, _z, _z)
- cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0_k0n1, b1_k0n1, _z, _z, _z, _z)
- a0 = cutlass.Float32(sAkk[rA + gid, cA + 8 + 2 * tid])
- a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + 8 + 2 * tid])
- a2 = cutlass.Float32(sAkk[rA + gid, cA + 8 + 2 * tid + 1])
- a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 8 + 2 * tid + 1])
- cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8(
- a0, a1, a2, a3, b0_k1n0, b1_k1n0, cn0_0, cn0_1, cn0_2, cn0_3
- )
- cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8(
- a0, a1, a2, a3, b0_k1n1, b1_k1n1, cn1_0, cn1_1, cn1_2, cn1_3
- )
- return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3
-
-
-@dsl_user_op
-def _store_neg_C(
- sAkk: cute.Tensor, br, bc, c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None
-):
- gid = lane_id // 4
- tid = lane_id % 4
- r = br * 16
- c = bc * 16
- sAkk[r + gid, c + 2 * tid] = -c0
- sAkk[r + gid, c + 2 * tid + 1] = -c1
- sAkk[r + gid + 8, c + 2 * tid] = -c2
- sAkk[r + gid + 8, c + 2 * tid + 1] = -c3
- sAkk[r + gid, c + 8 + 2 * tid] = -c4
- sAkk[r + gid, c + 8 + 2 * tid + 1] = -c5
- sAkk[r + gid + 8, c + 8 + 2 * tid] = -c6
- sAkk[r + gid + 8, c + 8 + 2 * tid + 1] = -c7
-
-
-@dsl_user_op
-def _shuffle_C_to_B(c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None):
- gid = lane_id // 4
- tid = lane_id % 4
- src_a = 8 * tid + gid // 2
- src_b = src_a + 4
- f_odd = cutlass.Float32(gid % 2)
- f_even = cutlass.Float32(1) - f_odd
- c0_a = cute.arch.shuffle_sync(c0, src_a)
- c1_a = cute.arch.shuffle_sync(c1, src_a)
- c2_a = cute.arch.shuffle_sync(c2, src_a)
- c3_a = cute.arch.shuffle_sync(c3, src_a)
- c4_a = cute.arch.shuffle_sync(c4, src_a)
- c5_a = cute.arch.shuffle_sync(c5, src_a)
- c6_a = cute.arch.shuffle_sync(c6, src_a)
- c7_a = cute.arch.shuffle_sync(c7, src_a)
- c0_b = cute.arch.shuffle_sync(c0, src_b)
- c1_b = cute.arch.shuffle_sync(c1, src_b)
- c2_b = cute.arch.shuffle_sync(c2, src_b)
- c3_b = cute.arch.shuffle_sync(c3, src_b)
- c4_b = cute.arch.shuffle_sync(c4, src_b)
- c5_b = cute.arch.shuffle_sync(c5, src_b)
- c6_b = cute.arch.shuffle_sync(c6, src_b)
- c7_b = cute.arch.shuffle_sync(c7, src_b)
- return (
- c0_a * f_even + c1_a * f_odd,
- c0_b * f_even + c1_b * f_odd,
- c2_a * f_even + c3_a * f_odd,
- c2_b * f_even + c3_b * f_odd,
- c4_a * f_even + c5_a * f_odd,
- c4_b * f_even + c5_b * f_odd,
- c6_a * f_even + c7_a * f_odd,
- c6_b * f_even + c7_b * f_odd,
- )
-
-
-@dsl_user_op
-def tma_store_fence(*, loc=None, ip=None):
- """Fence: wait for all outstanding TMA S2G stores to finish reading SMEM.
- Must be called after TMA store and before SMEM is reused (e.g., by K123 phase)."""
- llvm.inline_asm(
- T.i32(),
- [],
- "cp.async.bulk.commit_group; cp.async.bulk.wait_group.read 0; mov.u32 $0, 0;",
- "=r",
- has_side_effects=True,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
-
-
-@dsl_user_op
-def threadfence_gl(*, loc=None, ip=None):
- """Global memory fence — flush L1 writes to L2 for TMA visibility."""
- llvm.inline_asm(
- T.i32(),
- [],
- "membar.gl; mov.u32 $0, 0;",
- "=r",
- has_side_effects=True,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
-
-
-@dsl_user_op
-def _store_C_temp(
- sT: cute.Tensor, buf, c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None
-):
- gid = lane_id // 4
- tid = lane_id % 4
- sT[gid, 2 * tid, buf] = c0
- sT[gid, 2 * tid + 1, buf] = c1
- sT[gid + 8, 2 * tid, buf] = c2
- sT[gid + 8, 2 * tid + 1, buf] = c3
- sT[gid, 8 + 2 * tid, buf] = c4
- sT[gid, 8 + 2 * tid + 1, buf] = c5
- sT[gid + 8, 8 + 2 * tid, buf] = c6
- sT[gid + 8, 8 + 2 * tid + 1, buf] = c7
-
-
-@dsl_user_op
-def _load_C_temp(sT: cute.Tensor, buf, lane_id, *, loc=None, ip=None):
- gid = lane_id // 4
- tid = lane_id % 4
- return (
- cutlass.Float32(sT[gid, 2 * tid, buf]),
- cutlass.Float32(sT[gid, 2 * tid + 1, buf]),
- cutlass.Float32(sT[gid + 8, 2 * tid, buf]),
- cutlass.Float32(sT[gid + 8, 2 * tid + 1, buf]),
- cutlass.Float32(sT[gid, 8 + 2 * tid, buf]),
- cutlass.Float32(sT[gid, 8 + 2 * tid + 1, buf]),
- cutlass.Float32(sT[gid + 8, 8 + 2 * tid, buf]),
- cutlass.Float32(sT[gid + 8, 8 + 2 * tid + 1, buf]),
- )
-
-
-# =============================================================================
-# Fused K1234 Kernel
-# =============================================================================
-
-
-@cute.kernel
-def fused_k1234_kernel(
- # --- K4 MMA descriptors ---
- tiled_mma_kmn: cute.TiledMma,
- tiled_mma_mn_mn: cute.TiledMma,
- # --- K4 TMA load atoms (V only; S goes TMEM→SMEM via readout WG) ---
- tma_atom_v: cute.CopyAtom,
- mV_nkl: cute.Tensor,
- v_sl: cute.ComposedLayout,
- s_sl: cute.ComposedLayout, # sST SMEM layout (filled from TMEM, not TMA)
- # --- K4 Zone B swizzle layouts (for SMEM allocation + write views) ---
- ab_sl: cute.ComposedLayout,
- ks_sl: cute.ComposedLayout,
- qs_sl: cute.ComposedLayout,
- aqc_sl: cute.ComposedLayout,
- kg_sl: cute.ComposedLayout,
- # --- K4 TMA store ---
- tma_atom_o_st: cute.CopyAtom,
- mOo: cute.Tensor,
- store_sl: cute.ComposedLayout,
- # --- K4 readout/reinterpret layouts ---
- readout_k_sl: cute.ComposedLayout,
- nv_b_sl: cute.ComposedLayout,
- nv_a_sl: cute.ComposedLayout,
- # --- K123 TMA load atoms + tensors ---
- tma_atom_q_k123: cute.CopyAtom,
- tma_tensor_q_k123: cute.Tensor,
- tma_atom_k_k123: cute.CopyAtom,
- tma_tensor_k_k123: cute.Tensor,
- tma_atom_g_k123: cute.CopyAtom,
- tma_tensor_g_k123: cute.Tensor,
- # --- K123 SMEM layouts (passed from host) ---
- qk_smem_layout: cute.ComposedLayout,
- g_smem_layout,
- g_cumsum_layout,
- # --- K123 tiled copies ---
- tiled_copy_qk_k1,
- tiled_mma_k2,
- tiled_copy_mma_A,
- tiled_copy_mma_B,
- tiled_copy_Gcum_norm,
- tiled_copy_Gcum_gate,
- # --- K123 GMEM tensors (K123 outputs now go to Zone B SMEM, not GMEM) ---
- mA_log: cute.Tensor,
- mBeta: cute.Tensor,
- scale: cutlass.Float32,
- # --- K4 GMEM (state in TMEM; GMEM only for init/final store) ---
- mS_fp32: cute.Tensor, # [V, K, B*H] fp32 — initial state + final state output
- # --- dt_bias / safe_gate ---
- mDtBias: cute.Tensor,
- lower_bound: cutlass.Float32,
- HAS_BIAS: cutlass.Constexpr[int],
- USE_SAFE_GATE: cutlass.Constexpr[int],
- # --- Clock profiling (optional) ---
- mClocks: cute.Tensor, # [16, B*H] i64 — per-role clock profiling
- PROFILE_CLOCKS: cutlass.Constexpr[int],
- # --- Dimensions (dynamic — one compilation serves all shapes) ---
- num_chunks: cutlass.Int32,
- num_heads: cutlass.Int32,
- batch_size: cutlass.Int32,
-):
- # === Thread indices ===
- tidx, _, _ = cute.arch.thread_idx()
- warp_idx = cute.arch.warp_idx()
- warp_idx = cute.arch.make_warp_uniform(warp_idx)
- warpgroup_idx = cute.arch.make_warp_uniform(tidx // WG_SIZE)
- warpgroup_tidx = tidx % WG_SIZE
- lane_id = tidx % WARP_SIZE
-
- # ========================================================================
- # Per-warpgroup register reallocation (Blackwell SM100, CUDA 13.1).
- #
- # CRITICAL: setmaxnreg.inc is effectively BLOCKING on SM100. If CTAPOOL
- # doesn't have the full requested regs, INC spins forever (hang at
- # torch.cuda.synchronize). Total INC claims MUST be <= total DEC donations.
- #
- # Requirements:
- # (1) nvidia-cutlass-dsl-libs-cu13==4.4.2 installed (start_env.sh)
- # (2) CUDA_HOME = CUDA 13.1 toolkit (so 13.1 ptxas is used)
- # (3) DEC target >= 40 (else ptxas rejects with C7507)
- # (4) sum(DEC donations) >= sum(INC claims)
- #
- # Budget (20 warps, 640 threads, compiler default 96/thread):
- # WG0 (W0-3, 128t): dec(56) donates (96-56)*128 = 5120 regs
- # WG1 (W4-7, 128t): inc(104) claims (104-96)*128 = 1024 regs
- # WG2 (W8-11, 128t): default 96, no call
- # WG3+4 (W12-19, 256t): inc(112) claims (112-96)*256 = 4096 regs
- # Donate 5120 == Claim 5120 (exact balance).
- # ========================================================================
- if warpgroup_idx == 0:
- cute.arch.setmaxregister_decrease(56)
- elif warpgroup_idx == 1:
- cute.arch.setmaxregister_increase(104)
- elif warpgroup_idx >= 3:
- cute.arch.setmaxregister_increase(112)
-
- bid = cute.arch.block_idx()[0]
- i_b = bid // num_heads
- i_h = bid % num_heads
-
- # === K4 thr slices (created once, used every K4 phase) ===
- thr_kmn = tiled_mma_kmn.get_slice(0)
- thr_mn = tiled_mma_mn_mn.get_slice(0)
- dice = (None, None, None)
-
- # ========================================================================
- # SMEM ALLOCATION — Persistent + Aliased Layout (~224KB)
- # Persistent (0-112KB): K4 operands + sST + sV_ext
- # sQS/sKS: K4-only, NOT aliased by K123 (separate sQ_k123/sK_k123 in Zone A)
- # sST: persistent — K1 warps write decayed state bf16 during K123.
- # Zone A (112-224KB): K123 scratch / K4 readout (aliased)
- # During K123: sQ, sK, sG, sGcum, sAqk, sAkk, sTemp (~109KB)
- # During K4: sW, sNV, sO readout buffers (48KB)
- # sGkLast (512B): separate allocation, persists from K1 to K4
- # State S[V,K] fp32 lives in TMEM offset 256-383 (persistent across chunks).
- # ========================================================================
- smem = cutlass.utils.SmemAllocator()
- AL = 128
-
- # --- Persistent region (written by K123, read by K4) ---
- sAB = smem.allocate_tensor(mma_dtype, ab_sl.outer, AL, ab_sl.inner) # 8KB A_kk_inv
- sAQC = smem.allocate_tensor(mma_dtype, aqc_sl.outer, AL, aqc_sl.inner) # 8KB A_qk
- sKG = smem.allocate_tensor(mma_dtype, kg_sl.outer, AL, kg_sl.inner) # 16KB kg
- sQS = smem.allocate_tensor(mma_dtype, qs_sl.outer, AL, qs_sl.inner) # 16KB q_scaled (K4 only)
- sKS = smem.allocate_tensor(mma_dtype, ks_sl.outer, AL, ks_sl.inner) # 16KB k_scaled (K4 only)
- sST = smem.allocate_tensor(
- mma_dtype, s_sl.outer, AL, s_sl.inner
- ) # 32KB state bf16 (persistent)
- sV_ext = smem.allocate_tensor(mma_dtype, v_sl.outer, AL, v_sl.inner) # 16KB V (non-aliased)
- # Persistent total: 112KB
-
- # --- Zone A: K4 readout (aliases K123 scratch) ---
- sW = smem.allocate_tensor(mma_dtype, readout_k_sl.outer, AL, readout_k_sl.inner) # 16KB readout
- sNV = smem.allocate_tensor(
- mma_dtype, readout_k_sl.outer, AL, readout_k_sl.inner
- ) # 16KB readout
- sO = smem.allocate_tensor(mma_dtype, readout_k_sl.outer, AL, readout_k_sl.inner) # 16KB readout
- # K4 readout: 48KB. K123 scratch needs ~111KB, extra ~64KB allocated below.
- _smem_extra = smem.allocate_array(cutlass.Float32, 16384) # ~64KB for K123 overflow
-
- # --- K123 alias: sQ_k123/sK_k123 in Zone A (separate from sQS/sKS) ---
- # sQ_k123 aliases sW, sK_k123 aliases sNV during K123.
- # TMA loads Q/K here in K_SW128 swizzle layout. K1 writes KS/QS to persistent sQS/sKS.
- _za = sW.iterator # Zone A base (bf16 pointer)
- sQ_k123 = cute.make_tensor(
- cute.recast_ptr(_za, qk_smem_layout.inner, cutlass.BFloat16), qk_smem_layout.outer
- )
- sK_k123 = cute.make_tensor(
- cute.recast_ptr(_za + 8192, qk_smem_layout.inner, cutlass.BFloat16), qk_smem_layout.outer
- )
-
- # --- K123 alias: scratch tensors (overlap Zone A, after sQ/sK) ---
- # sG_k123 [64,128] bf16 at +32KB (+16384 bf16), aliases sO during K4
- sG_k123 = cute.make_tensor(cute.recast_ptr(_za + 16384, dtype=cutlass.BFloat16), g_smem_layout)
- # sGcum [64,136] fp32 at +48KB (+24576 bf16), 34816B
- sGcum = cute.make_tensor(cute.recast_ptr(_za + 24576, dtype=cutlass.Float32), g_cumsum_layout)
- # sPartialLast removed — 1-col-per-thread K1 doesn't need cross-warp shuffle
- # sAqk_k123 [16,168] bf16 at +85KB (+43040 bf16), 5376B
- aqk_tile_layout = cute.make_layout((BC, AQK_TILE_STRIDE), stride=(AQK_TILE_STRIDE, 1))
- sAqk_k123 = cute.make_tensor(
- cute.recast_ptr(_za + 43040, dtype=cutlass.BFloat16), aqk_tile_layout
- )
- # sAkk_k123 [64,72] fp32 at +91KB (+45728 bf16), 18432B
- akk_tile_layout = cute.make_layout((BT, AKK_STRIDE), stride=(AKK_STRIDE, 1))
- sAkk_k123 = cute.make_tensor(
- cute.recast_ptr(_za + 45728, dtype=cutlass.Float32), akk_tile_layout
- )
- # sTemp [16,24,2] fp32 at +110KB (+54944 bf16), 3072B
- temp_layout = cute.make_layout(
- (BC, TEMP_COLS, NUM_TEMPS), stride=(TEMP_COLS, 1, BC * TEMP_COLS)
- )
- sTemp = cute.make_tensor(cute.recast_ptr(_za + 54944, dtype=cutlass.Float32), temp_layout)
-
- # --- K4 reinterpret views ---
- sNV_b = cute.make_tensor(cute.recast_ptr(sNV.iterator, nv_b_sl.inner, mma_dtype), nv_b_sl.outer)
- sNV_a = cute.make_tensor(cute.recast_ptr(sNV.iterator, nv_a_sl.inner, mma_dtype), nv_a_sl.outer)
- sO_st = cute.make_tensor(
- cute.recast_ptr(sO.iterator, store_sl.inner, out_dtype), store_sl.outer
- )
-
- # sSU_fp32 removed — SU stays in TMEM/RMEM, no SMEM copy needed
-
- # gk_last: 128 fp32 per (b,h) chunk — persists from K1 to K4 within same chunk
- _sGkLast_buf = smem.allocate_array(cutlass.Float32, K_DIM) # 128 fp32 = 512 bytes
- sGkLast = cute.make_tensor(_sGkLast_buf, cute.make_layout((K_DIM,), stride=(1,)))
-
- # beta: 64 bf16 per chunk — loaded once in K1, read by K2/K3
- _sBeta_buf = smem.allocate_array(cutlass.BFloat16, BT) # 64 bf16 = 128 bytes
- sBeta = cute.make_tensor(_sBeta_buf, cute.make_layout((BT,), stride=(1,)))
-
- # ========================================================================
- # TMEM ALLOCATION (K4, warp 0 only)
- # ========================================================================
- tmem_smem = smem.allocate_array(cutlass.Int32, 1)
- if warp_idx == 0:
- cute.arch.alloc_tmem(512, tmem_smem)
-
- # ========================================================================
- # K123 MBARRIERS (single-stage, no double-buffering)
- # ========================================================================
- k123_tma_mbar = smem.allocate_array(cutlass.Int64, 1)
- k123_k1_done_mbar = smem.allocate_array(cutlass.Int64, 1)
- k123_mma_done_mbar = smem.allocate_array(cutlass.Int64, 1)
- # Asymmetric mbarriers for WG1 async (arrive = non-blocking signal, wait = blocking)
- mma6_done_mbar = smem.allocate_array(cutlass.Int64, 1) # MMA→WG1: MMA6 complete
- st_ready_mbar = smem.allocate_array(cutlass.Int64, 1) # WG1→MMA: sST bf16 ready
- gk_last_ready_mbar = smem.allocate_array(cutlass.Int64, 1) # K1→WG1: gkLast ready
- state_decayed_mbar = smem.allocate_array(cutlass.Int64, 1) # WG1→MMA: decayed state ready
- final_state_done_mbar = smem.allocate_array(cutlass.Int64, 1) # WG1→MMA: final state stored
-
- # K4 named barriers (symmetric — both sides arrive and wait)
- sW_ready_nbar = pipeline.NamedBarrier(4, WG_SIZE + WARP_SIZE)
- sNV_ready_nbar = pipeline.NamedBarrier(5, WG_SIZE + WARP_SIZE)
- store_nbar = pipeline.NamedBarrier(6, WG_SIZE + 2 * WARP_SIZE)
- phase_nbar = pipeline.NamedBarrier(7, THREADS) # all threads (WG1 now participates)
-
- elect_one = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1)
- wg_coop = pipeline.CooperativeGroup(pipeline.Agent.Thread, WG_SIZE)
-
- # K4 TMA pipelines
- def _make_tma_pipe(byte_count):
- ptr = smem.allocate_array(cutlass.Int64, 2)
- return pipeline.PipelineTmaUmma.create(
- barrier_storage=ptr,
- num_stages=1,
- producer_group=elect_one,
- consumer_group=elect_one,
- tx_count=byte_count,
- defer_sync=True,
- ).make_participants()
-
- # Only V still uses TMA pipeline (KS/QS now in persistent SMEM)
- v_bytes = cute.size_in_bytes(mma_dtype, v_sl)
- v_prod, v_cons = _make_tma_pipe(v_bytes)
-
- # K4 UMMA pipelines (MMA → readout)
- def _make_umma_pipe():
- ptr = smem.allocate_array(cutlass.Int64, 2)
- return pipeline.PipelineUmmaAsync.create(
- barrier_storage=ptr,
- num_stages=1,
- producer_group=elect_one,
- consumer_group=wg_coop,
- defer_sync=True,
- ).make_participants()
-
- w_prod, w_cons = _make_umma_pipe()
- nv_prod, nv_cons = _make_umma_pipe()
- o_prod, o_cons = _make_umma_pipe()
- # kv_acc_pipeline removed: barrier() between K123→K4 ensures state readout/decay done
-
- # ========================================================================
- # INIT
- # ========================================================================
- k123_tma_bytes = BT * K_DIM * 2 * 3 # q+k+g bf16
-
- if tidx == 0:
- cute.arch.mbarrier_init(k123_tma_mbar, 1)
- cute.arch.mbarrier_init(k123_k1_done_mbar, NUM_K1_WARPS * WARP_SIZE)
- cute.arch.mbarrier_init(k123_mma_done_mbar, NUM_MMA_WARPS * WARP_SIZE)
- cute.arch.mbarrier_init(mma6_done_mbar, WARP_SIZE) # MMA warp (W0, 32t) arrives
- cute.arch.mbarrier_init(st_ready_mbar, WG_SIZE) # WG1 (128t) arrives
- cute.arch.mbarrier_init(gk_last_ready_mbar, NUM_K1_WARPS * WARP_SIZE) # K1 (128t) arrives
- cute.arch.mbarrier_init(state_decayed_mbar, WG_SIZE) # WG1 (128t) arrives
- cute.arch.mbarrier_init(final_state_done_mbar, WG_SIZE) # WG1 (128t) arrives
- cute.arch.mbarrier_init_fence()
- cute.arch.barrier()
-
- # ========================================================================
- # TMEM tensors (K4, created once)
- # ========================================================================
- tmem_ptr = cute.arch.retrieve_tmem_ptr(cutlass.Int32, 16, tmem_smem)
-
- tCtW_shape = tiled_mma_kmn.partition_shape_C((M4, N4))
- tCtW_fake = tiled_mma_kmn.make_fragment_C(tCtW_shape)
- tCtW = cute.make_tensor(cute.recast_ptr(tmem_ptr + 0, dtype=acc_dtype), tCtW_fake.layout)
-
- tCtNV_shape = tiled_mma_kmn.partition_shape_C((M4, N4))
- tCtNV_fake = tiled_mma_kmn.make_fragment_C(tCtNV_shape)
- tCtNV = cute.make_tensor(cute.recast_ptr(tmem_ptr + 128, dtype=acc_dtype), tCtNV_fake.layout)
- tCtO = cute.make_tensor(cute.recast_ptr(tmem_ptr + 384, dtype=acc_dtype), tCtNV_fake.layout)
-
- tCtS_shape = tiled_mma_mn_mn.partition_shape_C((M6, N6))
- tCtS_fake = tiled_mma_mn_mn.make_fragment_C(tCtS_shape)
- # tCtS = cute.make_tensor(cute.recast_ptr(tmem_ptr + 128, dtype=acc_dtype), tCtS_fake.layout)
-
- # State [V=128, K=128] fp32 at TMEM offset 256 — persistent across chunks (GDN pattern)
- # Same MN-MN layout as SU (tCtS) so corresponding elements align for state update
- tCtState = cute.make_tensor(cute.recast_ptr(tmem_ptr + 256, dtype=acc_dtype), tCtS_fake.layout)
-
- # ========================================================================
- # WARP-SPECIALIZED INDEPENDENT LOOPS (GDN pattern)
- # Each role: per-role setup -> own chunk loop -> cleanup
- # (setmaxnreg deferred to Phase 2)
- # ========================================================================
-
- # ==============================================================
- # WG1 (warps 4-7): state readout + decay (K123) + W/NV/O readout (K4)
- # Both tasks in one WG — they never overlap (K123 vs K4 phases).
- # ==============================================================
- if warpgroup_idx == K4_READOUT_WG:
- # --- W/NV/O readout setup (M=64: Ld16x256bOp) ---
- tCtW_mn = transform_partitioned_tensor_layout(tCtW)
- tCtNV_mn = transform_partitioned_tensor_layout(tCtNV)
- tCtO_mn = transform_partitioned_tensor_layout(tCtO)
-
- atom_t2r = cute.make_copy_atom(tcgen05.Ld16x256bOp(tcgen05.Repetition(1)), acc_dtype)
- tiled_t2r = tcgen05.make_tmem_copy(atom_t2r, tCtW[(None, None), 0, 0])
- thr_t2r = tiled_t2r.get_slice(warpgroup_tidx)
-
- tTR_W = thr_t2r.partition_S(tCtW_mn)
- tTR_NV = thr_t2r.partition_S(tCtNV_mn)
- tTR_O = thr_t2r.partition_S(tCtO_mn)
-
- atom_r2s_k = sm100_utils.get_smem_store_op(
- utils.LayoutEnum.ROW_MAJOR, mma_dtype, acc_dtype, tiled_t2r
- )
- tiled_r2s_k = cute.make_tiled_copy_D(atom_r2s_k, tiled_t2r)
- thr_r2s_k = tiled_r2s_k.get_slice(warpgroup_tidx)
- tCsW = thr_r2s_k.partition_D(transform_partitioned_tensor_layout(sW))
- tCsO = thr_r2s_k.partition_D(transform_partitioned_tensor_layout(sO))
- tCsNV = thr_r2s_k.partition_D(transform_partitioned_tensor_layout(sNV))
-
- cId = cute.make_identity_tensor((M4, N4))
- tTR_cId = thr_t2r.partition_D(cId)
-
- # --- State TMEM setup (M=128, GDN pattern) ---
- cId_128 = cute.make_identity_tensor((M6, N6))
-
- # --- State TMEM copy atoms (M=128, GDN pattern) ---
- tCtState_mn = transform_partitioned_tensor_layout(tCtState)
-
- # TMEM->RMEM (Ld32x32bOp for M=128)
- atom_state_t2r = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition(32)), acc_dtype)
- tiled_state_t2r = tcgen05.make_tmem_copy(atom_state_t2r, tCtState[(None, None), 0, 0])
- thr_state_t2r = tiled_state_t2r.get_slice(warpgroup_tidx)
- tTR_tCtState = thr_state_t2r.partition_S(tCtState_mn)
- tTR_tCcState = thr_state_t2r.partition_D(cId_128)
- tRrState = cute.make_rmem_tensor_like(tTR_tCcState, acc_dtype) # fp32 RMEM
-
- # RMEM->TMEM (St32x32bOp for M=128, state write-back)
- atom_state_r2t = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition(32)), acc_dtype)
- tiled_state_r2t = tcgen05.make_tmem_copy(atom_state_r2t, tCtState[(None, None), 0, 0])
- thr_state_r2t = tiled_state_r2t.get_slice(warpgroup_tidx)
- tRT_tCtState = thr_state_r2t.partition_D(tCtState_mn)
-
- # RMEM bf16 -> swizzled sST SMEM (CopyUniversalOp, NO domain transpose)
- atom_state_r2s = cute.make_copy_atom(
- cute.nvgpu.CopyUniversalOp(), mma_dtype, num_bits_per_copy=16
- )
- tiled_state_r2s = cute.make_tiled_copy_D(atom_state_r2s, tiled_state_t2r)
- thr_state_r2s = tiled_state_r2s.get_slice(warpgroup_tidx)
- sST_mn_view = transform_partitioned_tensor_layout(sST)
- tCsState_inp = thr_state_r2s.partition_D(sST_mn_view)
- tRrState_bf16 = cute.make_rmem_tensor_like(tTR_tCcState, mma_dtype) # bf16 RMEM
- tCrState_bf16 = tiled_state_r2s.retile(tRrState_bf16)
-
- # GMEM fp32 -> RMEM fp32 (for initial/final state, runs once)
- atom_state_g2r = cute.make_copy_atom(
- cute.nvgpu.CopyUniversalOp(), acc_dtype, num_bits_per_copy=32
- )
- tiled_state_g2r = cute.make_tiled_copy_S(atom_state_g2r, tiled_state_r2t)
- thr_state_g2r = tiled_state_g2r.get_slice(warpgroup_tidx)
-
- # RMEM fp32 -> GMEM fp32 (for final state store, runs once)
- atom_state_r2g = cute.make_copy_atom(
- cute.nvgpu.CopyUniversalOp(), acc_dtype, num_bits_per_copy=32
- )
- tiled_state_r2g = cute.make_tiled_copy_D(atom_state_r2g, tiled_state_t2r)
- thr_state_r2g = tiled_state_r2g.get_slice(warpgroup_tidx)
-
- # --- Initial state load: GMEM fp32 -> RMEM -> state TMEM ---
- gS_init = mS_fp32[(None, None, bid)] # [V=128, K=128] fp32
- tGR_tCgState_in = thr_state_g2r.partition_S(gS_init)
- tGR_tCrState_in = thr_state_g2r.retile(tRrState)
- # 1. GMEM fp32 -> RMEM fp32 (coalesced CopyUniversalOp)
- cute.copy(tiled_state_g2r, tGR_tCgState_in, tGR_tCrState_in)
- # 2. RMEM fp32 -> state TMEM (St32x32bOp per sub-tile)
- num_state_subs = tRrState.shape[2]
- for sub in cutlass.range(num_state_subs):
- cute.copy(tiled_state_r2t, tRrState[None, 0, sub], tRT_tCtState[None, 0, sub])
- cute.arch.fence_view_async_tmem_store()
-
- # --- WG1 chunk loop ---
- clk_wg1_start = cutlass.Int64(0)
- if PROFILE_CLOCKS:
- clk_wg1_start = read_clock()
- for chunk_c in cutlass.range(num_chunks):
- num_state_subs = tRrState.shape[2]
- _sub_tile_size = cute.size(tRrState.shape[0])
-
- # ---- Part 1: State readout + decay ----
- # Wait previous MMA6 (chunk > 0). Chunk 0: state from init.
- if chunk_c > 0:
- cute.arch.mbarrier_wait(mma6_done_mbar, phase=(chunk_c - 1) % 2)
-
- # Step A: State TMEM -> RMEM(fp32) -> bf16 -> sST SMEM
- # Per-sub streaming: only 1 sub-tile in RMEM at a time (32 regs, not 128).
- # tRrState[sub] is overwritten each iteration — NOT held across subs.
- for sub in cutlass.range(num_state_subs):
- cute.copy(tiled_state_t2r, tTR_tCtState[None, 0, sub], tRrState[None, 0, sub])
- tRrState_bf16[None, 0, sub].store(tRrState[None, 0, sub].load().to(mma_dtype))
- cute.copy(
- tiled_state_r2s, tCrState_bf16[None, 0, sub], tCsState_inp[None, 0, sub, 0]
- )
- cute.arch.fence_view_async_shared()
- # Signal MMA: sST bf16 ready for MMA4/MMA3 (non-blocking)
- cute.arch.mbarrier_arrive(st_ready_mbar)
-
- # Step B: Wait K1 to finish gk_last (current chunk)
- cute.arch.mbarrier_wait(gk_last_ready_mbar, phase=chunk_c % 2)
-
- # Step C+D: Re-read state from TMEM, decay, write back.
- # tRrState was overwritten per-sub in Step A, so re-read is needed.
- # Cost: 4 extra TMEM reads (~40 cycles total) vs saving 96 spilled regs.
- for sub in cutlass.range(num_state_subs):
- # Re-read this sub-tile from TMEM (only 32 regs live)
- cute.copy(tiled_state_t2r, tTR_tCtState[None, 0, sub], tRrState[None, 0, sub])
- # Decay by gk_last
- for i in cutlass.range(_sub_tile_size):
- coord = tTR_tCcState[i, 0, sub]
- k_idx = coord[1] # N dimension = K in (V,K) convention
- gk_val = cutlass.Float32(sGkLast[k_idx])
- tRrState[i, 0, sub] = tRrState[i, 0, sub] * gk_val
- # Write decayed sub-tile back to TMEM immediately
- cute.copy(tiled_state_r2t, tRrState[None, 0, sub], tRT_tCtState[None, 0, sub])
- cute.arch.fence_view_async_tmem_store()
- # Signal MMA: decayed state in TMEM ready for MMA6 (non-blocking)
- cute.arch.mbarrier_arrive(state_decayed_mbar)
-
- # ---- K4 phase: W/NV/O readout ----
- phase_nbar.arrive_and_wait() # K123->K4
-
- tRrR = cute.make_rmem_tensor_like(tTR_cId, acc_dtype)
- tRrR_out = cute.make_rmem_tensor_like(tRrR, mma_dtype)
- tCrR_k = tiled_r2s_k.retile(tRrR_out)
- num_subs = tRrR.shape[2]
-
- # W readout
- wh = w_cons.wait_and_advance()
- for sub in cutlass.range(num_subs):
- cute.copy(tiled_t2r, tTR_W[None, 0, sub], tRrR[None, 0, sub])
- tRrR_out[None, 0, sub].store(tRrR[None, 0, sub].load().to(mma_dtype))
- cute.copy(tiled_r2s_k, tCrR_k[None, 0, sub], tCsW[None, 0, sub, 0])
- cute.arch.fence_view_async_tmem_load()
- wh.release()
- cute.arch.fence_view_async_shared()
- sW_ready_nbar.arrive_and_wait()
-
- # NV readout
- nvh = nv_cons.wait_and_advance()
- for sub in cutlass.range(num_subs):
- cute.copy(tiled_t2r, tTR_NV[None, 0, sub], tRrR[None, 0, sub])
- tRrR_out[None, 0, sub].store(tRrR[None, 0, sub].load().to(mma_dtype))
- cute.copy(tiled_r2s_k, tCrR_k[None, 0, sub], tCsNV[None, 0, sub, 0])
- cute.arch.fence_view_async_tmem_load()
- nvh.release()
- cute.arch.fence_view_async_shared()
- sNV_ready_nbar.arrive_and_wait()
-
- # O readout
- oh = o_cons.wait_and_advance()
- for sub in cutlass.range(num_subs):
- cute.copy(tiled_t2r, tTR_O[None, 0, sub], tRrR[None, 0, sub])
- tRrR_out[None, 0, sub].store(tRrR[None, 0, sub].load().to(mma_dtype))
- cute.copy(tiled_r2s_k, tCrR_k[None, 0, sub], tCsO[None, 0, sub, 0])
- cute.arch.fence_view_async_tmem_load()
- oh.release()
- cute.arch.fence_view_async_shared()
- store_nbar.arrive_and_wait()
-
- phase_nbar.arrive_and_wait() # K4->next
-
- # --- Wait last MMA6, final state store ---
- cute.arch.mbarrier_wait(mma6_done_mbar, phase=(num_chunks - 1) % 2)
-
- # TMEM -> RMEM fp32
- gS_out = mS_fp32[(None, None, bid)] # [V=128, K=128] fp32
- tGR_tCgState_out = thr_state_r2g.partition_D(gS_out)
- tGR_tCrState_out = thr_state_r2g.retile(tRrState)
- num_state_subs_final = tRrState.shape[2]
- for sub in cutlass.range(num_state_subs_final):
- cute.copy(tiled_state_t2r, tTR_tCtState[None, 0, sub], tRrState[None, 0, sub])
- cute.arch.fence_view_async_tmem_load()
- # RMEM fp32 -> GMEM fp32 (coalesced CopyUniversalOp)
- for sub in cutlass.range(num_state_subs_final):
- cute.copy(
- tiled_state_r2g, tGR_tCrState_out[None, 0, sub], tGR_tCgState_out[None, 0, sub]
- )
-
- # Signal MMA warp: safe to dealloc TMEM
- cute.arch.mbarrier_arrive(final_state_done_mbar)
-
- if PROFILE_CLOCKS:
- clk_wg1_end = read_clock()
- if warpgroup_tidx == 0:
- mClocks[8, bid] = clk_wg1_end - clk_wg1_start
-
- # ==============================================================
- # K1 (warps 8-11): TMA + gate activation + cumsum + KG/KS/QS
- # ==============================================================
- elif warp_idx >= K1_FIRST_WARP and warp_idx < K1_FIRST_WARP + NUM_K1_WARPS:
- # K1 per-warp setup — each thread handles 1 column, all 64 rows sequentially
- clk_k1_start = cutlass.Int64(0)
- if PROFILE_CLOCKS:
- clk_k1_start = read_clock()
- k1_warp = warp_idx - K1_FIRST_WARP
- my_col = k1_warp * WARP_SIZE + lane_id # 0-127, one column per thread
-
- # Per-head constants (invariant across chunks)
- exp_A = cute.exp(mA_log[i_h], fastmath=True)
- cumsum_scale = cutlass.Float32(RCP_LN2)
- rBias_val = cutlass.Float32(0.0)
- if HAS_BIAS:
- rBias_val = mDtBias[i_h, my_col].to(cutlass.Float32)
-
- # Zero sAQC once (upper triangle stays zero across all chunks)
- _aqc_zero = cutlass.BFloat16(0.0)
- for ri in cutlass.range(BC):
- aqc_row = k1_warp * BC + ri
- c0 = lane_id * 2
- c1 = lane_id * 2 + 1
- sAQC[aqc_row + (c0 % 16) * 64, 0, c0 // 16, 0] = _aqc_zero
- sAQC[aqc_row + (c1 % 16) * 64, 0, c1 // 16, 0] = _aqc_zero
- inv_internal_barrier()
-
- for chunk_c in cutlass.range(num_chunks):
- c_phase = chunk_c % 2
- chunk_start = i_b * num_chunks * BT + chunk_c * BT
- chunk_start_local = chunk_c * BT
-
- # TMA load q, k, g for this chunk
- if warp_idx == K1_FIRST_WARP:
- gQ_head = tma_tensor_q_k123[(None, None, i_h)]
- gK_head = tma_tensor_k_k123[(None, None, i_h)]
- gG_head = tma_tensor_g_k123[(None, None, i_h)]
-
- if lane_id == 0:
- cute.arch.mbarrier_expect_tx(k123_tma_mbar, k123_tma_bytes)
-
- sQ_s = sQ_k123
- gQ_s = cute.local_tile(
- cute.domain_offset((chunk_start, 0), gQ_head), (BT, K_DIM), (0, 0)
- )
- ts_q, tg_q = cpasync.tma_partition(
- tma_atom_q_k123,
- 0,
- cute.make_layout(1),
- cute.group_modes(sQ_s, 0, 2),
- cute.group_modes(gQ_s, 0, 2),
- )
- cute.copy(tma_atom_q_k123, tg_q, ts_q, tma_bar_ptr=k123_tma_mbar)
-
- sK_s = sK_k123
- gK_s = cute.local_tile(
- cute.domain_offset((chunk_start, 0), gK_head), (BT, K_DIM), (0, 0)
- )
- ts_k, tg_k = cpasync.tma_partition(
- tma_atom_k_k123,
- 0,
- cute.make_layout(1),
- cute.group_modes(sK_s, 0, 2),
- cute.group_modes(gK_s, 0, 2),
- )
- cute.copy(tma_atom_k_k123, tg_k, ts_k, tma_bar_ptr=k123_tma_mbar)
-
- sG_s = sG_k123
- gG_s = cute.local_tile(
- cute.domain_offset((chunk_start, 0), gG_head), (BT, K_DIM), (0, 0)
- )
- ts_g, tg_g = cpasync.tma_partition(
- tma_atom_g_k123,
- 0,
- cute.make_layout(1),
- cute.group_modes(sG_s, 0, 2),
- cute.group_modes(gG_s, 0, 2),
- )
- cute.copy(tma_atom_g_k123, tg_g, ts_g, tma_bar_ptr=k123_tma_mbar)
-
- if lane_id == 0:
- cute.arch.mbarrier_arrive(k123_tma_mbar)
-
- # Wait TMA
- cute.arch.mbarrier_wait(k123_tma_mbar, phase=c_phase)
-
- # Load beta[chunk] into SMEM (first K1 warp, 32 threads load 64 bf16 in 2 iters)
- if warp_idx == K1_FIRST_WARP:
- for bi in cutlass.range_constexpr(2):
- beta_t = bi * WARP_SIZE + lane_id # 0..63
- sBeta[beta_t] = mBeta[i_b, chunk_start_local + beta_t, i_h]
-
- # ============================================================
- # Fused gate activation + cumsum + KG/KS/QS scaling
- # Each thread processes 1 column across all 64 rows sequentially.
- # No rGact, no sPartialLast, no cross-warp shuffle, no barriers.
- # ============================================================
-
- # Fused Pass: gate activation → cumsum → sGcum (single pass, 64 rows)
- # 16 dynamic iterations × 4 constexpr unrolled = 64 rows
- running_sum = cutlass.Float32(0.0)
- for row_base in cutlass.range(BT // 4):
- for ri in cutlass.range_constexpr(4):
- row = row_base * 4 + ri
- g_val = sG_k123[row, my_col].to(cutlass.Float32)
- if HAS_BIAS:
- g_val = g_val + rBias_val
- g_activated = cutlass.Float32(0.0)
- if USE_SAFE_GATE:
- sigmoid_g = fast_rcp(
- cutlass.Float32(1.0) + cute.exp2(-exp_A * g_val * LOG2E, fastmath=True)
- )
- g_activated = lower_bound * sigmoid_g
- else:
- softplus_g = (
- cute.log2(
- cutlass.Float32(1.0) + cute.exp2(g_val * LOG2E, fastmath=True),
- fastmath=True,
- )
- * LN2
- )
- g_activated = -exp_A * softplus_g
- running_sum = running_sum + g_activated
- sGcum[row, my_col] = running_sum * cumsum_scale
-
- # Signal K2: sGcum ready (all 64 rows written)
- cute.arch.mbarrier_arrive(k123_k1_done_mbar)
-
- # KG/KS/QS scaling (reads sGcum + sK + sQ, writes Zone B)
- # 16 dynamic × 4 constexpr = 64 rows
- gk_last_cs = running_sum * cumsum_scale # total cumsum for this column
- for row_base in cutlass.range(BT // 4):
- for ri in cutlass.range_constexpr(4):
- row = row_base * 4 + ri
- cs = sGcum[row, my_col]
- k_val = sK_k123[row, my_col].to(cutlass.Float32)
- q_val = sQ_k123[row, my_col].to(cutlass.Float32)
- exp2_cs = cute.exp2(cs, fastmath=True)
- exp2_kg = cute.exp2(gk_last_cs - cs, fastmath=True)
- # Write directly to Zone B swizzled layout
- sKG[my_col + (row % 16) * 128, 0, row // 16, 0] = (k_val * exp2_kg).to(
- cutlass.BFloat16
- )
- sKS[my_col + (row % 16) * 128, 0, row // 16, 0] = (
- cutlass.Float32(-1.0) * k_val * exp2_cs
- ).to(cutlass.BFloat16)
- sQS[row + (my_col % 16) * 64, 0, my_col // 16, 0] = (
- q_val * exp2_cs * scale
- ).to(cutlass.BFloat16)
-
- # gkLast: all threads write their own column
- sGkLast[my_col] = cute.exp2(gk_last_cs, fastmath=True)
-
- # Signal readout WG: gk_last is in SMEM (K1 arrive, non-blocking)
- cute.arch.mbarrier_arrive(gk_last_ready_mbar)
-
- # ============================================================
- # K3 phase: wait K2, then inversion + store to Zone B
- # (merged into K1 warps — K1 is idle here, K3 reuses same 4 warps)
- # ============================================================
- store_warp = k1_warp # k1_warp 0-3 maps directly to store_warp 0-3
-
- cute.arch.mbarrier_wait(k123_mma_done_mbar, phase=c_phase)
-
- # Write lower-triangle tiles with actual A_qk values
- # (sAQC zeroed once before chunk loop — upper triangle stays zero)
- for tile_idx in cutlass.range_constexpr(NUM_TILES):
- i_q = _TILE_IQ[tile_idx]
- i_k = _TILE_IK[tile_idx]
- is_diag = _TILE_IQ[tile_idx] == _TILE_IK[tile_idx]
- aqk_col_base = tile_idx * BC
-
- for ri in cutlass.range(BC // NUM_STORE_WARPS):
- local_row = store_warp * (BC // NUM_STORE_WARPS) + ri
- if lane_id < BC:
- local_col = lane_id
- aqk_val = sAqk_k123[local_row, aqk_col_base + local_col]
- if is_diag and local_row < local_col:
- aqk_val = cutlass.BFloat16(0.0)
- aqc_m = i_q * BC + local_row
- aqc_k = i_k * BC + local_col
- sAQC[aqc_m + (aqc_k % 16) * 64, 0, aqc_k // 16, 0] = aqk_val
-
- # Akk inversion: 4 stages
- if store_warp == 0:
- _invert_diag(sAkk_k123, lane_id // 16, lane_id)
- if store_warp == 1:
- _invert_diag(sAkk_k123, 2 + lane_id // 16, lane_id)
- inv_internal_barrier()
-
- # Stage 2: Ai10, Ai21, Ai32
- if store_warp == 0:
- t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk_k123, 1, 1, 0, 1, lane_id)
- r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_B(
- sAkk_k123, 0, 0, t0, t2, t1, t3, t4, t6, t5, t7, lane_id
- )
- _store_neg_C(sAkk_k123, 1, 0, r0, r1, r2, r3, r4, r5, r6, r7, lane_id)
- if store_warp == 1:
- t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk_k123, 2, 2, 1, 2, lane_id)
- r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_B(
- sAkk_k123, 1, 1, t0, t2, t1, t3, t4, t6, t5, t7, lane_id
- )
- _store_neg_C(sAkk_k123, 2, 1, r0, r1, r2, r3, r4, r5, r6, r7, lane_id)
- if store_warp == 2:
- t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk_k123, 3, 3, 2, 3, lane_id)
- r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_B(
- sAkk_k123, 2, 2, t0, t2, t1, t3, t4, t6, t5, t7, lane_id
- )
- _store_neg_C(sAkk_k123, 3, 2, r0, r1, r2, r3, r4, r5, r6, r7, lane_id)
- inv_internal_barrier()
-
- # Stage 3: Ai20, Ai31
- _zz = cutlass.Float32(0.0)
- t0 = _zz
- t1 = _zz
- t2 = _zz
- t3 = _zz
- t4 = _zz
- t5 = _zz
- t6 = _zz
- t7 = _zz
- if store_warp == 0:
- t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk_k123, 0, 2, 0, 0, lane_id)
- if store_warp == 2:
- s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk_k123, 1, 2, 1, 0, lane_id)
- _store_C_temp(sTemp, 0, s0, s1, s2, s3, s4, s5, s6, s7, lane_id)
- if store_warp == 1:
- t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk_k123, 1, 3, 1, 1, lane_id)
- if store_warp == 3:
- s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk_k123, 2, 3, 2, 1, lane_id)
- _store_C_temp(sTemp, 1, s0, s1, s2, s3, s4, s5, s6, s7, lane_id)
- inv_internal_barrier()
- if store_warp == 0:
- e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 0, lane_id)
- t0 = t0 + e0
- t1 = t1 + e1
- t2 = t2 + e2
- t3 = t3 + e3
- t4 = t4 + e4
- t5 = t5 + e5
- t6 = t6 + e6
- t7 = t7 + e7
- sb = _shuffle_C_to_B(t0, t1, t2, t3, t4, t5, t6, t7, lane_id)
- r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_A(
- sAkk_k123, 2, 2, sb[0], sb[1], sb[4], sb[5], sb[2], sb[3], sb[6], sb[7], lane_id
- )
- _store_neg_C(sAkk_k123, 2, 0, r0, r1, r2, r3, r4, r5, r6, r7, lane_id)
- if store_warp == 1:
- e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 1, lane_id)
- t0 = t0 + e0
- t1 = t1 + e1
- t2 = t2 + e2
- t3 = t3 + e3
- t4 = t4 + e4
- t5 = t5 + e5
- t6 = t6 + e6
- t7 = t7 + e7
- sb = _shuffle_C_to_B(t0, t1, t2, t3, t4, t5, t6, t7, lane_id)
- r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_A(
- sAkk_k123, 3, 3, sb[0], sb[1], sb[4], sb[5], sb[2], sb[3], sb[6], sb[7], lane_id
- )
- _store_neg_C(sAkk_k123, 3, 1, r0, r1, r2, r3, r4, r5, r6, r7, lane_id)
- inv_internal_barrier()
-
- # Stage 4: Ai30
- t0 = _zz
- t1 = _zz
- t2 = _zz
- t3 = _zz
- t4 = _zz
- t5 = _zz
- t6 = _zz
- t7 = _zz
- if store_warp == 0:
- t0, t1, t2, t3, t4, t5, t6, t7 = _matmul_AB(sAkk_k123, 0, 3, 0, 0, lane_id)
- if store_warp == 1:
- s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk_k123, 1, 3, 1, 0, lane_id)
- _store_C_temp(sTemp, 0, s0, s1, s2, s3, s4, s5, s6, s7, lane_id)
- if store_warp == 2:
- s0, s1, s2, s3, s4, s5, s6, s7 = _matmul_AB(sAkk_k123, 2, 3, 2, 0, lane_id)
- _store_C_temp(sTemp, 1, s0, s1, s2, s3, s4, s5, s6, s7, lane_id)
- inv_internal_barrier()
- if store_warp == 0:
- e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 0, lane_id)
- t0 = t0 + e0
- t1 = t1 + e1
- t2 = t2 + e2
- t3 = t3 + e3
- t4 = t4 + e4
- t5 = t5 + e5
- t6 = t6 + e6
- t7 = t7 + e7
- e0, e1, e2, e3, e4, e5, e6, e7 = _load_C_temp(sTemp, 1, lane_id)
- t0 = t0 + e0
- t1 = t1 + e1
- t2 = t2 + e2
- t3 = t3 + e3
- t4 = t4 + e4
- t5 = t5 + e5
- t6 = t6 + e6
- t7 = t7 + e7
- sb = _shuffle_C_to_B(t0, t1, t2, t3, t4, t5, t6, t7, lane_id)
- r0, r1, r2, r3, r4, r5, r6, r7 = _chain_mma_A(
- sAkk_k123, 3, 3, sb[0], sb[1], sb[4], sb[5], sb[2], sb[3], sb[6], sb[7], lane_id
- )
- _store_neg_C(sAkk_k123, 3, 0, r0, r1, r2, r3, r4, r5, r6, r7, lane_id)
- inv_internal_barrier()
-
- # Write inverted Akk * beta to K4's swizzled SMEM (sAB in Zone B)
- inv_row_start = store_warp * 16
- for ri in cutlass.range(BC):
- inv_row = inv_row_start + ri
- c0 = lane_id * 2
- c1 = lane_id * 2 + 1
- beta_c0 = sBeta[c0].to(cutlass.Float32)
- beta_c1 = sBeta[c1].to(cutlass.Float32)
- v0 = (
- cutlass.Float32(sAkk_k123[inv_row, c0])
- * cutlass.Float32(inv_row >= c0)
- * beta_c0
- )
- v1 = (
- cutlass.Float32(sAkk_k123[inv_row, c1])
- * cutlass.Float32(inv_row >= c1)
- * beta_c1
- )
- sAB[inv_row + (c0 % 16) * 64, 0, c0 // 16, 0] = v0.to(cutlass.BFloat16)
- sAB[inv_row + (c1 % 16) * 64, 0, c1 // 16, 0] = v1.to(cutlass.BFloat16)
-
- # K1+K3 done, idle during K4 phase
- phase_nbar.arrive_and_wait() # K123->K4
- phase_nbar.arrive_and_wait() # K4->next
-
- if PROFILE_CLOCKS:
- clk_k1_end = read_clock()
- if k1_warp == 0 and lane_id == 0:
- mClocks[9, bid] = clk_k1_end - clk_k1_start
-
- # ==============================================================
- # K2 (warps 12-19): intra sub-chunk attention MMA
- # 8 warps: warps 0-5 do 1 tile each, warps 6-7 do 2 tiles each (row 3)
- # ==============================================================
- elif warp_idx >= K2_FIRST_WARP and warp_idx < K2_FIRST_WARP + NUM_MMA_WARPS:
- # K2 warp decode
- clk_k2_start = cutlass.Int64(0)
- if PROFILE_CLOCKS:
- clk_k2_start = read_clock()
- mma_warp = warp_idx - K2_FIRST_WARP
- # Tile A assignment (all 8 warps)
- my_i_q = cutlass.Int32(0)
- my_i_k = cutlass.Int32(0)
- if mma_warp < 1:
- my_i_q = cutlass.Int32(0)
- my_i_k = mma_warp
- elif mma_warp < 3:
- my_i_q = cutlass.Int32(1)
- my_i_k = mma_warp - 1
- elif mma_warp < 6:
- my_i_q = cutlass.Int32(2)
- my_i_k = mma_warp - 3
- else:
- # Warps 6-7: first tile of row 3 pair
- # W6: (3,0), W7: (3,2)
- my_i_q = cutlass.Int32(3)
- my_i_k = (mma_warp - 6) * 2
- # Tile B assignment (warps 6-7 only): i_k_b = i_k_a + 1
- # W6: (3,1), W7: (3,3)
- my_i_k_b = my_i_k + cutlass.Int32(1)
-
- for chunk_c in cutlass.range(num_chunks):
- c_phase = chunk_c % 2
-
- cute.arch.mbarrier_wait(k123_k1_done_mbar, phase=c_phase)
-
- # --- Tile A (all 8 warps) ---
- q_row_base = my_i_q * BC
- k_row_base = my_i_k * BC
- # tile_col_base maps to sAqk slot index:
- # warps 0-5: 1:1 mapping (tile slot = mma_warp)
- # warps 6-7: tile slot = 6+(mma_warp-6)*2 = 6 or 8
- tile_col_base = mma_warp * BC
- if mma_warp >= 6:
- tile_col_base = (6 + (mma_warp - 6) * 2) * BC
- akk_row_base = k_row_base
- akk_col_base = q_row_base
- norm_row = q_row_base
- if my_i_q == my_i_k:
- norm_row = q_row_base + cutlass.Int32(BC // 2)
-
- group_id = lane_id // 4
- tid_in_group = lane_id % 4
- row0, row1 = group_id, group_id + 8
-
- thr_mma = tiled_mma_k2.get_slice(lane_id)
- thr_copy_A = tiled_copy_mma_A.get_slice(lane_id)
- thr_copy_B = tiled_copy_mma_B.get_slice(lane_id)
- thr_copy_Gn = tiled_copy_Gcum_norm.get_slice(tid_in_group)
- thr_copy_Ggate = tiled_copy_Gcum_gate.get_slice(lane_id)
-
- beta_row0 = sBeta[q_row_base + row0].to(cutlass.Float32)
- beta_row1 = sBeta[q_row_base + row1].to(cutlass.Float32)
-
- _z = cutlass.Float32(0.0)
- acc_aqk_n0_0, acc_aqk_n0_1, acc_aqk_n0_2, acc_aqk_n0_3 = _z, _z, _z, _z
- acc_aqk_n1_0, acc_aqk_n1_1, acc_aqk_n1_2, acc_aqk_n1_3 = _z, _z, _z, _z
- acc_akk_n0_0, acc_akk_n0_1, acc_akk_n0_2, acc_akk_n0_3 = _z, _z, _z, _z
- acc_akk_n1_0, acc_akk_n1_1, acc_akk_n1_2, acc_akk_n1_3 = _z, _z, _z, _z
-
- for k_block in cutlass.range_constexpr(NUM_MMA_K_TILES):
- sQ_tile = cute.local_tile(sQ_k123, tiler=(16, 8), coord=(my_i_q, k_block))
- tCrQ = tiled_mma_k2.make_fragment_A(thr_mma.partition_A(sQ_tile))
- cute.copy(
- tiled_copy_mma_A, thr_copy_A.partition_S(sQ_tile), thr_copy_A.retile(tCrQ)
- )
-
- sKq_tile = cute.local_tile(sK_k123, tiler=(16, 8), coord=(my_i_q, k_block))
- tCrKq = tiled_mma_k2.make_fragment_A(thr_mma.partition_A(sKq_tile))
- cute.copy(
- tiled_copy_mma_A, thr_copy_A.partition_S(sKq_tile), thr_copy_A.retile(tCrKq)
- )
-
- sGn_tile = cute.local_tile(sGcum, tiler=(1, 8), coord=(norm_row, k_block))
- tCsGn = thr_copy_Gn.partition_S(sGn_tile)
- tCrGn = cute.make_fragment_like(tCsGn, cutlass.Float32)
- cute.copy(tiled_copy_Gcum_norm, tCsGn, thr_copy_Gn.retile(tCrGn))
- g_norm_0 = tCrGn[0]
- g_norm_1 = tCrGn[1]
-
- sGq_tile = cute.local_tile(sGcum, tiler=(16, 8), coord=(my_i_q, k_block))
- tCrGq = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGq_tile))
- cute.copy(
- tiled_copy_Gcum_gate,
- thr_copy_Ggate.partition_S(sGq_tile),
- thr_copy_Ggate.retile(tCrGq),
- )
- gate_q_0 = cute.exp2(tCrGq[0] - g_norm_0, fastmath=True)
- gate_q_1 = cute.exp2(tCrGq[1] - g_norm_1, fastmath=True)
- gate_q_2 = cute.exp2(tCrGq[2] - g_norm_0, fastmath=True)
- gate_q_3 = cute.exp2(tCrGq[3] - g_norm_1, fastmath=True)
-
- qa0 = tCrQ[0].to(cutlass.Float32) * gate_q_0
- qa1 = tCrQ[2].to(cutlass.Float32) * gate_q_2
- qa2 = tCrQ[1].to(cutlass.Float32) * gate_q_1
- qa3 = tCrQ[3].to(cutlass.Float32) * gate_q_3
- ka0 = tCrKq[0].to(cutlass.Float32) * gate_q_0
- ka1 = tCrKq[2].to(cutlass.Float32) * gate_q_2
- ka2 = tCrKq[1].to(cutlass.Float32) * gate_q_1
- ka3 = tCrKq[3].to(cutlass.Float32) * gate_q_3
-
- sK_tile_n0 = cute.local_tile(sK_k123, tiler=(8, 8), coord=(my_i_k * 2, k_block))
- tCrK_n0 = tiled_mma_k2.make_fragment_B(thr_mma.partition_B(sK_tile_n0))
- cute.copy(
- tiled_copy_mma_B, thr_copy_B.partition_S(sK_tile_n0), thr_copy_B.retile(tCrK_n0)
- )
-
- sK_tile_n1 = cute.local_tile(sK_k123, tiler=(8, 8), coord=(my_i_k * 2 + 1, k_block))
- tCrK_n1 = tiled_mma_k2.make_fragment_B(thr_mma.partition_B(sK_tile_n1))
- cute.copy(
- tiled_copy_mma_B, thr_copy_B.partition_S(sK_tile_n1), thr_copy_B.retile(tCrK_n1)
- )
-
- sGk_tile = cute.local_tile(sGcum, tiler=(16, 8), coord=(my_i_k, k_block))
- tCrGk = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGk_tile))
- cute.copy(
- tiled_copy_Gcum_gate,
- thr_copy_Ggate.partition_S(sGk_tile),
- thr_copy_Ggate.retile(tCrGk),
- )
- gk_n0_0 = cute.exp2(g_norm_0 - tCrGk[0], fastmath=True)
- gk_n0_1 = cute.exp2(g_norm_1 - tCrGk[1], fastmath=True)
- k_n0_b0 = tCrK_n0[0].to(cutlass.Float32) * gk_n0_0
- k_n0_b1 = tCrK_n0[1].to(cutlass.Float32) * gk_n0_1
- gk_n1_0 = cute.exp2(g_norm_0 - tCrGk[2], fastmath=True)
- gk_n1_1 = cute.exp2(g_norm_1 - tCrGk[3], fastmath=True)
- k_n1_b0 = tCrK_n1[0].to(cutlass.Float32) * gk_n1_0
- k_n1_b1 = tCrK_n1[1].to(cutlass.Float32) * gk_n1_1
-
- acc_aqk_n0_0, acc_aqk_n0_1, acc_aqk_n0_2, acc_aqk_n0_3 = mma_tf32_m16n8k8(
- qa0,
- qa1,
- qa2,
- qa3,
- k_n0_b0,
- k_n0_b1,
- acc_aqk_n0_0,
- acc_aqk_n0_1,
- acc_aqk_n0_2,
- acc_aqk_n0_3,
- )
- acc_aqk_n1_0, acc_aqk_n1_1, acc_aqk_n1_2, acc_aqk_n1_3 = mma_tf32_m16n8k8(
- qa0,
- qa1,
- qa2,
- qa3,
- k_n1_b0,
- k_n1_b1,
- acc_aqk_n1_0,
- acc_aqk_n1_1,
- acc_aqk_n1_2,
- acc_aqk_n1_3,
- )
- acc_akk_n0_0, acc_akk_n0_1, acc_akk_n0_2, acc_akk_n0_3 = mma_tf32_m16n8k8(
- ka0,
- ka1,
- ka2,
- ka3,
- k_n0_b0,
- k_n0_b1,
- acc_akk_n0_0,
- acc_akk_n0_1,
- acc_akk_n0_2,
- acc_akk_n0_3,
- )
- acc_akk_n1_0, acc_akk_n1_1, acc_akk_n1_2, acc_akk_n1_3 = mma_tf32_m16n8k8(
- ka0,
- ka1,
- ka2,
- ka3,
- k_n1_b0,
- k_n1_b1,
- acc_akk_n1_0,
- acc_akk_n1_1,
- acc_akk_n1_2,
- acc_akk_n1_3,
- )
-
- # Store tile A
- col0, col1 = tid_in_group * 2, tid_in_group * 2 + 1
- col2, col3 = 8 + tid_in_group * 2, 8 + tid_in_group * 2 + 1
-
- sAqk_k123[row0, tile_col_base + col0] = (acc_aqk_n0_0 * scale).to(cutlass.BFloat16)
- sAqk_k123[row0, tile_col_base + col1] = (acc_aqk_n0_1 * scale).to(cutlass.BFloat16)
- sAqk_k123[row1, tile_col_base + col0] = (acc_aqk_n0_2 * scale).to(cutlass.BFloat16)
- sAqk_k123[row1, tile_col_base + col1] = (acc_aqk_n0_3 * scale).to(cutlass.BFloat16)
- sAqk_k123[row0, tile_col_base + col2] = (acc_aqk_n1_0 * scale).to(cutlass.BFloat16)
- sAqk_k123[row0, tile_col_base + col3] = (acc_aqk_n1_1 * scale).to(cutlass.BFloat16)
- sAqk_k123[row1, tile_col_base + col2] = (acc_aqk_n1_2 * scale).to(cutlass.BFloat16)
- sAqk_k123[row1, tile_col_base + col3] = (acc_aqk_n1_3 * scale).to(cutlass.BFloat16)
-
- sAkk_k123[akk_row_base + row0, akk_col_base + col0] = acc_akk_n0_0 * beta_row0
- sAkk_k123[akk_row_base + row0, akk_col_base + col1] = acc_akk_n0_1 * beta_row0
- sAkk_k123[akk_row_base + row1, akk_col_base + col0] = acc_akk_n0_2 * beta_row1
- sAkk_k123[akk_row_base + row1, akk_col_base + col1] = acc_akk_n0_3 * beta_row1
- sAkk_k123[akk_row_base + row0, akk_col_base + col2] = acc_akk_n1_0 * beta_row0
- sAkk_k123[akk_row_base + row0, akk_col_base + col3] = acc_akk_n1_1 * beta_row0
- sAkk_k123[akk_row_base + row1, akk_col_base + col2] = acc_akk_n1_2 * beta_row1
- sAkk_k123[akk_row_base + row1, akk_col_base + col3] = acc_akk_n1_3 * beta_row1
-
- # --- Tile B (warps 6-7 only): second tile of row 3 pair ---
- # W6: (3,1), W7: (3,3). Accumulators reused (zero extra reg pressure).
- if mma_warp >= 6:
- k_row_base_b = my_i_k_b * BC
- tile_col_base_b = tile_col_base + BC # tile B = tile A + 1 slot
- akk_row_base_b = k_row_base_b
- norm_row_b = q_row_base
- if my_i_q == my_i_k_b:
- norm_row_b = q_row_base + cutlass.Int32(BC // 2)
-
- acc_aqk_n0_0, acc_aqk_n0_1, acc_aqk_n0_2, acc_aqk_n0_3 = _z, _z, _z, _z
- acc_aqk_n1_0, acc_aqk_n1_1, acc_aqk_n1_2, acc_aqk_n1_3 = _z, _z, _z, _z
- acc_akk_n0_0, acc_akk_n0_1, acc_akk_n0_2, acc_akk_n0_3 = _z, _z, _z, _z
- acc_akk_n1_0, acc_akk_n1_1, acc_akk_n1_2, acc_akk_n1_3 = _z, _z, _z, _z
-
- for k_block in cutlass.range_constexpr(NUM_MMA_K_TILES):
- sQ_tile = cute.local_tile(sQ_k123, tiler=(16, 8), coord=(my_i_q, k_block))
- tCrQ = tiled_mma_k2.make_fragment_A(thr_mma.partition_A(sQ_tile))
- cute.copy(
- tiled_copy_mma_A, thr_copy_A.partition_S(sQ_tile), thr_copy_A.retile(tCrQ)
- )
-
- sKq_tile = cute.local_tile(sK_k123, tiler=(16, 8), coord=(my_i_q, k_block))
- tCrKq = tiled_mma_k2.make_fragment_A(thr_mma.partition_A(sKq_tile))
- cute.copy(
- tiled_copy_mma_A, thr_copy_A.partition_S(sKq_tile), thr_copy_A.retile(tCrKq)
- )
-
- sGn_tile = cute.local_tile(sGcum, tiler=(1, 8), coord=(norm_row_b, k_block))
- tCsGn = thr_copy_Gn.partition_S(sGn_tile)
- tCrGn = cute.make_fragment_like(tCsGn, cutlass.Float32)
- cute.copy(tiled_copy_Gcum_norm, tCsGn, thr_copy_Gn.retile(tCrGn))
- g_norm_0 = tCrGn[0]
- g_norm_1 = tCrGn[1]
-
- sGq_tile = cute.local_tile(sGcum, tiler=(16, 8), coord=(my_i_q, k_block))
- tCrGq = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGq_tile))
- cute.copy(
- tiled_copy_Gcum_gate,
- thr_copy_Ggate.partition_S(sGq_tile),
- thr_copy_Ggate.retile(tCrGq),
- )
- gate_q_0 = cute.exp2(tCrGq[0] - g_norm_0, fastmath=True)
- gate_q_1 = cute.exp2(tCrGq[1] - g_norm_1, fastmath=True)
- gate_q_2 = cute.exp2(tCrGq[2] - g_norm_0, fastmath=True)
- gate_q_3 = cute.exp2(tCrGq[3] - g_norm_1, fastmath=True)
-
- qa0 = tCrQ[0].to(cutlass.Float32) * gate_q_0
- qa1 = tCrQ[2].to(cutlass.Float32) * gate_q_2
- qa2 = tCrQ[1].to(cutlass.Float32) * gate_q_1
- qa3 = tCrQ[3].to(cutlass.Float32) * gate_q_3
- ka0 = tCrKq[0].to(cutlass.Float32) * gate_q_0
- ka1 = tCrKq[2].to(cutlass.Float32) * gate_q_2
- ka2 = tCrKq[1].to(cutlass.Float32) * gate_q_1
- ka3 = tCrKq[3].to(cutlass.Float32) * gate_q_3
-
- sK_tile_n0 = cute.local_tile(
- sK_k123, tiler=(8, 8), coord=(my_i_k_b * 2, k_block)
- )
- tCrK_n0 = tiled_mma_k2.make_fragment_B(thr_mma.partition_B(sK_tile_n0))
- cute.copy(
- tiled_copy_mma_B,
- thr_copy_B.partition_S(sK_tile_n0),
- thr_copy_B.retile(tCrK_n0),
- )
-
- sK_tile_n1 = cute.local_tile(
- sK_k123, tiler=(8, 8), coord=(my_i_k_b * 2 + 1, k_block)
- )
- tCrK_n1 = tiled_mma_k2.make_fragment_B(thr_mma.partition_B(sK_tile_n1))
- cute.copy(
- tiled_copy_mma_B,
- thr_copy_B.partition_S(sK_tile_n1),
- thr_copy_B.retile(tCrK_n1),
- )
-
- sGk_tile = cute.local_tile(sGcum, tiler=(16, 8), coord=(my_i_k_b, k_block))
- tCrGk = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGk_tile))
- cute.copy(
- tiled_copy_Gcum_gate,
- thr_copy_Ggate.partition_S(sGk_tile),
- thr_copy_Ggate.retile(tCrGk),
- )
- gk_n0_0 = cute.exp2(g_norm_0 - tCrGk[0], fastmath=True)
- gk_n0_1 = cute.exp2(g_norm_1 - tCrGk[1], fastmath=True)
- k_n0_b0 = tCrK_n0[0].to(cutlass.Float32) * gk_n0_0
- k_n0_b1 = tCrK_n0[1].to(cutlass.Float32) * gk_n0_1
- gk_n1_0 = cute.exp2(g_norm_0 - tCrGk[2], fastmath=True)
- gk_n1_1 = cute.exp2(g_norm_1 - tCrGk[3], fastmath=True)
- k_n1_b0 = tCrK_n1[0].to(cutlass.Float32) * gk_n1_0
- k_n1_b1 = tCrK_n1[1].to(cutlass.Float32) * gk_n1_1
-
- acc_aqk_n0_0, acc_aqk_n0_1, acc_aqk_n0_2, acc_aqk_n0_3 = mma_tf32_m16n8k8(
- qa0,
- qa1,
- qa2,
- qa3,
- k_n0_b0,
- k_n0_b1,
- acc_aqk_n0_0,
- acc_aqk_n0_1,
- acc_aqk_n0_2,
- acc_aqk_n0_3,
- )
- acc_aqk_n1_0, acc_aqk_n1_1, acc_aqk_n1_2, acc_aqk_n1_3 = mma_tf32_m16n8k8(
- qa0,
- qa1,
- qa2,
- qa3,
- k_n1_b0,
- k_n1_b1,
- acc_aqk_n1_0,
- acc_aqk_n1_1,
- acc_aqk_n1_2,
- acc_aqk_n1_3,
- )
- acc_akk_n0_0, acc_akk_n0_1, acc_akk_n0_2, acc_akk_n0_3 = mma_tf32_m16n8k8(
- ka0,
- ka1,
- ka2,
- ka3,
- k_n0_b0,
- k_n0_b1,
- acc_akk_n0_0,
- acc_akk_n0_1,
- acc_akk_n0_2,
- acc_akk_n0_3,
- )
- acc_akk_n1_0, acc_akk_n1_1, acc_akk_n1_2, acc_akk_n1_3 = mma_tf32_m16n8k8(
- ka0,
- ka1,
- ka2,
- ka3,
- k_n1_b0,
- k_n1_b1,
- acc_akk_n1_0,
- acc_akk_n1_1,
- acc_akk_n1_2,
- acc_akk_n1_3,
- )
-
- # Store tile B
- sAqk_k123[row0, tile_col_base_b + col0] = (acc_aqk_n0_0 * scale).to(
- cutlass.BFloat16
- )
- sAqk_k123[row0, tile_col_base_b + col1] = (acc_aqk_n0_1 * scale).to(
- cutlass.BFloat16
- )
- sAqk_k123[row1, tile_col_base_b + col0] = (acc_aqk_n0_2 * scale).to(
- cutlass.BFloat16
- )
- sAqk_k123[row1, tile_col_base_b + col1] = (acc_aqk_n0_3 * scale).to(
- cutlass.BFloat16
- )
- sAqk_k123[row0, tile_col_base_b + col2] = (acc_aqk_n1_0 * scale).to(
- cutlass.BFloat16
- )
- sAqk_k123[row0, tile_col_base_b + col3] = (acc_aqk_n1_1 * scale).to(
- cutlass.BFloat16
- )
- sAqk_k123[row1, tile_col_base_b + col2] = (acc_aqk_n1_2 * scale).to(
- cutlass.BFloat16
- )
- sAqk_k123[row1, tile_col_base_b + col3] = (acc_aqk_n1_3 * scale).to(
- cutlass.BFloat16
- )
-
- sAkk_k123[akk_row_base_b + row0, akk_col_base + col0] = acc_akk_n0_0 * beta_row0
- sAkk_k123[akk_row_base_b + row0, akk_col_base + col1] = acc_akk_n0_1 * beta_row0
- sAkk_k123[akk_row_base_b + row1, akk_col_base + col0] = acc_akk_n0_2 * beta_row1
- sAkk_k123[akk_row_base_b + row1, akk_col_base + col1] = acc_akk_n0_3 * beta_row1
- sAkk_k123[akk_row_base_b + row0, akk_col_base + col2] = acc_akk_n1_0 * beta_row0
- sAkk_k123[akk_row_base_b + row0, akk_col_base + col3] = acc_akk_n1_1 * beta_row0
- sAkk_k123[akk_row_base_b + row1, akk_col_base + col2] = acc_akk_n1_2 * beta_row1
- sAkk_k123[akk_row_base_b + row1, akk_col_base + col3] = acc_akk_n1_3 * beta_row1
-
- cute.arch.mbarrier_arrive(k123_mma_done_mbar)
-
- # K2 idle during K4 phase
- phase_nbar.arrive_and_wait() # K123->K4
- phase_nbar.arrive_and_wait() # K4->next
-
- if PROFILE_CLOCKS:
- clk_k2_end = read_clock()
- if mma_warp == 0 and lane_id == 0:
- mClocks[10, bid] = clk_k2_end - clk_k2_start
-
- # ==============================================================
- # MMA warp (warp 0): 6 MMAs per chunk
- # ==============================================================
- elif warp_idx == K4_MMA_WARP:
- mc = (0, 0, 0, 0)
- ml = cute.make_layout((1, 1, 1, 1))
-
- # Clock accumulators (thread 0 only, i64)
- clk_k123_acc = cutlass.Int64(0)
- clk_k4_acc = cutlass.Int64(0)
- clk0 = cutlass.Int64(0)
- clk1 = cutlass.Int64(0)
- clk2 = cutlass.Int64(0)
- clk_mma1_acc = cutlass.Int64(0)
- clk_mma2_acc = cutlass.Int64(0)
- clk_wait_s_mma4_acc = cutlass.Int64(0)
- clk_wait_w_mma3_acc = cutlass.Int64(0)
- clk_wait_nv_mma5_acc = cutlass.Int64(0)
- clk_mma6_drain_acc = cutlass.Int64(0)
- clk_k4a = cutlass.Int64(0)
- clk_k4b = cutlass.Int64(0)
-
- for chunk_c in cutlass.range(num_chunks):
- chunk_start = i_b * num_chunks * BT + chunk_c * BT
-
- if PROFILE_CLOCKS:
- clk0 = read_clock()
-
- # Wait for K123 to finish
- phase_nbar.arrive_and_wait() # K123->K4
-
- if PROFILE_CLOCKS:
- clk1 = read_clock()
- clk_k123_acc = clk_k123_acc + (clk1 - clk0)
-
- if DEBUG_K4_LEVEL == 0:
- pass # Skip K4 entirely for debugging
- else:
- fA_kmn = thr_kmn.make_fragment_A(sAB)
- fB_ks = thr_kmn.make_fragment_B(sKS)
- fB_v = thr_kmn.make_fragment_B(sV_ext)
- fA_w = thr_kmn.make_fragment_A(sW)
- fB_s = thr_kmn.make_fragment_B(sST)
- fA_q = thr_kmn.make_fragment_A(sQS)
- fA_aqc = thr_kmn.make_fragment_A(sAQC)
- fB_nv = thr_kmn.make_fragment_B(sNV_b)
- fA_nv_mn = thr_mn.make_fragment_A(sNV_a)
- fB_kg = thr_mn.make_fragment_B(sKG)
-
- if PROFILE_CLOCKS:
- clk_k4a = read_clock()
-
- # MMA1: W = AB @ KS
- w_h = w_prod.acquire_and_advance()
- tiled_mma_kmn.set(Field.ACCUMULATE, False)
- for k in cutlass.range_constexpr(cute.size(sKS.shape[2])):
- cute.gemm(
- tiled_mma_kmn,
- tCtW,
- fA_kmn[dice + (0,)][None, None, k],
- fB_ks[dice + (0,)][None, None, k],
- tCtW,
- )
- if k == 0:
- tiled_mma_kmn.set(Field.ACCUMULATE, True)
- w_h.commit()
-
- if PROFILE_CLOCKS:
- clk_k4b = read_clock()
- clk_mma1_acc = clk_mma1_acc + (clk_k4b - clk_k4a)
- clk_k4a = clk_k4b
-
- # MMA2: U = AB @ V
- vh = v_cons.wait_and_advance()
- tiled_mma_kmn.set(Field.ACCUMULATE, False)
- for k in cutlass.range_constexpr(cute.size(sV_ext.shape[2])):
- cute.gemm(
- tiled_mma_kmn,
- tCtNV,
- fA_kmn[dice + (0,)][None, None, k],
- fB_v[dice + (vh.index,)][None, None, k],
- tCtNV,
- )
- if k == 0:
- tiled_mma_kmn.set(Field.ACCUMULATE, True)
- vh.release()
-
- if PROFILE_CLOCKS:
- clk_k4b = read_clock()
- clk_mma2_acc = clk_mma2_acc + (clk_k4b - clk_k4a)
- clk_k4a = clk_k4b
-
- # Wait WG1: sST bf16 ready in SMEM
- cute.arch.mbarrier_wait(st_ready_mbar, phase=chunk_c % 2)
-
- # MMA4: OI = QS @ S
- tiled_mma_kmn.set(Field.ACCUMULATE, False)
- for k in cutlass.range_constexpr(cute.size(sST.shape[2])):
- cute.gemm(
- tiled_mma_kmn,
- tCtO,
- fA_q[dice + (0,)][None, None, k],
- fB_s[dice + (0,)][None, None, k],
- tCtO,
- )
- if k == 0:
- tiled_mma_kmn.set(Field.ACCUMULATE, True)
-
- if PROFILE_CLOCKS:
- clk_k4b = read_clock()
- clk_wait_s_mma4_acc = clk_wait_s_mma4_acc + (clk_k4b - clk_k4a)
- clk_k4a = clk_k4b
-
- # Wait W readout
- sW_ready_nbar.arrive_and_wait()
-
- # MMA3: NV += sW @ S (accumulate into U)
- nv_h = nv_prod.acquire_and_advance()
- tiled_mma_kmn.set(Field.ACCUMULATE, True)
- for k in cutlass.range_constexpr(cute.size(sST.shape[2])):
- cute.gemm(
- tiled_mma_kmn,
- tCtNV,
- fA_w[dice + (0,)][None, None, k],
- fB_s[dice + (0,)][None, None, k],
- tCtNV,
- )
- nv_h.commit()
-
- if PROFILE_CLOCKS:
- clk_k4b = read_clock()
- clk_wait_w_mma3_acc = clk_wait_w_mma3_acc + (clk_k4b - clk_k4a)
- clk_k4a = clk_k4b
-
- # Wait NV readout
- sNV_ready_nbar.arrive_and_wait()
-
- # MMA5: O += AQC @ NV
- o_h = o_prod.acquire_and_advance()
- if not DEBUG_SKIP_MMA5:
- tiled_mma_kmn.set(Field.ACCUMULATE, True)
- for k in cutlass.range_constexpr(cute.size(sNV_b.shape[2])):
- cute.gemm(
- tiled_mma_kmn,
- tCtO,
- fA_aqc[dice + (0,)][None, None, k],
- fB_nv[dice + (0,)][None, None, k],
- tCtO,
- )
- o_h.commit()
-
- if PROFILE_CLOCKS:
- clk_k4b = read_clock()
- clk_wait_nv_mma5_acc = clk_wait_nv_mma5_acc + (clk_k4b - clk_k4a)
- clk_k4a = clk_k4b
-
- # Wait WG1: decayed state written back to TMEM
- cute.arch.mbarrier_wait(state_decayed_mbar, phase=chunk_c % 2)
-
- # MMA6: accumulate SU = NV^T @ KG onto decayed state in TMEM@256
- tiled_mma_mn_mn.set(Field.ACCUMULATE, True)
- for k in cutlass.range_constexpr(cute.size(sKG.shape[2])):
- cute.gemm(
- tiled_mma_mn_mn,
- tCtState,
- fA_nv_mn[dice + (0,)][None, None, k],
- fB_kg[dice + (0,)][None, None, k],
- tCtState,
- )
-
- # Drain W/NV/O UMMA pipelines, then sync with TMA+Readout
- w_prod.tail()
- nv_prod.tail()
- o_prod.tail()
- store_nbar.arrive_and_wait()
-
- # Signal WG1: MMA6 done (all chunks, including last)
- cute.arch.mbarrier_arrive(mma6_done_mbar)
-
- if PROFILE_CLOCKS:
- clk_k4b = read_clock()
- clk_mma6_drain_acc = clk_mma6_drain_acc + (clk_k4b - clk_k4a)
-
- # K4->next chunk
- phase_nbar.arrive_and_wait()
-
- if PROFILE_CLOCKS:
- clk2 = read_clock()
- clk_k4_acc = clk_k4_acc + (clk2 - clk1)
-
- # Store accumulated clocks (thread 0 = warp 0 = MMA warp)
- if PROFILE_CLOCKS:
- if tidx == 0:
- mClocks[0, bid] = clk_k123_acc
- mClocks[1, bid] = clk_k4_acc
- mClocks[2, bid] = clk_mma1_acc
- mClocks[3, bid] = clk_mma2_acc
- mClocks[4, bid] = clk_wait_s_mma4_acc
- mClocks[5, bid] = clk_wait_w_mma3_acc
- mClocks[6, bid] = clk_wait_nv_mma5_acc
- mClocks[7, bid] = clk_mma6_drain_acc
-
- # Wait for WG1 final state store, then dealloc TMEM
- cute.arch.mbarrier_wait(final_state_done_mbar, phase=0)
- cute.arch.relinquish_tmem_alloc_permit()
- cute.arch.dealloc_tmem(tmem_ptr, 512)
-
- # ==============================================================
- # TMA warp (warp 2): V early load + O store
- # ==============================================================
- elif warp_idx == K4_TMA_WARP:
- mc = (0, 0, 0, 0)
- ml = cute.make_layout((1, 1, 1, 1))
- clk_tma_start = cutlass.Int64(0)
- if PROFILE_CLOCKS:
- clk_tma_start = read_clock()
-
- for chunk_c in cutlass.range(num_chunks):
- chunk_start = i_b * num_chunks * BT + chunk_c * BT
-
- # V early load (overlaps K123 phase)
- gV_h = mV_nkl[(None, None, i_h)]
- gV = cute.local_tile(
- cute.domain_offset((0, chunk_start), gV_h), tiler=(N4, K4_K), coord=(0, 0)
- )
- tBsV, tBgV = cpasync.tma_partition(
- tma_atom_v,
- mc,
- ml,
- cute.group_modes(sV_ext, 0, 3),
- cute.group_modes(thr_kmn.partition_B(gV), 0, 3),
- )
- vh = v_prod.acquire_and_advance()
- cute.copy(tma_atom_v, tBgV, tBsV[None, vh.index], tma_bar_ptr=vh.barrier)
-
- # Wait for K123 to finish
- phase_nbar.arrive_and_wait() # K123->K4
-
- # Store O[c] to GMEM
- gO_h = mOo[(None, None, i_h)] # [T_total, V_dim]
-
- # Wait for MMA + Readout to finish (O is in SMEM)
- store_nbar.arrive_and_wait()
-
- gOo_c = cute.local_tile(
- cute.domain_offset((chunk_start, 0), gO_h), tiler=(M4, N4), coord=(0, 0)
- )
- sOt, gOt = cpasync.tma_partition(
- tma_atom_o_st, mc, ml, cute.group_modes(sO_st, 0, 2), cute.group_modes(gOo_c, 0, 2)
- )
- cute.copy(tma_atom_o_st, sOt, gOt)
-
- # Fence: wait for TMA to finish reading sO from SMEM
- tma_store_fence()
-
- # K4->next chunk
- phase_nbar.arrive_and_wait()
-
- # Drain V pipeline
- v_prod.tail()
-
- if PROFILE_CLOCKS:
- clk_tma_end = read_clock()
- if lane_id == 0:
- mClocks[12, bid] = clk_tma_end - clk_tma_start
-
- # ==============================================================
- # Idle warps (1, 3): phase barriers only
- # ==============================================================
- else:
- for chunk_c in cutlass.range(num_chunks):
- phase_nbar.arrive_and_wait() # K123->K4
- phase_nbar.arrive_and_wait() # K4->next
-
-
-# =============================================================================
-# Host Function
-# =============================================================================
-
-
-def make_host_fn(has_bias=False, use_safe_gate=False, profile_clocks=False):
- """Create the host function for the fused K1234 kernel.
-
- B, H, NC are all dynamic — one compilation serves all shapes.
- Only mode flags (has_bias, use_safe_gate, profile_clocks) specialize at compile time.
- """
- _HAS_BIAS = 1 if has_bias else 0
- _USE_SAFE_GATE = 1 if use_safe_gate else 0
- _PROFILE_CLOCKS = 1 if profile_clocks else 0
-
- tile1 = (M4, N4, K4_K)
- tile3 = (M4, N4, K4_K3)
- tile6 = (M6, N6, K6)
-
- @cute.jit
- def host_fn(
- # K123 raw inputs
- mQ,
- mK,
- mG,
- mA_log,
- mBeta,
- scale_val,
- # K4 inputs (raw pointers — host_fn creates 3D views internally)
- mV_in, # [B, T, H, V_dim] bf16
- mO_out, # [B, T, H, V_dim] bf16 output
- mS_fp32, # [B*H, K, V] fp32 state (init + final output)
- # dt_bias / safe_gate
- mDtBias, # [H, K] fp32 (or dummy [1,1] if no bias)
- lower_bound_val, # float (0.0 if unused)
- # Clock profiling
- mClocks, # [2, B*H] i64 (or dummy if not profiling)
- # Runtime dimensions (dynamic — no recompilation needed)
- num_chunks: cutlass.Int32,
- num_heads: cutlass.Int32,
- batch_size: cutlass.Int32,
- # Launch stream — runtime argument; launching on the DSL default
- # stream races with the executor's non-blocking execution stream.
- stream: cuda.CUstream,
- ):
- # Derive GMEM shapes from runtime dimensions
- T_total = batch_size * num_chunks * BT
- s_row = num_heads * K_DIM
- s_col = 1
- s_h = K_DIM
- BH = batch_size * num_heads
-
- # --- Common 3D layout for [T_total, K_DIM, H] tensors ---
- view_layout_3d = cute.make_layout((T_total, K_DIM, num_heads), stride=(s_row, s_col, s_h))
-
- # --- K4 MMA setup ---
- mma_kmn = sm100_utils.make_trivial_tiled_mma(
- mma_dtype,
- OperandMajorMode.K,
- OperandMajorMode.MN,
- acc_dtype,
- tcgen05.CtaGroup.ONE,
- (M4, N4),
- OperandSource.SMEM,
- )
- mma_mn_mn = sm100_utils.make_trivial_tiled_mma(
- mma_dtype,
- OperandMajorMode.MN,
- OperandMajorMode.MN,
- acc_dtype,
- tcgen05.CtaGroup.ONE,
- (M6, N6),
- OperandSource.SMEM,
- )
-
- sl_ab = sm100_utils.make_smem_layout_a(mma_kmn, tile1, mma_dtype, 1)
- sl_ks = sm100_utils.make_smem_layout_b(mma_kmn, tile1, mma_dtype, 1)
- sl_v = sm100_utils.make_smem_layout_b(mma_kmn, tile1, mma_dtype, 1)
- sl_s = sm100_utils.make_smem_layout_b(mma_kmn, tile3, mma_dtype, 1)
- sl_qs = sm100_utils.make_smem_layout_a(mma_kmn, tile3, mma_dtype, 1)
- sl_aqc = sm100_utils.make_smem_layout_a(mma_kmn, tile1, mma_dtype, 1)
- sl_kg = sm100_utils.make_smem_layout_b(mma_mn_mn, tile6, mma_dtype, 1)
-
- sl_readout_k = sm100_utils.make_smem_layout_a(mma_kmn, tile3, mma_dtype, 1)
- sl_nv_b = sm100_utils.make_smem_layout_b(mma_kmn, tile1, mma_dtype, 1)
- sl_nv_a = sm100_utils.make_smem_layout_a(mma_mn_mn, tile6, mma_dtype, 1)
-
- # --- K4 3D GMEM views (V, O, S_fp32 — AB/KS/QS/AQC/KG in Zone B SMEM) ---
- b_view_k4 = cute.make_layout(
- (K_DIM, T_total, num_heads), stride=(1, num_heads * K_DIM, K_DIM)
- )
- mV_k4 = cute.make_tensor(mV_in.iterator, b_view_k4)
-
- # O output: [T_total, V_dim, H] strides (H*V_dim, 1, V_dim) — same as A K_DIM view
- mO_k4 = cute.make_tensor(mO_out.iterator, view_layout_3d)
-
- # State S_fp32: [V, K, B*H] view — mode 0=V(stride 1), mode 1=K(stride V_dim)
- # Matches TMEM state [V,K] convention (M=V, N=K from tiled_mma_mn_mn)
- s_fp32_vk_view = cute.make_layout((K_DIM, K_DIM, BH), stride=(1, K_DIM, K_DIM * K_DIM))
- mS_fp32_vk = cute.make_tensor(mS_fp32.iterator, s_fp32_vk_view)
-
- tma_ld = cpasync.CopyBulkTensorTileG2SOp()
- # Only V loaded via TMA (S now in TMEM, not TMA)
- ta_v, mVt = cute.nvgpu.make_tiled_tma_atom_B(
- tma_ld, mV_k4, cute.select(sl_v, mode=[0, 1, 2]), tile1, mma_kmn
- )
-
- # TMA store for O (3D view)
- sk = sm100_utils.get_smem_layout_atom_ab(OperandMajorMode.K, out_dtype, (M4, N4))
- sl_store = cute.tile_to_shape(
- sm100_utils.make_smem_layout_atom(sk, out_dtype), (M4, N4), order=(0, 1)
- )
- tma_st = cpasync.CopyBulkTensorTileS2GOp()
- ta_o, mOo = cpasync.make_tiled_tma_atom(tma_st, mO_k4, sl_store, (M4, N4))
-
- # --- K123 TMA setup ---
- mQ_view = cute.make_tensor(mQ.iterator, view_layout_3d)
- mK_view = cute.make_tensor(mK.iterator, view_layout_3d)
- mG_view = cute.make_tensor(mG.iterator, view_layout_3d)
-
- smem_atom_qk = tcgen05.make_smem_layout_atom(
- tcgen05.SmemLayoutAtomKind.K_SW128, cutlass.BFloat16
- )
- qk_smem_2d = cute.tile_to_shape(smem_atom_qk, (BT, K_DIM), order=(0, 1))
-
- g_smem_2d = cute.make_layout((BT, K_DIM), stride=(K_DIM, 1))
-
- tma_op_k123 = cpasync.CopyBulkTensorTileG2SOp(cpasync.CtaGroup.ONE)
- ta_q_k123, tt_q = cpasync.make_tiled_tma_atom(
- tma_op_k123, mQ_view, qk_smem_2d, cute.product_each(qk_smem_2d.shape), num_multicast=1
- )
- ta_k_k123, tt_k = cpasync.make_tiled_tma_atom(
- tma_op_k123, mK_view, qk_smem_2d, cute.product_each(qk_smem_2d.shape), num_multicast=1
- )
- ta_g_k123, tt_g = cpasync.make_tiled_tma_atom(
- tma_op_k123, mG_view, g_smem_2d, cute.product_each(g_smem_2d.shape), num_multicast=1
- )
-
- g_cumsum_layout = cute.make_layout((BT, K_STRIDE), stride=(K_STRIDE, 1))
-
- # K123 tiled copies
- copy_atom_qk_k1 = cute.make_copy_atom(
- cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=64
- )
- tiled_copy_qk_k1 = cute.make_tiled_copy_tv(
- copy_atom_qk_k1,
- thr_layout=cute.make_layout((1, 32)),
- val_layout=cute.make_layout((1, 4)),
- )
-
- # K123 output v2 views removed — k_scaled/q_scaled/kg/gk_last now write to SMEM
-
- mma_op_k2 = cute.nvgpu.warp.MmaF16BF16Op(cutlass.BFloat16, cutlass.Float32, (16, 8, 8))
- tiled_mma_k2 = cute.make_tiled_mma(
- mma_op_k2, cute.make_layout((1, 1, 1)), permutation_mnk=(16, 8, 8)
- )
- tiled_copy_mma_A = cute.make_tiled_copy_A(
- cute.make_copy_atom(cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 2), cutlass.BFloat16),
- tiled_mma_k2,
- )
- tiled_copy_mma_B = cute.make_tiled_copy_B(
- cute.make_copy_atom(cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 1), cutlass.BFloat16),
- tiled_mma_k2,
- )
- copy_atom_Gcum = cute.make_copy_atom(
- cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=64
- )
- tiled_copy_Gcum_norm = cute.make_tiled_copy_tv(
- copy_atom_Gcum, thr_layout=cute.make_layout((1, 4)), val_layout=cute.make_layout((1, 2))
- )
- tiled_copy_Gcum_gate = cute.make_tiled_copy_C(copy_atom_Gcum, tiled_mma_k2)
-
- fused_k1234_kernel(
- mma_kmn,
- mma_mn_mn,
- # K4 TMA: V only (S in TMEM, not TMA)
- ta_v,
- mVt,
- sl_v,
- sl_s, # sST SMEM layout (filled from TMEM, not TMA)
- # Zone B swizzle layouts
- sl_ab,
- sl_ks,
- sl_qs,
- sl_aqc,
- sl_kg,
- # K4 TMA store + readout
- ta_o,
- mOo,
- sl_store,
- sl_readout_k,
- sl_nv_b,
- sl_nv_a,
- # K123 TMA
- ta_q_k123,
- tt_q,
- ta_k_k123,
- tt_k,
- ta_g_k123,
- tt_g,
- qk_smem_2d,
- g_smem_2d,
- g_cumsum_layout,
- tiled_copy_qk_k1,
- tiled_mma_k2,
- tiled_copy_mma_A,
- tiled_copy_mma_B,
- tiled_copy_Gcum_norm,
- tiled_copy_Gcum_gate,
- # K123 GMEM
- mA_log,
- mBeta,
- scale_val,
- # K4 GMEM: S_fp32 [V,K,BH] for initial/final state
- mS_fp32_vk,
- mDtBias,
- lower_bound_val,
- _HAS_BIAS,
- _USE_SAFE_GATE,
- mClocks,
- _PROFILE_CLOCKS,
- num_chunks,
- num_heads,
- batch_size,
- ).launch(
- grid=(BH, 1, 1),
- block=(THREADS, 1, 1),
- smem=225 * 1024,
- stream=stream,
- ) # Force high SMEM to prevent >1 block per SM (TMEM conflict)
-
- return host_fn
diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/fused_prefill_conv.py b/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/fused_prefill_conv.py
new file mode 100644
index 000000000000..f5105feb8252
--- /dev/null
+++ b/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/fused_prefill_conv.py
@@ -0,0 +1,1280 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Fused KDA prefill short convolution + SiLU + indexed state update + Q/K/V layout.
+
+One CuTe DSL kernel. The packed QKV source is token-major and the requested
+outputs are token-major [T, H*128] with the same channel order as each source
+section, so no transpose is needed: each thread owns 8 consecutive channels
+(one 16-byte vector) and walks a run of tokens, carrying the 3-token causal
+history in registers.
+
+Alignment: vector source loads require a 16-byte-aligned base and a row
+stride divisible by 8 BF16 elements. Misaligned sources use elementwise
+loads of channels [cc, cc+8), including history and final-state loads.
+Rounding a misaligned pointer down can read uninitialized storage before
+the tensor view or in row padding, even when it stays inside the allocation.
+
+SiLU: the accumulator is h = x/2 (pre-halved weights, exact). The fast form
+silu(x) = h + h*tanh.approx(h) has relative error <= 2^-11 * |tanh(h)| / (1 + tanh(h)),
+which is <= 1.3% for h >= -2 but grows without bound as h -> -inf (cancellation).
+Each lane keeps the minimum h of the current 8-token group; one warp vote per group
+flags groups where some h < -2, and flagged groups are recomputed after the loop with
+x*e/(1+e), e = exp(-|x|) (ex2/rcp, subnormal-preserving), stored over the fast results
+by the same thread in program order.
+
+cp.async ring (aligned sources): per lane, the prologue packs ASYNC_GROUP adjacent
+tokens into each commit group (empty copies are omitted). Each token pack waits until
+at most D/ASYNC_GROUP-1 groups remain, consumes its distinct slots, refills those slots
+for tokens i+D..., then commits the refills as one group. Thus pack p retires group p,
+including every token it consumes. A lane refills a slot only after its own lds_v4 of
+that slot in program order; partial tails still commit one final group.
+
+Launch: one-warp CTAs on grid (channel blocks, runs, sequences). Every CTA
+grid-strides over the runs of its sequence with the run count taken from the
+device-side length (starts), so max_seqlen only sizes grid.y (parallelism), never
+which tokens are processed: an understated or overstated bound still covers every
+token. grid.y is also capped at 65535. Only when B > 65535 does a shape-selected
+specialization grid-stride over sequences as well. A max_seqlen that is a CUDA tensor
+is never read on the host (T is used as the bound).
+
+State pool: run 0 of each sequence (one warp per 256 channels) owns the slot's rows for
+its channels: it reads the old history, then writes the final three raw inputs after
+its own loop (same thread, program order). With a dense [C, 3] slot layout, 16-byte
+aligned base and slot stride % 8 == 0, each lane moves its 8 channels x 3 taps as
+3 x 16-byte loads/stores; otherwise elementwise.
+
+Slots: -1 skips the state pool (raw copy); -pool <= slot < -1 index the pool from the
+end, as in the reference's Python indexing.
+
+Caller contract (stated by the problem: "Slot -1 bypasses ... Other slots are unique and
+valid"; not checked on device): slots other than -1 are unique after wrapping and lie in
+[-pool, pool). Duplicate slots are outside that contract: the reference's sequential
+semantics would need a cross-sequence ordering of state reads and writes that one
+launch does not provide.
+"""
+
+from collections.abc import Callable
+from functools import cache
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+import torch
+from cutlass._mlir.dialects import llvm
+from cutlass.cute.runtime import make_fake_tensor
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+WARPS = 4
+THREADS = WARPS * 32
+VEC = 8 # channels per lane
+WARP_CH = 32 * VEC # channels per warp
+WIDTH = 4
+ASYNC_GROUP = 4 # adjacent token copies per cp.async commit/wait group
+SLOT = 32 * 16 # bytes per ring slot: one 16-byte granule per lane
+# Fast SiLU only while every accumulator h = x/2 of the warp step is >= this.
+FAST_MIN_H = -2.0
+H_INIT = 3.0e38 # running-min start value
+NEG_2LOG2E = -2.8853900817779268 # -2*log2(e)
+
+
+def _asm(res, args, text, cons, side=False, loc=None, ip=None):
+ return llvm.inline_asm(
+ res,
+ args,
+ text,
+ cons,
+ has_side_effects=side,
+ is_align_stack=False,
+ asm_dialect=llvm.AsmDialect.AD_ATT,
+ loc=loc,
+ ip=ip,
+ )
+
+
+@dsl_user_op
+def ld_v4(addr, *, loc=None, ip=None):
+ out = _asm(
+ llvm.StructType.get_literal([T.i32()] * 4),
+ [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)],
+ "ld.global.nc.v4.b32 {$0, $1, $2, $3}, [$4];",
+ "=r,=r,=r,=r,l",
+ loc=loc,
+ ip=ip,
+ )
+ return [cutlass.Int32(llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip)) for i in range(4)]
+
+
+@dsl_user_op
+def ld_v4_first(addr, *, loc=None, ip=None):
+ """ld.global.nc.v4 kept in program order with the cp.async issue (volatile asm), so the
+ halo rows are requested before the ring prologue."""
+ out = _asm(
+ llvm.StructType.get_literal([T.i32()] * 4),
+ [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)],
+ "ld.global.nc.v4.b32 {$0, $1, $2, $3}, [$4];",
+ "=r,=r,=r,=r,l",
+ side=True,
+ loc=loc,
+ ip=ip,
+ )
+ return [cutlass.Int32(llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip)) for i in range(4)]
+
+
+@dsl_user_op
+def st_v4(addr, a, b, c, d, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [
+ cutlass.Int64(addr).ir_value(loc=loc, ip=ip),
+ cutlass.Int32(a).ir_value(loc=loc, ip=ip),
+ cutlass.Int32(b).ir_value(loc=loc, ip=ip),
+ cutlass.Int32(c).ir_value(loc=loc, ip=ip),
+ cutlass.Int32(d).ir_value(loc=loc, ip=ip),
+ ],
+ "st.global.v4.b32 [$1], {$2, $3, $4, $5};\n\tmov.b32 $0, 0;",
+ "=r,l,r,r,r,r",
+ side=True,
+ loc=loc,
+ ip=ip,
+ )
+
+
+@dsl_user_op
+def ld_v4_rw(addr, *, loc=None, ip=None):
+ """Coherent 16-byte load (for the state pool, which this kernel also writes)."""
+ out = _asm(
+ llvm.StructType.get_literal([T.i32()] * 4),
+ [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)],
+ "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];",
+ "=r,=r,=r,=r,l",
+ side=True,
+ loc=loc,
+ ip=ip,
+ )
+ return [cutlass.Int32(llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip)) for i in range(4)]
+
+
+@dsl_user_op
+def prmt(a, b, sel: int, *, loc=None, ip=None):
+ """Byte permute of the pair (a: bytes 0-3, b: bytes 4-7) with an immediate selector."""
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [cutlass.Int32(a).ir_value(loc=loc, ip=ip), cutlass.Int32(b).ir_value(loc=loc, ip=ip)],
+ f"prmt.b32 $0, $1, $2, {sel:#x};",
+ "=r,r,r",
+ loc=loc,
+ ip=ip,
+ )
+ )
+
+
+@dsl_user_op
+def unpack_bf16x2(w, *, loc=None, ip=None):
+ out = _asm(
+ llvm.StructType.get_literal([T.f32(), T.f32()]),
+ [cutlass.Int32(w).ir_value(loc=loc, ip=ip)],
+ "{\n\t.reg .b32 t;\n\tshl.b32 t, $2, 16;\n\tmov.b32 $0, t;\n\t"
+ "and.b32 t, $2, 0xFFFF0000;\n\tmov.b32 $1, t;\n\t}",
+ "=f,=f,r",
+ loc=loc,
+ ip=ip,
+ )
+ return (
+ cutlass.Float32(llvm.extractvalue(T.f32(), out, [0], loc=loc, ip=ip)),
+ cutlass.Float32(llvm.extractvalue(T.f32(), out, [1], loc=loc, ip=ip)),
+ )
+
+
+@dsl_user_op
+def pack_bf16x2(lo, hi, *, loc=None, ip=None):
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [
+ cutlass.Float32(lo).ir_value(loc=loc, ip=ip),
+ cutlass.Float32(hi).ir_value(loc=loc, ip=ip),
+ ],
+ "cvt.rn.bf16x2.f32 $0, $2, $1;",
+ "=r,f,f",
+ loc=loc,
+ ip=ip,
+ )
+ )
+
+
+@dsl_user_op
+def tanh_approx(x, *, loc=None, ip=None):
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [cutlass.Float32(x).ir_value(loc=loc, ip=ip)],
+ "tanh.approx.f32 $0, $1;",
+ "=f,f",
+ loc=loc,
+ ip=ip,
+ )
+ )
+
+
+@dsl_user_op
+def fma_rn(a, b, c, *, loc=None, ip=None):
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [
+ cutlass.Float32(a).ir_value(loc=loc, ip=ip),
+ cutlass.Float32(b).ir_value(loc=loc, ip=ip),
+ cutlass.Float32(c).ir_value(loc=loc, ip=ip),
+ ],
+ "fma.rn.f32 $0, $1, $2, $3;",
+ "=f,f,f,f",
+ loc=loc,
+ ip=ip,
+ )
+ )
+
+
+@dsl_user_op
+def cp_async16(saddr, gaddr, *, loc=None, ip=None):
+ _asm(
+ T.i32(),
+ [
+ cutlass.Int32(saddr).ir_value(loc=loc, ip=ip),
+ cutlass.Int64(gaddr).ir_value(loc=loc, ip=ip),
+ ],
+ "cp.async.cg.shared.global [$1], [$2], 16;\n\tmov.b32 $0, 0;",
+ "=r,r,l",
+ side=True,
+ loc=loc,
+ ip=ip,
+ )
+
+
+@dsl_user_op
+def lds_v4(saddr, *, loc=None, ip=None):
+ out = _asm(
+ llvm.StructType.get_literal([T.i32()] * 4),
+ [cutlass.Int32(saddr).ir_value(loc=loc, ip=ip)],
+ "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];",
+ "=r,=r,=r,=r,r",
+ side=True,
+ loc=loc,
+ ip=ip,
+ )
+ return [cutlass.Int32(llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip)) for i in range(4)]
+
+
+@dsl_user_op
+def smem_u32(ptr, *, loc=None, ip=None):
+ return cutlass.Int32(
+ _asm(
+ T.i32(),
+ [cutlass.Int64(ptr).ir_value(loc=loc, ip=ip)],
+ "cvta.to.shared.u32 $0, $1;",
+ "=r,l",
+ loc=loc,
+ ip=ip,
+ )
+ )
+
+
+@dsl_user_op
+def bf_lo(w, *, loc=None, ip=None):
+ """FP32 value of the low BF16 half of a packed word."""
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [cutlass.Int32(w).ir_value(loc=loc, ip=ip)],
+ "shl.b32 $0, $1, 16;",
+ "=f,r",
+ loc=loc,
+ ip=ip,
+ )
+ )
+
+
+@dsl_user_op
+def bf_hi(w, *, loc=None, ip=None):
+ """FP32 value of the high BF16 half of a packed word."""
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [cutlass.Int32(w).ir_value(loc=loc, ip=ip)],
+ "and.b32 $0, $1, 0xFFFF0000;",
+ "=f,r",
+ loc=loc,
+ ip=ip,
+ )
+ )
+
+
+def _pair_op(text, n_in, loc, ip, args):
+ regs = ", ".join(f"p{i}" for i in range(n_in))
+ body = "{\n\t.reg .b64 " + regs + ", d;\n\t"
+ for i in range(n_in):
+ body += f"mov.b64 p{i}, {{${2 + 2 * i}, ${3 + 2 * i}}};\n\t"
+ body += text.format(*(f"p{i}" for i in range(n_in))) + "\n\tmov.b64 {$0, $1}, d;\n\t}"
+ out = _asm(
+ llvm.StructType.get_literal([T.f32(), T.f32()]),
+ [cutlass.Float32(a).ir_value(loc=loc, ip=ip) for a in args],
+ body,
+ "=f,=f," + ",".join(["f"] * (2 * n_in)),
+ loc=loc,
+ ip=ip,
+ )
+ return (
+ cutlass.Float32(llvm.extractvalue(T.f32(), out, [0], loc=loc, ip=ip)),
+ cutlass.Float32(llvm.extractvalue(T.f32(), out, [1], loc=loc, ip=ip)),
+ )
+
+
+@dsl_user_op
+def mul2(a0, a1, b0, b1, *, loc=None, ip=None):
+ return _pair_op("mul.rn.f32x2 d, {0}, {1};", 2, loc, ip, [a0, a1, b0, b1])
+
+
+@dsl_user_op
+def fma2(a0, a1, b0, b1, c0, c1, *, loc=None, ip=None):
+ return _pair_op("fma.rn.f32x2 d, {0}, {1}, {2};", 3, loc, ip, [a0, a1, b0, b1, c0, c1])
+
+
+@dsl_user_op
+def fmin3(a, b, c, *, loc=None, ip=None):
+ return cutlass.Float32(
+ _asm(
+ T.f32(),
+ [cutlass.Float32(x).ir_value(loc=loc, ip=ip) for x in (a, b, c)],
+ "min.f32 $0, $1, $2, $3;",
+ "=f,f,f,f",
+ loc=loc,
+ ip=ip,
+ )
+ )
+
+
+@dsl_user_op
+def silu2_exact(h0, h1, *, loc=None, ip=None):
+ """For h = x/2, silu(x) on a pair: e = exp(-|x|) in (0, 1], so 1 + e in (1, 2].
+
+ x >= 0: x / (1 + e); x < 0: (x * e) / (1 + e). The reciprocal never sees a
+ large or subnormal argument and e keeps subnormals (no FTZ on ex2), so deep
+ negative outputs keep their relative accuracy.
+ """
+ out = _asm(
+ llvm.StructType.get_literal([T.f32(), T.f32()]),
+ [cutlass.Float32(a).ir_value(loc=loc, ip=ip) for a in (h0, h1, NEG_2LOG2E)],
+ "{\n\t.reg .f32 a<2>, e<2>, r<2>, x<2>, n<2>;\n\t.reg .pred p<2>;\n\t"
+ "abs.f32 a0, $2;\n\tabs.f32 a1, $3;\n\t"
+ "mul.rn.f32 a0, a0, $4;\n\tmul.rn.f32 a1, a1, $4;\n\t"
+ "ex2.approx.f32 e0, a0;\n\tex2.approx.f32 e1, a1;\n\t"
+ "add.rn.f32 r0, e0, 0f3F800000;\n\tadd.rn.f32 r1, e1, 0f3F800000;\n\t"
+ "rcp.approx.ftz.f32 r0, r0;\n\trcp.approx.ftz.f32 r1, r1;\n\t"
+ "add.rn.f32 x0, $2, $2;\n\tadd.rn.f32 x1, $3, $3;\n\t"
+ "mul.rn.f32 n0, x0, e0;\n\tmul.rn.f32 n1, x1, e1;\n\t"
+ "setp.lt.f32 p0, $2, 0f00000000;\n\tsetp.lt.f32 p1, $3, 0f00000000;\n\t"
+ "selp.f32 n0, n0, x0, p0;\n\tselp.f32 n1, n1, x1, p1;\n\t"
+ "mul.rn.f32 $0, n0, r0;\n\tmul.rn.f32 $1, n1, r1;\n\t}",
+ "=f,=f,f,f,f",
+ loc=loc,
+ ip=ip,
+ )
+ return (
+ cutlass.Float32(llvm.extractvalue(T.f32(), out, [0], loc=loc, ip=ip)),
+ cutlass.Float32(llvm.extractvalue(T.f32(), out, [1], loc=loc, ip=ip)),
+ )
+
+
+class KdaConv:
+ def __init__(
+ self,
+ heads,
+ tokens_per_warp,
+ unroll,
+ src_vec,
+ w_vec,
+ out_vec,
+ row_stride,
+ state_stride,
+ seq_loop,
+ depth=8,
+ warps=4,
+ st_vec=False,
+ ):
+ self.W = warps
+ self.MB = FORCE_MB
+ self.heads = heads
+ self.P = heads * 128
+ self.C = 3 * heads * 128
+ self.L = tokens_per_warp
+ self.U = unroll
+ self.src_vec = src_vec
+ self.w_vec = w_vec
+ self.out_vec = out_vec
+ self.SR = row_stride
+ self.SS = state_stride
+ self.seq_loop = seq_loop
+ self.st_vec = st_vec
+ self.D = depth
+ self.U = 8
+ self.RUN = warps * tokens_per_warp # tokens per CTA (power of two)
+ self.RUN_LOG2 = self.RUN.bit_length() - 1
+ assert self.RUN == 1 << self.RUN_LOG2
+ assert self.D % ASYNC_GROUP == 0 and self.U % ASYNC_GROUP == 0
+
+ @cute.jit
+ def loadw(self, src, src_base, row, cc, xw: cute.Tensor, off: cutlass.Constexpr):
+ """Load source channels [cc, cc+8) of `row` as 4 packed BF16x2 words into xw[off:off+4]."""
+ if cutlass.const_expr(self.src_vec):
+ w = ld_v4(src_base + (cutlass.Int64(row) * self.SR + cc) * 2)
+ for j in cutlass.range_constexpr(4):
+ xw[off + j] = w[j]
+ else:
+ for j in cutlass.range_constexpr(4):
+ xw[off + j] = pack_bf16x2(
+ cutlass.Float32(src[cutlass.Int64(row), cc + 2 * j]),
+ cutlass.Float32(src[cutlass.Int64(row), cc + 2 * j + 1]),
+ )
+
+ @cute.jit
+ def load8(self, src, src_base, row, cc, xs: cute.Tensor, off: cutlass.Constexpr):
+ """Load source channels [cc, cc+8) of `row` as FP32 into xs[off:off+8]."""
+ xw = cute.make_rmem_tensor(cute.make_layout(4), cutlass.Int32)
+ self.loadw(src, src_base, row, cc, xw, 0)
+ for j in cutlass.range_constexpr(4):
+ xlo, xhi = unpack_bf16x2(xw[j])
+ xs[off + 2 * j] = xlo
+ xs[off + 2 * j + 1] = xhi
+
+ @cute.jit
+ def issue(self, src_base, row, cc, sl):
+ """cp.async the 16-byte granule holding channels [cc, cc+8) of `row` into sl."""
+ cp_async16(sl, src_base + (cutlass.Int64(row) * self.SR + cc) * 2)
+
+ @cute.jit
+ def storew(self, q, k, v, sec, row, inner, yw: cute.Tensor):
+ """Elementwise store of 8 channels for non-contiguous or misaligned q/k/v."""
+ h = inner // 128
+ d = inner - h * 128
+ row64 = cutlass.Int64(row)
+ for j in cutlass.range_constexpr(4):
+ olo, ohi = unpack_bf16x2(yw[j])
+ for e in cutlass.range_constexpr(2):
+ val = cutlass.BFloat16(olo)
+ if cutlass.const_expr(e == 1):
+ val = cutlass.BFloat16(ohi)
+ if sec == 0:
+ q[0, row64, h, d + 2 * j + e] = val
+ elif sec == 1:
+ k[0, row64, h, d + 2 * j + e] = val
+ else:
+ v[0, row64, h, d + 2 * j + e] = val
+
+ @cute.kernel
+ def kernel(self, src, weight, state, starts, slots, has_init, q, k, v):
+ tid, _, _ = cute.arch.thread_idx()
+ b0, b1, b2 = cute.arch.block_idx()
+ _, gdy, gdz = cute.arch.grid_dim()
+ smem = cutlass.utils.SmemAllocator()
+ ring = smem.allocate_tensor(cutlass.Int8, cute.make_layout(self.W * self.D * SLOT), 16)
+ ring_warp = cutlass.Int32(ring.iterator.toint()) + (tid // 32) * (self.D * SLOT)
+ if cutlass.const_expr(self.seq_loop):
+ # B > 65535 (grid.z limit): grid-stride over sequences too.
+ for seq in cutlass.range(b2, slots.shape[0], gdz, unroll=1):
+ self.runs(
+ seq,
+ b0,
+ b1,
+ gdy,
+ tid,
+ ring_warp,
+ src,
+ weight,
+ state,
+ starts,
+ slots,
+ has_init,
+ q,
+ k,
+ v,
+ )
+ else:
+ self.runs(
+ b2,
+ b0,
+ b1,
+ gdy,
+ tid,
+ ring_warp,
+ src,
+ weight,
+ state,
+ starts,
+ slots,
+ has_init,
+ q,
+ k,
+ v,
+ )
+
+ @cute.jit
+ def runs(
+ self, seq, b0, b1, gdy, tid, ring_warp, src, weight, state, starts, slots, has_init, q, k, v
+ ):
+ """All runs b1, b1 + gdy, ... of sequence `seq` for channel block b0.
+
+ Run b1 is straight-line (tile() skips it when it starts past the sequence end,
+ except run 0, which owns the state update, also for empty sequences). The cold
+ loop covers runs past grid.y, using the run count from the device-side length:
+ it is empty whenever max_seqlen bounds every length, and otherwise (understated
+ bound, or the 65535 cap) it still processes every token. Each tile drains its
+ ring before returning.
+ """
+ begin, length, slot, init = self.meta(seq, state, starts, slots, has_init)
+ # Only this bit stays live across the hot tile; the cold path reloads the metadata.
+ more = b1 + gdy < ((length + (self.RUN - 1)) >> self.RUN_LOG2)
+ self.tile(b1, b0, tid, ring_warp, begin, length, slot, init, src, weight, state, q, k, v)
+ if more:
+ begin, length, slot, init = self.meta(seq, state, starts, slots, has_init)
+ nrun = (length + (self.RUN - 1)) >> self.RUN_LOG2
+ for bx in cutlass.range(b1 + gdy, nrun, gdy, unroll=1):
+ self.tile(
+ bx, b0, tid, ring_warp, begin, length, slot, init, src, weight, state, q, k, v
+ )
+
+ @cute.jit
+ def meta(self, seq, state, starts, slots, has_init):
+ begin = starts[seq]
+ length = starts[seq + 1] - begin
+ slot = slots[seq]
+ # Issued with starts/slots so the old-state load is one round trip from launch.
+ init = has_init[seq]
+ # Python indexing, as in the reference: -pool <= slot < -1 counts from the end.
+ if slot < -1:
+ slot = slot + state.shape[0]
+ return begin, length, slot, init
+
+ @cute.jit
+ def tile(self, bx, by, tid, ring_warp, begin, length, slot, init, src, weight, state, q, k, v):
+ """One warp: run t0 = (bx * W + warp) * L of one sequence, channel block `by`."""
+ C, P, L, U = self.C, self.P, self.L, self.U
+ lane = tid % 32
+ warp = tid // 32
+ state_slot = cute.make_tensor(
+ state.iterator + cutlass.Int64(slot) * self.SS[0],
+ cute.make_layout((C, WIDTH - 1), stride=(self.SS[1], self.SS[2])),
+ )
+ t0 = (bx * self.W + warp) * L
+ c = by * WARP_CH + lane * VEC
+ if cutlass.const_expr(C % WARP_CH == 0):
+ cvalid = True
+ cc = c
+ else:
+ cvalid = c < C
+ cc = cutlass.min(c, C - VEC)
+ sec = cc // P
+ inner = cc - sec * P
+ src_base = src.iterator.toint()
+ ob = q.iterator.toint()
+ if sec == 1:
+ ob = k.iterator.toint()
+ if sec == 2:
+ ob = v.iterator.toint()
+
+ D = self.D
+
+ if t0 < length or t0 == 0:
+ count = cutlass.max(cutlass.min(length - t0, L), 0)
+ row0 = begin + t0
+ ring_lane = ring_warp + lane * 16
+ # Halo rows t0-3..t0-1 (needed by token 0) are requested before the ring prologue.
+ hw = cute.make_rmem_tensor(cute.make_layout(4 * (WIDTH - 1)), cutlass.Int32)
+ if cutlass.const_expr(self.src_vec):
+ if t0 != 0:
+ for j in cutlass.range_constexpr(WIDTH - 1):
+ w = ld_v4_first(
+ src_base + (cutlass.Int64(row0 - (WIDTH - 1) + j) * self.SR + cc) * 2
+ )
+ for e in cutlass.range_constexpr(4):
+ hw[4 * j + e] = w[e]
+ if cutlass.const_expr(self.src_vec):
+ for pg in cutlass.range_constexpr(D // ASYNC_GROUP):
+ for pe in cutlass.range_constexpr(ASYNC_GROUP):
+ d = ASYNC_GROUP * pg + pe
+ if d < count:
+ self.issue(src_base, row0 + d, cc, ring_lane + d * SLOT)
+ cute.arch.cp_async_commit_group()
+ wt = cute.make_rmem_tensor(cute.make_layout(WIDTH * VEC), cutlass.Float32)
+ hist = cute.make_rmem_tensor(cute.make_layout((WIDTH - 1) * VEC), cutlass.Float32)
+ xs = cute.make_rmem_tensor(cute.make_layout(VEC), cutlass.Float32)
+ fin = cute.make_rmem_tensor(cute.make_layout(4 * (WIDTH - 1)), cutlass.Int32)
+ # Weights: weight[cc + i, w] -> wt[w * 8 + i]
+ if cutlass.const_expr(self.w_vec):
+ wbase = weight.iterator.toint() + cutlass.Int64(cc) * (WIDTH * 2)
+ for j in cutlass.range_constexpr(4):
+ ww = ld_v4(wbase + 16 * j)
+ for jj in cutlass.range_constexpr(4):
+ ch = 2 * j + jj // 2
+ tap = 2 * (jj % 2)
+ wlo, whi = unpack_bf16x2(ww[jj])
+ wt[tap * VEC + ch] = wlo * 0.5
+ wt[(tap + 1) * VEC + ch] = whi * 0.5
+ else:
+ for i in cutlass.range_constexpr(VEC):
+ for w in cutlass.range_constexpr(WIDTH):
+ wt[w * VEC + i] = cutlass.Float32(weight[cc + i, w]) * 0.5
+
+ # Causal history for token t0: tokens t0-3 .. t0-1.
+ if t0 == 0:
+ for j in cutlass.range_constexpr((WIDTH - 1) * VEC):
+ hist[j] = cutlass.Float32(0)
+ if slot >= 0:
+ if init:
+ if cutlass.const_expr(self.st_vec):
+ # The lane's 8 channels x 3 taps are 48 contiguous, 16-byte-aligned
+ # bytes: element e = 3 * i + j is half e % 2 of word e // 2.
+ sw = []
+ for m in cutlass.range_constexpr(3):
+ sw.extend(ld_v4_rw(self.state_addr(state, slot, cc) + 16 * m))
+ for i in cutlass.range_constexpr(VEC):
+ for j in cutlass.range_constexpr(WIDTH - 1):
+ e = 3 * i + j
+ if cutlass.const_expr(e % 2 == 0):
+ hist[j * VEC + i] = bf_lo(sw[e // 2])
+ else:
+ hist[j * VEC + i] = bf_hi(sw[e // 2])
+ else:
+ for i in cutlass.range_constexpr(VEC):
+ for j in cutlass.range_constexpr(WIDTH - 1):
+ hist[j * VEC + i] = cutlass.Float32(state_slot[cc + i, j])
+ # New state = last three raw inputs, older ones from the
+ # initial history. Gathered (packed BF16x2) now so the loads
+ # overlap the main loop; written after the loop.
+ for j in cutlass.range_constexpr(WIDTH - 1):
+ rel = length - (WIDTH - 1) + j
+ if rel >= 0:
+ self.loadw(src, src_base, begin + rel, cc, fin, 4 * j)
+ else:
+ hidx = rel + (WIDTH - 1)
+ for i in cutlass.range_constexpr(VEC):
+ val = hist[i]
+ if hidx == 1:
+ val = hist[VEC + i]
+ if hidx == 2:
+ val = hist[2 * VEC + i]
+ xs[i] = val
+ for w in cutlass.range_constexpr(4):
+ fin[4 * j + w] = pack_bf16x2(xs[2 * w], xs[2 * w + 1])
+ else:
+ if cutlass.const_expr(self.src_vec):
+ for j in cutlass.range_constexpr(WIDTH - 1):
+ for e in cutlass.range_constexpr(4):
+ xlo, xhi = unpack_bf16x2(hw[4 * j + e])
+ hist[j * VEC + 2 * e] = xlo
+ hist[j * VEC + 2 * e + 1] = xhi
+ else:
+ for j in cutlass.range_constexpr(WIDTH - 1):
+ self.load8(src, src_base, begin + t0 - (WIDTH - 1) + j, cc, hist, j * VEC)
+
+ yw = cute.make_rmem_tensor(cute.make_layout(4), cutlass.Int32)
+ gbase = src_base + (cutlass.Int64(row0) * self.SR + cc) * 2
+ sbase = ob + (cutlass.Int64(row0) * P + inner) * 2
+ # hm: running min of the SiLU arguments h of the current token group.
+ # flag bit g: some lane of the warp had h < FAST_MIN_H in group g
+ # (bit ngroup: the tail). Flagged groups are recomputed after the loop.
+ hm = cute.make_rmem_tensor(cute.make_layout(1), cutlass.Float32)
+ flag = cute.make_rmem_tensor(cute.make_layout(1), cutlass.Int32)
+ flag[0] = cutlass.Int32(0)
+ if slot >= 0:
+ ngroup = count // U
+ for g in cutlass.range(ngroup, unroll=1):
+ hm[0] = cutlass.Float32(H_INIT)
+ # Group base addresses: per-token offsets below are immediates.
+ gg = gbase + cutlass.Int64(g * U) * (self.SR * 2)
+ sg = sbase + cutlass.Int64(g * U) * (P * 2)
+ for p in cutlass.range_constexpr(U // ASYNC_GROUP):
+ if cutlass.const_expr(self.src_vec):
+ cute.arch.cp_async_wait_group(D // ASYNC_GROUP - 1)
+ for pe in cutlass.range_constexpr(ASYNC_GROUP):
+ u = ASYNC_GROUP * p + pe
+ self.step(
+ g * U + u,
+ u,
+ True,
+ src,
+ src_base,
+ row0,
+ cc,
+ ring_lane,
+ count,
+ wt,
+ hist,
+ yw,
+ hm,
+ cvalid,
+ q,
+ k,
+ v,
+ sec,
+ inner,
+ gg,
+ sg,
+ )
+ if cutlass.const_expr(self.src_vec):
+ cute.arch.cp_async_commit_group()
+ deep = cute.arch.vote_any_sync(hm[0] < FAST_MIN_H)
+ flag[0] = flag[0] | (cutlass.Int32(deep) << g)
+ hm[0] = cutlass.Float32(H_INIT)
+ first = ngroup * U
+ npack = (count - first + (ASYNC_GROUP - 1)) // ASYNC_GROUP
+ for p in cutlass.range(npack, unroll=1):
+ i = first + ASYNC_GROUP * p
+ if cutlass.const_expr(self.src_vec):
+ cute.arch.cp_async_wait_group(D // ASYNC_GROUP - 1)
+ gg = gbase + cutlass.Int64(i) * (self.SR * 2)
+ sg = sbase + cutlass.Int64(i) * (P * 2)
+ self.step(
+ i,
+ 0,
+ True,
+ src,
+ src_base,
+ row0,
+ cc,
+ ring_lane,
+ count,
+ wt,
+ hist,
+ yw,
+ hm,
+ cvalid,
+ q,
+ k,
+ v,
+ sec,
+ inner,
+ gg,
+ sg,
+ )
+ for pe in cutlass.range_constexpr(1, ASYNC_GROUP):
+ if i + pe < count:
+ self.step(
+ i + pe,
+ pe,
+ True,
+ src,
+ src_base,
+ row0,
+ cc,
+ ring_lane,
+ count,
+ wt,
+ hist,
+ yw,
+ hm,
+ cvalid,
+ q,
+ k,
+ v,
+ sec,
+ inner,
+ gg,
+ sg,
+ )
+ if cutlass.const_expr(self.src_vec):
+ cute.arch.cp_async_commit_group()
+ deep = cute.arch.vote_any_sync(hm[0] < FAST_MIN_H)
+ flag[0] = flag[0] | (cutlass.Int32(deep) << ngroup)
+ # Rare, warp-uniform: exact SiLU for flagged groups. hist is dead here.
+ if flag[0] != 0:
+ for g in cutlass.range(ngroup + 1, unroll=1):
+ if (flag[0] >> g) & 1:
+ first = g * U
+ self.redo(
+ first,
+ cutlass.min(count - first, U),
+ src,
+ src_base,
+ begin,
+ t0,
+ init,
+ state_slot,
+ row0,
+ cc,
+ wt,
+ hist,
+ yw,
+ cvalid,
+ q,
+ k,
+ v,
+ sec,
+ inner,
+ sbase,
+ )
+ else:
+ npack = (count + (ASYNC_GROUP - 1)) // ASYNC_GROUP
+ for p in cutlass.range(npack, unroll=1):
+ i = ASYNC_GROUP * p
+ if cutlass.const_expr(self.src_vec):
+ cute.arch.cp_async_wait_group(D // ASYNC_GROUP - 1)
+ gg = gbase + cutlass.Int64(i) * (self.SR * 2)
+ sg = sbase + cutlass.Int64(i) * (P * 2)
+ self.step(
+ i,
+ 0,
+ False,
+ src,
+ src_base,
+ row0,
+ cc,
+ ring_lane,
+ count,
+ wt,
+ hist,
+ yw,
+ hm,
+ cvalid,
+ q,
+ k,
+ v,
+ sec,
+ inner,
+ gg,
+ sg,
+ )
+ for pe in cutlass.range_constexpr(1, ASYNC_GROUP):
+ if i + pe < count:
+ self.step(
+ i + pe,
+ pe,
+ False,
+ src,
+ src_base,
+ row0,
+ cc,
+ ring_lane,
+ count,
+ wt,
+ hist,
+ yw,
+ hm,
+ cvalid,
+ q,
+ k,
+ v,
+ sec,
+ inner,
+ gg,
+ sg,
+ )
+ if cutlass.const_expr(self.src_vec):
+ cute.arch.cp_async_commit_group()
+ if cutlass.const_expr(self.src_vec):
+ cute.arch.cp_async_wait_group(0)
+ if t0 == 0:
+ if slot >= 0:
+ if cvalid:
+ if cutlass.const_expr(self.st_vec):
+ # Element e = 3 * i + j (channel i, tap j) is half i % 2 of
+ # fin[4 * j + i // 2]; pack element pairs, then 3 x 16-byte stores.
+ ow = []
+ for mw in cutlass.range_constexpr(12):
+ srcw = []
+ sel = 0
+ for hh in cutlass.range_constexpr(2):
+ e = 2 * mw + hh
+ i, j = e // 3, e % 3
+ srcw.append(fin[4 * j + i // 2])
+ b = 2 * (i % 2) + 4 * hh
+ sel |= (b | ((b + 1) << 4)) << (8 * hh)
+ ow.append(prmt(srcw[0], srcw[1], sel))
+ sa = self.state_addr(state, slot, cc)
+ for m in cutlass.range_constexpr(3):
+ st_v4(
+ sa + 16 * m,
+ ow[4 * m],
+ ow[4 * m + 1],
+ ow[4 * m + 2],
+ ow[4 * m + 3],
+ )
+ else:
+ for j in cutlass.range_constexpr(WIDTH - 1):
+ for w in cutlass.range_constexpr(4):
+ slo, shi = unpack_bf16x2(fin[4 * j + w])
+ state_slot[cc + 2 * w, j] = cutlass.BFloat16(slo)
+ state_slot[cc + 2 * w + 1, j] = cutlass.BFloat16(shi)
+
+ @cute.jit
+ def state_addr(self, state, slot, cc):
+ """Byte address of state[slot, cc, 0] (vector path: 16-byte aligned by the host check)."""
+ return state.iterator.toint() + (cutlass.Int64(slot) * self.SS[0] + cc * self.SS[1]) * 2
+
+ @cute.jit
+ def step(
+ self,
+ i,
+ u: cutlass.Constexpr,
+ conv: cutlass.Constexpr,
+ src,
+ src_base,
+ row0,
+ cc,
+ ring_lane,
+ count,
+ wt: cute.Tensor,
+ hist: cute.Tensor,
+ yw: cute.Tensor,
+ hm: cute.Tensor,
+ cvalid,
+ q,
+ k,
+ v,
+ sec,
+ inner,
+ gb,
+ sb,
+ ):
+ """Consume token i of this warp's run: ring slot -> conv+SiLU (or raw copy) -> store.
+
+ gb/sb are the source/output addresses of token i - u (u is a compile-time offset)."""
+ D = self.D
+ sl = ring_lane + (i % D) * SLOT
+ xt = cute.make_rmem_tensor(cute.make_layout(4), cutlass.Int32)
+ if cutlass.const_expr(self.src_vec):
+ w4 = lds_v4(sl)
+ for j in cutlass.range_constexpr(4):
+ xt[j] = w4[j]
+ else:
+ self.loadw(src, src_base, row0 + i, cc, xt, 0)
+ if cutlass.const_expr(conv):
+ xs = []
+ for e2 in cutlass.range_constexpr(4):
+ xs.append(bf_lo(xt[e2]))
+ xs.append(bf_hi(xt[e2]))
+ hs = []
+ for j in cutlass.range_constexpr(4):
+ a, b = 2 * j, 2 * j + 1
+ acc0, acc1 = mul2(wt[a], wt[b], hist[a], hist[b])
+ acc0, acc1 = fma2(
+ wt[VEC + a], wt[VEC + b], hist[VEC + a], hist[VEC + b], acc0, acc1
+ )
+ acc0, acc1 = fma2(
+ wt[2 * VEC + a],
+ wt[2 * VEC + b],
+ hist[2 * VEC + a],
+ hist[2 * VEC + b],
+ acc0,
+ acc1,
+ )
+ acc0, acc1 = fma2(wt[3 * VEC + a], wt[3 * VEC + b], xs[a], xs[b], acc0, acc1)
+ hs.append(acc0)
+ hs.append(acc1)
+ for e3 in cutlass.range_constexpr(2):
+ ch = a + e3
+ hist[ch] = hist[VEC + ch]
+ hist[VEC + ch] = hist[2 * VEC + ch]
+ hist[2 * VEC + ch] = xs[ch]
+ # acc = conv/2 (weights pre-halved): silu(2h) = h + h * tanh(h)
+ for j in cutlass.range_constexpr(4):
+ h0, h1 = hs[2 * j], hs[2 * j + 1]
+ y0, y1 = fma2(h0, h1, tanh_approx(h0), tanh_approx(h1), h0, h1)
+ yw[j] = pack_bf16x2(y0, y1)
+ # That form cancels for strongly negative h: track the group minimum
+ # (the caller recomputes the group exactly if it is below the bound).
+ hm[0] = fmin3(
+ fmin3(hs[0], hs[1], hs[2]), fmin3(hs[3], hs[4], hs[5]), fmin3(hs[6], hs[7], hm[0])
+ )
+ else:
+ for j in cutlass.range_constexpr(4):
+ yw[j] = xt[j]
+ # Refill this slot in program order after its same-thread lds_v4. The
+ # caller commits one adjacent-token pack of refills as one cp.async group.
+ if cutlass.const_expr(self.src_vec):
+ if i + D < count:
+ cp_async16(sl, gb + (u + D) * (self.SR * 2))
+ if cvalid:
+ if cutlass.const_expr(self.out_vec):
+ st_v4(sb + u * (self.P * 2), yw[0], yw[1], yw[2], yw[3])
+ else:
+ self.storew(q, k, v, sec, row0 + i, inner, yw)
+
+ @cute.jit
+ def redo(
+ self,
+ first,
+ n,
+ src,
+ src_base,
+ begin,
+ t0,
+ init,
+ state_slot,
+ row0,
+ cc,
+ wt: cute.Tensor,
+ hx: cute.Tensor,
+ yr: cute.Tensor,
+ cvalid,
+ q,
+ k,
+ v,
+ sec,
+ inner,
+ sbase,
+ ):
+ """Rare path: recompute tokens [first, first+n) of this run with exact SiLU and store
+ over the fast results (same thread, same addresses, program order).
+
+ Inputs are re-read from global memory (just read, so L2-resident); the history
+ before `first` comes from src or, at the sequence start, from the initial state,
+ which this thread only overwrites after this. hx is scratch (the caller's dead
+ rolling-history registers).
+ """
+ for j in cutlass.range_constexpr(WIDTH - 1):
+ r = t0 + first - (WIDTH - 1) + j # sequence-relative token index
+ if r >= 0:
+ self.load8(src, src_base, begin + r, cc, hx, j * VEC)
+ else:
+ for i in cutlass.range_constexpr(VEC):
+ hx[j * VEC + i] = cutlass.Float32(0)
+ if init:
+ for i in cutlass.range_constexpr(VEC):
+ hx[j * VEC + i] = cutlass.Float32(state_slot[cc + i, r + (WIDTH - 1)])
+ xr = cute.make_rmem_tensor(cute.make_layout(VEC), cutlass.Float32)
+ for u in cutlass.range(n, unroll=1):
+ self.load8(src, src_base, row0 + first + u, cc, xr, 0)
+ for j in cutlass.range_constexpr(4):
+ a, b = 2 * j, 2 * j + 1
+ acc0, acc1 = mul2(wt[a], wt[b], hx[a], hx[b])
+ acc0, acc1 = fma2(wt[VEC + a], wt[VEC + b], hx[VEC + a], hx[VEC + b], acc0, acc1)
+ acc0, acc1 = fma2(
+ wt[2 * VEC + a], wt[2 * VEC + b], hx[2 * VEC + a], hx[2 * VEC + b], acc0, acc1
+ )
+ acc0, acc1 = fma2(wt[3 * VEC + a], wt[3 * VEC + b], xr[a], xr[b], acc0, acc1)
+ y0, y1 = silu2_exact(acc0, acc1)
+ yr[j] = pack_bf16x2(y0, y1)
+ for e3 in cutlass.range_constexpr(2):
+ ch = a + e3
+ hx[ch] = hx[VEC + ch]
+ hx[VEC + ch] = hx[2 * VEC + ch]
+ hx[2 * VEC + ch] = xr[ch]
+ self.store(first + u, cvalid, q, k, v, sec, row0, inner, sbase, yr)
+
+ @cute.jit
+ def store(self, i, cvalid, q, k, v, sec, row0, inner, sbase, yw: cute.Tensor):
+ if cvalid:
+ if cutlass.const_expr(self.out_vec):
+ st_v4(sbase + cutlass.Int64(i) * (self.P * 2), yw[0], yw[1], yw[2], yw[3])
+ else:
+ self.storew(q, k, v, sec, row0 + i, inner, yw)
+
+ @cute.jit
+ def __call__(
+ self,
+ src,
+ weight,
+ state,
+ starts,
+ slots,
+ has_init,
+ q,
+ k,
+ v,
+ grid_y: cutlass.Int32,
+ stream: cuda.CUstream,
+ ):
+ channel_blocks = (self.C + WARP_CH - 1) // WARP_CH
+ grid = (channel_blocks, grid_y, min(slots.shape[0], MAX_GRID_YZ))
+ self.kernel(src, weight, state, starts, slots, has_init, q, k, v).launch(
+ grid=grid, block=(self.W * 32, 1, 1), min_blocks_per_mp=self.MB, stream=stream
+ )
+
+
+MAX_GRID_YZ = 65535 # grid.y/z launch limit; the kernel grid-strides past it
+FORCE_L = 32
+FORCE_D = 8
+FORCE_W = 1
+FORCE_MB = 16
+
+
+def _pick_tokens(heads, max_seqlen, batch):
+ # Enough warps to fill the machine; longer runs amortise the 3-token halo.
+ lanes_c = (3 * heads * 128 + WARP_CH - 1) // WARP_CH
+ for L in (32, 16):
+ warps = lanes_c * batch * ((max_seqlen + L - 1) // L)
+ if warps >= 160 * 48:
+ return L
+ return 8
+
+
+def _align(t):
+ return t.data_ptr() % 16
+
+
+def _layout_key(t, *, dynamic_dim=None, dynamic_stride=None):
+ """Keep only layout properties that specialize the convolution."""
+ shape = tuple(None if i == dynamic_dim else n for i, n in enumerate(t.shape))
+ stride = tuple(None if i == dynamic_stride else s for i, s in enumerate(t.stride()))
+ return (shape, stride, t.dtype)
+
+
+@cache
+def compile_prefill_conv(
+ device: torch.device,
+ layouts: tuple[tuple[tuple[int | None, ...], tuple[int | None, ...], torch.dtype], ...],
+ tokens_per_warp: int,
+ unroll: int,
+ depth: int,
+ warps: int,
+ min_blocks_per_mp: int,
+ src_vec: bool,
+ w_vec: bool,
+ out_vec: bool,
+ st_vec: bool,
+ seq_loop: bool,
+) -> Callable:
+ """Compile and cache a convolution layout without retaining caller-owned tensors."""
+ with torch.cuda.device(device):
+ if torch.cuda.is_current_stream_capturing():
+ raise RuntimeError("Warm KDA convolution layout before CUDA graph capture")
+ dtypes = {
+ torch.bfloat16: (cutlass.BFloat16, 2),
+ torch.int32: (cutlass.Int32, 4),
+ torch.bool: (cutlass.Boolean, 1),
+ }
+ tokens, batch, pool, offsets = (cute.sym_int() for _ in range(4))
+ # starts has B + 1 entries; its length is validated by the caller.
+ extents = (tokens, None, pool, offsets, batch, batch, tokens, tokens, tokens)
+ args = [
+ make_fake_tensor(
+ dtypes[dtype][0],
+ tuple(extent if n is None else n for n in shape),
+ stride=tuple(cute.sym_int() if s is None else s for s in stride),
+ assumed_align=dtypes[dtype][1],
+ )
+ for (shape, stride, dtype), extent in zip(layouts, extents)
+ ]
+ op = KdaConv(
+ layouts[1][0][0] // 384,
+ tokens_per_warp,
+ unroll,
+ src_vec,
+ w_vec,
+ out_vec,
+ layouts[0][1][0],
+ layouts[2][1],
+ seq_loop,
+ depth,
+ warps,
+ st_vec,
+ )
+ op.MB = min_blocks_per_mp
+ stream = cuda.CUstream(torch.cuda.current_stream(device).cuda_stream)
+ # Grid size and stream remain runtime arguments, not cache-key inputs.
+ return cute.compile(op, *args, cutlass.Int32(1), stream, options="--enable-tvm-ffi")
+
+
+def fused_prefill_conv(src, weight, conv_state, starts, slots, has_init, max_seqlen, q, k, v):
+ """Convolve raw packed QKV into contiguous, unnormalized BF16 Q/K/V.
+
+ Starts must be monotone, with unique nonnegative slots in the pool. The
+ scheduler prepares both arrays once per batch, before layer execution.
+ Outputs and the state pool must not alias the immutable inputs.
+ """
+ tensors = (src, weight, conv_state, starts, slots, has_init, q, k, v)
+ if any(t.device != src.device for t in tensors) or src.device.type != "cuda":
+ raise ValueError("KDA convolution requires tensors on one CUDA device")
+ if any(t.dtype != torch.bfloat16 for t in (src, weight, conv_state, q, k, v)):
+ raise ValueError("KDA convolution requires BF16 inputs, outputs and state")
+ channels = weight.shape[0]
+ if (
+ weight.shape != (channels, 4)
+ or channels % 384
+ or src.ndim != 2
+ or src.shape[1] < channels
+ or conv_state.ndim != 3
+ or conv_state.shape[1:] != (channels, 3)
+ ):
+ raise ValueError(
+ "KDA convolution requires width 4, head dimension 128 and three state taps"
+ )
+ shape = (1, src.shape[0], channels // 384, 128)
+ if any(t.shape != shape or not t.is_contiguous() for t in (q, k, v)):
+ raise ValueError("KDA convolution outputs must be contiguous [1, T, H, 128]")
+ if (
+ starts.dtype != torch.int32
+ or starts.shape != (slots.numel() + 1,)
+ or slots.dtype != torch.int32
+ or slots.ndim != 1
+ or has_init.dtype != torch.bool
+ or has_init.shape != slots.shape
+ or any(not t.is_contiguous() for t in (starts, slots, has_init))
+ ):
+ raise ValueError("Prepare contiguous int32 starts/slots and bool initial-state flags")
+ # The scalar fallback uses CuTe layout arithmetic with int32 extents.
+ # Reject overflowing layouts before launch rather than wrap an address.
+ if any(
+ sum((n - 1) * stride for n, stride in zip(t.shape, t.stride())) >= 2**31 for t in tensors
+ ):
+ raise ValueError("KDA convolution tensor offsets must fit int32")
+ writable = (conv_state, q, k, v)
+ for i, output in enumerate(writable):
+ for other in (src, weight, starts, slots, has_init, *writable[:i]):
+ if output.untyped_storage().data_ptr() == other.untyped_storage().data_ptr():
+ raise ValueError("KDA convolution outputs/state must use disjoint storage")
+ heads = weight.shape[0] // 384
+ batch = slots.shape[0]
+ if batch == 0 or heads == 0:
+ return
+ if isinstance(max_seqlen, torch.Tensor) and max_seqlen.device.type != "cpu":
+ # Reading a device bound would sync (and break graph capture); T bounds every length.
+ max_seqlen = src.shape[0]
+ max_seqlen = int(max_seqlen)
+ L = FORCE_L or _pick_tokens(heads, max(max_seqlen, 1), batch)
+ D = FORCE_D
+ W = FORCE_W or 4
+ U = min(8, L)
+ sr = src.stride(0)
+ # Vector (cp.async) source path only for a 16-byte-aligned base and row stride.
+ src_vec = src.stride(1) == 1 and sr % 8 == 0 and _align(src) == 0
+ w_vec = _align(weight) == 0 and weight.is_contiguous()
+ out_vec = all(_align(t) == 0 and t.is_contiguous() for t in (q, k, v))
+ # 16-byte state rows per lane (8 channels x 3 taps) need a dense [C, 3] slot layout,
+ # a 16-byte-aligned base and a slot stride divisible by 8 elements.
+ ss = conv_state.stride()
+ st_vec = ss[1] == 3 and ss[2] == 1 and ss[0] % 8 == 0 and _align(conv_state) == 0
+ stream = cuda.CUstream(torch.cuda.current_stream(src.device).cuda_stream)
+ # grid.y only sets parallelism (CTAs loop over the device-side run count): no
+ # sequence is longer than T, so a looser bound only adds idle CTAs and an
+ # understated one only makes CTAs loop.
+ seqlen_bound = max(1, min(max_seqlen, src.shape[0]))
+ run_tokens = L * W
+ grid_y = min((seqlen_bound + run_tokens - 1) // run_tokens, MAX_GRID_YZ)
+ seq_loop = batch > MAX_GRID_YZ # shape-only choice
+ # Token count, batch size and pool capacity are runtime extents. The
+ # singleton leading output stride also varies with the token count.
+ layouts = (
+ _layout_key(src, dynamic_dim=0),
+ _layout_key(weight),
+ _layout_key(conv_state, dynamic_dim=0),
+ _layout_key(starts, dynamic_dim=0),
+ _layout_key(slots, dynamic_dim=0),
+ _layout_key(has_init, dynamic_dim=0),
+ *(_layout_key(t, dynamic_dim=1, dynamic_stride=0) for t in (q, k, v)),
+ )
+ fn = compile_prefill_conv(
+ src.device,
+ layouts,
+ L,
+ U,
+ D,
+ W,
+ FORCE_MB,
+ src_vec,
+ w_vec,
+ out_vec,
+ st_vec,
+ seq_loop,
+ )
+ fn(src, weight, conv_state, starts, slots, has_init, q, k, v, grid_y, stream)
diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/k4_persistent.py b/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/k4_persistent.py
deleted file mode 100644
index 019c09efe0ff..000000000000
--- a/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/k4_persistent.py
+++ /dev/null
@@ -1,1508 +0,0 @@
-# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""K4-only fused KDA kernel: persistent variant with varlen support.
-
-Persistent scheduler (GDNTileScheduler) with 3D tensor layout, TensorMapManager
-for per-tile TMA descriptor updates, and domain_offset+flat_divide for per-chunk
-addressing. Supports variable-length sequences via cu_seqlens.
-
-Precondition: every sequence in cu_seqlens must have length > 0. A zero-length
-sequence yields num_chunks = 0, so its scheduler tile skips the chunk loop and
-never commits the per-chunk mbarriers the other warp groups wait on (deadlock).
-The prefill runtime never produces zero-length sequences, so this is not
-checked on the device or in the launch path (cu_seqlens lives on the GPU; a
-host-side check would force a sync on the hot path).
-
-K4 chunk loop with 6 MMAs per chunk:
- MMA1: W = AB @ KS (K-MN, K=64)
- MMA2: U = AB @ V (K-MN, K=64)
- MMA3: NV = U - W_bf16 @ S (SS-mode: A=SMEM sO, B=SMEM sST, K=128, accumulate;
- W is negated in its readout, FLA: b_v = u - w @ h)
- MMA4: OI = QS @ S (K-MN, K=128)
- MMA5: O = OI + AQC @ NV (K-MN, K=64, accumulate)
- MMA6: State += NV^T @ KG (MN-MN, K=64, accumulate on decayed state)
-
-Execution order: MMA1->MMA2->MMA4->MMA3->MMA5->MMA6
-
-Warp assignment (3 warpgroups, 12 warps, 384 threads):
- WG0 (W0-3): W0=MMA issue, W2=TMA load/store, W1/W3=idle
- WG1 (W4-7): State 2-pass: TMEM->bf16->sST, gk(SMEM)->decay->TMEM
- WG2 (W8-11): GDN readout: W/NV/O TMEM->bf16->SMEM
-
-TMA pipelines (all PipelineTmaUmma, 1-stage):
- - KS, V, QS: prefetch c+1 while MMA uses c
- - AB, AQC, KG: 1-stage with explicit consumer release (no prefetch)
- - TMA warp decoupled from O readout (no store_nbar wait)
-
-State management:
- - tCtState [128,128] fp32 @ TMEM offset 256: persistent
- - 2-pass per chunk: bf16->sST early, gk decay->TMEM late
- - gk_last preloaded to SMEM (coalesced), read in Pass 2
- - single state_ready_mbar after both passes (TMEM no-overlap constraint)
-"""
-
-import cuda.bindings.driver as cuda
-import cutlass
-import cutlass.cute as cute
-import cutlass.cutlass_dsl as _dsl_mod
-import cutlass.pipeline as pipeline
-import cutlass.utils as utils
-import cutlass.utils.blackwell_helpers as sm100_utils
-from cutlass._mlir import ir
-from cutlass._mlir.dialects import llvm
-from cutlass.cute.nvgpu import cpasync, tcgen05
-from cutlass.cute.nvgpu.tcgen05 import Field, OperandMajorMode, OperandSource
-from cutlass.cutlass_dsl import Int32, T, dsl_user_op
-from cutlass.utils import TensorMapManager, TensorMapUpdateMode
-
-# Compatibility shim, global effect: flashinfer's gated_delta_net_tile_scheduler
-# imports CuteExperimentalDSL from cutlass.cutlass_dsl, which older CUTLASS DSL
-# releases do not define. Install an inert placeholder so the import below
-# succeeds; other modules in the process see the same attribute.
-if not hasattr(_dsl_mod, "CuteExperimentalDSL"):
-
- class _DummyExperimentalDSL:
- jit = None
- kernel = None
- compile = None
-
- _dsl_mod.CuteExperimentalDSL = _DummyExperimentalDSL
-
-from flashinfer.gdn_kernels.blackwell.gated_delta_net_tile_scheduler import (
- GDNTileScheduler,
- GDNTileSchedulerParams,
-)
-
-SB = 16
-AKK_PAD = 8
-AKK_STRIDE = 64 + AKK_PAD # 72
-TEMP_COLS = SB + AKK_PAD # 24
-NUM_TEMPS = 2
-
-
-# ===========================================================================
-# Inverse dsl_user_op functions (TF32 MMA m16n8k8, barrier 7)
-# ===========================================================================
-@dsl_user_op
-def mma_tf32_m16n8k8(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, *, loc=None, ip=None):
- a0b = llvm.bitcast(T.i32(), a0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- a1b = llvm.bitcast(T.i32(), a1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- a2b = llvm.bitcast(T.i32(), a2.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- a3b = llvm.bitcast(T.i32(), a3.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- b0b = llvm.bitcast(T.i32(), b0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- b1b = llvm.bitcast(T.i32(), b1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
- result = llvm.inline_asm(
- ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"),
- [
- a0b,
- a1b,
- a2b,
- a3b,
- b0b,
- b1b,
- c0.ir_value(loc=loc, ip=ip),
- c1.ir_value(loc=loc, ip=ip),
- c2.ir_value(loc=loc, ip=ip),
- c3.ir_value(loc=loc, ip=ip),
- ],
- """{
- mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32
- {$0, $1, $2, $3},
- {$4, $5, $6, $7},
- {$8, $9},
- {$10, $11, $12, $13};
- }""",
- "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
- has_side_effects=True,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
- d0 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip))
- d1 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip))
- d2 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip))
- d3 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip))
- return d0, d1, d2, d3
-
-
-@dsl_user_op
-def inv_internal_barrier(*, loc=None, ip=None):
- llvm.inline_asm(
- T.i32(),
- [],
- "bar.sync 7, 128; mov.u32 $0, 0;",
- "=r",
- has_side_effects=True,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
-
-
-@dsl_user_op
-def _invert_diag(sAkk: cute.Tensor, block_rc, lane_id, *, loc=None, ip=None):
- my_row = lane_id % 16
- halfwarp_base = (lane_id // 16) * 16
- r_off = block_rc * 16
- c_off = block_rc * 16
- rInv = cute.make_rmem_tensor(cute.make_layout((16,), stride=(1,)), cutlass.Float32)
- rInv[0] = cutlass.Float32(1.0)
- for x in range(1, 16):
- rInv[x] = cutlass.Float32(0.0)
- for d in range(1, 16):
- col_d = my_row - d
- valid = cutlass.Float32(col_d >= 0)
- a_val = cutlass.Float32(sAkk[r_off + my_row, c_off + col_d]) * valid
- acc = cutlass.Float32(0.0)
- for j in range(1, d):
- a_re = cutlass.Float32(sAkk[r_off + my_row, c_off + my_row - (d - j)])
- inv_shfl = cute.arch.shuffle_sync(rInv[j], halfwarp_base + my_row - d + j)
- acc = acc + a_re * inv_shfl
- rInv[d] = (-a_val - acc) * valid
- rInv[0] = cutlass.Float32(1.0)
- sAkk[r_off + my_row, c_off + my_row] = rInv[0]
- for d in range(1, 16):
- sAkk[r_off + my_row, c_off + (my_row + 16 - d) % 16] = rInv[d] * cutlass.Float32(
- my_row >= d
- )
-
-
-@dsl_user_op
-def _matmul_AB_inv(sAkk: cute.Tensor, br_A, bc_A, br_B, bc_B, lane_id, *, loc=None, ip=None):
- gid = lane_id // 4
- tid = lane_id % 4
- _z = cutlass.Float32(0.0)
- rA = br_A * 16
- cA = bc_A * 16
- rB = br_B * 16
- cB = bc_B * 16
- a0 = cutlass.Float32(sAkk[rA + gid, cA + 2 * tid])
- a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + 2 * tid])
- a2 = cutlass.Float32(sAkk[rA + gid, cA + 2 * tid + 1])
- a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 2 * tid + 1])
- b0n0 = cutlass.Float32(sAkk[rB + 2 * tid, cB + gid])
- b1n0 = cutlass.Float32(sAkk[rB + 2 * tid + 1, cB + gid])
- b0n1 = cutlass.Float32(sAkk[rB + 2 * tid, cB + 8 + gid])
- b1n1 = cutlass.Float32(sAkk[rB + 2 * tid + 1, cB + 8 + gid])
- cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0n0, b1n0, _z, _z, _z, _z)
- cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0n1, b1n1, _z, _z, _z, _z)
- a0 = cutlass.Float32(sAkk[rA + gid, cA + 8 + 2 * tid])
- a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + 8 + 2 * tid])
- a2 = cutlass.Float32(sAkk[rA + gid, cA + 8 + 2 * tid + 1])
- a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 8 + 2 * tid + 1])
- b0n0 = cutlass.Float32(sAkk[rB + 8 + 2 * tid, cB + gid])
- b1n0 = cutlass.Float32(sAkk[rB + 8 + 2 * tid + 1, cB + gid])
- b0n1 = cutlass.Float32(sAkk[rB + 8 + 2 * tid, cB + 8 + gid])
- b1n1 = cutlass.Float32(sAkk[rB + 8 + 2 * tid + 1, cB + 8 + gid])
- cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8(
- a0, a1, a2, a3, b0n0, b1n0, cn0_0, cn0_1, cn0_2, cn0_3
- )
- cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8(
- a0, a1, a2, a3, b0n1, b1n1, cn1_0, cn1_1, cn1_2, cn1_3
- )
- return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3
-
-
-@dsl_user_op
-def _chain_mma_B_inv(
- sAkk: cute.Tensor,
- br_B,
- bc_B,
- a0k0,
- a1k0,
- a2k0,
- a3k0,
- a0k1,
- a1k1,
- a2k1,
- a3k1,
- lane_id,
- *,
- loc=None,
- ip=None,
-):
- gid = lane_id // 4
- tid = lane_id % 4
- _z = cutlass.Float32(0.0)
- rB = br_B * 16
- cB = bc_B * 16
- b0n0 = cutlass.Float32(sAkk[rB + 2 * tid, cB + gid])
- b1n0 = cutlass.Float32(sAkk[rB + 2 * tid + 1, cB + gid])
- b0n1 = cutlass.Float32(sAkk[rB + 2 * tid, cB + 8 + gid])
- b1n1 = cutlass.Float32(sAkk[rB + 2 * tid + 1, cB + 8 + gid])
- cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8(
- a0k0, a1k0, a2k0, a3k0, b0n0, b1n0, _z, _z, _z, _z
- )
- cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8(
- a0k0, a1k0, a2k0, a3k0, b0n1, b1n1, _z, _z, _z, _z
- )
- b0n0 = cutlass.Float32(sAkk[rB + 8 + 2 * tid, cB + gid])
- b1n0 = cutlass.Float32(sAkk[rB + 8 + 2 * tid + 1, cB + gid])
- b0n1 = cutlass.Float32(sAkk[rB + 8 + 2 * tid, cB + 8 + gid])
- b1n1 = cutlass.Float32(sAkk[rB + 8 + 2 * tid + 1, cB + 8 + gid])
- cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8(
- a0k1, a1k1, a2k1, a3k1, b0n0, b1n0, cn0_0, cn0_1, cn0_2, cn0_3
- )
- cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8(
- a0k1, a1k1, a2k1, a3k1, b0n1, b1n1, cn1_0, cn1_1, cn1_2, cn1_3
- )
- return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3
-
-
-@dsl_user_op
-def _chain_mma_A_inv(
- sAkk: cute.Tensor,
- br_A,
- bc_A,
- b0_k0n0,
- b1_k0n0,
- b0_k0n1,
- b1_k0n1,
- b0_k1n0,
- b1_k1n0,
- b0_k1n1,
- b1_k1n1,
- lane_id,
- *,
- loc=None,
- ip=None,
-):
- gid = lane_id // 4
- tid = lane_id % 4
- _z = cutlass.Float32(0.0)
- rA = br_A * 16
- cA = bc_A * 16
- a0 = cutlass.Float32(sAkk[rA + gid, cA + 2 * tid])
- a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + 2 * tid])
- a2 = cutlass.Float32(sAkk[rA + gid, cA + 2 * tid + 1])
- a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 2 * tid + 1])
- cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0_k0n0, b1_k0n0, _z, _z, _z, _z)
- cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0_k0n1, b1_k0n1, _z, _z, _z, _z)
- a0 = cutlass.Float32(sAkk[rA + gid, cA + 8 + 2 * tid])
- a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + 8 + 2 * tid])
- a2 = cutlass.Float32(sAkk[rA + gid, cA + 8 + 2 * tid + 1])
- a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 8 + 2 * tid + 1])
- cn0_0, cn0_1, cn0_2, cn0_3 = mma_tf32_m16n8k8(
- a0, a1, a2, a3, b0_k1n0, b1_k1n0, cn0_0, cn0_1, cn0_2, cn0_3
- )
- cn1_0, cn1_1, cn1_2, cn1_3 = mma_tf32_m16n8k8(
- a0, a1, a2, a3, b0_k1n1, b1_k1n1, cn1_0, cn1_1, cn1_2, cn1_3
- )
- return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3
-
-
-@dsl_user_op
-def _store_neg_C_inv(
- sAkk: cute.Tensor, br, bc, c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None
-):
- gid = lane_id // 4
- tid = lane_id % 4
- r = br * 16
- c = bc * 16
- sAkk[r + gid, c + 2 * tid] = -c0
- sAkk[r + gid, c + 2 * tid + 1] = -c1
- sAkk[r + gid + 8, c + 2 * tid] = -c2
- sAkk[r + gid + 8, c + 2 * tid + 1] = -c3
- sAkk[r + gid, c + 8 + 2 * tid] = -c4
- sAkk[r + gid, c + 8 + 2 * tid + 1] = -c5
- sAkk[r + gid + 8, c + 8 + 2 * tid] = -c6
- sAkk[r + gid + 8, c + 8 + 2 * tid + 1] = -c7
-
-
-@dsl_user_op
-def _shuffle_C_to_B_inv(c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None):
- gid = lane_id // 4
- tid = lane_id % 4
- src_a = 8 * tid + gid // 2
- src_b = src_a + 4
- f_odd = cutlass.Float32(gid % 2)
- f_even = cutlass.Float32(1) - f_odd
- c0_a = cute.arch.shuffle_sync(c0, src_a)
- c1_a = cute.arch.shuffle_sync(c1, src_a)
- c2_a = cute.arch.shuffle_sync(c2, src_a)
- c3_a = cute.arch.shuffle_sync(c3, src_a)
- c4_a = cute.arch.shuffle_sync(c4, src_a)
- c5_a = cute.arch.shuffle_sync(c5, src_a)
- c6_a = cute.arch.shuffle_sync(c6, src_a)
- c7_a = cute.arch.shuffle_sync(c7, src_a)
- c0_b = cute.arch.shuffle_sync(c0, src_b)
- c1_b = cute.arch.shuffle_sync(c1, src_b)
- c2_b = cute.arch.shuffle_sync(c2, src_b)
- c3_b = cute.arch.shuffle_sync(c3, src_b)
- c4_b = cute.arch.shuffle_sync(c4, src_b)
- c5_b = cute.arch.shuffle_sync(c5, src_b)
- c6_b = cute.arch.shuffle_sync(c6, src_b)
- c7_b = cute.arch.shuffle_sync(c7, src_b)
- b0_00 = c0_a * f_even + c1_a * f_odd
- b1_00 = c0_b * f_even + c1_b * f_odd
- b0_10 = c2_a * f_even + c3_a * f_odd
- b1_10 = c2_b * f_even + c3_b * f_odd
- b0_01 = c4_a * f_even + c5_a * f_odd
- b1_01 = c4_b * f_even + c5_b * f_odd
- b0_11 = c6_a * f_even + c7_a * f_odd
- b1_11 = c6_b * f_even + c7_b * f_odd
- return b0_00, b1_00, b0_10, b1_10, b0_01, b1_01, b0_11, b1_11
-
-
-@dsl_user_op
-def _store_C_temp_inv(
- sT: cute.Tensor, buf, c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None
-):
- gid = lane_id // 4
- tid = lane_id % 4
- sT[gid, 2 * tid, buf] = c0
- sT[gid, 2 * tid + 1, buf] = c1
- sT[gid + 8, 2 * tid, buf] = c2
- sT[gid + 8, 2 * tid + 1, buf] = c3
- sT[gid, 8 + 2 * tid, buf] = c4
- sT[gid, 8 + 2 * tid + 1, buf] = c5
- sT[gid + 8, 8 + 2 * tid, buf] = c6
- sT[gid + 8, 8 + 2 * tid + 1, buf] = c7
-
-
-@dsl_user_op
-def _load_C_temp_inv(sT: cute.Tensor, buf, lane_id, *, loc=None, ip=None):
- gid = lane_id // 4
- tid = lane_id % 4
- c0 = cutlass.Float32(sT[gid, 2 * tid, buf])
- c1 = cutlass.Float32(sT[gid, 2 * tid + 1, buf])
- c2 = cutlass.Float32(sT[gid + 8, 2 * tid, buf])
- c3 = cutlass.Float32(sT[gid + 8, 2 * tid + 1, buf])
- c4 = cutlass.Float32(sT[gid, 8 + 2 * tid, buf])
- c5 = cutlass.Float32(sT[gid, 8 + 2 * tid + 1, buf])
- c6 = cutlass.Float32(sT[gid + 8, 8 + 2 * tid, buf])
- c7 = cutlass.Float32(sT[gid + 8, 8 + 2 * tid + 1, buf])
- return c0, c1, c2, c3, c4, c5, c6, c7
-
-
-def transform_partitioned_tensor_layout(tensor):
- layout = tensor.layout
- stored_layout = layout
- if isinstance(stored_layout, cute.ComposedLayout):
- layout = layout.outer
- shape = layout.shape
- stride = layout.stride
- new_shape = ((shape[0][0], shape[1]), (shape[0][1], shape[2]), *shape[3:])
- new_stride = ((stride[0][0], stride[1]), (stride[0][1], stride[2]), *stride[3:])
- new_layout = cute.make_layout(shape=new_shape, stride=new_stride)
- if isinstance(stored_layout, cute.ComposedLayout):
- new_layout = cute.make_composed_layout(
- stored_layout.inner, stored_layout.offset, new_layout
- )
- return cute.make_tensor(tensor.iterator, new_layout)
-
-
-mma_dtype = cutlass.BFloat16
-acc_dtype = cutlass.Float32
-out_dtype = cutlass.BFloat16
-
-M = 64
-N = 128
-K = 64
-K3 = 128
-M6 = 128
-N6 = 128
-K6 = 64
-
-threads_per_cta = 384
-warp_threads = 32
-warpgroup_threads = 128
-
-BYTES_PER_TENSORMAP = 128
-NUM_TENSORMAPS = 7 # a, b, v, q, aqc, kg, o
-
-MMA_WARP = 0
-O_STORE_WARP = 1
-TMA_WARP = 2
-STATE_WG = 1
-READOUT_WG = 2
-
-NUM_REGS_WG0 = 40
-NUM_REGS_WG1 = 232
-NUM_REGS_WG2 = 232
-MAX_REGS = 168
-
-try:
- from tensorrt_llm.logger import logger as _logger
-except ImportError: # standalone kernel use outside the full package
- import logging as _logging
-
- _logger = _logging.getLogger(__name__)
-
-# Monkey-patch, GLOBAL effect: CuTeDSL._get_pipeline is class-level state, so
-# every CuTe DSL kernel compiled in this process after this import (not just
-# K4) gets ptxas --uumn (unified uniform register allocation, needed here to
-# keep WG0 under its 40-register budget). The DSL exposes no per-compile
-# ptx-options hook at the cute.compile call sites, hence the patch. The extra
-# flag is benign for the other kernels in this package. Idempotent: skips
-# pipelines that already carry ptx-options.
-try:
- from cutlass.cutlass_dsl.cutlass import CuTeDSL as _CuTeDSL
-
- _orig_get_pipeline = _CuTeDSL._get_pipeline
- _patch_applied = False
-
- def _patched_get_pipeline(self, _pipeline):
- global _patch_applied
- result = _orig_get_pipeline(self, _pipeline)
- if result and "ptx-options=" not in result:
- if "cubin-format=bin" in result:
- result = result.replace("cubin-format=bin", "cubin-format=bin ptx-options='--uumn'")
- _patch_applied = True
- else:
- _logger.warning(
- f"k4_persistent ptx-options patch: 'cubin-format=bin' not found "
- f"in pipeline: {result[:200]}"
- )
- elif result and "ptx-options=" in result:
- _logger.debug("k4_persistent ptx-options patch: ptx-options already present")
- _patch_applied = True
- return result
-
- _CuTeDSL._get_pipeline = _patched_get_pipeline
-# Workaround for a CuTe DSL parser bug: the nvidia-cutlass-dsl 4.5.0 AST
-# preprocessor cannot parse tuple except handlers anywhere in a kernel
-# module ("'Tuple' object has no attribute 'id'"), which breaks every
-# cute.compile of this file. Keep a single Exception clause until the DSL
-# pin moves past the bug; the isinstance re-raise below narrows it to the
-# expected failure kinds without putting a tuple in the except clause.
-except Exception as e:
- if not isinstance(e, (AttributeError, ImportError)):
- # Not a known-benign "DSL not present / API surface moved" case:
- # swallowing it would defer the blow-up to the first K4 compile
- # (ptxas over WG0's 40-register budget without --uumn) with no
- # trail back to this patch. Fail here, where the cause is visible.
- raise
- _logger.warning(f"k4_persistent ptx-options patch failed: {e}")
- _patch_applied = False
-
-
-@cute.kernel
-def k4_persistent_kernel(
- tiled_mma_kmn: cute.TiledMma,
- tiled_mma_mn_mn: cute.TiledMma,
- tma_a,
- a_sl: cute.ComposedLayout,
- tma_b,
- b_sl: cute.ComposedLayout,
- tma_v,
- v_sl: cute.ComposedLayout,
- s_sl: cute.ComposedLayout,
- tma_q,
- q_sl: cute.ComposedLayout,
- tma_aqc,
- aqc_sl: cute.ComposedLayout,
- tma_kg,
- kg_sl: cute.ComposedLayout,
- tma_o,
- store_sl: cute.ComposedLayout,
- readout_k_sl: cute.ComposedLayout,
- nv_b_sl: cute.ComposedLayout,
- nv_a_sl: cute.ComposedLayout,
- kg_a_sl: cute.ComposedLayout,
- nv_b_mn_sl: cute.ComposedLayout,
- mGkLastExp: cute.Tensor,
- mS_fp32: cute.Tensor,
- state_indices: cute.Tensor,
- cu_seqlens: cute.Tensor,
- chunk_offsets: cute.Tensor,
- mA: cute.Tensor,
- mB: cute.Tensor,
- mV_g: cute.Tensor,
- mQ: cute.Tensor,
- mAQC: cute.Tensor,
- mKG: cute.Tensor,
- mO: cute.Tensor,
- tensormap_workspace: cute.Tensor,
- scheduler_params: GDNTileSchedulerParams,
-):
- bidx, bidy, bidz = cute.arch.block_idx()
- grid_dim = cute.arch.grid_dim()
- tidx, _, _ = cute.arch.thread_idx()
- warp_idx = cute.arch.make_warp_uniform(tidx // warp_threads)
- warpgroup_idx = cute.arch.make_warp_uniform(tidx // warpgroup_threads)
- warpgroup_tidx = tidx % warpgroup_threads
- thr_kmn = tiled_mma_kmn.get_slice(0)
- thr_mn = tiled_mma_mn_mn.get_slice(0)
- dice = (None, None, None)
-
- cta_linear_idx = bidz * grid_dim[1] * grid_dim[0] + bidy * grid_dim[0] + bidx
- tensormap_manager = TensorMapManager(TensorMapUpdateMode.GMEM, BYTES_PER_TENSORMAP)
- tm_ws = cute.make_tensor(
- tensormap_workspace.iterator,
- cute.make_layout(
- (grid_dim[0] * grid_dim[1] * grid_dim[2], NUM_TENSORMAPS, BYTES_PER_TENSORMAP),
- stride=(NUM_TENSORMAPS * BYTES_PER_TENSORMAP, BYTES_PER_TENSORMAP, 1),
- ),
- )
- tm_a_ptr = tensormap_manager.get_tensormap_ptr(tm_ws[(cta_linear_idx, 0, None)].iterator)
- tm_b_ptr = tensormap_manager.get_tensormap_ptr(tm_ws[(cta_linear_idx, 1, None)].iterator)
- tm_v_ptr = tensormap_manager.get_tensormap_ptr(tm_ws[(cta_linear_idx, 2, None)].iterator)
- tm_q_ptr = tensormap_manager.get_tensormap_ptr(tm_ws[(cta_linear_idx, 3, None)].iterator)
- tm_aqc_ptr = tensormap_manager.get_tensormap_ptr(tm_ws[(cta_linear_idx, 4, None)].iterator)
- tm_kg_ptr = tensormap_manager.get_tensormap_ptr(tm_ws[(cta_linear_idx, 5, None)].iterator)
- tm_o_ptr = tensormap_manager.get_tensormap_ptr(tm_ws[(cta_linear_idx, 6, None)].iterator)
-
- smem = cutlass.utils.SmemAllocator()
- AL = 128
- sA = smem.allocate_tensor(mma_dtype, a_sl.outer, AL, a_sl.inner)
- sB = smem.allocate_tensor(mma_dtype, b_sl.outer, AL, b_sl.inner)
- sV = smem.allocate_tensor(mma_dtype, v_sl.outer, AL, v_sl.inner)
- sST = smem.allocate_tensor(mma_dtype, s_sl.outer, AL, s_sl.inner)
- sQ = smem.allocate_tensor(mma_dtype, q_sl.outer, AL, q_sl.inner)
- sAQC = smem.allocate_tensor(mma_dtype, aqc_sl.outer, AL, aqc_sl.inner)
- sKG = smem.allocate_tensor(mma_dtype, kg_sl.outer, AL, kg_sl.inner)
- sNV = smem.allocate_tensor(mma_dtype, readout_k_sl.outer, AL, readout_k_sl.inner)
- sO = smem.allocate_tensor(mma_dtype, readout_k_sl.outer, AL, readout_k_sl.inner)
- sO_out = smem.allocate_tensor(mma_dtype, readout_k_sl.outer, AL, readout_k_sl.inner)
- sGk_buf = smem.allocate_array(acc_dtype, N)
-
- sNV_b = cute.make_tensor(cute.recast_ptr(sNV.iterator, nv_b_sl.inner, mma_dtype), nv_b_sl.outer)
- # sNV_a = cute.make_tensor(cute.recast_ptr(sNV.iterator, nv_a_sl.inner, mma_dtype), nv_a_sl.outer)
- sKG_a = cute.make_tensor(cute.recast_ptr(sKG.iterator, kg_a_sl.inner, mma_dtype), kg_a_sl.outer)
- sNV_b_mn = cute.make_tensor(
- cute.recast_ptr(sNV.iterator, nv_b_mn_sl.inner, mma_dtype), nv_b_mn_sl.outer
- )
- sO_out_st = cute.make_tensor(
- cute.recast_ptr(sO_out.iterator, store_sl.inner, out_dtype), store_sl.outer
- )
-
- tmem_smem = smem.allocate_array(cutlass.Int32, 1)
- if warp_idx == 0:
- cute.arch.alloc_tmem(512, tmem_smem)
-
- sNV_ready_nbar = pipeline.NamedBarrier(3, warpgroup_threads + warp_threads)
- sW_ready_nbar = pipeline.NamedBarrier(2, warpgroup_threads + warp_threads)
- gk_load_nbar = pipeline.NamedBarrier(5, warpgroup_threads)
-
- elect_one = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1)
- wg_coop = pipeline.CooperativeGroup(pipeline.Agent.Thread, warpgroup_threads)
- warp_coop = pipeline.CooperativeGroup(pipeline.Agent.Thread, warp_threads)
-
- mma6_done_mbar = smem.allocate_array(cutlass.Int64, 1)
- gmem_done_mbar = smem.allocate_array(cutlass.Int64, 1)
- state_ready_mbar = smem.allocate_array(cutlass.Int64, 1)
- final_state_done_mbar = smem.allocate_array(cutlass.Int64, 1)
-
- if warp_idx == 0:
- with cute.arch.elect_one():
- cute.arch.mbarrier_init(mma6_done_mbar, warp_threads)
- cute.arch.mbarrier_init(gmem_done_mbar, warpgroup_threads)
- cute.arch.mbarrier_init(state_ready_mbar, warpgroup_threads)
- cute.arch.mbarrier_init(final_state_done_mbar, warpgroup_threads)
-
- def _make_tma_pipe(total_byte_count, num_stages=1):
- ptr = smem.allocate_array(cutlass.Int64, 2 * num_stages)
- return pipeline.PipelineTmaUmma.create(
- barrier_storage=ptr,
- num_stages=num_stages,
- producer_group=elect_one,
- consumer_group=elect_one,
- tx_count=total_byte_count // num_stages,
- defer_sync=True,
- ).make_participants()
-
- b_prod, b_cons = _make_tma_pipe(cute.size_in_bytes(mma_dtype, b_sl), num_stages=1)
- v_prod, v_cons = _make_tma_pipe(cute.size_in_bytes(mma_dtype, v_sl), num_stages=1)
- q_prod, q_cons = _make_tma_pipe(cute.size_in_bytes(mma_dtype, q_sl), num_stages=1)
- a_prod, a_cons = _make_tma_pipe(cute.size_in_bytes(mma_dtype, a_sl), num_stages=1)
- aqc_prod, aqc_cons = _make_tma_pipe(cute.size_in_bytes(mma_dtype, aqc_sl), num_stages=1)
- kg_prod, kg_cons = _make_tma_pipe(cute.size_in_bytes(mma_dtype, kg_sl), num_stages=1)
-
- def _make_umma_pipe():
- ptr = smem.allocate_array(cutlass.Int64, 2)
- return pipeline.PipelineUmmaAsync.create(
- barrier_storage=ptr,
- num_stages=1,
- producer_group=elect_one,
- consumer_group=wg_coop,
- defer_sync=True,
- ).make_participants()
-
- w_prod, w_cons = _make_umma_pipe()
- nv_prod, nv_cons = _make_umma_pipe()
- o_prod, o_cons = _make_umma_pipe()
-
- o_store_mbar = smem.allocate_array(cutlass.Int64, 2)
- o_store_prod, o_store_cons = pipeline.PipelineAsync.create(
- barrier_storage=o_store_mbar,
- num_stages=1,
- producer_group=wg_coop,
- consumer_group=warp_coop,
- defer_sync=True,
- ).make_participants()
-
- cute.arch.sync_threads()
-
- tmem_ptr = cute.arch.retrieve_tmem_ptr(cutlass.Int32, 16, tmem_smem)
-
- tCtW_shape = tiled_mma_kmn.partition_shape_C((M, N))
- tCtW_fake = tiled_mma_kmn.make_fragment_C(tCtW_shape)
- tCtW = cute.make_tensor(cute.recast_ptr(tmem_ptr + 0, dtype=acc_dtype), tCtW_fake.layout)
-
- tCtNV_shape = tiled_mma_kmn.partition_shape_C((M, N))
- tCtNV_fake = tiled_mma_kmn.make_fragment_C(tCtNV_shape)
- tCtNV = cute.make_tensor(cute.recast_ptr(tmem_ptr + 128, dtype=acc_dtype), tCtNV_fake.layout)
- tCtO = cute.make_tensor(cute.recast_ptr(tmem_ptr + 384, dtype=acc_dtype), tCtNV_fake.layout)
-
- tCtS_shape = tiled_mma_mn_mn.partition_shape_C((M6, N6))
- tCtS_fake = tiled_mma_mn_mn.make_fragment_C(tCtS_shape)
- tCtState = cute.make_tensor(cute.recast_ptr(tmem_ptr + 256, dtype=acc_dtype), tCtS_fake.layout)
-
- # PDL: setup is done; now wait for upstream akk_inv to commit before reading
- # gmem inputs (TMA loads, gk_last_exp reads). PDL allows the entire setup phase
- # above (~tmem alloc, smem layout, pipeline init) to overlap with akk_inv's tail.
- cute.arch.griddepcontrol_wait()
-
- # ==== WG1: State readout + decay ====
- if warpgroup_idx == STATE_WG:
- cId_128 = cute.make_identity_tensor((M6, N6))
- tCtState_mn = transform_partitioned_tensor_layout(tCtState)
-
- atom_state_t2r = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition(32)), acc_dtype)
- tiled_state_t2r = tcgen05.make_tmem_copy(atom_state_t2r, tCtState[(None, None), 0, 0])
- thr_state_t2r = tiled_state_t2r.get_slice(warpgroup_tidx)
- tTR_tCtState = thr_state_t2r.partition_S(tCtState_mn)
- tTR_tCcState = thr_state_t2r.partition_D(cId_128)
- tRrState = cute.make_rmem_tensor_like(tTR_tCcState, acc_dtype)
-
- atom_state_r2t = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition(32)), acc_dtype)
- tiled_state_r2t = tcgen05.make_tmem_copy(atom_state_r2t, tCtState[(None, None), 0, 0])
- thr_state_r2t = tiled_state_r2t.get_slice(warpgroup_tidx)
- tRT_tCtState = thr_state_r2t.partition_D(tCtState_mn)
-
- atom_state_g2r = cute.make_copy_atom(
- cute.nvgpu.CopyUniversalOp(), acc_dtype, num_bits_per_copy=32
- )
- tiled_state_g2r = cute.make_tiled_copy_S(atom_state_g2r, tiled_state_r2t)
- thr_state_g2r = tiled_state_g2r.get_slice(warpgroup_tidx)
-
- atom_state_r2g = cute.make_copy_atom(
- cute.nvgpu.CopyUniversalOp(), acc_dtype, num_bits_per_copy=32
- )
- tiled_state_r2g = cute.make_tiled_copy_D(atom_state_r2g, tiled_state_t2r)
- thr_state_r2g = tiled_state_r2g.get_slice(warpgroup_tidx)
-
- atom_state_r2s = cute.make_copy_atom(
- cute.nvgpu.CopyUniversalOp(), mma_dtype, num_bits_per_copy=128
- )
- tiled_state_r2s = cute.make_tiled_copy_D(atom_state_r2s, tiled_state_t2r)
- thr_state_r2s = tiled_state_r2s.get_slice(warpgroup_tidx)
- sST_vk_view = transform_partitioned_tensor_layout(sST)
- sST_kv_view = cute.make_tensor(
- sST.iterator, cute.select(sST_vk_view.layout, mode=[1, 0, 2])
- )
- tCsState_inp = thr_state_r2s.partition_D(sST_kv_view)
- tRrState_bf16 = cute.make_rmem_tensor_like(tTR_tCcState, mma_dtype)
- tCrState_bf16 = tiled_state_r2s.retile(tRrState_bf16)
-
- sGk = cute.make_tensor(sGk_buf, cute.make_layout(((N,), (1,))))
-
- scheduler = GDNTileScheduler.create(scheduler_params, (bidx, bidy, bidz), grid_dim)
- work = scheduler.initial_work_tile_info()
- global_chunk = Int32(0)
-
- while work.is_valid_tile:
- batch_idx, head_idx, _ = work.tile_idx
- state_idx = state_indices[batch_idx]
- batch_start = cu_seqlens[batch_idx]
- batch_end = cu_seqlens[batch_idx + 1]
- num_chunks = cute.ceil_div(batch_end - batch_start, M)
- # chunk_base from cumulative chunk-offsets array (handles non-64-aligned varlen).
- # batch_start // M only works when all seq lengths are multiples of M.
- chunk_base = chunk_offsets[batch_idx]
-
- state_vk = mS_fp32[state_idx, head_idx, None, None]
- state_kv = cute.make_tensor(
- state_vk.iterator,
- cute.make_layout((M6, N6), stride=(1, M6)),
- )
- gS_init = cute.flat_divide(state_kv, (M6, N6))[None, None, 0, 0]
- tGR_tCgState_in = thr_state_g2r.partition_S(gS_init)
- tGR_tCrState_in = thr_state_g2r.retile(tRrState)
- cute.copy(tiled_state_g2r, tGR_tCgState_in, tGR_tCrState_in)
- num_state_subs = tRrState.shape[2]
- for sub in cutlass.range(num_state_subs):
- cute.copy(tiled_state_r2t, tRrState[None, 0, sub], tRT_tCtState[None, 0, sub])
- cute.arch.fence_view_async_tmem_store()
-
- sGk[warpgroup_tidx] = cutlass.Float32(mGkLastExp[chunk_base, head_idx, warpgroup_tidx])
- gk_load_nbar.arrive_and_wait()
-
- for chunk_c in cutlass.range(num_chunks):
- num_state_subs = tRrState.shape[2]
- _sub_tile_size = cute.size(tRrState.shape[0])
-
- if chunk_c > 0:
- cute.arch.mbarrier_wait(mma6_done_mbar, phase=(global_chunk - 1) % 2)
- gk_load_nbar.arrive_and_wait()
-
- cute.copy(tiled_state_t2r, tTR_tCtState, tRrState)
- tRrState_bf16.store(tRrState.load().to(mma_dtype))
- cute.copy(tiled_state_r2s, tCrState_bf16, tCsState_inp[None, None, None, 0])
- cute.arch.fence_view_async_shared()
- cute.arch.mbarrier_arrive(gmem_done_mbar)
-
- for sub in cutlass.range(num_state_subs):
- cute.copy(tiled_state_t2r, tTR_tCtState[None, 0, sub], tRrState[None, 0, sub])
- cute.arch.fence_view_async_tmem_load()
- for sub in cutlass.range(num_state_subs):
- for i in cutlass.range(_sub_tile_size):
- coord = tTR_tCcState[i, 0, sub]
- # FIX: probe shows kernel state TMEM[m=K, n=V] (m corresponds to K-dim);
- # OLD K4 / fla wants K-axis decay (state[k, v] *= gk[k]) → use coord[0]=m=K-idx.
- k_idx = coord[0]
- gk_val = cutlass.Float32(sGk[k_idx])
- tRrState[i, 0, sub] = tRrState[i, 0, sub] * gk_val
- cute.copy(tiled_state_r2t, tRrState[None, 0, sub], tRT_tCtState[None, 0, sub])
- cute.arch.fence_view_async_tmem_store()
- cute.arch.mbarrier_arrive(state_ready_mbar)
-
- if chunk_c + 1 < num_chunks:
- sGk[warpgroup_tidx] = cutlass.Float32(
- mGkLastExp[chunk_base + chunk_c + 1, head_idx, warpgroup_tidx]
- )
- global_chunk = global_chunk + 1
-
- cute.arch.mbarrier_wait(mma6_done_mbar, phase=(global_chunk - 1) % 2)
- num_state_subs_final = tRrState.shape[2]
- for sub in cutlass.range(num_state_subs_final):
- cute.copy(tiled_state_t2r, tTR_tCtState[None, 0, sub], tRrState[None, 0, sub])
- cute.arch.fence_view_async_tmem_load()
- state_vk = mS_fp32[state_idx, head_idx, None, None]
- state_kv = cute.make_tensor(
- state_vk.iterator,
- cute.make_layout((M6, N6), stride=(1, M6)),
- )
- gS_out = cute.flat_divide(state_kv, (M6, N6))[None, None, 0, 0]
- tGR_tCgState_out = thr_state_r2g.partition_D(gS_out)
- tGR_tCrState_out = thr_state_r2g.retile(tRrState)
- for sub in cutlass.range(num_state_subs_final):
- cute.copy(
- tiled_state_r2g,
- tGR_tCrState_out[None, 0, sub],
- tGR_tCgState_out[None, 0, sub],
- )
-
- cute.arch.mbarrier_arrive(final_state_done_mbar)
-
- scheduler.advance_to_next_work()
- work = scheduler.get_current_work()
-
- # ==== MMA warp (warp 0): 6 MMAs ====
- elif warp_idx == MMA_WARP:
- fA_kmn = thr_kmn.make_fragment_A(sA)
- fB_ks = thr_kmn.make_fragment_B(sB)
- fB_v = thr_kmn.make_fragment_B(sV)
- fB_s = thr_kmn.make_fragment_B(sST)
- fA_q = thr_kmn.make_fragment_A(sQ)
- fA_aqc = thr_kmn.make_fragment_A(sAQC)
- fB_nv = thr_kmn.make_fragment_B(sNV_b)
- fA_kg = thr_mn.make_fragment_A(sKG_a)
- fB_nv_mn = thr_mn.make_fragment_B(sNV_b_mn)
- fA_w = thr_kmn.make_fragment_A(sO)
-
- scheduler = GDNTileScheduler.create(scheduler_params, (bidx, bidy, bidz), grid_dim)
- work = scheduler.initial_work_tile_info()
- global_chunk = Int32(0)
- tile_count = Int32(0)
-
- while work.is_valid_tile:
- batch_idx_m, head_idx_m, _ = work.tile_idx
- num_chunks_m = cute.ceil_div(cu_seqlens[batch_idx_m + 1] - cu_seqlens[batch_idx_m], M)
-
- for chunk_c in cutlass.range(num_chunks_m):
- c_phase = global_chunk % 2
-
- ah = a_cons.wait_and_advance()
- bh = b_cons.wait_and_advance()
- w_h = w_prod.acquire_and_advance()
- tiled_mma_kmn.set(Field.ACCUMULATE, False)
- for k in cutlass.range_constexpr(cute.size(sB.shape[2])):
- cute.gemm(
- tiled_mma_kmn,
- tCtW,
- fA_kmn[dice + (0,)][None, None, k],
- fB_ks[dice + (bh.index,)][None, None, k],
- tCtW,
- )
- if k == 0:
- tiled_mma_kmn.set(Field.ACCUMULATE, True)
- bh.release()
- w_h.commit()
-
- vh = v_cons.wait_and_advance()
- tiled_mma_kmn.set(Field.ACCUMULATE, False)
- for k in cutlass.range_constexpr(cute.size(sV.shape[2])):
- cute.gemm(
- tiled_mma_kmn,
- tCtNV,
- fA_kmn[dice + (0,)][None, None, k],
- fB_v[dice + (vh.index,)][None, None, k],
- tCtNV,
- )
- if k == 0:
- tiled_mma_kmn.set(Field.ACCUMULATE, True)
- vh.release()
- ah.release()
-
- cute.arch.mbarrier_wait(gmem_done_mbar, phase=c_phase)
- sW_ready_nbar.arrive_and_wait()
-
- nv_h = nv_prod.acquire_and_advance()
- tiled_mma_kmn.set(Field.ACCUMULATE, True)
- for k in cutlass.range_constexpr(cute.size(sST.shape[2])):
- cute.gemm(
- tiled_mma_kmn,
- tCtNV,
- fA_w[dice + (0,)][None, None, k],
- fB_s[dice + (0,)][None, None, k],
- tCtNV,
- )
- nv_h.commit()
-
- qh = q_cons.wait_and_advance()
- tiled_mma_kmn.set(Field.ACCUMULATE, False)
- for k in cutlass.range_constexpr(cute.size(sST.shape[2])):
- cute.gemm(
- tiled_mma_kmn,
- tCtO,
- fA_q[dice + (qh.index,)][None, None, k],
- fB_s[dice + (0,)][None, None, k],
- tCtO,
- )
- if k == 0:
- tiled_mma_kmn.set(Field.ACCUMULATE, True)
- qh.release()
-
- sNV_ready_nbar.arrive_and_wait()
-
- aqch = aqc_cons.wait_and_advance()
- o_h = o_prod.acquire_and_advance()
- tiled_mma_kmn.set(Field.ACCUMULATE, True)
- for k in cutlass.range_constexpr(cute.size(sNV_b.shape[2])):
- cute.gemm(
- tiled_mma_kmn,
- tCtO,
- fA_aqc[dice + (0,)][None, None, k],
- fB_nv[dice + (0,)][None, None, k],
- tCtO,
- )
- o_h.commit()
- aqch.release()
-
- cute.arch.mbarrier_wait(state_ready_mbar, phase=c_phase)
- kgh = kg_cons.wait_and_advance()
- tiled_mma_mn_mn.set(Field.ACCUMULATE, True)
- for k in cutlass.range_constexpr(cute.size(sKG_a.shape[2])):
- cute.gemm(
- tiled_mma_mn_mn,
- tCtState,
- fA_kg[dice + (0,)][None, None, k],
- fB_nv_mn[dice + (0,)][None, None, k],
- tCtState,
- )
- kgh.release()
-
- w_prod.tail()
- nv_prod.tail()
- o_prod.tail()
- tcgen05.commit(mma6_done_mbar)
- global_chunk = global_chunk + 1
-
- cute.arch.mbarrier_wait(final_state_done_mbar, phase=tile_count % 2)
- tile_count = tile_count + 1
-
- scheduler.advance_to_next_work()
- work = scheduler.get_current_work()
-
- cute.arch.relinquish_tmem_alloc_permit()
- cute.arch.dealloc_tmem(tmem_ptr, 512)
-
- # ==== TMA warp (warp 2) ====
- elif warp_idx == TMA_WARP:
- cta_layout = cute.make_layout(1)
-
- scheduler = GDNTileScheduler.create(scheduler_params, (bidx, bidy, bidz), grid_dim)
- work = scheduler.initial_work_tile_info()
-
- if work.is_valid_tile:
- tensormap_manager.init_tensormap_from_atom(tma_a[0], tm_a_ptr, TMA_WARP)
- tensormap_manager.init_tensormap_from_atom(tma_b[0], tm_b_ptr, TMA_WARP)
- tensormap_manager.init_tensormap_from_atom(tma_v[0], tm_v_ptr, TMA_WARP)
- tensormap_manager.init_tensormap_from_atom(tma_q[0], tm_q_ptr, TMA_WARP)
- tensormap_manager.init_tensormap_from_atom(tma_aqc[0], tm_aqc_ptr, TMA_WARP)
- tensormap_manager.init_tensormap_from_atom(tma_kg[0], tm_kg_ptr, TMA_WARP)
- tensormap_manager.fence_tensormap_initialization()
-
- while work.is_valid_tile:
- batch_idx_t, head_idx_t, _ = work.tile_idx
- batch_start_t = cu_seqlens[batch_idx_t]
- batch_end_t = cu_seqlens[batch_idx_t + 1]
- num_chunks_t = cute.ceil_div(batch_end_t - batch_start_t, M)
-
- bounded_a = cute.make_tensor(
- mA.iterator,
- cute.make_layout(
- (batch_end_t, mA.shape[1], mA.shape[2]),
- stride=(mA.stride[0], mA.stride[1], mA.stride[2]),
- ),
- )
- bounded_b = cute.make_tensor(
- mB.iterator,
- cute.make_layout(
- (mB.shape[0], batch_end_t, mB.shape[2]),
- stride=(mB.stride[0], mB.stride[1], mB.stride[2]),
- ),
- )
- bounded_v = cute.make_tensor(
- mV_g.iterator,
- cute.make_layout(
- (mV_g.shape[0], batch_end_t, mV_g.shape[2]),
- stride=(mV_g.stride[0], mV_g.stride[1], mV_g.stride[2]),
- ),
- )
- bounded_q = cute.make_tensor(
- mQ.iterator,
- cute.make_layout(
- (batch_end_t, mQ.shape[1], mQ.shape[2]),
- stride=(mQ.stride[0], mQ.stride[1], mQ.stride[2]),
- ),
- )
- bounded_aqc = cute.make_tensor(
- mAQC.iterator,
- cute.make_layout(
- (batch_end_t, mAQC.shape[1], mAQC.shape[2]),
- stride=(mAQC.stride[0], mAQC.stride[1], mAQC.stride[2]),
- ),
- )
- bounded_kg = cute.make_tensor(
- mKG.iterator,
- cute.make_layout(
- (mKG.shape[0], batch_end_t, mKG.shape[2]),
- stride=(mKG.stride[0], mKG.stride[1], mKG.stride[2]),
- ),
- )
- tensormap_manager.update_tensormap(
- (bounded_a, bounded_b, bounded_v, bounded_q, bounded_aqc, bounded_kg),
- (tma_a[0], tma_b[0], tma_v[0], tma_q[0], tma_aqc[0], tma_kg[0]),
- (tm_a_ptr, tm_b_ptr, tm_v_ptr, tm_q_ptr, tm_aqc_ptr, tm_kg_ptr),
- TMA_WARP,
- (None, None, None, None, None, None),
- )
-
- for chunk_c in cutlass.range(num_chunks_t):
- chunk_offset = batch_start_t + chunk_c * M
-
- # AB: A-operand (T, K, H) → domain_offset on (tokens, 0)
- mA_c = cute.domain_offset(
- (chunk_offset, Int32(0)), tma_a[1][None, None, head_idx_t]
- )
- gA = cute.flat_divide(mA_c, (M, K))
- tCgA = thr_kmn.partition_A(gA)
- tAsA, tAgA = cpasync.tma_partition(
- tma_a[0],
- 0,
- cta_layout,
- cute.group_modes(sA, 0, 3),
- cute.group_modes(tCgA, 0, 3),
- )
- ah = a_prod.acquire_and_advance()
- if chunk_c == 0:
- tensormap_manager.fence_tensormap_update(tm_a_ptr)
- cute.copy(
- tma_a[0],
- tAgA[(None, 0, 0)],
- tAsA[(None, ah.index)],
- tma_bar_ptr=ah.barrier,
- tma_desc_ptr=tensormap_manager.get_tensormap_ptr(
- tm_a_ptr, cute.AddressSpace.generic
- ),
- )
-
- # KS: B-operand (N, T, H) → domain_offset on (0, tokens)
- mB_c = cute.domain_offset(
- (Int32(0), chunk_offset), tma_b[1][None, None, head_idx_t]
- )
- gB_c = cute.flat_divide(mB_c, (N, K))
- tCgB = thr_kmn.partition_B(gB_c)
- tBsB, tBgB = cpasync.tma_partition(
- tma_b[0],
- 0,
- cta_layout,
- cute.group_modes(sB, 0, 3),
- cute.group_modes(tCgB, 0, 3),
- )
- bh = b_prod.acquire_and_advance()
- if chunk_c == 0:
- tensormap_manager.fence_tensormap_update(tm_b_ptr)
- cute.copy(
- tma_b[0],
- tBgB[(None, 0, 0)],
- tBsB[(None, bh.index)],
- tma_bar_ptr=bh.barrier,
- tma_desc_ptr=tensormap_manager.get_tensormap_ptr(
- tm_b_ptr, cute.AddressSpace.generic
- ),
- )
-
- # V: B-operand (N, T, H) → domain_offset on (0, tokens)
- mV_c = cute.domain_offset(
- (Int32(0), chunk_offset), tma_v[1][None, None, head_idx_t]
- )
- gV_c = cute.flat_divide(mV_c, (N, K))
- tCgV = thr_kmn.partition_B(gV_c)
- tBsV, tBgV = cpasync.tma_partition(
- tma_v[0],
- 0,
- cta_layout,
- cute.group_modes(sV, 0, 3),
- cute.group_modes(tCgV, 0, 3),
- )
- vh = v_prod.acquire_and_advance()
- if chunk_c == 0:
- tensormap_manager.fence_tensormap_update(tm_v_ptr)
- cute.copy(
- tma_v[0],
- tBgV[(None, 0, 0)],
- tBsV[(None, vh.index)],
- tma_bar_ptr=vh.barrier,
- tma_desc_ptr=tensormap_manager.get_tensormap_ptr(
- tm_v_ptr, cute.AddressSpace.generic
- ),
- )
-
- # QS: A-operand (T, N, H) → domain_offset on (tokens, 0)
- mQ_c = cute.domain_offset(
- (chunk_offset, Int32(0)), tma_q[1][None, None, head_idx_t]
- )
- gQ_c = cute.flat_divide(mQ_c, (M, K3))
- tCgQ = thr_kmn.partition_A(gQ_c)
- tAsQ, tAgQ = cpasync.tma_partition(
- tma_q[0],
- 0,
- cta_layout,
- cute.group_modes(sQ, 0, 3),
- cute.group_modes(tCgQ, 0, 3),
- )
- qh = q_prod.acquire_and_advance()
- if chunk_c == 0:
- tensormap_manager.fence_tensormap_update(tm_q_ptr)
- cute.copy(
- tma_q[0],
- tAgQ[(None, 0, 0)],
- tAsQ[(None, qh.index)],
- tma_bar_ptr=qh.barrier,
- tma_desc_ptr=tensormap_manager.get_tensormap_ptr(
- tm_q_ptr, cute.AddressSpace.generic
- ),
- )
-
- # AQC: A-operand (T, K, H) → domain_offset on (tokens, 0)
- mAQC_c = cute.domain_offset(
- (chunk_offset, Int32(0)), tma_aqc[1][None, None, head_idx_t]
- )
- gAQC_c = cute.flat_divide(mAQC_c, (M, K))
- tCgAQC = thr_kmn.partition_A(gAQC_c)
- tAsAQC, tAgAQC = cpasync.tma_partition(
- tma_aqc[0],
- 0,
- cta_layout,
- cute.group_modes(sAQC, 0, 3),
- cute.group_modes(tCgAQC, 0, 3),
- )
- aqch = aqc_prod.acquire_and_advance()
- if chunk_c == 0:
- tensormap_manager.fence_tensormap_update(tm_aqc_ptr)
- cute.copy(
- tma_aqc[0],
- tAgAQC[(None, 0, 0)],
- tAsAQC[(None, aqch.index)],
- tma_bar_ptr=aqch.barrier,
- tma_desc_ptr=tensormap_manager.get_tensormap_ptr(
- tm_aqc_ptr, cute.AddressSpace.generic
- ),
- )
-
- # KG: B-operand (N6, T, H) → domain_offset on (0, tokens)
- mKG_c = cute.domain_offset(
- (Int32(0), chunk_offset), tma_kg[1][None, None, head_idx_t]
- )
- gKG_c = cute.flat_divide(mKG_c, (N6, K6))
- tCgKG = thr_mn.partition_B(gKG_c)
- tBsKG, tBgKG = cpasync.tma_partition(
- tma_kg[0],
- 0,
- cta_layout,
- cute.group_modes(sKG, 0, 3),
- cute.group_modes(tCgKG, 0, 3),
- )
- kgh = kg_prod.acquire_and_advance()
- if chunk_c == 0:
- tensormap_manager.fence_tensormap_update(tm_kg_ptr)
- cute.copy(
- tma_kg[0],
- tBgKG[(None, 0, 0)],
- tBsKG[(None, kgh.index)],
- tma_bar_ptr=kgh.barrier,
- tma_desc_ptr=tensormap_manager.get_tensormap_ptr(
- tm_kg_ptr, cute.AddressSpace.generic
- ),
- )
-
- scheduler.advance_to_next_work()
- work = scheduler.get_current_work()
-
- # ==== Warp 1: O TMA store ====
- elif warp_idx == O_STORE_WARP:
- cta_layout_o = cute.make_layout(1)
-
- scheduler = GDNTileScheduler.create(scheduler_params, (bidx, bidy, bidz), grid_dim)
- work = scheduler.initial_work_tile_info()
-
- if work.is_valid_tile:
- tensormap_manager.init_tensormap_from_atom(tma_o[0], tm_o_ptr, O_STORE_WARP)
- tensormap_manager.fence_tensormap_initialization()
-
- while work.is_valid_tile:
- batch_idx_o, head_idx_o, _ = work.tile_idx
- batch_start_o = cu_seqlens[batch_idx_o]
- batch_end_o = cu_seqlens[batch_idx_o + 1]
- num_chunks_o = cute.ceil_div(batch_end_o - batch_start_o, M)
-
- bounded_o = cute.make_tensor(
- mO.iterator,
- cute.make_layout(
- (batch_end_o, mO.shape[1], mO.shape[2]),
- stride=(mO.stride[0], mO.stride[1], mO.stride[2]),
- ),
- )
- tensormap_manager.update_tensormap(
- (bounded_o,), (tma_o[0],), (tm_o_ptr,), O_STORE_WARP, (None,)
- )
- tensormap_manager.fence_tensormap_update(tm_o_ptr)
-
- for chunk_c in cutlass.range(num_chunks_o):
- os_h = o_store_cons.wait_and_advance()
- chunk_offset_o = batch_start_o + chunk_c * M
- mO_c = cute.domain_offset(
- (chunk_offset_o, Int32(0)), tma_o[1][None, None, head_idx_o]
- )
- gOo = cute.flat_divide(mO_c, (M, N))
- sOt, gOt = cpasync.tma_partition(
- tma_o[0],
- 0,
- cta_layout_o,
- cute.group_modes(sO_out_st, 0, 2),
- cute.group_modes(gOo, 0, 2),
- )
- cute.copy(
- tma_o[0],
- sOt[None],
- gOt[(None, 0, 0)],
- tma_desc_ptr=tensormap_manager.get_tensormap_ptr(
- tm_o_ptr, cute.AddressSpace.generic
- ),
- )
- cute.arch.cp_async_bulk_commit_group()
- cute.arch.cp_async_bulk_wait_group(0, read=True)
- os_h.release()
-
- scheduler.advance_to_next_work()
- work = scheduler.get_current_work()
-
- # ==== WG2: W/NV/O readout (SS-mode, no Phase 2) ====
- elif warpgroup_idx == READOUT_WG:
- tCtW_mn = transform_partitioned_tensor_layout(tCtW)
- tCtNV_mn = transform_partitioned_tensor_layout(tCtNV)
- tCtO_mn = transform_partitioned_tensor_layout(tCtO)
-
- atom_t2r = cute.make_copy_atom(tcgen05.Ld16x256bOp(tcgen05.Repetition(1)), acc_dtype)
- tiled_t2r = tcgen05.make_tmem_copy(atom_t2r, tCtW[(None, None), 0, 0])
- thr_t2r = tiled_t2r.get_slice(warpgroup_tidx)
-
- tTR_W = thr_t2r.partition_S(tCtW_mn)
- tTR_NV = thr_t2r.partition_S(tCtNV_mn)
- tTR_O = thr_t2r.partition_S(tCtO_mn)
-
- atom_r2s_k = sm100_utils.get_smem_store_op(
- utils.LayoutEnum.ROW_MAJOR, mma_dtype, acc_dtype, tiled_t2r
- )
- tiled_r2s_k = cute.make_tiled_copy_D(atom_r2s_k, tiled_t2r)
- thr_r2s_k = tiled_r2s_k.get_slice(warpgroup_tidx)
- tCsO = thr_r2s_k.partition_D(transform_partitioned_tensor_layout(sO))
- tCsO_out = thr_r2s_k.partition_D(transform_partitioned_tensor_layout(sO_out))
- tCsNV = thr_r2s_k.partition_D(transform_partitioned_tensor_layout(sNV))
-
- cId = cute.make_identity_tensor((M, N))
- tTR_cId = thr_t2r.partition_D(cId)
-
- scheduler = GDNTileScheduler.create(scheduler_params, (bidx, bidy, bidz), grid_dim)
- work = scheduler.initial_work_tile_info()
- global_chunk = Int32(0)
-
- while work.is_valid_tile:
- batch_idx_r, head_idx_r, _ = work.tile_idx
- num_chunks_r = cute.ceil_div(cu_seqlens[batch_idx_r + 1] - cu_seqlens[batch_idx_r], M)
-
- for chunk_c in cutlass.range(num_chunks_r):
- c_phase = global_chunk % 2
- tRrR = cute.make_rmem_tensor_like(tTR_cId, acc_dtype)
- tRrR_out = cute.make_rmem_tensor_like(tRrR, mma_dtype)
- tCrR_k = tiled_r2s_k.retile(tRrR_out)
- num_subs = tRrR.shape[2]
-
- wh = w_cons.wait_and_advance()
- for sub in cutlass.range(num_subs):
- cute.copy(tiled_t2r, tTR_W[None, 0, sub], tRrR[None, 0, sub])
- # Negate W in its readout so MMA3 accumulates NV = U - W @ S
- # (FLA: b_v = u - w @ h). Folded into the bf16 cast — free.
- tRrR_out[None, 0, sub].store((-(tRrR[None, 0, sub].load())).to(mma_dtype))
- cute.copy(tiled_r2s_k, tCrR_k[None, 0, sub], tCsO[None, 0, sub, 0])
- cute.arch.fence_view_async_tmem_load()
- wh.release()
- cute.arch.fence_view_async_shared()
- sW_ready_nbar.arrive_and_wait()
-
- nvh = nv_cons.wait_and_advance()
- for sub in cutlass.range(num_subs):
- cute.copy(tiled_t2r, tTR_NV[None, 0, sub], tRrR[None, 0, sub])
- tRrR_out[None, 0, sub].store(tRrR[None, 0, sub].load().to(mma_dtype))
- cute.copy(tiled_r2s_k, tCrR_k[None, 0, sub], tCsNV[None, 0, sub, 0])
- cute.arch.fence_view_async_tmem_load()
- nvh.release()
- cute.arch.fence_view_async_shared()
- sNV_ready_nbar.arrive_and_wait()
-
- os_h = o_store_prod.acquire_and_advance()
-
- oh = o_cons.wait_and_advance()
- for sub in cutlass.range(num_subs):
- cute.copy(tiled_t2r, tTR_O[None, 0, sub], tRrR[None, 0, sub])
- tRrR_out[None, 0, sub].store(tRrR[None, 0, sub].load().to(mma_dtype))
- cute.copy(tiled_r2s_k, tCrR_k[None, 0, sub], tCsO_out[None, 0, sub, 0])
- cute.arch.fence_view_async_tmem_load()
- oh.release()
- cute.arch.fence_view_async_shared()
- os_h.commit()
- global_chunk = global_chunk + 1
-
- scheduler.advance_to_next_work()
- work = scheduler.get_current_work()
- o_store_prod.tail()
-
-
-def make_host_fn(num_sm=148):
- """Create one K4 host function for all runtime context batch sizes.
-
- ``num_seqs`` is a runtime scalar. ``H`` and the token extent come from
- dynamic tensor layouts, while ``num_sm`` is the fixed persistent-grid
- ceiling. The scheduler computes ``min(num_seqs * H, num_sm)`` at launch.
- """
- _num_sm = num_sm
-
- @cute.jit
- def host_fn(
- a_raw: cute.Tensor,
- b_raw: cute.Tensor,
- v_raw: cute.Tensor,
- q_raw: cute.Tensor,
- aqc_raw: cute.Tensor,
- kg_raw: cute.Tensor,
- o_raw: cute.Tensor,
- gk_last_exp: cute.Tensor,
- s_fp32: cute.Tensor,
- state_indices: cute.Tensor,
- cu_seqlens: cute.Tensor,
- chunk_offsets: cute.Tensor,
- tm_workspace: cute.Tensor,
- num_seqs: cutlass.Int32,
- # Launch stream — runtime argument; launching on the DSL default
- # stream races with the executor's non-blocking execution stream.
- # Same stream as akk_inv, so the PDL chain stays intact.
- stream: cuda.CUstream,
- ):
- # Raw tensors from from_dlpack have PyTorch layout (T, H, dim)
- # stride (H*dim, dim, 1). Reshape to 3D for TMA:
- # A-operands: (T, dim, H), stride (dim*H, 1, dim)
- # B-operands: (dim, T, H), stride (1, dim*H, dim)
- T_tok = a_raw.shape[0]
- H = a_raw.shape[1]
-
- a = cute.make_tensor(
- a_raw.iterator,
- cute.make_layout(
- (T_tok, a_raw.shape[2], H),
- stride=(a_raw.stride[0], a_raw.stride[2], a_raw.stride[1]),
- ),
- )
- q = cute.make_tensor(
- q_raw.iterator,
- cute.make_layout(
- (T_tok, q_raw.shape[2], H),
- stride=(q_raw.stride[0], q_raw.stride[2], q_raw.stride[1]),
- ),
- )
- aqc = cute.make_tensor(
- aqc_raw.iterator,
- cute.make_layout(
- (T_tok, aqc_raw.shape[2], H),
- stride=(aqc_raw.stride[0], aqc_raw.stride[2], aqc_raw.stride[1]),
- ),
- )
- o_out = cute.make_tensor(
- o_raw.iterator,
- cute.make_layout(
- (T_tok, o_raw.shape[2], H),
- stride=(o_raw.stride[0], o_raw.stride[2], o_raw.stride[1]),
- ),
- )
-
- b = cute.make_tensor(
- b_raw.iterator,
- cute.make_layout(
- (b_raw.shape[2], T_tok, H),
- stride=(b_raw.stride[2], b_raw.stride[0], b_raw.stride[1]),
- ),
- )
- v = cute.make_tensor(
- v_raw.iterator,
- cute.make_layout(
- (v_raw.shape[2], T_tok, H),
- stride=(v_raw.stride[2], v_raw.stride[0], v_raw.stride[1]),
- ),
- )
- kg = cute.make_tensor(
- kg_raw.iterator,
- cute.make_layout(
- (kg_raw.shape[2], T_tok, H),
- stride=(kg_raw.stride[2], kg_raw.stride[0], kg_raw.stride[1]),
- ),
- )
-
- tile1 = (M, N, K)
- tile3 = (M, N, K3)
- tile6 = (M6, N6, K6)
-
- mma_kmn = sm100_utils.make_trivial_tiled_mma(
- mma_dtype,
- OperandMajorMode.K,
- OperandMajorMode.MN,
- acc_dtype,
- tcgen05.CtaGroup.ONE,
- (M, N),
- OperandSource.SMEM,
- )
- mma_mn_mn = sm100_utils.make_trivial_tiled_mma(
- mma_dtype,
- OperandMajorMode.MN,
- OperandMajorMode.MN,
- acc_dtype,
- tcgen05.CtaGroup.ONE,
- (M6, N6),
- OperandSource.SMEM,
- )
- sl_a = sm100_utils.make_smem_layout_a(mma_kmn, tile1, mma_dtype, 1)
- sl_b = sm100_utils.make_smem_layout_b(mma_kmn, tile1, mma_dtype, 1)
- sl_v = sm100_utils.make_smem_layout_b(mma_kmn, tile1, mma_dtype, 1)
- sl_s = sm100_utils.make_smem_layout_b(mma_kmn, tile3, mma_dtype, 1)
- sl_q = sm100_utils.make_smem_layout_a(mma_kmn, tile3, mma_dtype, 1)
- sl_aqc = sm100_utils.make_smem_layout_a(mma_kmn, tile1, mma_dtype, 1)
- sl_kg = sm100_utils.make_smem_layout_b(mma_mn_mn, tile6, mma_dtype, 1)
-
- sl_readout_k = sm100_utils.make_smem_layout_a(mma_kmn, tile3, mma_dtype, 1)
- sl_nv_b = sm100_utils.make_smem_layout_b(mma_kmn, tile1, mma_dtype, 1)
- sl_nv_a = sm100_utils.make_smem_layout_a(mma_mn_mn, tile6, mma_dtype, 1)
- sl_kg_a = sm100_utils.make_smem_layout_a(mma_mn_mn, tile6, mma_dtype, 1)
- sl_nv_b_mn = sm100_utils.make_smem_layout_b(mma_mn_mn, tile6, mma_dtype, 1)
-
- tma_ld = cpasync.CopyBulkTensorTileG2SOp()
-
- ta_a = cute.nvgpu.make_tiled_tma_atom_A(
- tma_ld, a, cute.select(sl_a, mode=[0, 1, 2]), tile1, mma_kmn
- )
- ta_b = cute.nvgpu.make_tiled_tma_atom_B(
- tma_ld, b, cute.select(sl_b, mode=[0, 1, 2]), tile1, mma_kmn
- )
- ta_v = cute.nvgpu.make_tiled_tma_atom_B(
- tma_ld, v, cute.select(sl_v, mode=[0, 1, 2]), tile1, mma_kmn
- )
- ta_q = cute.nvgpu.make_tiled_tma_atom_A(
- tma_ld, q, cute.select(sl_q, mode=[0, 1, 2]), tile3, mma_kmn
- )
- ta_aqc = cute.nvgpu.make_tiled_tma_atom_A(
- tma_ld, aqc, cute.select(sl_aqc, mode=[0, 1, 2]), tile1, mma_kmn
- )
- ta_kg = cute.nvgpu.make_tiled_tma_atom_B(
- tma_ld, kg, cute.select(sl_kg, mode=[0, 1, 2]), tile6, mma_mn_mn
- )
- sk = sm100_utils.get_smem_layout_atom_ab(OperandMajorMode.K, out_dtype, (M, N))
- sl_store = cute.tile_to_shape(
- sm100_utils.make_smem_layout_atom(sk, out_dtype), (M, N), order=(0, 1)
- )
- tma_st = cpasync.CopyBulkTensorTileS2GOp()
- ta_o = cpasync.make_tiled_tma_atom(tma_st, o_out, sl_store, (M, N))
-
- scheduler_params = GDNTileSchedulerParams(
- num_seqs=num_seqs,
- num_q_heads=H,
- num_v_heads=H,
- is_GQA=False,
- is_persistent=True,
- )
- grid_shape = GDNTileScheduler.get_grid_shape(scheduler_params, _num_sm)
-
- k4_persistent_kernel(
- mma_kmn,
- mma_mn_mn,
- ta_a,
- sl_a,
- ta_b,
- sl_b,
- ta_v,
- sl_v,
- sl_s,
- ta_q,
- sl_q,
- ta_aqc,
- sl_aqc,
- ta_kg,
- sl_kg,
- ta_o,
- sl_store,
- sl_readout_k,
- sl_nv_b,
- sl_nv_a,
- sl_kg_a,
- sl_nv_b_mn,
- gk_last_exp,
- s_fp32,
- state_indices,
- cu_seqlens,
- chunk_offsets,
- a,
- b,
- v,
- q,
- aqc,
- kg,
- o_out,
- tm_workspace,
- scheduler_params,
- ).launch(grid=grid_shape, block=(threads_per_cta, 1, 1), use_pdl=True, stream=stream)
-
- return host_fn
diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_chunked_bt16.py b/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_chunked_bt16.py
new file mode 100644
index 000000000000..38903d429cbc
--- /dev/null
+++ b/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_chunked_bt16.py
@@ -0,0 +1,10836 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""KDA chunked BT=16 forward kernel using direct CUTLASS primitives.
+
+The current body implements the chunk-size 16 KDA schedule:
+
+ load q/k/v/gate
+ L2-normalize q/k
+ exp2(g), exp2(-g), stage final-token exp2(g) as exp2(g_last)
+ auxiliary MMA warp: kk/qk/blockwise inverse + apply beta
+ tcgen05-MMA: state*k/state*q/new-v/kv-update/qkv
+ store o
+
+Inputs use a packed variable-length ABI: q/k/v/gate/BF16 beta logits
+have a singleton batch dimension and `cu_seqlens` partitions their token
+dimension. Initial and final states retain one batch entry per logical sequence
+and may use BF16 or FP32 storage. State presence (initial / final / periodic
+checkpoints) is a compile-time specialization, but the returned callable
+DERIVES it per call from which state tensors are actually passed (`None` vs a
+real tensor) and transparently uses the matching build -- the `has_state_*`
+compile flags only choose which build is compiled eagerly. The gate input dtype
+is a compile-time parameter `gate_dtype`; it DEFAULTS to BF16, matching
+FlashKDA's published ABI ("g: Gate before activation, bf16, shape [B,T,H,K]"),
+so every tensor of the default ABI agrees with FlashKDA (q/k/v/beta/out BF16,
+A_log/dt_bias FP32). FP32 gates remain fully supported by passing
+`gate_dtype=cutlass.Float32`; that build is bit-for-bit the historical kernel.
+Only the MEMORY FORMAT of the gate changes: a 16-bit raw logit is widened to
+FP32 on the SMEM read and the whole activation/prefix/exp2 chain stays FP32.
+The real branch point is 32-bit vs 16-bit, not one dtype per path: under FP32
+the `exp2(g_prefix)` exchange tile stays fused inside the raw-gate SMEM ring
+(which is FP32 and therefore big enough), while a 16-bit gate shrinks that ring
+to a 16-bit landing buffer and gives the FP32 exchange its own 4-deep ring
+(`gate_exchange_smem`). Scale is a runtime FP32 argument. The
+torch reference below is the numerical contract used by the verify path.
+
+On top of the plain engine kernel this file adds a two-kernel decomposition of
+the same math. A chunk-parallel factor-prep kernel (kernel 1) computes the
+rank-16 chunk factors from q/k/gate/beta and writes them, pre-permuted, to a
+user-allocated workspace; a serial chain kernel (kernel 2, DV2 split: two CTAs
+per sequence-head, each owning half the value dimension) runs the state
+recursion, reading those factors by TMA. Decomp always runs PREP-FIRST on one
+stream -- k1 completes its full grid, then k2 reads a workspace k1 has fully
+written, so the kernel boundary is the only synchronization. A single host
+`compile()` returns one callable that routes the plain engine or the
+decomposition per launch shape by one occupancy rule (decomp iff its doubled k2
+grid fits one wave: `n_seq * heads * 2 <= sm_count`); the decomp route reads an
+opaque workspace whose byte size is queried with `workspace_size()` (0 for the
+engine route). All device kernels keep the packed varlen ABI.
+
+Optimization notes (each expanded at its definition site):
+
+* Beta TMA transport -- token-major beta makes the per-chunk beta read a
+ strided gather (16 lanes at `heads * 2B`; up to 16 distinct 32B sectors where
+ a head-major layout touches 1, ~1us per extra sector at H96). A 2D TMA
+ descriptor over the same memory moves the gather onto the TMA engine, where
+ it is free; `g = 8 / gcd(heads, 8)` pair-packs rows so the descriptor stride
+ stays a legal 16B multiple for EVERY head count (H12 = the TP8 shard of H96
+ uses g=2). See the constants block above `tma_stage_load_inputs`.
+* Prep-first decomposition -- a co-resident overlap of k1/k2 was measured and
+ removed: it forces a per-chunk flag/fence ring onto both kernels and caps
+ k1's grid to the leftover SMs, which together cost more than the concurrency
+ returns. Sequential prep runs its full grid with zero cross-kernel
+ synchronization. See the route rule in `host_unified`.
+* K1_CPC = 4 -- each prep CTA walks four contiguous chunks (measured ladder);
+ amortizes launch/preamble and enables mid-chunk prefetch of the next chunk's
+ raw tiles.
+* Recurrent state lives in TMEM as an fp32 accumulator across the whole chunk
+ walk; the final store, the optional periodic checkpoints, and the DV2 halves
+ all drain the same accumulator (`tcgen05_store_final_state_tmem`).
+
+State checkpoints (opt-in): `compile(has_state_ckpt=True)` returns a callable
+taking `state_ckpt`, `checkpoint_cu_starts`, and `ckpt_interval` (a positive
+multiple of BT). Without `checkpoint_state_indices`, each sequence stores its
+initial state first, followed by states at interior interval boundaries. With
+indices, completed boundaries (including an aligned final boundary, excluding
+the initial state) are written directly to the selected rows of `state_ckpt`.
+
+Reproducing the results::
+
+ # single-shape correctness vs the torch MMA reference below
+ python kda_chunked_bt16.py --batch 1 --heads 32 --seqlen 8192
+
+ # timing protocol used for the reported numbers: engine route =
+ # CUPTI/Kineto device-time sum; decomp route = back-to-back wall clock
+ # with one synchronize at the end (per-launch syncs both add ~10us of
+ # launch overhead per sample and serialize the two kernels' tail).
+"""
+
+import argparse
+import weakref
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Callable, Protocol, cast
+
+import torch
+import cuda.bindings.driver as cuda_driver
+
+import cutlass
+import cutlass.experimental.cuda as cuda
+import cutlass.experimental.primitives as prims
+import cutlass.cute as cute
+from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream
+
+
+LOG2_E: float = 1.4426950408889634
+DEFAULT_GATE_LOWER_BOUND: float = -5.0
+
+
+class CompiledKDA(Protocol):
+ """Callable returned by ``compile`` with its workspace query attached."""
+
+ workspace_size: Callable[..., int]
+ workspace_size_from_total_chunks: Callable[..., int]
+
+ def __call__(self, *args, **kwargs) -> None: ...
+
+
+def packed_f32x2_binary(
+ op: Callable,
+ lhs: tuple[cutlass.Float32, cutlass.Float32],
+ rhs: tuple[cutlass.Float32, cutlass.Float32],
+) -> tuple[cutlass.Float32, cutlass.Float32]:
+ """Apply a CUTLASS packed-FP32 primitive to two scalar pairs."""
+
+ lhs_vec = cutlass.Vector.from_elements(lhs, cutlass.Float32)
+ rhs_vec = cutlass.Vector.from_elements(rhs, cutlass.Float32)
+ result = op(lhs_vec, rhs_vec, ftz=False, rnd="rn")
+ return cutlass.Float32(result[0]), cutlass.Float32(result[1])
+
+
+def fadd2(lhs, rhs):
+ return packed_f32x2_binary(prims.add_packed_f32x2, lhs, rhs)
+
+
+def fmul2(lhs, rhs):
+ return packed_f32x2_binary(prims.mul_packed_f32x2, lhs, rhs)
+
+
+@cute.jit
+def movmatrix_b16(value: cutlass.Int32) -> cutlass.Int32:
+ """Transpose one packed m8n8 b16 register fragment."""
+
+ return prims.inline_ptx_hl(
+ "movmatrix.sync.aligned.m8n8.trans.b16 {$w0}, {$r0};",
+ write_only_types=[cutlass.Int32],
+ read_only_args=[value],
+ )
+
+
+@cute.jit
+def mul_f16x2(value: cutlass.Int32, scale: cutlass.Int32) -> cutlass.Int32:
+ """Multiply two packed FP16 pairs."""
+
+ return prims.inline_ptx_hl(
+ "mul.f16x2 {$w0}, {$r0}, {$r1};",
+ write_only_types=[cutlass.Int32],
+ read_only_args=[value, scale],
+ )
+
+
+@cute.jit
+def sub_b16x2_input_dtype(
+ lhs: cutlass.Int32,
+ rhs: cutlass.Int32,
+ input_dtype: cutlass.Constexpr,
+) -> cutlass.Int32:
+ """Subtract two packed pairs using the compile-time input dtype."""
+
+ if cutlass.const_expr(input_dtype is cutlass.BFloat16):
+ return prims.inline_ptx_hl(
+ "sub.bf16x2 {$w0}, {$r0}, {$r1};",
+ write_only_types=[cutlass.Int32],
+ read_only_args=[lhs, rhs],
+ )
+ return prims.inline_ptx_hl(
+ "sub.f16x2 {$w0}, {$r0}, {$r1};",
+ write_only_types=[cutlass.Int32],
+ read_only_args=[lhs, rhs],
+ )
+
+
+@cute.jit
+def mul_b16x2_input_dtype(
+ lhs: cutlass.Int32,
+ rhs: cutlass.Int32,
+ input_dtype: cutlass.Constexpr,
+) -> cutlass.Int32:
+ """Multiply two packed pairs using the compile-time input dtype."""
+
+ if cutlass.const_expr(input_dtype is cutlass.BFloat16):
+ return prims.mul_bf16x2(lhs, rhs)
+ return mul_f16x2(lhs, rhs)
+
+
+@cute.jit
+def safe_gate_log2_increment_prehalved(
+ half_scaled_gate: cutlass.Float32,
+ SAFE_GATE: cutlass.Constexpr,
+ GATE_SCALE_LOG2: cutlass.Constexpr,
+) -> cutlass.Float32:
+ """Apply the safe-gate transform after folding its exact half scales."""
+
+ if cutlass.const_expr(SAFE_GATE):
+ half_scale = cutlass.Float32(GATE_SCALE_LOG2 * 0.5)
+ tanh_value = cute.math.tanh(half_scaled_gate, approx=True)
+ return tanh_value * half_scale + half_scale
+ return half_scaled_gate
+
+
+@cute.jit
+def softplus_log2_f32(value: cutlass.Float32) -> cutlass.Float32:
+ """Compute ``softplus(value) * log2(e)`` with predicated MUFU ops.
+
+ Evaluates the FLA non-safe activation directly in the log2 domain so the
+ prefix scan never leaves base-2 units. Matches Triton's ``x > 20``
+ overflow guard: above the threshold ``softplus(x) ~= x`` and the result is
+ just ``x * log2(e)`` (the ``@!p`` ops are skipped).
+ """
+
+ return prims.inline_ptx_hl(
+ """
+ {
+ .reg .pred p;
+ setp.gt.f32 p, {$r0}, 20.0;
+ mul.f32 {$w0}, {$r0}, 1.4426950408889634;
+ @!p ex2.approx.ftz.f32 {$w0}, {$w0};
+ @!p add.f32 {$w0}, {$w0}, 1.0;
+ @!p lg2.approx.ftz.f32 {$w0}, {$w0};
+ }
+ """,
+ write_only_types=[cutlass.Float32],
+ read_only_args=[value],
+ )
+
+
+@cute.jit
+def pack_input_b16x2_to_i32(
+ value0: cutlass.Float32,
+ value1: cutlass.Float32,
+ input_dtype: cutlass.Constexpr,
+):
+ """Pack two FP32 values through the compile-time input 16-bit dtype."""
+
+ return (
+ cutlass.Vector.from_elements(
+ (value0, value1),
+ cutlass.Float32,
+ )
+ .to(input_dtype)
+ .bitcast(cutlass.Int32)[0]
+ )
+
+
+# ---------------------------------------------------------------------------
+# Fixed KDA BT=16 shape
+# ---------------------------------------------------------------------------
+
+BT: int = 16
+DK: int = 128
+DV: int = 128
+DEFAULT_SEQLEN: int = 8192
+DEFAULT_BATCH: int = 1
+DEFAULT_HEADS: int = 32
+DEFAULT_SCALE: float = 1.0 / (DK**0.5)
+L2_NORM_EPS: float = 1.0e-12
+VERIFY_RTOL: float = 1.0e-2
+VERIFY_ATOL: float = 1.0e-3
+
+THREADS_PER_WARP: int = 32
+THREADS_PER_CTA: int = 16 * THREADS_PER_WARP
+CG0_GROUP_COUNT: int = 2
+CG0_WARPS_PER_GROUP: int = 4
+CG0_THREADS_PER_GROUP: int = CG0_WARPS_PER_GROUP * THREADS_PER_WARP
+NBAR_CG0_GROUP0_ID: int = 1
+TMEM_USER_WARP_COUNT: int = 5
+TMEM_USER_THREADS: int = TMEM_USER_WARP_COUNT * THREADS_PER_WARP
+NBAR_TMEM_LIFECYCLE_ID: int = 2
+NBAR_CG0_GROUP1_ID: int = 3
+# Third CG0 producer group (kernel 1 runs warps 8-11 as group 2).
+NBAR_CG0_GROUP2_ID: int = 4
+KDA_CG0_REGS: int = 160
+KDA_CG1_REGS: int = 136
+KDA_SERVICE_REGS: int = 56
+TCGEN05_VALID_ALLOC_COLS: tuple[int, ...] = (32, 64, 128, 256, 512)
+TCGEN05_F16_K_ATOM: int = 16
+TCGEN05_F16_ELEM_BYTES: int = 2
+TCGEN05_F16_A_TMEM_PAIR_XOR: int = 4
+TCGEN05_SW128_BYTES: int = 128
+TCGEN05_SW128_K_PHASES_PER_SLICE: int = 4
+# `state` and decay operands use K-box-major SW128 staging. K-major F16 SMEM
+# descriptors use a 16B leading offset and 1024B stride; the decay store applies
+# a row-group key xor so tcgen05 B reads logical `[DK, BT]` correctly.
+TCGEN05_STATE_K_B_LEADING_BYTES: int = 16
+TCGEN05_STATE_K_B_STRIDE_BYTES: int = 1024
+TCGEN05_STATE_K_B_K_STEP_BYTES: int = TCGEN05_F16_K_ATOM * TCGEN05_F16_ELEM_BYTES
+TCGEN05_SW32_BT_HALF_XOR: int = BT // 2
+PAIRWISE_SW32_ROW_STRIDE: int = BT
+PAIRWISE_SW32_COL_XOR: int = TCGEN05_SW32_BT_HALF_XOR
+PAIRWISE_SW32_TILE_ELEMS: int = BT * BT
+TCGEN05_TMEM_LOAD_COLS: int = BT
+TCGEN05_STATE_INPUT_LOAD_COLS: int = 16
+TCGEN05_STATE_INPUT_PACKED_COLS: int = TCGEN05_STATE_INPUT_LOAD_COLS // 2
+TCGEN05_STATE_K_TMEM_ROW_BLOCKS: int = DV // THREADS_PER_WARP
+
+
+def _tcgen05_accumulator_tmem_cols(n_dim: int) -> int:
+ """Return TMEM columns for an FP32 `[128, n_dim]` accumulator tile."""
+
+ if n_dim <= 0:
+ raise ValueError(f"n_dim must be positive, got {n_dim}")
+ if n_dim % 8 != 0:
+ raise ValueError(f"n_dim must be a multiple of 8, got {n_dim}")
+ return n_dim
+
+
+def _tcgen05_f16_input_tmem_cols(k_dim: int) -> int:
+ """Return TMEM columns for an F16 `[128, k_dim]` A-input staging tile."""
+
+ if k_dim <= 0:
+ raise ValueError(f"k_dim must be positive, got {k_dim}")
+ if k_dim % 2 != 0:
+ raise ValueError(f"k_dim must be even for packed F16 TMEM, got {k_dim}")
+ return k_dim // 2
+
+
+def _tcgen05_allocation_tmem_cols(required_cols: int) -> int:
+ """Round required TMEM columns up to a tcgen05.alloc-supported count."""
+
+ for alloc_cols in TCGEN05_VALID_ALLOC_COLS:
+ if required_cols <= alloc_cols:
+ return alloc_cols
+ raise ValueError(f"required_cols must be <= 512, got {required_cols}")
+
+
+# Value-side TMEM layout for the BT=16 KDA schedule. Gate TMEM is intentionally
+# omitted until that path is implemented. The state slot holds the recurrent VK
+# state across chunks and is reused in-place as the final_state accumulator late
+# in each chunk. The two shared_input slots stage F16 A operands, while
+# shared_acc is the runtime-selected accumulator pool used by state*k and update.
+#
+KDA_TMEM_N16_ACC_COLS: int = _tcgen05_accumulator_tmem_cols(BT)
+KDA_TMEM_N128_ACC_COLS: int = _tcgen05_accumulator_tmem_cols(DK)
+KDA_TMEM_STATE_COLS: int = KDA_TMEM_N128_ACC_COLS
+KDA_TMEM_STATE_AS_INPUT_COLS: int = _tcgen05_f16_input_tmem_cols(DK)
+KDA_TMEM_SHARED_INPUT_COLS: int = _tcgen05_f16_input_tmem_cols(BT)
+KDA_TMEM_SHARED_INPUT_STAGE_COUNT: int = 2
+KDA_TMEM_QSTATE_ACC_STAGE_COUNT: int = 2
+KDA_TMEM_SHARED_ACC_STAGE_COUNT: int = 2
+
+
+KDA_TMEM_STATE_COL_OFFSET: int = 0
+KDA_TMEM_FINAL_STATE_ACC_COL_OFFSET: int = KDA_TMEM_STATE_COL_OFFSET
+KDA_TMEM_STATE_AS_INPUT_COL_OFFSET: int = (
+ KDA_TMEM_STATE_COL_OFFSET + KDA_TMEM_STATE_COLS
+)
+KDA_TMEM_SHARED_INPUT_COL_OFFSET: int = (
+ KDA_TMEM_STATE_AS_INPUT_COL_OFFSET + KDA_TMEM_STATE_AS_INPUT_COLS
+)
+KDA_TMEM_QSTATE_ACC_COL_OFFSET: int = (
+ KDA_TMEM_SHARED_INPUT_COL_OFFSET
+ + KDA_TMEM_SHARED_INPUT_STAGE_COUNT * KDA_TMEM_SHARED_INPUT_COLS
+)
+KDA_TMEM_SHARED_ACC_COL_OFFSET: int = (
+ KDA_TMEM_QSTATE_ACC_COL_OFFSET + KDA_TMEM_N16_ACC_COLS
+)
+KDA_TMEM_QSTATE_ACC_STAGE1_COL_OFFSET: int = (
+ KDA_TMEM_SHARED_ACC_COL_OFFSET
+ + KDA_TMEM_SHARED_ACC_STAGE_COUNT * KDA_TMEM_N16_ACC_COLS
+)
+KDA_TMEM_QSTATE_ACC_STAGE_STRIDE_COLS: int = (
+ KDA_TMEM_QSTATE_ACC_STAGE1_COL_OFFSET - KDA_TMEM_QSTATE_ACC_COL_OFFSET
+)
+
+
+KDA_TMEM_LAYOUT_COLS: int = (
+ KDA_TMEM_QSTATE_ACC_STAGE1_COL_OFFSET + KDA_TMEM_N16_ACC_COLS
+)
+KDA_TMEM_ALLOC_COLS: int = _tcgen05_allocation_tmem_cols(KDA_TMEM_LAYOUT_COLS)
+
+
+@cute.jit
+def cta_sync() -> None:
+ """Synchronize all threads in the CTA."""
+
+ prims.barrier_cta_sync(0, thread_count=THREADS_PER_CTA)
+
+
+@cute.jit
+def cg0_sync(cg0_group_id) -> None:
+ """Synchronize one four-warp CG0 producer group."""
+
+ if cg0_group_id == 0:
+ prims.barrier_cta_sync(
+ NBAR_CG0_GROUP0_ID,
+ thread_count=CG0_THREADS_PER_GROUP,
+ )
+ elif cg0_group_id == 1:
+ prims.barrier_cta_sync(
+ NBAR_CG0_GROUP1_ID,
+ thread_count=CG0_THREADS_PER_GROUP,
+ )
+ else:
+ prims.barrier_cta_sync(
+ NBAR_CG0_GROUP2_ID,
+ thread_count=CG0_THREADS_PER_GROUP,
+ )
+
+
+@cute.jit
+def tmem_user_sync() -> None:
+ """Named barrier for CG1 plus the tcgen05 warp during TMEM lifecycle setup."""
+
+ prims.barrier_cta_sync(NBAR_TMEM_LIFECYCLE_ID, thread_count=TMEM_USER_THREADS)
+
+
+@cute.jit
+def is_compute_group0_warp(warp_idx) -> cutlass.Boolean:
+ """Return whether this warp belongs to CG0 preprocessing."""
+
+ return (warp_idx >= ROLES.compute_group0_first) & (
+ warp_idx <= ROLES.compute_group0_last
+ )
+
+
+@cute.jit
+def is_compute_group1_warp(warp_idx) -> cutlass.Boolean:
+ """Return whether this warp belongs to CG1 value/final-state work."""
+
+ return (warp_idx >= ROLES.compute_group1_first) & (
+ warp_idx <= ROLES.compute_group1_last
+ )
+
+
+@cute.jit
+def is_tmem_user_warp(warp_idx) -> cutlass.Boolean:
+ """Return whether this warp needs the allocated TMEM base pointer."""
+
+ return is_compute_group1_warp(warp_idx) | (warp_idx == ROLES.tcgen05_mma)
+
+
+@cute.jit
+def is_service_warpgroup(warp_idx) -> cutlass.Boolean:
+ """Return whether this warp belongs to the non-CG0/CG1 service warpgroup."""
+
+ return (warp_idx >= ROLES.super_mma) & (warp_idx <= ROLES.epilogue)
+
+
+@cute.jit
+def warp_group_sum_8(value: cutlass.Float32) -> cutlass.Float32:
+ """Reduce independent 8-lane row groups inside one warp."""
+
+ value = value + cutlass.Float32(
+ prims.shfl_sync(cute.arch.FULL_MASK, value, 4, 0x1F, prims.Shfl.BFLY)
+ )
+ value = value + cutlass.Float32(
+ prims.shfl_sync(cute.arch.FULL_MASK, value, 2, 0x1F, prims.Shfl.BFLY)
+ )
+ return value + cutlass.Float32(
+ prims.shfl_sync(cute.arch.FULL_MASK, value, 1, 0x1F, prims.Shfl.BFLY)
+ )
+
+
+@cute.jit
+def mma_input_dtype(
+ value: cutlass.Float32,
+ input_dtype: cutlass.Constexpr,
+) -> cutlass.Float32:
+ """Round a scalar through the compile-time input dtype for MMA reuse."""
+
+ return value.to(input_dtype).to(cutlass.Float32)
+
+
+@cute.jit
+def pairwise_eye(row_coord, col_coord) -> cutlass.Float32:
+ """Return the 16x16 identity value used by inverse product factors."""
+
+ return cutlass.Float32(1.0) if row_coord == col_coord else cutlass.Float32(0.0)
+
+
+@cute.jit
+def pairwise_sw32_smem_index(offset: cutlass.Constexpr, row_coord, col_coord):
+ """Return the SW32 physical index for a logical pairwise tile element."""
+
+ storage_col_coord = col_coord ^ PAIRWISE_SW32_COL_XOR
+ return offset + tcgen05_swizzle_32b_elem_index(
+ row_coord * PAIRWISE_SW32_ROW_STRIDE + storage_col_coord,
+ TCGEN05_F16_ELEM_BYTES,
+ )
+
+
+@cute.jit
+def pairwise_stmatrix_m8n8x4_ptr(
+ pairwise_smem,
+ offset: cutlass.Constexpr,
+ lane,
+):
+ """Return the row-start pointer for a 16x16 F16 pairwise STSM store."""
+
+ matrix_id = lane // 8
+ row_coord = lane & 7
+ col_coord = cutlass.Int32(0)
+ if matrix_id & 1:
+ row_coord = row_coord + cutlass.Int32(8)
+ if matrix_id >= 2:
+ col_coord = cutlass.Int32(8)
+ return pairwise_smem.subview(
+ pairwise_sw32_smem_index(offset, row_coord, col_coord)
+ ).data_ptr()
+
+
+@cute.jit
+def tcgen05_swizzle_128b_elem_index(
+ linear_elem_idx,
+ elem_bytes: cutlass.Constexpr,
+ rows: cutlass.Constexpr,
+):
+ """Return the K-box-major SW128 physical element index."""
+
+ elems_per_128b = 128 // elem_bytes
+ row_coord = linear_elem_idx // DK
+ col_coord = linear_elem_idx - row_coord * DK
+ slice_coord = col_coord // elems_per_128b
+ col_in_slice = col_coord - slice_coord * elems_per_128b
+ slice_linear_idx = row_coord * elems_per_128b + col_in_slice
+ byte_offset = slice_linear_idx * elem_bytes
+ swizzle_mask = ((byte_offset >> 7) & 0x7) << 4
+ return slice_coord * rows * elems_per_128b + (
+ (byte_offset ^ swizzle_mask) // elem_bytes
+ )
+
+
+@cute.jit
+def tcgen05_swizzle_32b_elem_index(
+ linear_elem_idx,
+ elem_bytes: cutlass.Constexpr,
+):
+ """Return the SW32 physical element index for a row-major logical tile."""
+
+ byte_offset = linear_elem_idx * elem_bytes
+ # SW32 is CuTe Swizzle<1,4,3>: xor address bit 7 into bit 4.
+ swizzle_mask = ((byte_offset >> 7) & 0x1) << 4
+ return (byte_offset ^ swizzle_mask) // elem_bytes
+
+
+@cute.jit
+def tcgen05_decay_b_key_storage_dim_runtime(token_coord, key_dim):
+ """Return the runtime key coordinate for tcgen05 SW128 decay operands.
+
+ Only the constant half-atom interleave applies. The former extra
+ token-dependent XOR (``(token & 2) * K_ATOM``) was matched by the warp
+ ldmatrix reader (self-consistent, so the intra-chunk path stayed clean)
+ but NOT by the tcgen05 state MMA, which reads this tile through a
+ standard-layout SMEM descriptor: tokens with ``t % 4 in {2, 3}`` read a
+ decay row displaced by 32 key channels every chunk. The error scales
+ with state persistence (large-negative ``dt_bias``), which is why random
+ small-gate tests never caught it.
+ """
+
+ key_mask = cutlass.Int32(TCGEN05_F16_K_ATOM // 2)
+ return key_dim ^ key_mask
+
+
+# Token-major staged output tiles consumed by the epilogue role. RHS/update
+# inputs are now packed by CG1 directly into shared_input TMEM, and value-side B
+# operands that need tcgen05-specific layouts live in distinct
+# `tcgen05_*_smem` buffers.
+O_STAGE_COUNT: int = 2
+O_STAGE_COLS: int = DV
+O_OUT_OFFSET: int = 0
+O_ELEM_BYTES: int = TCGEN05_F16_ELEM_BYTES
+O_TMA_SWIZZLE_BYTES: int = 128
+O_TMA_SWIZZLE_ELEMS: int = O_TMA_SWIZZLE_BYTES // O_ELEM_BYTES
+O_TMA_SWIZZLE_GROUP_BYTES: int = 16
+O_TMA_SWIZZLE_GROUP_ELEMS: int = O_TMA_SWIZZLE_GROUP_BYTES // O_ELEM_BYTES
+O_TMA_SWIZZLE_ROW_MASK: int = (O_TMA_SWIZZLE_ELEMS // O_TMA_SWIZZLE_GROUP_ELEMS) - 1
+O_TMA_SEGMENTS: int = DV // O_TMA_SWIZZLE_ELEMS
+O_TMA_SWIZZLE_ALIGNMENT_BYTES: int = O_TMA_SWIZZLE_BYTES * (
+ O_TMA_SWIZZLE_ELEMS // O_TMA_SWIZZLE_GROUP_ELEMS
+)
+O_SMEM_STAGE_SIZE: int = BT * O_STAGE_COLS
+O_SMEM_TILE_SIZE: int = O_STAGE_COUNT * O_SMEM_STAGE_SIZE
+
+TCGEN05_VALUE_PAIRWISE_B_LEADING_BYTES: int = 16
+TCGEN05_VALUE_PAIRWISE_B_STRIDE_BYTES: int = 8 * BT * TCGEN05_F16_ELEM_BYTES
+TCGEN05_FINAL_STATE_B_N_GROUP_ELEMS: int = TCGEN05_SW128_BYTES // TCGEN05_F16_ELEM_BYTES
+TCGEN05_FINAL_STATE_B_LEADING_BYTES: int = (
+ BT * TCGEN05_FINAL_STATE_B_N_GROUP_ELEMS * TCGEN05_F16_ELEM_BYTES
+)
+TCGEN05_FINAL_STATE_B_STRIDE_BYTES: int = (
+ 8 * TCGEN05_FINAL_STATE_B_N_GROUP_ELEMS * TCGEN05_F16_ELEM_BYTES
+)
+TCGEN05_FINAL_STATE_TMEM_LOAD_COLS: int = 32
+DECAY_STAGE_COUNT: int = 2
+Q_K_RESTORE_READY_STAGE_COUNT: int = 3
+TCGEN05_K_DECAY_STAGE_SIZE: int = DK * BT
+TCGEN05_Q_DECAY_STAGE_SIZE: int = DK * BT
+TCGEN05_K_RESTORE_STAGE_SIZE: int = DK * BT
+TCGEN05_K_DECAY_SMEM_TILE_SIZE: int = DECAY_STAGE_COUNT * TCGEN05_K_DECAY_STAGE_SIZE
+TCGEN05_Q_DECAY_SMEM_TILE_SIZE: int = DECAY_STAGE_COUNT * TCGEN05_Q_DECAY_STAGE_SIZE
+TCGEN05_K_RESTORE_SMEM_TILE_SIZE: int = DECAY_STAGE_COUNT * TCGEN05_K_RESTORE_STAGE_SIZE
+
+# Per-chunk q/k/v/gate/beta inputs staged by the TMA/load role. "raw" means
+# pre-normalization/pre-decay input; the 16-bit raw tiles still use SW128.
+RAW_STAGE_COUNT: int = 8
+# Depth of the TMA-completion mbarrier ring in the ENGINE-CLASS kernels
+# (engine / decomp-k1 prep; consumer-direct-wait). Chunk c's
+# transaction is issued against tma_mbar[c % D] and CONSUMERS wait that
+# slot's parity directly (no raw_ready relay), so D must equal the raw SMEM
+# ring depth: consumers reuse the raw-slot index and flip phase on the same
+# 8-chunk wrap. Warp 14 never waits completions; raw_consumed alone
+# throttles issue to <= D chunks in flight. (The decomp k2 chain kernels
+# keep their own relay ring, K2_TMA_MBAR_STAGE_COUNT below.)
+TMA_MBAR_STAGE_COUNT: int = RAW_STAGE_COUNT
+if TMA_MBAR_STAGE_COUNT != RAW_STAGE_COUNT:
+ raise ValueError("consumer-direct TMA ring must match the raw SMEM ring depth")
+
+RAW_F16_TMA_SWIZZLE_BYTES: int = 128
+RAW_F16_TMA_SWIZZLE_ELEMS: int = RAW_F16_TMA_SWIZZLE_BYTES // TCGEN05_F16_ELEM_BYTES
+RAW_F16_TMA_SWIZZLE_GROUP_BYTES: int = 16
+RAW_F16_TMA_SWIZZLE_GROUP_ELEMS: int = (
+ RAW_F16_TMA_SWIZZLE_GROUP_BYTES // TCGEN05_F16_ELEM_BYTES
+)
+RAW_F16_TMA_SWIZZLE_ROW_MASK: int = (
+ RAW_F16_TMA_SWIZZLE_ELEMS // RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+) - 1
+RAW_F16_TMA_SEGMENTS: int = DK // RAW_F16_TMA_SWIZZLE_ELEMS
+RAW_F16_TMA_SEGMENT_ELEMS: int = BT * RAW_F16_TMA_SWIZZLE_ELEMS
+RAW_F16_TMA_SWIZZLE_ALIGNMENT_BYTES: int = RAW_F16_TMA_SWIZZLE_BYTES * (
+ RAW_F16_TMA_SWIZZLE_ELEMS // RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+)
+RAW_F16_TAIL_ZERO_LANES: int = DK // RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+
+RAW_F32_ELEM_BYTES: int = 4
+RAW_F32_TMA_SWIZZLE_BYTES: int = 128
+RAW_F32_TMA_SWIZZLE_ELEMS: int = RAW_F32_TMA_SWIZZLE_BYTES // RAW_F32_ELEM_BYTES
+RAW_F32_TMA_SWIZZLE_GROUP_BYTES: int = 16
+RAW_F32_TMA_SWIZZLE_GROUP_ELEMS: int = (
+ RAW_F32_TMA_SWIZZLE_GROUP_BYTES // RAW_F32_ELEM_BYTES
+)
+CG0_TOKEN_ROWS_PER_WARP: int = BT // CG0_WARPS_PER_GROUP
+RAW_F32_TMA_SWIZZLE_ROW_MASK: int = (
+ RAW_F32_TMA_SWIZZLE_ELEMS // RAW_F32_TMA_SWIZZLE_GROUP_ELEMS
+) - 1
+RAW_F32_TMA_SEGMENTS: int = DK // RAW_F32_TMA_SWIZZLE_ELEMS
+RAW_F32_TMA_SEGMENT_ELEMS: int = BT * RAW_F32_TMA_SWIZZLE_ELEMS
+RAW_F32_TMA_SWIZZLE_ALIGNMENT_BYTES: int = RAW_F32_TMA_SWIZZLE_BYTES * (
+ RAW_F32_TMA_SWIZZLE_ELEMS // RAW_F32_TMA_SWIZZLE_GROUP_ELEMS
+)
+RAW_F32_TAIL_ZERO_LANES: int = DK // RAW_F32_TMA_SWIZZLE_GROUP_ELEMS
+
+RAW_Q_STAGE_SIZE: int = DK * BT
+RAW_K_STAGE_SIZE: int = DK * BT
+RAW_V_STAGE_SIZE: int = DV * BT
+RAW_GATE_STAGE_SIZE: int = DK * BT
+# FP32 gate-prefix exchange ring. Sizes are ELEMENT counts, so the raw-gate
+# ring above is dtype-agnostic (its byte size follows `gate_dtype`); this ring
+# is always FP32 and is only ALLOCATED on the 16-bit gate path -- with an FP32
+# gate the exchange keeps aliasing the raw-gate stage exactly as it always did.
+#
+# CG0 writes exp2(g_prefix) for one chunk, then CG0 (materialize + pack_rescale)
+# and CG1 (publish_projection_then_rescale) read it in the SAME chunk. There is
+# NO consumed mbarrier for this tile (it used to be covered by living inside the
+# 8-deep raw ring under `raw_consumed`), so the depth alone must order producer
+# chunk c+D against the readers of chunk c. D=2 does NOT close: CG0 group g's
+# write of c+2 and CG1's read of c are both gated only by
+# k_restore_consumed(c-1). D=4 closes through k_restore_consumed_l(c+1) ->
+# update_ready(c+1) -> CG1 chunk c+1 -> k_restore_consumed(c) -> update_ready(c)
+# -> CG0 group g's pack_rescale(c), and the two CG0 groups own disjoint stage
+# parities.
+GATE_EXCHANGE_STAGE_COUNT: int = 4
+GATE_EXCHANGE_STAGE_SIZE: int = DK * BT
+RAW_BETA_STAGE_SIZE: int = BT
+RAW_DT_BIAS_STAGE_SIZE: int = DK
+RAW_DT_BIAS_A_LOG_EXP_OFFSET: int = RAW_DT_BIAS_STAGE_SIZE
+RAW_Q_SMEM_TILE_SIZE: int = RAW_STAGE_COUNT * RAW_Q_STAGE_SIZE
+RAW_K_SMEM_TILE_SIZE: int = RAW_STAGE_COUNT * RAW_K_STAGE_SIZE
+RAW_V_SMEM_TILE_SIZE: int = RAW_STAGE_COUNT * RAW_V_STAGE_SIZE
+RAW_GATE_SMEM_TILE_SIZE: int = RAW_STAGE_COUNT * RAW_GATE_STAGE_SIZE
+GATE_EXCHANGE_SMEM_TILE_SIZE: int = GATE_EXCHANGE_STAGE_COUNT * GATE_EXCHANGE_STAGE_SIZE
+RAW_BETA_SMEM_TILE_SIZE: int = RAW_STAGE_COUNT * RAW_BETA_STAGE_SIZE
+RAW_DT_BIAS_SMEM_TILE_SIZE: int = RAW_DT_BIAS_STAGE_SIZE + 1
+
+
+# The public gate ABI is BF16. Float32 preserves compatibility with callers of
+# the historical kernel; other dtypes are intentionally rejected so they do not
+# silently multiply the compile-specialization matrix.
+GATE_DTYPES_SUPPORTED: tuple = (cutlass.BFloat16, cutlass.Float32)
+
+
+def gate_dtype_is_f32(gate_dtype) -> bool:
+ """True when the gate rides the historical FP32 memory format."""
+
+ return gate_dtype is cutlass.Float32
+
+
+def validate_gate_dtype(gate_dtype):
+ """Reject any gate dtype the kernel has no lowering for."""
+
+ if gate_dtype not in GATE_DTYPES_SUPPORTED:
+ names = ", ".join(getattr(d, "__name__", str(d)) for d in GATE_DTYPES_SUPPORTED)
+ raise TypeError(
+ "KDA gate_dtype must be one of {"
+ + names
+ + "}, got "
+ + str(getattr(gate_dtype, "__name__", gate_dtype))
+ )
+ return gate_dtype
+
+
+K_INV_STAGE_SIZE: int = BT * DK
+K_INV_SMEM_TILE_SIZE: int = DECAY_STAGE_COUNT * K_INV_STAGE_SIZE
+PAIRWISE_SMEM_QK_OFFSET: int = 0
+PAIRWISE_SMEM_AINV_OFFSET: int = PAIRWISE_SMEM_QK_OFFSET + PAIRWISE_SW32_TILE_ELEMS
+PAIRWISE_SMEM_STAGE_SIZE: int = PAIRWISE_SMEM_AINV_OFFSET + PAIRWISE_SW32_TILE_ELEMS
+PAIRWISE_STAGE_COUNT: int = 2
+PAIRWISE_SMEM_TILE_SIZE: int = PAIRWISE_STAGE_COUNT * PAIRWISE_SMEM_STAGE_SIZE
+
+SUPER_MMA_ATOM_N: int = 8
+SUPER_MMA_ATOM_K: int = 16
+SUPER_MMA_K_BLOCKS: int = DK // SUPER_MMA_ATOM_K
+SUPER_MMA_ACCUMULATORS_PER_LANE: int = 4
+
+
+@dataclass(frozen=True)
+class WarpRoles:
+ """Warp assignment for the BT=16 fully fused KDA kernel."""
+
+ compute_group0_first: int = 0
+ compute_group0_last: int = 7
+ compute_group1_first: int = 8
+ compute_group1_last: int = 11
+ super_mma: int = 12
+ tcgen05_mma: int = 13
+ tma_load: int = 14
+ epilogue: int = 15
+
+
+ROLES = WarpRoles()
+
+
+DTYPE_MAP: dict[type, torch.dtype] = {
+ cutlass.Float16: torch.float16,
+ cutlass.BFloat16: torch.bfloat16,
+ cutlass.Float32: torch.float32,
+}
+
+CLI_DTYPES: dict[str, type] = {
+ "fp16": cutlass.Float16,
+ "bf16": cutlass.BFloat16,
+}
+
+CLI_STATE_DTYPES: dict[str, type] = {
+ "bf16": cutlass.BFloat16,
+ "fp32": cutlass.Float32,
+}
+
+# Gate ABI dtype selectable from the CLI harness. BF16 is the default (see the
+# module docstring); FP32 selects the historical build.
+CLI_GATE_DTYPES: dict[str, type] = {
+ "bf16": cutlass.BFloat16,
+ "fp32": cutlass.Float32,
+}
+
+
+# ---------------------------------------------------------------------------
+# Torch reference: current numerical contract
+# ---------------------------------------------------------------------------
+
+
+_LN_2: float = 0.6931471805599453
+
+
+def _torch_exp2(x: torch.Tensor) -> torch.Tensor:
+ """``torch.exp2`` via the ``exp`` identity.
+
+ ``torch.exp2`` lowers through the CUDA jiterator (a runtime NVRTC
+ compile), which fails when the pip NVRTC build predates the running
+ device's architecture; ``torch.exp`` is precompiled ATen and never
+ touches NVRTC. The identity costs at most ~1 ulp on the FP32 reference
+ path, far below the kernel comparison tolerances.
+ """
+
+ return torch.exp(x * _LN_2)
+
+
+def _torch_l2_normalize_qk(
+ q: torch.Tensor, k: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """L2-normalize q/k on the head dimension (rsqrt of the squared sum)."""
+
+ return (
+ torch.nn.functional.normalize(q.float(), p=2.0, dim=-1, eps=L2_NORM_EPS),
+ torch.nn.functional.normalize(k.float(), p=2.0, dim=-1, eps=L2_NORM_EPS),
+ )
+
+
+def _torch_mma_input(x: torch.Tensor, mma_dtype: torch.dtype) -> torch.Tensor:
+ """Round tensors as 16-bit MMA inputs while keeping FP32 accumulation sites."""
+
+ if mma_dtype not in (torch.float16, torch.bfloat16):
+ raise TypeError(f"mma_dtype must be fp16 or bf16, got {mma_dtype}")
+ return x.to(mma_dtype).float()
+
+
+def _torch_mma_input_mul(
+ lhs: torch.Tensor,
+ rhs: torch.Tensor,
+ mma_dtype: torch.dtype,
+) -> torch.Tensor:
+ """Model packed 16-bit multiply used when materializing decay operands."""
+
+ return _torch_mma_input(
+ _torch_mma_input(lhs, mma_dtype) * _torch_mma_input(rhs, mma_dtype),
+ mma_dtype,
+ )
+
+
+def _torch_safe_gate_log2_increment(
+ raw_gate: torch.Tensor,
+ a_log: torch.Tensor | None,
+ dt_bias: torch.Tensor | None,
+ safe_gate: bool,
+ gate_lower_bound: float,
+) -> torch.Tensor:
+ """Apply FLA-compatible KDA gate preprocessing in log2-domain units.
+
+ safe_gate=True -> the bounded sigmoid gate;
+ safe_gate=False -> FLA's unbounded ``-exp(A_log) * softplus(g + dt_bias)``
+ activation, both returned as base-2 (log2) decay increments.
+ """
+
+ if a_log is None:
+ a_log = torch.zeros(
+ raw_gate.shape[1], dtype=torch.float32, device=raw_gate.device
+ )
+ if dt_bias is None:
+ dt_bias = torch.zeros(
+ raw_gate.shape[1],
+ raw_gate.shape[-1],
+ dtype=torch.float32,
+ device=raw_gate.device,
+ )
+ if a_log.shape != (raw_gate.shape[1],):
+ raise ValueError(f"a_log must be [H], got {a_log.shape}")
+ if dt_bias.shape != (raw_gate.shape[1], raw_gate.shape[-1]):
+ raise ValueError(
+ f"dt_bias must be [H,{raw_gate.shape[-1]}], got {dt_bias.shape}"
+ )
+
+ a_log_exp = _torch_exp2(a_log.float() * LOG2_E).view(1, raw_gate.shape[1], 1, 1)
+ biased_gate = raw_gate.float() + dt_bias.float().view(
+ 1, raw_gate.shape[1], 1, raw_gate.shape[-1]
+ )
+ if safe_gate:
+ sigmoid = 0.5 * torch.tanh(0.5 * a_log_exp * biased_gate) + 0.5
+ return (gate_lower_bound * LOG2_E) * sigmoid
+ return -a_log_exp * torch.nn.functional.softplus(biased_gate) * LOG2_E
+
+
+def _state_vk_to_kv(state_vk: torch.Tensor) -> torch.Tensor:
+ """Interpret a VK `[B, H, DV, DK]` state tensor with the reference KV axes."""
+
+ if state_vk.shape[-2:] != (DV, DK):
+ raise ValueError(f"VK state must end with [{DV},{DK}], got {state_vk.shape}")
+ return state_vk.transpose(-1, -2)
+
+
+def _state_kv_to_vk(state_kv: torch.Tensor) -> torch.Tensor:
+ """Interpret a KV `[B, H, DK, DV]` reference state with kernel VK axes."""
+
+ if state_kv.shape[-2:] != (DK, DV):
+ raise ValueError(f"KV state must end with [{DK},{DV}], got {state_kv.shape}")
+ return state_kv.transpose(-1, -2)
+
+
+def _blockwise_inverse_unit_lower_bt16_mma(
+ strict_lower: torch.Tensor,
+ mma_dtype: torch.dtype,
+) -> torch.Tensor:
+ """BT=16 blockwise inverse mirroring the device auxiliary-MMA stage.
+
+ Hierarchical (GDN-style) block inverse expressed entirely as MMAs:
+
+ D = blockdiag(L11, L22) (strict-lower 8x8 diagonal blocks)
+ Binv = (I - D)(I + D^2)(I + D^4) (exact: blocks nilpotent at 8)
+ A^-1 = Binv - Binv @ A21hat @ Binv (exact: (Binv @ A21hat)^2 = 0)
+
+ Chain and combine operands are FP16 (FP32 accumulate) regardless of the
+ kernel input dtype, matching the device stage; only the staged A_inv
+ tile is rounded to `mma_dtype`.
+ """
+
+ if strict_lower.shape[-2:] != (BT, BT):
+ raise ValueError(
+ f"strict_lower must end with [{BT},{BT}], got {strict_lower.shape}"
+ )
+
+ chain_dtype = torch.float16
+ strict_lower = strict_lower.float()
+ eye = torch.eye(BT, dtype=torch.float32, device=strict_lower.device)
+
+ d = torch.zeros_like(strict_lower)
+ d[..., :8, :8] = strict_lower[..., :8, :8]
+ d[..., 8:, 8:] = strict_lower[..., 8:, 8:]
+ d16 = _torch_mma_input(d, chain_dtype)
+
+ d2 = _torch_mma_input(d16 @ d16, chain_dtype)
+ d4 = _torch_mma_input(d2 @ d2, chain_dtype)
+
+ b01_lhs = _torch_mma_input(eye - d, chain_dtype)
+ b01_rhs = _torch_mma_input(eye + d2, chain_dtype)
+ b01 = _torch_mma_input(b01_lhs @ b01_rhs, chain_dtype)
+ binv = b01 @ _torch_mma_input(eye + d4, chain_dtype) # FP32 accumulator
+
+ a21hat = torch.zeros_like(strict_lower)
+ a21hat[..., 8:, :8] = strict_lower[..., 8:, :8]
+ t1 = _torch_mma_input(binv, chain_dtype) @ _torch_mma_input(a21hat, chain_dtype)
+ x21 = -(_torch_mma_input(t1, chain_dtype) @ _torch_mma_input(binv, chain_dtype))
+
+ ainv = binv.clone()
+ ainv[..., 8:, :8] = x21[..., 8:, :8]
+ return _torch_mma_input(ainv, mma_dtype)
+
+
+@dataclass(frozen=True)
+class PairwiseF16MmaTiles:
+ """Input-dtype-rounded tiles owned by the KDA auxiliary-MMA warp."""
+
+ kk: torch.Tensor
+ qk: torch.Tensor
+ strict_lower: torch.Tensor
+ a_inv: torch.Tensor
+
+
+@cute.jit
+def raw_f16_s128_smem_index(token_coord, dim):
+ """Return the physical s128 SMEM index for raw F16 q/k/v staging."""
+
+ segment = dim // RAW_F16_TMA_SWIZZLE_ELEMS
+ segment_dim = dim - segment * RAW_F16_TMA_SWIZZLE_ELEMS
+ col_group = segment_dim // RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ col_in_group = segment_dim - col_group * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ row_swizzle = token_coord & RAW_F16_TMA_SWIZZLE_ROW_MASK
+ return (
+ segment * RAW_F16_TMA_SEGMENT_ELEMS
+ + token_coord * RAW_F16_TMA_SWIZZLE_ELEMS
+ + ((col_group ^ row_swizzle) * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS)
+ + col_in_group
+ )
+
+
+@cute.jit
+def raw_f32_s128_smem_index(token_coord, dim):
+ """Return the physical s128 SMEM index for raw F32 gate staging."""
+
+ segment = dim // RAW_F32_TMA_SWIZZLE_ELEMS
+ segment_dim = dim - segment * RAW_F32_TMA_SWIZZLE_ELEMS
+ col_group = segment_dim // RAW_F32_TMA_SWIZZLE_GROUP_ELEMS
+ col_in_group = segment_dim - col_group * RAW_F32_TMA_SWIZZLE_GROUP_ELEMS
+ row_swizzle = token_coord & RAW_F32_TMA_SWIZZLE_ROW_MASK
+ return (
+ segment * RAW_F32_TMA_SEGMENT_ELEMS
+ + token_coord * RAW_F32_TMA_SWIZZLE_ELEMS
+ + ((col_group ^ row_swizzle) * RAW_F32_TMA_SWIZZLE_GROUP_ELEMS)
+ + col_in_group
+ )
+
+
+def raw_f32_exchange_smem_index(token_coord, dim):
+ """Return a conflict-free SMEM index for the CG0 gate-prefix exchange.
+
+ The TMA s128 swizzle XORs banks only within one 128-byte segment, so the
+ vectorized prefix re-reads hit the same bank group once per segment
+ (2-way conflicts). The exchange scratch is FP32 and either reuses the
+ raw-gate stage after its raw values are consumed (FP32 gate) or is a
+ dedicated `gate_exchange_smem` buffer (16-bit gate), so both sides of the
+ exchange can add a segment XOR that spreads segments across bank groups;
+ column-wise prefix stores stay conflict-free because the segment term is
+ constant per warp.
+
+ NOTE the layout math stays on the RAW_F32 s128 geometry on BOTH gate paths:
+ this buffer is FP32 and independent of the TMA landing layout.
+ """
+
+ segment = dim // RAW_F32_TMA_SWIZZLE_ELEMS
+ segment_dim = dim - segment * RAW_F32_TMA_SWIZZLE_ELEMS
+ col_group = segment_dim // RAW_F32_TMA_SWIZZLE_GROUP_ELEMS
+ col_in_group = segment_dim - col_group * RAW_F32_TMA_SWIZZLE_GROUP_ELEMS
+ row_swizzle = token_coord & RAW_F32_TMA_SWIZZLE_ROW_MASK
+ return (
+ segment * RAW_F32_TMA_SEGMENT_ELEMS
+ + token_coord * RAW_F32_TMA_SWIZZLE_ELEMS
+ + ((col_group ^ row_swizzle ^ segment) * RAW_F32_TMA_SWIZZLE_GROUP_ELEMS)
+ + col_in_group
+ )
+
+
+@cute.jit
+def k_inv_s128_smem_index(token_coord, key_dim):
+ """Return the physical s128 SMEM index for the K-inverse auxiliary-MMA RHS."""
+
+ return raw_f16_s128_smem_index(token_coord, key_dim)
+
+
+@cute.jit
+def o_smem_swizzle_128b_elem_index(
+ o_stage_base,
+ value_dim,
+ token_coord,
+):
+ """Return the physical W128 SMEM index for one staged output element."""
+
+ segment = value_dim // O_TMA_SWIZZLE_ELEMS
+ segment_value_dim = value_dim - segment * O_TMA_SWIZZLE_ELEMS
+ col_group = segment_value_dim // O_TMA_SWIZZLE_GROUP_ELEMS
+ col_in_group = segment_value_dim - col_group * O_TMA_SWIZZLE_GROUP_ELEMS
+ row_swizzle = token_coord & O_TMA_SWIZZLE_ROW_MASK
+ return (
+ o_stage_base
+ + O_OUT_OFFSET
+ + segment * BT * O_TMA_SWIZZLE_ELEMS
+ + token_coord * O_TMA_SWIZZLE_ELEMS
+ + ((col_group ^ row_swizzle) * O_TMA_SWIZZLE_GROUP_ELEMS)
+ + col_in_group
+ )
+
+
+@cute.jit
+def o_smem_stmatrix_128b_ptr(
+ o_smem,
+ o_stage_base,
+ value_dim_base,
+ lane,
+):
+ """Return the per-lane W128 row-start pointer for one 16x16 STSM.T tile."""
+
+ matrix_id = lane // 8
+ row_in_matrix = lane & 7
+ token_block = matrix_id // 2
+ value_block = matrix_id & 1
+ token_coord = token_block * 8 + row_in_matrix
+ value_dim = value_dim_base + value_block * 8
+ smem_idx = o_smem_swizzle_128b_elem_index(
+ o_stage_base,
+ value_dim,
+ token_coord,
+ )
+ return o_smem.subview(smem_idx).data_ptr()
+
+
+@cute.jit
+def raw_v_ldmatrix_trans_ptr(raw_v_smem, value_dim_base, lane):
+ """Return the per-lane row-start pointer for raw V `ldmatrix.x4.trans`."""
+
+ matrix_id = lane // 8
+ row_in_matrix = lane & 7
+ token_block = matrix_id // 2
+ value_block = matrix_id & 1
+ token_coord = token_block * 8 + row_in_matrix
+ value_dim = value_dim_base + value_block * 8
+ smem_idx = raw_f16_s128_smem_index(token_coord, value_dim)
+ return raw_v_smem.subview(smem_idx).data_ptr()
+
+
+def _kda_pairwise_mma_tiles(
+ q_decay: torch.Tensor,
+ k_decay: torch.Tensor,
+ k_inv: torch.Tensor,
+ beta_chunk: torch.Tensor,
+ mma_dtype: torch.dtype,
+) -> PairwiseF16MmaTiles:
+ """Reference contract for the BT=16 KK/QK/inverse auxiliary-MMA path."""
+
+ expected_tile_shape = (BT, DK)
+ if q_decay.shape[-2:] != expected_tile_shape:
+ raise ValueError(
+ f"q_decay must end with {expected_tile_shape}, got {q_decay.shape}"
+ )
+ if k_decay.shape != q_decay.shape:
+ raise ValueError(f"k_decay must match q_decay, got {k_decay.shape}")
+ if k_inv.shape != q_decay.shape:
+ raise ValueError(f"k_inv must match q_decay, got {k_inv.shape}")
+ if beta_chunk.shape != q_decay.shape[:-1]:
+ raise ValueError(f"beta_chunk must be [...,{BT}], got {beta_chunk.shape}")
+
+ tril_strict = torch.tril(
+ torch.ones(BT, BT, dtype=torch.bool, device=q_decay.device), diagonal=-1
+ )
+ tril_inclusive = torch.tril(
+ torch.ones(BT, BT, dtype=torch.bool, device=q_decay.device), diagonal=0
+ )
+
+ kk = _torch_mma_input(
+ torch.einsum("...ik,...jk->...ij", k_decay, k_inv),
+ mma_dtype,
+ )
+ qk_full = _torch_mma_input(
+ torch.einsum("...ik,...jk->...ij", q_decay, k_inv),
+ mma_dtype,
+ )
+ kk_lower = torch.where(tril_strict, kk, torch.zeros_like(kk))
+ qk = _torch_mma_input(
+ torch.where(tril_inclusive, qk_full, torch.zeros_like(qk_full)),
+ mma_dtype,
+ )
+ strict_lower = _torch_mma_input(kk_lower * beta_chunk[..., :, None], mma_dtype)
+ a_inv = _blockwise_inverse_unit_lower_bt16_mma(strict_lower, mma_dtype)
+ return PairwiseF16MmaTiles(
+ kk=kk,
+ qk=qk,
+ strict_lower=strict_lower,
+ a_inv=a_inv,
+ )
+
+
+def _pad_chunk_to_bt(value: torch.Tensor, valid_tokens: int) -> torch.Tensor:
+ """Pad a `[B,H,T,...]` chunk to BT tokens with trailing zeros."""
+
+ if valid_tokens == BT:
+ return value
+ pad_shape = list(value.shape)
+ pad_shape[2] = BT - valid_tokens
+ return torch.cat([value, value.new_zeros(pad_shape)], dim=2)
+
+
+def _pad_beta_chunk_to_bt(value: torch.Tensor, valid_tokens: int) -> torch.Tensor:
+ """Pad a `[B,H,T]` beta chunk to BT tokens with trailing zeros."""
+
+ if valid_tokens == BT:
+ return value
+ pad_shape = list(value.shape)
+ pad_shape[2] = BT - valid_tokens
+ return torch.cat([value, value.new_zeros(pad_shape)], dim=2)
+
+
+def kda_chunked_mma_reference(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ raw_gate: torch.Tensor,
+ beta_logits: torch.Tensor,
+ initial_state: torch.Tensor | None = None,
+ scale: float | None = None,
+ mma_dtype: torch.dtype | None = None,
+ a_log: torch.Tensor | None = None,
+ dt_bias: torch.Tensor | None = None,
+ safe_gate: bool = True,
+ gate_lower_bound: float = DEFAULT_GATE_LOWER_BOUND,
+ cu_seqlens: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """Chunk KDA reference with dtype-rounded operands at MMA input boundaries."""
+
+ if cu_seqlens is not None:
+ if q.shape[0] != 1:
+ raise ValueError(f"packed q batch must be 1, got {q.shape[0]}")
+ boundaries = cu_seqlens.tolist()
+ num_sequences = len(boundaries) - 1
+ if initial_state is not None and initial_state.shape[0] != num_sequences:
+ raise ValueError("packed initial_state must have one row per sequence")
+ packed_out = torch.empty_like(v)
+ final_states: list[torch.Tensor] = []
+ for sequence_idx, (start, end) in enumerate(
+ zip(boundaries[:-1], boundaries[1:], strict=True)
+ ):
+ sequence_out, sequence_state = kda_chunked_mma_reference(
+ q[:, :, start:end, :],
+ k[:, :, start:end, :],
+ v[:, :, start:end, :],
+ raw_gate[:, :, start:end, :],
+ beta_logits[:, :, start:end],
+ None
+ if initial_state is None
+ else initial_state[sequence_idx : sequence_idx + 1],
+ scale,
+ mma_dtype,
+ a_log,
+ dt_bias,
+ safe_gate,
+ gate_lower_bound,
+ )
+ packed_out[:, :, start:end, :] = sequence_out
+ final_states.append(sequence_state)
+ return packed_out, torch.cat(final_states, dim=0)
+
+ if q.shape != k.shape:
+ raise ValueError(
+ f"q and k must have the same shape, got {q.shape} and {k.shape}"
+ )
+ if v.shape[:-1] != q.shape[:-1] or v.shape[-1] != DV:
+ raise ValueError(f"v must be [B,H,T,{DV}], got {v.shape}")
+ if raw_gate.shape != q.shape:
+ raise ValueError(f"raw_gate must be [B,H,T,{DK}], got {raw_gate.shape}")
+ if beta_logits.shape != q.shape[:-1]:
+ raise ValueError(f"beta_logits must be [B,H,T], got {beta_logits.shape}")
+
+ batch, heads, seqlen, d_k = q.shape
+ if d_k != DK:
+ raise ValueError(f"DK must be {DK}, got {d_k}")
+ if scale is None:
+ scale = 1.0 / (DK**0.5)
+ if mma_dtype is None:
+ mma_dtype = q.dtype
+
+ state = (
+ torch.zeros(batch, heads, DK, DV, dtype=torch.float32, device=q.device)
+ if initial_state is None
+ else initial_state.float().clone()
+ )
+ out = torch.empty(batch, heads, seqlen, DV, dtype=q.dtype, device=q.device)
+
+ q_f, k_f = _torch_l2_normalize_qk(q, k)
+ v_f = v.float()
+ gate_f = _torch_safe_gate_log2_increment(
+ raw_gate,
+ a_log,
+ dt_bias,
+ safe_gate,
+ gate_lower_bound,
+ )
+ beta_f = torch.tanh(beta_logits.float() * 0.5) * 0.5 + 0.5
+ for chunk_start in range(0, seqlen, BT):
+ chunk_end = min(chunk_start + BT, seqlen)
+ valid_tokens = chunk_end - chunk_start
+ q_chunk = _pad_chunk_to_bt(q_f[:, :, chunk_start:chunk_end, :], valid_tokens)
+ k_chunk = _pad_chunk_to_bt(k_f[:, :, chunk_start:chunk_end, :], valid_tokens)
+ v_chunk = _pad_chunk_to_bt(v_f[:, :, chunk_start:chunk_end, :], valid_tokens)
+ gate_chunk = _pad_chunk_to_bt(
+ gate_f[:, :, chunk_start:chunk_end, :],
+ valid_tokens,
+ )
+ beta_chunk = _pad_beta_chunk_to_bt(
+ beta_f[:, :, chunk_start:chunk_end],
+ valid_tokens,
+ )
+
+ g_prefix = torch.cumsum(gate_chunk, dim=2)
+ exp_g = _torch_exp2(g_prefix)
+ exp_neg_g = _torch_exp2(-g_prefix)
+ exp_g_last = _torch_exp2(g_prefix[:, :, -1, :])
+
+ q_decay = _torch_mma_input_mul(q_chunk, exp_g, mma_dtype)
+ k_decay = _torch_mma_input_mul(k_chunk, exp_g, mma_dtype)
+ k_inv = _torch_mma_input(k_chunk * exp_neg_g, mma_dtype)
+ k_restore = _torch_mma_input_mul(
+ k_inv,
+ exp_g_last[:, :, None, :],
+ mma_dtype,
+ )
+
+ pairwise_tiles = _kda_pairwise_mma_tiles(
+ q_decay,
+ k_decay,
+ k_inv,
+ beta_chunk,
+ mma_dtype,
+ )
+
+ state_mma = _torch_mma_input(state, mma_dtype)
+ state_k = torch.einsum("bhtk,bhkv->bhtv", k_decay, state_mma)
+ rhs = _torch_mma_input(
+ beta_chunk[:, :, :, None] * (v_chunk - state_k),
+ mma_dtype,
+ )
+ update = torch.einsum("bhij,bhjv->bhiv", pairwise_tiles.a_inv, rhs)
+ update_mma = _torch_mma_input(update, mma_dtype)
+
+ state_q = torch.einsum("bhtk,bhkv->bhtv", q_decay, state_mma)
+ intra = torch.einsum("bhij,bhjv->bhiv", pairwise_tiles.qk, update_mma)
+ out[:, :, chunk_start:chunk_end, :] = ((state_q + intra) * scale).to(q.dtype)[
+ :, :, :valid_tokens, :
+ ]
+
+ state = state * exp_g_last[:, :, :, None] + torch.einsum(
+ "bhtk,bhtv->bhkv", k_restore, update_mma
+ )
+
+ return out, state
+
+
+# ---------------------------------------------------------------------------
+# Device-stage helpers. These keep each warp group's numerical contract
+# explicit while the heavy MMA/TMEM lowering lands stage by stage.
+# ---------------------------------------------------------------------------
+
+
+# --- beta TMA transport ------------------------------------------------------
+# beta arrives by 2D TMA instead of a per-lane strided LSU gather (which cost
+# ~1.07us per extra 32B sector at H96: 16 lanes x heads*2B stride = 16 sectors
+# vs 1). One descriptor family, runtime-selected by g = 8/gcd(heads, 8):
+# g == 1 (heads % 8 == 0): view (heads, T) stride (1, heads); box = (8, BT)
+# -- 16B rows, the CTA fetches only its 8-head group. Verified H96/64/32,
+# H96-T8192 engine 450.6 -> 434.2us (beats head-major 438.0; NCU: the
+# +786K-sector LSU signature vanishes exactly).
+# g > 1: PAIR-PACKED view (g*heads, ceil(T/g)) over the same memory -- one
+# packed row = g consecutive tokens x all heads, so the row stride
+# g*heads*2B is a legal 16B multiple even when heads % 8 != 0 (H12 = TP8
+# of H96: g=2, 48B). Box = one packed row x (BT/g + 1) rows; the +1
+# parity row lets a chunk whose start token is not g-aligned round DOWN
+# to the packed grain and consume at offset (start % g). Verified H12
+# bitwise incl. odd-bos varlen; requires heads <= 14 when heads % 8 != 0
+# (SMEM stage cap + Int8 box dims), which covers every TP split of
+# production head counts.
+# The descriptor needs beta.data_ptr() 16B-aligned: guaranteed for allocator-
+# fresh tensors (512B granule) and checked loudly in the wrapper; only an
+# odd-token-offset *view* at g > 1 can violate it (see test_beta_alignment).
+BETA_TILE_STAGE_ELEMS: int = 256 # bf16; >= max(8*BT group, heads*(BT+g) pair)
+BETA_TILE_STAGE_COUNT: int = 8
+BETA_TMA_LOOKAHEAD: int = 4
+
+
+@cute.jit
+def tma_stage_load_inputs(
+ tma_desc_q: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_k: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_v: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_gate: cutlass.GridConstant[cuda.TensorMap],
+ beta_tile_stage,
+ beta_g,
+ heads32,
+ raw_q_smem,
+ raw_k_smem,
+ raw_v_smem,
+ raw_gate_smem,
+ raw_beta_smem,
+ sequence_start,
+ head_idx,
+ lane,
+ chunk_start,
+ seqlen,
+ tma_mbar,
+ tma_tx_bytes: cutlass.Constexpr,
+ gate_dtype: cutlass.Constexpr,
+) -> None:
+ """Issue tensor-operand TMA + sigmoid(beta logits) for one chunk (no wait).
+
+ Consumers wait the chunk's tma_mbar slot parity directly, so the
+ beta scalar store MUST precede the mbarrier_arrive_expect_tx: the
+ arrive's release (after the warp sync) is the only edge that publishes
+ beta to the consumers' acquire (there is no raw_ready relay any more).
+ """
+
+ global_chunk_start = sequence_start + chunk_start
+ if lane < BT:
+ token_idx = chunk_start + lane
+ beta_value = cutlass.Float32(0.0)
+ if token_idx < seqlen:
+ # The tile is a contiguous [token][head-slot] slab staged by TMA:
+ # g == 1 -> 8-head group rows (slot = head % 8, token = lane);
+ # g > 1 -> whole packed rows from token (start - start % g)
+ # (slot = head, token = start % g + lane).
+ b_idx = lane * cutlass.Int32(8) + head_idx % cutlass.Int32(8)
+ if beta_g > cutlass.Int32(1):
+ b_idx = (sequence_start % beta_g + lane) * heads32 + head_idx
+ beta_logit = beta_tile_stage[b_idx].to(cutlass.Float32)
+ half = cutlass.Float32(0.5)
+ beta_value = cute.math.tanh(beta_logit * half, approx=True) * half + half
+ raw_beta_smem[lane] = beta_value
+ prims.bar_warp_sync(cute.arch.FULL_MASK)
+ if prims.elect_sync():
+ prims.mbarrier_arrive_expect_tx(tma_mbar, tma_tx_bytes)
+ if prims.elect_sync():
+ for segment in cutlass.range_constexpr(RAW_F16_TMA_SEGMENTS):
+ tma_coord = (
+ cutlass.Int32(segment * RAW_F16_TMA_SWIZZLE_ELEMS),
+ global_chunk_start,
+ head_idx,
+ cutlass.Int32(0),
+ )
+ smem_offset = segment * RAW_F16_TMA_SEGMENT_ELEMS
+ prims.cp_async_bulk_tensor_shared_cta_global(
+ raw_q_smem.subview(smem_offset),
+ tma_desc_q.get_ptr(),
+ tma_coord,
+ tma_mbar,
+ )
+ prims.cp_async_bulk_tensor_shared_cta_global(
+ raw_k_smem.subview(smem_offset),
+ tma_desc_k.get_ptr(),
+ tma_coord,
+ tma_mbar,
+ )
+ prims.cp_async_bulk_tensor_shared_cta_global(
+ raw_v_smem.subview(smem_offset),
+ tma_desc_v.get_ptr(),
+ tma_coord,
+ tma_mbar,
+ )
+ if cutlass.const_expr(not gate_dtype_is_f32(gate_dtype)):
+ # 16-bit gate: it rides the SAME 16-bit s128 box/swizzle family
+ # as q/k/v, so it folds into this loop (2 x 128 B segments
+ # instead of the FP32 path's 4 x 128 B segments).
+ prims.cp_async_bulk_tensor_shared_cta_global(
+ raw_gate_smem.subview(smem_offset),
+ tma_desc_gate.get_ptr(),
+ tma_coord,
+ tma_mbar,
+ )
+
+ if cutlass.const_expr(gate_dtype_is_f32(gate_dtype)):
+ for segment in cutlass.range_constexpr(RAW_F32_TMA_SEGMENTS):
+ tma_coord = (
+ cutlass.Int32(segment * RAW_F32_TMA_SWIZZLE_ELEMS),
+ global_chunk_start,
+ head_idx,
+ cutlass.Int32(0),
+ )
+ smem_offset = segment * RAW_F32_TMA_SEGMENT_ELEMS
+ prims.cp_async_bulk_tensor_shared_cta_global(
+ raw_gate_smem.subview(smem_offset),
+ tma_desc_gate.get_ptr(),
+ tma_coord,
+ tma_mbar,
+ )
+
+
+@cute.jit
+def tma_transfer_wait(tma_mbar, tma_phase) -> None:
+ """Spin until one tma_mbar ring slot's expect_tx transaction completes."""
+
+ while not prims.mbarrier_wait_parity(
+ tma_mbar,
+ tma_phase,
+ prims.MBarrierWait.TRY,
+ ):
+ pass
+
+
+@cute.jit
+def cg0_zero_tail_raw_operands(
+ raw_q_smem,
+ raw_k_smem,
+ raw_v_smem,
+ raw_gate_smem,
+ lane,
+ chunk_start,
+ seqlen,
+ input_dtype: cutlass.Constexpr,
+ gate_dtype: cutlass.Constexpr,
+) -> None:
+ """Zero padded tail rows after TMA has produced the raw input tile.
+
+ `gate_dtype` is carried separately from `input_dtype` because the gate has
+ its own ABI dtype: a 16-bit gate shares the q/k/v s128 geometry (and can be
+ zeroed in the same lane block), an FP32 gate keeps its own wider geometry.
+ """
+
+ raw_q_ptr = raw_q_smem.data_ptr()
+ raw_k_ptr = raw_k_smem.data_ptr()
+ raw_v_ptr = raw_v_smem.data_ptr()
+ raw_gate_ptr = raw_gate_smem.data_ptr()
+ f16_zero = input_dtype(0.0)
+ f16_zero_vec = cutlass.Vector.from_elements(
+ (
+ f16_zero,
+ f16_zero,
+ f16_zero,
+ f16_zero,
+ f16_zero,
+ f16_zero,
+ f16_zero,
+ f16_zero,
+ ),
+ input_dtype,
+ )
+ if cutlass.const_expr(gate_dtype_is_f32(gate_dtype)):
+ f32_zero = cutlass.Float32(0.0)
+ gate_zero_vec = cutlass.Vector.from_elements(
+ (f32_zero, f32_zero, f32_zero, f32_zero),
+ cutlass.Float32,
+ )
+ else:
+ gate_zero = gate_dtype(0.0)
+ gate_zero_vec = cutlass.Vector.from_elements(
+ (
+ gate_zero,
+ gate_zero,
+ gate_zero,
+ gate_zero,
+ gate_zero,
+ gate_zero,
+ gate_zero,
+ gate_zero,
+ ),
+ gate_dtype,
+ )
+ for row in cutlass.range_constexpr(BT):
+ token_idx = chunk_start + cutlass.Int32(row)
+ if token_idx >= seqlen:
+ if lane < RAW_F16_TAIL_ZERO_LANES:
+ f16_dim_base = lane * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ f16_idx = raw_f16_s128_smem_index(row, f16_dim_base)
+ (raw_q_ptr + f16_idx).store(
+ f16_zero_vec,
+ alignment=RAW_F16_TMA_SWIZZLE_GROUP_BYTES,
+ )
+ (raw_k_ptr + f16_idx).store(
+ f16_zero_vec,
+ alignment=RAW_F16_TMA_SWIZZLE_GROUP_BYTES,
+ )
+ (raw_v_ptr + f16_idx).store(
+ f16_zero_vec,
+ alignment=RAW_F16_TMA_SWIZZLE_GROUP_BYTES,
+ )
+ if cutlass.const_expr(not gate_dtype_is_f32(gate_dtype)):
+ (raw_gate_ptr + f16_idx).store(
+ gate_zero_vec,
+ alignment=RAW_F16_TMA_SWIZZLE_GROUP_BYTES,
+ )
+ if cutlass.const_expr(gate_dtype_is_f32(gate_dtype)):
+ if lane < RAW_F32_TAIL_ZERO_LANES:
+ f32_dim_base = lane * RAW_F32_TMA_SWIZZLE_GROUP_ELEMS
+ f32_idx = raw_f32_s128_smem_index(row, f32_dim_base)
+ (raw_gate_ptr + f32_idx).store(
+ gate_zero_vec,
+ alignment=RAW_F32_TMA_SWIZZLE_GROUP_BYTES,
+ )
+
+
+@cute.jit
+def cg0_materialize_decay_operands(
+ raw_q_smem,
+ raw_k_smem,
+ raw_gate_smem,
+ gate_exchange_smem,
+ a_log_exp,
+ dt_bias_value,
+ k_inv_smem,
+ tcgen05_k_decay_smem,
+ tcgen05_q_decay_smem,
+ tcgen05_k_restore_smem,
+ cg0_k_ready_stage_mbar,
+ cg0_k_half_ready_stage_mbar,
+ diag_ready_stage_mbar,
+ operand_smem_consumed_stage_mbar,
+ k_restore_consumed_stage_mbar,
+ chunk,
+ seqlen,
+ input_dtype: cutlass.Constexpr,
+ gate_dtype: cutlass.Constexpr,
+ SAFE_GATE: cutlass.Constexpr,
+ GATE_SCALE_LOG2: cutlass.Constexpr,
+ FULL_CHUNKS: cutlass.Constexpr,
+ cg0_group_id,
+ cg0_local_warp,
+ lane,
+) -> None:
+ """Materialize safe-gate KDA decay operands for one key dimension.
+
+ `FULL_CHUNKS` is the engine peel's per-call selector (True for a guard-free
+ interior chunk, False for the peeled partial tail), NOT a compile-time
+ specialization: the guard-free interior skips the masked-tail fixup below.
+ """
+
+ row_group_start = cg0_local_warp * CG0_TOKEN_ROWS_PER_WARP
+ lane_row_group = lane // RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ lane_in_row_group = lane - lane_row_group * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ decay_row = row_group_start + lane_row_group
+
+ raw_q_ptr = raw_q_smem.data_ptr()
+ raw_k_ptr = raw_k_smem.data_ptr()
+ # The FP32 gate-prefix exchange buffer. With an FP32 gate the caller
+ # passes the raw-gate stage itself (the historical aliasing); with a 16-bit
+ # gate it is the dedicated `gate_exchange_smem` tile.
+ g_prefix_ptr = gate_exchange_smem.data_ptr()
+ k_inv_ptr = k_inv_smem.data_ptr()
+ tcgen05_k_decay_ptr = tcgen05_k_decay_smem.data_ptr()
+ tcgen05_q_decay_ptr = tcgen05_q_decay_smem.data_ptr()
+ tcgen05_k_restore_ptr = tcgen05_k_restore_smem.data_ptr()
+
+ prefix_dim = cg0_local_warp * THREADS_PER_WARP + lane
+ g_prefix_regs = cutlass.Array(
+ cutlass.Float32,
+ BT,
+ alignment=16,
+ )
+ if cutlass.const_expr(SAFE_GATE):
+ # Fold tanh's exact 0.5 scale into the chunk-uniform coefficient.
+ a_log_exp_half = a_log_exp * cutlass.Float32(0.5)
+ for row_pair in cutlass.range_constexpr(BT // 2):
+ row0 = row_pair * 2
+ row1 = row0 + 1
+ # Only the MEMORY FORMAT of the raw gate depends on gate_dtype:
+ # a 16-bit gate lands in the q/k s128 geometry and is widened to
+ # FP32 right here -- ALL gate arithmetic below stays FP32.
+ if cutlass.const_expr(gate_dtype_is_f32(gate_dtype)):
+ prefix_idx0 = raw_f32_s128_smem_index(row0, prefix_dim)
+ prefix_idx1 = raw_f32_s128_smem_index(row1, prefix_dim)
+ gate0 = raw_gate_smem[prefix_idx0]
+ gate1 = raw_gate_smem[prefix_idx1]
+ else:
+ prefix_idx0 = raw_f16_s128_smem_index(row0, prefix_dim)
+ prefix_idx1 = raw_f16_s128_smem_index(row1, prefix_dim)
+ gate0 = raw_gate_smem[prefix_idx0].to(cutlass.Float32)
+ gate1 = raw_gate_smem[prefix_idx1].to(cutlass.Float32)
+ gate0 = a_log_exp_half * (gate0 + dt_bias_value)
+ gate1 = a_log_exp_half * (gate1 + dt_bias_value)
+ gate0 = safe_gate_log2_increment_prehalved(
+ gate0,
+ SAFE_GATE,
+ GATE_SCALE_LOG2,
+ )
+ gate1 = safe_gate_log2_increment_prehalved(
+ gate1,
+ SAFE_GATE,
+ GATE_SCALE_LOG2,
+ )
+ gate_pair = cutlass.Vector.from_elements((gate0, gate1), cutlass.Float32)
+ g_prefix_regs[row0] = gate_pair[0]
+ g_prefix_regs[row1] = gate_pair[1]
+ else:
+ # FLA-compatible non-safe gate: accept raw gate logits + A_log +
+ # dt_bias and compute the log2-domain decay increment in-kernel as
+ # ``-exp(A_log) * softplus(raw_gate + dt_bias) * log2(e)``. a_log_exp
+ # already carries ``exp2(A_log * log2(e)) == exp(A_log)`` and the
+ # softplus is evaluated in base-2 units (softplus_log2_f32).
+ for row in cutlass.range_constexpr(BT):
+ if cutlass.const_expr(gate_dtype_is_f32(gate_dtype)):
+ prefix_idx = raw_f32_s128_smem_index(row, prefix_dim)
+ gate = raw_gate_smem[prefix_idx]
+ else:
+ prefix_idx = raw_f16_s128_smem_index(row, prefix_dim)
+ gate = raw_gate_smem[prefix_idx].to(cutlass.Float32)
+ gate = -a_log_exp * softplus_log2_f32(gate + dt_bias_value)
+ g_prefix_regs[row] = gate
+
+ if cutlass.const_expr(not FULL_CHUNKS):
+ # tail-block peeling: interior (full) chunks run the lean
+ # unconditional gate stream above; only a sequence's genuinely
+ # partial tail chunk (warp-uniform runtime condition, one compare
+ # per chunk) pays the masked zeroing fixup below.
+ tail_valid_rows = seqlen - chunk * cutlass.Int32(BT)
+ if tail_valid_rows < cutlass.Int32(BT):
+ tail_mask_pt = cutlass.vector.create_mask([BT], [tail_valid_rows])
+ for row_pair_pt in cutlass.range_constexpr(BT // 2):
+ row0_pt = row_pair_pt * 2
+ row1_pt = row0_pt + 1
+ gate_pair_pt = cutlass.Vector.from_elements(
+ (g_prefix_regs[row0_pt], g_prefix_regs[row1_pt]),
+ cutlass.Float32,
+ )
+ gate_pair_pt = cutlass.vector.where(
+ tail_mask_pt[row0_pt : row1_pt + 1], gate_pair_pt, 0.0
+ )
+ g_prefix_regs[row0_pt] = gate_pair_pt[0]
+ g_prefix_regs[row1_pt] = gate_pair_pt[1]
+
+ prefix_acc = cutlass.Float32(0.0)
+ for row_pair in cutlass.range_constexpr(BT // 2):
+ row0 = row_pair * 2
+ row1 = row0 + 1
+ # The scalar scan has the same dependency depth as packed fadd2 but
+ # avoids the register copy needed to construct its first input pair.
+ prefix0 = prefix_acc + g_prefix_regs[row0]
+ prefix1 = prefix0 + g_prefix_regs[row1]
+ g_prefix_regs[row0] = prefix0
+ g_prefix_regs[row1] = prefix1
+ prefix_acc = prefix1
+
+ for row in cutlass.range_constexpr(BT):
+ g_prefix_regs[row] = cute.math.exp2(g_prefix_regs[row], fastmath=True)
+
+ for row in cutlass.range_constexpr(BT):
+ prefix_idx = raw_f32_exchange_smem_index(row, prefix_dim)
+ gate_exchange_smem[prefix_idx] = g_prefix_regs[row]
+
+ cg0_sync(cg0_group_id)
+ diag_ready_arrive(diag_ready_stage_mbar)
+
+ k_inv_regs = cutlass.Array(
+ input_dtype,
+ 2 * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=16,
+ )
+ k_restore_all_regs = cutlass.Array(
+ input_dtype,
+ 2 * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=16,
+ )
+ raw_q_regs = cutlass.Array(
+ cutlass.Float32,
+ 2 * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=16,
+ )
+ raw_k_regs = cutlass.Array(
+ cutlass.Float32,
+ 2 * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=16,
+ )
+ q_sum_sq = cutlass.Float32(0.0)
+ k_sum_sq = cutlass.Float32(0.0)
+ for dim_half in cutlass.range_constexpr(2):
+ dim_base = (
+ dim_half * (DK // 2) + lane_in_row_group * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ )
+ reg_base = dim_half * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ raw_f16_idx = raw_f16_s128_smem_index(decay_row, dim_base)
+ raw_q_vec = (raw_q_ptr + raw_f16_idx).load(
+ count=RAW_F16_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=RAW_F16_TMA_SWIZZLE_GROUP_BYTES,
+ )
+ raw_k_vec = (raw_k_ptr + raw_f16_idx).load(
+ count=RAW_F16_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=RAW_F16_TMA_SWIZZLE_GROUP_BYTES,
+ )
+ raw_q_vec_f32 = raw_q_vec.to(cutlass.Float32)
+ raw_k_vec_f32 = raw_k_vec.to(cutlass.Float32)
+ for dim_offset in cutlass.range_constexpr(RAW_F16_TMA_SWIZZLE_GROUP_ELEMS):
+ q_val = raw_q_vec_f32[dim_offset]
+ k_val = raw_k_vec_f32[dim_offset]
+ raw_q_regs[reg_base + dim_offset] = q_val
+ raw_k_regs[reg_base + dim_offset] = k_val
+ q_sum_sq = q_sum_sq + q_val * q_val
+ k_sum_sq = k_sum_sq + k_val * k_val
+
+ q_sum_sq = warp_group_sum_8(q_sum_sq)
+ k_sum_sq = warp_group_sum_8(k_sum_sq)
+ norm_floor_sq = cutlass.Float32(L2_NORM_EPS * L2_NORM_EPS)
+ q_inv_norm = cute.math.rsqrt(
+ cute.math.max(q_sum_sq, norm_floor_sq, ftz=True),
+ fastmath=True,
+ )
+ k_inv_norm = cute.math.rsqrt(
+ cute.math.max(k_sum_sq, norm_floor_sq, ftz=True),
+ fastmath=True,
+ )
+
+ exp_g_regs = cutlass.Array(
+ cutlass.Float32,
+ 2 * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=16,
+ )
+ exp_g_last_regs = cutlass.Array(
+ cutlass.Float32,
+ 2 * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=16,
+ )
+ for dim_half in cutlass.range_constexpr(2):
+ dim_base = (
+ dim_half * (DK // 2) + lane_in_row_group * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ )
+ reg_base = dim_half * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ exp_neg_g_regs = cutlass.Array(
+ cutlass.Float32,
+ RAW_F16_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=16,
+ )
+ for f32_group in cutlass.range_constexpr(
+ RAW_F16_TMA_SWIZZLE_GROUP_ELEMS // RAW_F32_TMA_SWIZZLE_GROUP_ELEMS
+ ):
+ f32_dim_base = dim_base + f32_group * RAW_F32_TMA_SWIZZLE_GROUP_ELEMS
+ g_prefix_idx = raw_f32_exchange_smem_index(decay_row, f32_dim_base)
+ exp_g_vec = (g_prefix_ptr + g_prefix_idx).load(
+ count=RAW_F32_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=RAW_F32_TMA_SWIZZLE_GROUP_BYTES,
+ )
+ exp_g_last_idx = raw_f32_exchange_smem_index(BT - 1, f32_dim_base)
+ exp_g_last_vec = (g_prefix_ptr + exp_g_last_idx).load(
+ count=RAW_F32_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=RAW_F32_TMA_SWIZZLE_GROUP_BYTES,
+ )
+ half_reg_base = f32_group * RAW_F32_TMA_SWIZZLE_GROUP_ELEMS
+ f32_reg_base = reg_base + half_reg_base
+ exp_g_regs[f32_reg_base] = exp_g_vec[0]
+ exp_g_regs[f32_reg_base + 1] = exp_g_vec[1]
+ exp_g_regs[f32_reg_base + 2] = exp_g_vec[2]
+ exp_g_regs[f32_reg_base + 3] = exp_g_vec[3]
+ exp_neg_g_regs[half_reg_base] = cute.math.rcp(
+ exp_g_vec[0], approx=True, ftz=True
+ )
+ exp_neg_g_regs[half_reg_base + 1] = cute.math.rcp(
+ exp_g_vec[1], approx=True, ftz=True
+ )
+ exp_neg_g_regs[half_reg_base + 2] = cute.math.rcp(
+ exp_g_vec[2], approx=True, ftz=True
+ )
+ exp_neg_g_regs[half_reg_base + 3] = cute.math.rcp(
+ exp_g_vec[3], approx=True, ftz=True
+ )
+ exp_g_last_regs[f32_reg_base] = exp_g_last_vec[0]
+ exp_g_last_regs[f32_reg_base + 1] = exp_g_last_vec[1]
+ exp_g_last_regs[f32_reg_base + 2] = exp_g_last_vec[2]
+ exp_g_last_regs[f32_reg_base + 3] = exp_g_last_vec[3]
+
+ k_decay_vec_regs = cutlass.Array(
+ input_dtype,
+ RAW_F16_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=RAW_F16_TMA_SWIZZLE_GROUP_BYTES,
+ )
+ for pair_idx in cutlass.range_constexpr(RAW_F16_TMA_SWIZZLE_GROUP_ELEMS // 2):
+ dim0 = pair_idx * 2
+ dim1 = dim0 + 1
+ raw_reg_idx0 = reg_base + dim0
+ raw_reg_idx1 = reg_base + dim1
+ k_value0, k_value1 = fmul2(
+ (raw_k_regs[raw_reg_idx0], raw_k_regs[raw_reg_idx1]),
+ (k_inv_norm, k_inv_norm),
+ )
+ k_decay0, k_decay1 = fmul2(
+ (k_value0, k_value1),
+ (exp_g_regs[raw_reg_idx0], exp_g_regs[raw_reg_idx1]),
+ )
+ k_inv0, k_inv1 = fmul2(
+ (k_value0, k_value1),
+ (exp_neg_g_regs[dim0], exp_neg_g_regs[dim1]),
+ )
+ k_inv_regs[reg_base + dim0] = k_inv0.to(input_dtype)
+ k_inv_regs[reg_base + dim1] = k_inv1.to(input_dtype)
+ k_restore0, k_restore1 = fmul2(
+ (k_inv0, k_inv1),
+ (
+ exp_g_last_regs[reg_base + dim0],
+ exp_g_last_regs[reg_base + dim1],
+ ),
+ )
+ k_restore_all_regs[reg_base + dim0] = k_restore0.to(input_dtype)
+ k_restore_all_regs[reg_base + dim1] = k_restore1.to(input_dtype)
+ k_decay_vec_regs[dim0] = k_decay0.to(input_dtype)
+ k_decay_vec_regs[dim1] = k_decay1.to(input_dtype)
+
+ k_inv_vec = cutlass.Vector.from_elements(
+ (
+ k_inv_regs[reg_base],
+ k_inv_regs[reg_base + 1],
+ k_inv_regs[reg_base + 2],
+ k_inv_regs[reg_base + 3],
+ k_inv_regs[reg_base + 4],
+ k_inv_regs[reg_base + 5],
+ k_inv_regs[reg_base + 6],
+ k_inv_regs[reg_base + 7],
+ ),
+ input_dtype,
+ )
+ k_decay_vec = cutlass.Vector.from_elements(
+ (
+ k_decay_vec_regs[0],
+ k_decay_vec_regs[1],
+ k_decay_vec_regs[2],
+ k_decay_vec_regs[3],
+ k_decay_vec_regs[4],
+ k_decay_vec_regs[5],
+ k_decay_vec_regs[6],
+ k_decay_vec_regs[7],
+ ),
+ input_dtype,
+ )
+ if cutlass.const_expr(dim_half == 0):
+ operand_smem_consumed_phase = ((chunk // DECAY_STAGE_COUNT) + 1) % 2
+ operand_smem_consumed_wait(
+ operand_smem_consumed_stage_mbar,
+ operand_smem_consumed_phase,
+ )
+ k_inv_swizzled_idx = k_inv_s128_smem_index(decay_row, dim_base)
+ (k_inv_ptr + k_inv_swizzled_idx).store(
+ k_inv_vec,
+ alignment=RAW_F16_TMA_SWIZZLE_GROUP_BYTES,
+ )
+ decay_storage_dim_base = tcgen05_decay_b_key_storage_dim_runtime(
+ decay_row,
+ dim_base,
+ )
+ decay_linear_idx_base = decay_row * DK + decay_storage_dim_base
+ decay_swizzled_idx_base = tcgen05_swizzle_128b_elem_index(
+ decay_linear_idx_base,
+ TCGEN05_F16_ELEM_BYTES,
+ BT,
+ )
+ (tcgen05_k_decay_ptr + decay_swizzled_idx_base).store(
+ k_decay_vec,
+ alignment=RAW_F16_TMA_SWIZZLE_GROUP_BYTES,
+ )
+ if cutlass.const_expr(dim_half == 0):
+ # Publish the first half-DK of k_inv/k_decay early: warp 12's
+ # KK K-blocks 0..3 only read key dims [0, 64).
+ cg0_k_ready_arrive(cg0_k_half_ready_stage_mbar)
+ cg0_k_ready_arrive(cg0_k_ready_stage_mbar)
+
+ for dim_half in cutlass.range_constexpr(2):
+ dim_base = (
+ dim_half * (DK // 2) + lane_in_row_group * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ )
+ reg_base = dim_half * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ q_decay_vec_regs = cutlass.Array(
+ input_dtype,
+ RAW_F16_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=RAW_F16_TMA_SWIZZLE_GROUP_BYTES,
+ )
+ for pair_idx in cutlass.range_constexpr(RAW_F16_TMA_SWIZZLE_GROUP_ELEMS // 2):
+ dim0 = pair_idx * 2
+ dim1 = dim0 + 1
+ raw_reg_idx0 = reg_base + dim0
+ raw_reg_idx1 = reg_base + dim1
+ q_value0, q_value1 = fmul2(
+ (raw_q_regs[raw_reg_idx0], raw_q_regs[raw_reg_idx1]),
+ (q_inv_norm, q_inv_norm),
+ )
+ q_decay0, q_decay1 = fmul2(
+ (q_value0, q_value1),
+ (exp_g_regs[raw_reg_idx0], exp_g_regs[raw_reg_idx1]),
+ )
+ q_decay_vec_regs[dim0] = q_decay0.to(input_dtype)
+ q_decay_vec_regs[dim1] = q_decay1.to(input_dtype)
+
+ q_decay_vec = cutlass.Vector.from_elements(
+ (
+ q_decay_vec_regs[0],
+ q_decay_vec_regs[1],
+ q_decay_vec_regs[2],
+ q_decay_vec_regs[3],
+ q_decay_vec_regs[4],
+ q_decay_vec_regs[5],
+ q_decay_vec_regs[6],
+ q_decay_vec_regs[7],
+ ),
+ input_dtype,
+ )
+ decay_storage_dim_base = tcgen05_decay_b_key_storage_dim_runtime(
+ decay_row,
+ dim_base,
+ )
+ decay_linear_idx_base = decay_row * DK + decay_storage_dim_base
+ decay_swizzled_idx_base = tcgen05_swizzle_128b_elem_index(
+ decay_linear_idx_base,
+ TCGEN05_F16_ELEM_BYTES,
+ BT,
+ )
+ (tcgen05_q_decay_ptr + decay_swizzled_idx_base).store(
+ q_decay_vec,
+ alignment=RAW_F16_TMA_SWIZZLE_GROUP_BYTES,
+ )
+
+ k_restore_consumed_phase = ((chunk // DECAY_STAGE_COUNT) + 1) % 2
+ operand_smem_consumed_wait(
+ k_restore_consumed_stage_mbar,
+ k_restore_consumed_phase,
+ )
+
+ for dim_half in cutlass.range_constexpr(2):
+ dim_base = (
+ dim_half * (DK // 2) + lane_in_row_group * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ )
+ reg_base = dim_half * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ storage_row = decay_row ^ TCGEN05_SW32_BT_HALF_XOR
+ k_restore_idx = raw_f16_s128_smem_index(storage_row, dim_base)
+ k_restore_vec = cutlass.Vector.from_elements(
+ (
+ k_restore_all_regs[reg_base],
+ k_restore_all_regs[reg_base + 1],
+ k_restore_all_regs[reg_base + 2],
+ k_restore_all_regs[reg_base + 3],
+ k_restore_all_regs[reg_base + 4],
+ k_restore_all_regs[reg_base + 5],
+ k_restore_all_regs[reg_base + 6],
+ k_restore_all_regs[reg_base + 7],
+ ),
+ input_dtype,
+ )
+ (tcgen05_k_restore_ptr + k_restore_idx).store(
+ k_restore_vec,
+ alignment=RAW_F16_TMA_SWIZZLE_GROUP_BYTES,
+ )
+
+
+@cute.jit
+def ptx_mma_m16n8k16_b16_f32(
+ a0,
+ a1,
+ a2,
+ a3,
+ b0,
+ b1,
+ c0,
+ c1,
+ c2,
+ c3,
+ input_dtype: cutlass.Constexpr,
+):
+ """Issue `mma.sync.aligned.m16n8k16.row.col.f32.{f16|bf16}.{f16|bf16}.f32`."""
+
+ if cutlass.const_expr(
+ input_dtype != cutlass.Float16 and input_dtype != cutlass.BFloat16
+ ):
+ raise TypeError(f"Invalid auxiliary-MMA input dtype: {input_dtype}")
+ input_tag = "f16" if cutlass.const_expr(input_dtype == cutlass.Float16) else "bf16"
+
+ return cute.arch.inline_ptx(
+ f"mma.sync.aligned.m16n8k16.row.col.f32.{input_tag}.{input_tag}.f32"
+ " {$0,$1,$2,$3}, {$4,$5,$6,$7}, {$8,$9}, {$10,$11,$12,$13};",
+ write_only_types=[
+ cutlass.Float32,
+ cutlass.Float32,
+ cutlass.Float32,
+ cutlass.Float32,
+ ],
+ read_only_args=[a0, a1, a2, a3, b0, b1, c0, c1, c2, c3],
+ )
+
+
+@cute.jit
+def super_mma_accumulator_row(lane, accum_idx: cutlass.Constexpr) -> cutlass.Int32:
+ """Row coordinate for one `m16n8k16` accumulator element."""
+
+ row = lane // 4
+ if cutlass.const_expr(accum_idx >= 2):
+ row = row + cutlass.Int32(8)
+ return row
+
+
+@cute.jit
+def super_mma_accumulator_col(
+ lane,
+ n_block: cutlass.Constexpr,
+ accum_idx: cutlass.Constexpr,
+) -> cutlass.Int32:
+ """Column coordinate for one `m16n8k16` accumulator element."""
+
+ col = n_block * SUPER_MMA_ATOM_N + 2 * (lane % 4)
+ if cutlass.const_expr((accum_idx % 2) == 1):
+ col = col + cutlass.Int32(1)
+ return col
+
+
+@cute.jit
+def super_mma_store_pairwise_tile_stmatrix_x4(
+ pairwise_smem,
+ dst_offset: cutlass.Constexpr,
+ lane,
+ n0_acc,
+ n1_acc,
+ input_dtype: cutlass.Constexpr,
+) -> None:
+ """Store one 16x16 pairwise tile through `stmatrix.m8n8.x4.b16`."""
+
+ prims.stmatrix(
+ pairwise_stmatrix_m8n8x4_ptr(pairwise_smem, dst_offset, lane),
+ [
+ pack_input_b16x2_to_i32(n0_acc[0], n0_acc[1], input_dtype),
+ pack_input_b16x2_to_i32(n0_acc[2], n0_acc[3], input_dtype),
+ pack_input_b16x2_to_i32(n1_acc[0], n1_acc[1], input_dtype),
+ pack_input_b16x2_to_i32(n1_acc[2], n1_acc[3], input_dtype),
+ ],
+ prims.MMALayout.ROW,
+ shape=prims.StoreShape.M8N8,
+ )
+
+
+@cute.jit
+def super_mma_strict_lower_beta_value(
+ value: cutlass.Float32,
+ row_coord,
+ col_coord,
+ beta_scale: cutlass.Float32,
+) -> cutlass.Float32:
+ """Apply `tril(x, -1) * beta[row]` to one pairwise accumulator value."""
+
+ lower = value if row_coord > col_coord else cutlass.Float32(0.0)
+ return lower * beta_scale
+
+
+@cute.jit
+def super_mma_build_l_fragment(
+ raw_beta_smem,
+ lane,
+ n0_acc,
+ n1_acc,
+ l_frag,
+ input_dtype: cutlass.Constexpr,
+) -> None:
+ """Build the packed `L = beta * tril(KK, -1)` registers into `l_frag`."""
+
+ row_lo = super_mma_accumulator_row(lane, 0)
+ row_hi = super_mma_accumulator_row(lane, 2)
+ n0_col0 = super_mma_accumulator_col(lane, 0, 0)
+ n0_col1 = super_mma_accumulator_col(lane, 0, 1)
+ n0_col2 = super_mma_accumulator_col(lane, 0, 2)
+ n0_col3 = super_mma_accumulator_col(lane, 0, 3)
+ n1_col0 = super_mma_accumulator_col(lane, 1, 0)
+ n1_col1 = super_mma_accumulator_col(lane, 1, 1)
+ n1_col2 = super_mma_accumulator_col(lane, 1, 2)
+ n1_col3 = super_mma_accumulator_col(lane, 1, 3)
+ beta_lo = raw_beta_smem[row_lo].to(cutlass.Float32)
+ beta_hi = raw_beta_smem[row_hi].to(cutlass.Float32)
+
+ l_frag[0] = pack_input_b16x2_to_i32(
+ super_mma_strict_lower_beta_value(n0_acc[0], row_lo, n0_col0, beta_lo),
+ super_mma_strict_lower_beta_value(n0_acc[1], row_lo, n0_col1, beta_lo),
+ input_dtype,
+ )
+ l_frag[1] = pack_input_b16x2_to_i32(
+ super_mma_strict_lower_beta_value(n0_acc[2], row_hi, n0_col2, beta_hi),
+ super_mma_strict_lower_beta_value(n0_acc[3], row_hi, n0_col3, beta_hi),
+ input_dtype,
+ )
+ l_frag[2] = pack_input_b16x2_to_i32(
+ super_mma_strict_lower_beta_value(n1_acc[0], row_lo, n1_col0, beta_lo),
+ super_mma_strict_lower_beta_value(n1_acc[1], row_lo, n1_col1, beta_lo),
+ input_dtype,
+ )
+ l_frag[3] = pack_input_b16x2_to_i32(
+ super_mma_strict_lower_beta_value(n1_acc[2], row_hi, n1_col2, beta_hi),
+ super_mma_strict_lower_beta_value(n1_acc[3], row_hi, n1_col3, beta_hi),
+ input_dtype,
+ )
+
+
+@cute.jit
+def super_mma_sw128_decay_ldmatrix_index(
+ row_coord,
+ col_offset,
+ k_block: cutlass.Constexpr,
+):
+ """Return the SW128 row-segment index used by the auxiliary-MMA warp.
+
+ The decay tile is stored in the tcgen05 B-operand K-box-major layout with
+ only the constant half-atom interleave; the tcgen05 descriptor reads the
+ standard SW128 layout, so no row-dependent storage K-block xor exists (the
+ former ``(row & 2) * K_ATOM`` term here matched a write-side xor that the
+ tcgen05 state MMA never saw — both are removed together). ldmatrix reads
+ the 16B row segment at that physical K-block.
+ """
+
+ key_mask = cutlass.Int32(TCGEN05_F16_K_ATOM // 2)
+ logical_key = k_block * TCGEN05_F16_K_ATOM + col_offset
+ storage_key = logical_key ^ key_mask
+ elems_per_128b = cutlass.Int32(TCGEN05_SW128_BYTES // TCGEN05_F16_ELEM_BYTES)
+ storage_slice = storage_key // elems_per_128b
+ key_in_slice = storage_key - storage_slice * elems_per_128b
+ storage_phase = key_in_slice // cutlass.Int32(TCGEN05_F16_K_ATOM)
+ storage_col_offset = key_in_slice - storage_phase * cutlass.Int32(
+ TCGEN05_F16_K_ATOM
+ )
+ row_byte = row_coord * cutlass.Int32(TCGEN05_SW128_BYTES)
+ phase_byte = storage_phase * cutlass.Int32(
+ TCGEN05_F16_K_ATOM * TCGEN05_F16_ELEM_BYTES
+ )
+ col_byte = storage_col_offset * cutlass.Int32(TCGEN05_F16_ELEM_BYTES)
+ byte_in_slice = row_byte + phase_byte + col_byte
+ swizzle_mask = (row_coord & cutlass.Int32(7)) << 4
+ elems_per_slice = cutlass.Int32(BT) * elems_per_128b
+ return storage_slice * elems_per_slice + (
+ (byte_in_slice ^ swizzle_mask) // cutlass.Int32(TCGEN05_F16_ELEM_BYTES)
+ )
+
+
+@cute.jit
+def super_mma_load_decay_lhs_fragment(
+ tcgen05_decay_smem,
+ lane,
+ k_block: cutlass.Constexpr,
+):
+ """Load the A operand fragment for one `m16n8k16` K phase.
+
+ A is logically `[BT, DK]` and physically stored in the tcgen05 SW128 decay
+ operand layout. The register order matches the direct-Q fragment order used
+ by the CUTLASS primitives FMHA examples: rows 0/8 crossed with K cols 0/8.
+ """
+
+ lane_div8 = lane // 8
+ lane_mod8 = lane % 8
+ row_offset = cutlass.Int32(8) if (lane_div8 % 2) else cutlass.Int32(0)
+ col_offset = cutlass.Int32(8) if (lane_div8 // 2) else cutlass.Int32(0)
+ row_coord = lane_mod8 + row_offset
+ swizzled_idx = super_mma_sw128_decay_ldmatrix_index(
+ row_coord,
+ col_offset,
+ k_block,
+ )
+ ptr = tcgen05_decay_smem.subview(swizzled_idx)
+ return prims.ldmatrix(ptr.data_ptr(), 4, prims.MMALayout.ROW)
+
+
+@cute.jit
+def super_mma_load_k_inv_rhs_fragment(
+ k_inv_smem,
+ lane,
+ k_block: cutlass.Constexpr,
+):
+ """Load the B operand fragment for one `m16n8k16` K phase.
+
+ B is stored token-major `[BT, DK]`, which is column-major for the logical
+ `[DK, BT]` MMA RHS. The `ldmatrix.x4` return is split into two N fragments:
+ registers 0/1 for columns 0..7, and registers 2/3 for columns 8..15.
+ """
+
+ lane_div8 = lane // 8
+ lane_div16 = lane // 16
+ lane_mod8 = lane % 8
+ row_offset = cutlass.Int32(8) if lane_div16 else cutlass.Int32(0)
+ col_offset = cutlass.Int32(8) if (lane_div8 % 2) else cutlass.Int32(0)
+ row_coord = lane_mod8 + row_offset
+ col_coord = k_block * SUPER_MMA_ATOM_K + col_offset
+ ptr = k_inv_smem.subview(k_inv_s128_smem_index(row_coord, col_coord))
+ return prims.ldmatrix(ptr.data_ptr(), 4, prims.MMALayout.ROW)
+
+
+@cute.jit
+def super_mma_qk_causal_value(
+ value: cutlass.Float32,
+ lane,
+ n_block: cutlass.Constexpr,
+ accum_idx: cutlass.Constexpr,
+) -> cutlass.Float32:
+ """Zero an accumulator outside the inclusive-lower QK tile."""
+
+ row_coord = super_mma_accumulator_row(lane, accum_idx)
+ col_coord = super_mma_accumulator_col(lane, n_block, accum_idx)
+ return value if row_coord >= col_coord else cutlass.Float32(0.0)
+
+
+@cute.jit
+def super_mma_stage_kk_blocks(
+ tcgen05_k_decay_smem,
+ k_inv_smem,
+ lane,
+ input_dtype: cutlass.Constexpr,
+ k_block_lo: cutlass.Constexpr,
+ k_block_hi: cutlass.Constexpr,
+ kk_n0_acc,
+ kk_n1_acc,
+):
+ """Accumulate KK m16n8k16 K-blocks [k_block_lo, k_block_hi).
+
+ The KK product is split at the half-DK boundary so warp 12 can run
+ K-blocks 0..3 on the `cg0_k_half_ready` arrival (CG0's dim_half==0
+ stores cover key dims [0, 64): k_inv segment 0 and, because the decay
+ storage-key xor mask is 8 or 40, k_decay SW128 slice 0) while CG0 is
+ still staging the second half.
+ """
+
+ for k_block_off in cutlass.range_constexpr(k_block_hi - k_block_lo):
+ k_block = k_block_lo + k_block_off
+ rhs_vec = super_mma_load_k_inv_rhs_fragment(
+ k_inv_smem,
+ lane,
+ k_block,
+ )
+ kk_lhs_vec = super_mma_load_decay_lhs_fragment(
+ tcgen05_k_decay_smem,
+ lane,
+ k_block,
+ )
+
+ kk_n0_d0, kk_n0_d1, kk_n0_d2, kk_n0_d3 = ptx_mma_m16n8k16_b16_f32(
+ kk_lhs_vec[0],
+ kk_lhs_vec[1],
+ kk_lhs_vec[2],
+ kk_lhs_vec[3],
+ rhs_vec[0],
+ rhs_vec[1],
+ kk_n0_acc[0],
+ kk_n0_acc[1],
+ kk_n0_acc[2],
+ kk_n0_acc[3],
+ input_dtype,
+ )
+ kk_n0_acc[0] = kk_n0_d0
+ kk_n0_acc[1] = kk_n0_d1
+ kk_n0_acc[2] = kk_n0_d2
+ kk_n0_acc[3] = kk_n0_d3
+ kk_n1_d0, kk_n1_d1, kk_n1_d2, kk_n1_d3 = ptx_mma_m16n8k16_b16_f32(
+ kk_lhs_vec[0],
+ kk_lhs_vec[1],
+ kk_lhs_vec[2],
+ kk_lhs_vec[3],
+ rhs_vec[2],
+ rhs_vec[3],
+ kk_n1_acc[0],
+ kk_n1_acc[1],
+ kk_n1_acc[2],
+ kk_n1_acc[3],
+ input_dtype,
+ )
+ kk_n1_acc[0] = kk_n1_d0
+ kk_n1_acc[1] = kk_n1_d1
+ kk_n1_acc[2] = kk_n1_d2
+ kk_n1_acc[3] = kk_n1_d3
+
+
+@cute.jit
+def super_mma_stage_qk(
+ tcgen05_q_decay_smem,
+ k_inv_smem,
+ pairwise_smem,
+ lane,
+ input_dtype: cutlass.Constexpr,
+) -> None:
+ """Produce the causal QK tile consumed by the qkv tcgen05 MMA."""
+
+ qk_n0_acc = cutlass.Array(
+ cutlass.Float32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ qk_n1_acc = cutlass.Array(
+ cutlass.Float32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ for accum_idx in cutlass.range_constexpr(SUPER_MMA_ACCUMULATORS_PER_LANE):
+ qk_n0_acc[accum_idx] = cutlass.Float32(0.0)
+ qk_n1_acc[accum_idx] = cutlass.Float32(0.0)
+
+ for k_block in cutlass.range_constexpr(SUPER_MMA_K_BLOCKS):
+ rhs_vec = super_mma_load_k_inv_rhs_fragment(
+ k_inv_smem,
+ lane,
+ k_block,
+ )
+ qk_lhs_vec = super_mma_load_decay_lhs_fragment(
+ tcgen05_q_decay_smem,
+ lane,
+ k_block,
+ )
+
+ qk_n0_d0, qk_n0_d1, qk_n0_d2, qk_n0_d3 = ptx_mma_m16n8k16_b16_f32(
+ qk_lhs_vec[0],
+ qk_lhs_vec[1],
+ qk_lhs_vec[2],
+ qk_lhs_vec[3],
+ rhs_vec[0],
+ rhs_vec[1],
+ qk_n0_acc[0],
+ qk_n0_acc[1],
+ qk_n0_acc[2],
+ qk_n0_acc[3],
+ input_dtype,
+ )
+ qk_n0_acc[0] = qk_n0_d0
+ qk_n0_acc[1] = qk_n0_d1
+ qk_n0_acc[2] = qk_n0_d2
+ qk_n0_acc[3] = qk_n0_d3
+ qk_n1_d0, qk_n1_d1, qk_n1_d2, qk_n1_d3 = ptx_mma_m16n8k16_b16_f32(
+ qk_lhs_vec[0],
+ qk_lhs_vec[1],
+ qk_lhs_vec[2],
+ qk_lhs_vec[3],
+ rhs_vec[2],
+ rhs_vec[3],
+ qk_n1_acc[0],
+ qk_n1_acc[1],
+ qk_n1_acc[2],
+ qk_n1_acc[3],
+ input_dtype,
+ )
+ qk_n1_acc[0] = qk_n1_d0
+ qk_n1_acc[1] = qk_n1_d1
+ qk_n1_acc[2] = qk_n1_d2
+ qk_n1_acc[3] = qk_n1_d3
+
+ qk_n0_acc[0] = super_mma_qk_causal_value(qk_n0_acc[0], lane, 0, 0)
+ qk_n0_acc[1] = super_mma_qk_causal_value(qk_n0_acc[1], lane, 0, 1)
+ qk_n0_acc[2] = super_mma_qk_causal_value(qk_n0_acc[2], lane, 0, 2)
+ qk_n0_acc[3] = super_mma_qk_causal_value(qk_n0_acc[3], lane, 0, 3)
+ qk_n1_acc[0] = super_mma_qk_causal_value(qk_n1_acc[0], lane, 1, 0)
+ qk_n1_acc[1] = super_mma_qk_causal_value(qk_n1_acc[1], lane, 1, 1)
+ qk_n1_acc[2] = super_mma_qk_causal_value(qk_n1_acc[2], lane, 1, 2)
+ qk_n1_acc[3] = super_mma_qk_causal_value(qk_n1_acc[3], lane, 1, 3)
+
+ super_mma_store_pairwise_tile_stmatrix_x4(
+ pairwise_smem,
+ PAIRWISE_SMEM_QK_OFFSET,
+ lane,
+ qk_n0_acc,
+ qk_n1_acc,
+ input_dtype,
+ )
+
+
+@cute.jit
+def super_mma_pairwise_product_from_ab_regs(
+ lhs_frag,
+ rhs_frag,
+ n0_out,
+ n1_out,
+ input_dtype: cutlass.Constexpr,
+) -> None:
+ """Compute `lhs @ rhs` accumulators from packed A-layout fragments.
+
+ Results land in the caller-provided `n0_out`/`n1_out` accumulator arrays.
+ """
+
+ rhs_b0 = movmatrix_b16(rhs_frag[0])
+ rhs_b1 = movmatrix_b16(rhs_frag[1])
+ rhs_b2 = movmatrix_b16(rhs_frag[2])
+ rhs_b3 = movmatrix_b16(rhs_frag[3])
+ zero_acc = cutlass.Array(
+ cutlass.Float32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ for accum_idx in cutlass.range_constexpr(SUPER_MMA_ACCUMULATORS_PER_LANE):
+ zero_acc[accum_idx] = cutlass.Float32(0.0)
+ n0_out[0], n0_out[1], n0_out[2], n0_out[3] = ptx_mma_m16n8k16_b16_f32(
+ lhs_frag[0],
+ lhs_frag[1],
+ lhs_frag[2],
+ lhs_frag[3],
+ rhs_b0,
+ rhs_b1,
+ zero_acc[0],
+ zero_acc[1],
+ zero_acc[2],
+ zero_acc[3],
+ input_dtype,
+ )
+ n1_out[0], n1_out[1], n1_out[2], n1_out[3] = ptx_mma_m16n8k16_b16_f32(
+ lhs_frag[0],
+ lhs_frag[1],
+ lhs_frag[2],
+ lhs_frag[3],
+ rhs_b2,
+ rhs_b3,
+ zero_acc[0],
+ zero_acc[1],
+ zero_acc[2],
+ zero_acc[3],
+ input_dtype,
+ )
+
+
+@cute.jit
+def super_mma_initial_inverse_from_fragment(
+ l_values,
+ lane,
+ n_block: cutlass.Constexpr,
+ accum_idx: cutlass.Constexpr,
+) -> cutlass.Float32:
+ """Return one accumulator element of `I - L` from packed L registers."""
+
+ row_coord = super_mma_accumulator_row(lane, accum_idx)
+ col_coord = super_mma_accumulator_col(lane, n_block, accum_idx)
+ l_value_idx: cutlass.Constexpr[int] = (
+ n_block * SUPER_MMA_ACCUMULATORS_PER_LANE + accum_idx
+ )
+ return pairwise_eye(row_coord, col_coord) - l_values[l_value_idx]
+
+
+@cute.jit
+def super_mma_pack_pairwise_accumulator(
+ n0_acc,
+ n1_acc,
+ out_frag,
+ input_dtype: cutlass.Constexpr,
+) -> None:
+ """Round one 16x16 pairwise accumulator tile into an A-layout b16 fragment.
+
+ The packed registers land in the caller-provided `out_frag` array.
+ """
+
+ out_frag[0] = pack_input_b16x2_to_i32(n0_acc[0], n0_acc[1], input_dtype)
+ out_frag[1] = pack_input_b16x2_to_i32(n0_acc[2], n0_acc[3], input_dtype)
+ out_frag[2] = pack_input_b16x2_to_i32(n1_acc[0], n1_acc[1], input_dtype)
+ out_frag[3] = pack_input_b16x2_to_i32(n1_acc[2], n1_acc[3], input_dtype)
+
+
+@cute.jit
+def super_mma_square_pairwise_fragment(
+ src_frag,
+ out_frag,
+ input_dtype: cutlass.Constexpr,
+) -> None:
+ """Pack b16 registers for `src @ src` into `out_frag` without SMEM roundtrip."""
+
+ n0_acc = cutlass.Array(
+ cutlass.Float32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ n1_acc = cutlass.Array(
+ cutlass.Float32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ super_mma_pairwise_product_from_ab_regs(
+ src_frag,
+ src_frag,
+ n0_acc,
+ n1_acc,
+ input_dtype,
+ )
+ super_mma_pack_pairwise_accumulator(
+ n0_acc,
+ n1_acc,
+ out_frag,
+ input_dtype,
+ )
+
+
+@cute.jit
+def super_mma_update_inverse_with_power_regs(
+ rhs_frag,
+ n0_acc,
+ n1_acc,
+ input_dtype: cutlass.Constexpr,
+) -> None:
+ """Update `inv += inv @ Lpow` in place, keeping `inv` in registers."""
+
+ a_frag = cutlass.Array(
+ cutlass.Int32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ super_mma_pack_pairwise_accumulator(
+ n0_acc,
+ n1_acc,
+ a_frag,
+ input_dtype,
+ )
+ p0_acc = cutlass.Array(
+ cutlass.Float32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ p1_acc = cutlass.Array(
+ cutlass.Float32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ super_mma_pairwise_product_from_ab_regs(
+ a_frag,
+ rhs_frag,
+ p0_acc,
+ p1_acc,
+ input_dtype,
+ )
+
+ for accum_idx in cutlass.range_constexpr(SUPER_MMA_ACCUMULATORS_PER_LANE):
+ n0_acc[accum_idx] = (
+ mma_input_dtype(n0_acc[accum_idx], input_dtype) + p0_acc[accum_idx]
+ )
+ for accum_idx in cutlass.range_constexpr(SUPER_MMA_ACCUMULATORS_PER_LANE):
+ n1_acc[accum_idx] = (
+ mma_input_dtype(n1_acc[accum_idx], input_dtype) + p1_acc[accum_idx]
+ )
+
+
+@cute.jit
+def super_mma_stage_blockwise_inverse(
+ pairwise_smem,
+ lane,
+ l_frag,
+ input_dtype: cutlass.Constexpr,
+) -> None:
+ """Produce `A_inv` via the blockwise inverse, register-carried, all-MMA.
+
+ Hierarchical (GDN-style) block inverse of the unit-lower `A = I + L`,
+ expressed entirely as pairwise tensor-core MMA operations:
+
+ D = blockdiag(L11, L22) (strict-lower 8x8 diagonal blocks)
+ Binv = (I - D)(I + D^2)(I + D^4) -- EXACT: each 8x8 unit-lower
+ block is nilpotent at index 8
+ A^-1 = Binv - Binv @ A21hat @ Binv, A21hat = [[0,0],[A21,0]]
+ -- EXACT: (Binv @ A21hat)^2 = 0
+
+ Numerics: with near-identical keys and high beta the 16x16 Neumann
+ chain's intermediates L^2/L^4/L^8 grow combinatorially large and 16-bit
+ rounding destroys the alternating-sign cancellation that keeps the true
+ inverse bounded (geometric state growth on repeated tokens). The
+ block-diagonal chain never materializes them — D^2/D^4 stay O(10) — so
+ the result sits at the exact-inverse floor: adversarial d_state vs the
+ fp32 pipeline drops from 0.33-0.55 (16x16 fp16 Neumann) to 0.05-0.07.
+ Chain and combine operands are FP16 with FP32 accumulation, never bf16;
+ only the final staged tile is rounded to the input dtype, matching a
+ bf16 INV store.
+
+ Cost: 12 MMA operations (2 squares + 2 updates + 2 combine products) — the same
+ count as the 16x16 Neumann chain with one fewer pack/round/add update
+ stage; measured engine-kernel time is ~0.4% BELOW the Neumann baseline
+ (issue-bound auxiliary-MMA warp region, so avoiding non-MMA work is what
+ matters).
+ """
+
+ l_vec = cutlass.Vector.from_elements(
+ (l_frag[0], l_frag[1], l_frag[2], l_frag[3]),
+ cutlass.Int32,
+ )
+ l_values = l_vec.bitcast(input_dtype).to(cutlass.Float32)
+
+ zero_pair = cutlass.Int32(0)
+ # D fragment = the L fragment with the A21 quadrant zeroed, repacked to
+ # FP16. Slot 1 holds rows r+8 / cols c,c+1 (A21); slot 2 (rows r,
+ # cols c+8,c+9) is already the zero upper-right quadrant of the
+ # strictly-lower L.
+ d_frag = cutlass.Array(
+ cutlass.Int32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ d_frag[0] = pack_input_b16x2_to_i32(l_values[0], l_values[1], cutlass.Float16)
+ d_frag[1] = zero_pair
+ d_frag[2] = zero_pair
+ d_frag[3] = pack_input_b16x2_to_i32(l_values[6], l_values[7], cutlass.Float16)
+
+ # Accumulators start at I - D (= I - L with the A21 quadrant zeroed).
+ n0_acc = cutlass.Array(
+ cutlass.Float32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ n1_acc = cutlass.Array(
+ cutlass.Float32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ for accum_idx in cutlass.range_constexpr(SUPER_MMA_ACCUMULATORS_PER_LANE):
+ n0_acc[accum_idx] = super_mma_initial_inverse_from_fragment(
+ l_values,
+ lane,
+ 0,
+ accum_idx,
+ )
+ n1_acc[accum_idx] = super_mma_initial_inverse_from_fragment(
+ l_values,
+ lane,
+ 1,
+ accum_idx,
+ )
+ n0_acc[2] = cutlass.Float32(0.0)
+ n0_acc[3] = cutlass.Float32(0.0)
+
+ d2_frag = cutlass.Array(
+ cutlass.Int32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ super_mma_square_pairwise_fragment(
+ d_frag,
+ d2_frag,
+ cutlass.Float16,
+ )
+ super_mma_update_inverse_with_power_regs(
+ d2_frag,
+ n0_acc,
+ n1_acc,
+ cutlass.Float16,
+ )
+
+ d4_frag = cutlass.Array(
+ cutlass.Int32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ super_mma_square_pairwise_fragment(
+ d2_frag,
+ d4_frag,
+ cutlass.Float16,
+ )
+ super_mma_update_inverse_with_power_regs(
+ d4_frag,
+ n0_acc,
+ n1_acc,
+ cutlass.Float16,
+ )
+
+ # The accumulators now hold Binv: A11inv in n0[0:2], A22inv in n1[2:4];
+ # the off-diagonal quadrants (n0[2:4], n1[0:2]) stay exactly zero
+ # through the chain (block-diagonal times block-diagonal).
+ binv_frag = cutlass.Array(
+ cutlass.Int32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ binv_frag[0] = pack_input_b16x2_to_i32(n0_acc[0], n0_acc[1], cutlass.Float16)
+ binv_frag[1] = zero_pair
+ binv_frag[2] = zero_pair
+ binv_frag[3] = pack_input_b16x2_to_i32(n1_acc[2], n1_acc[3], cutlass.Float16)
+
+ a21_frag = cutlass.Array(
+ cutlass.Int32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ a21_frag[0] = zero_pair
+ a21_frag[1] = pack_input_b16x2_to_i32(l_values[2], l_values[3], cutlass.Float16)
+ a21_frag[2] = zero_pair
+ a21_frag[3] = zero_pair
+
+ # MMA pair 1: T1 = Binv @ A21hat = [[0,0],[A22inv A21, 0]]
+ t1_n0_acc = cutlass.Array(
+ cutlass.Float32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ t1_n1_acc = cutlass.Array(
+ cutlass.Float32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ super_mma_pairwise_product_from_ab_regs(
+ binv_frag,
+ a21_frag,
+ t1_n0_acc,
+ t1_n1_acc,
+ cutlass.Float16,
+ )
+
+ # MMA pair 2: X21hat = (-T1) @ Binv = [[0,0],[-A22inv A21 A11inv, 0]]
+ t1n_frag = cutlass.Array(
+ cutlass.Int32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ t1n_frag[0] = zero_pair
+ t1n_frag[1] = pack_input_b16x2_to_i32(
+ cutlass.Float32(0.0) - t1_n0_acc[2],
+ cutlass.Float32(0.0) - t1_n0_acc[3],
+ cutlass.Float16,
+ )
+ t1n_frag[2] = zero_pair
+ t1n_frag[3] = zero_pair
+ x21_n0_acc = cutlass.Array(
+ cutlass.Float32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ x21_n1_acc = cutlass.Array(
+ cutlass.Float32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ super_mma_pairwise_product_from_ab_regs(
+ t1n_frag,
+ binv_frag,
+ x21_n0_acc,
+ x21_n1_acc,
+ cutlass.Float16,
+ )
+
+ # A_inv = Binv with the A21 quadrant replaced by X21.
+ n0_acc[2] = x21_n0_acc[2]
+ n0_acc[3] = x21_n0_acc[3]
+
+ super_mma_store_pairwise_tile_stmatrix_x4(
+ pairwise_smem,
+ PAIRWISE_SMEM_AINV_OFFSET,
+ lane,
+ n0_acc,
+ n1_acc,
+ input_dtype,
+ )
+
+
+@cute.jit
+def super_mma_stage_pairwise_pipeline(
+ tcgen05_k_decay_smem,
+ k_inv_smem,
+ pairwise_smem,
+ raw_beta_smem,
+ cg0_k_ready_stage_mbar,
+ cg0_k_ready_phase,
+ lane,
+ input_dtype: cutlass.Constexpr,
+) -> None:
+ """Run the KK/L/inverse sequence inside the auxiliary-MMA warp.
+
+ The caller has already waited `cg0_k_half_ready` (first half-DK of
+ k_inv/k_decay staged); the full `cg0_k_ready` wait happens here, between
+ the two KK half-products.
+ """
+
+ kk_n0_acc = cutlass.Array(
+ cutlass.Float32, SUPER_MMA_ACCUMULATORS_PER_LANE, alignment=16
+ )
+ kk_n1_acc = cutlass.Array(
+ cutlass.Float32, SUPER_MMA_ACCUMULATORS_PER_LANE, alignment=16
+ )
+ for acc_idx in cutlass.range_constexpr(SUPER_MMA_ACCUMULATORS_PER_LANE):
+ kk_n0_acc[acc_idx] = cutlass.Float32(0.0)
+ kk_n1_acc[acc_idx] = cutlass.Float32(0.0)
+ super_mma_stage_kk_blocks(
+ tcgen05_k_decay_smem,
+ k_inv_smem,
+ lane,
+ input_dtype,
+ 0,
+ SUPER_MMA_K_BLOCKS // 2,
+ kk_n0_acc,
+ kk_n1_acc,
+ )
+ cg0_k_ready_wait(
+ cg0_k_ready_stage_mbar,
+ cg0_k_ready_phase,
+ )
+ super_mma_stage_kk_blocks(
+ tcgen05_k_decay_smem,
+ k_inv_smem,
+ lane,
+ input_dtype,
+ SUPER_MMA_K_BLOCKS // 2,
+ SUPER_MMA_K_BLOCKS,
+ kk_n0_acc,
+ kk_n1_acc,
+ )
+ l_frag = cutlass.Array(cutlass.Int32, SUPER_MMA_ACCUMULATORS_PER_LANE, alignment=16)
+ super_mma_build_l_fragment(
+ raw_beta_smem,
+ lane,
+ kk_n0_acc,
+ kk_n1_acc,
+ l_frag,
+ input_dtype,
+ )
+ super_mma_stage_blockwise_inverse(
+ pairwise_smem,
+ lane,
+ l_frag,
+ input_dtype,
+ )
+
+
+@cute.jit
+def pairwise_ready_arrive(pairwise_ready_mbar) -> None:
+ """Signal that the auxiliary-MMA pairwise/A_inv workspace is ready."""
+
+ prims.fence_proxy(
+ prims.Proxy.ASYNC_SHARED,
+ space=prims.SharedSpace.shared_cta,
+ )
+ if prims.elect_sync():
+ prims.mbarrier_arrive(pairwise_ready_mbar)
+
+
+@cute.jit
+def pairwise_ready_wait(pairwise_ready_mbar, pairwise_ready_phase):
+ """Wait for the auxiliary-MMA pairwise/A_inv workspace."""
+
+ while not prims.mbarrier_wait_parity(
+ pairwise_ready_mbar,
+ pairwise_ready_phase,
+ prims.MBarrierWait.TRY,
+ ):
+ pass
+ return pairwise_ready_phase ^ cutlass.Int32(1)
+
+
+@cute.jit
+def pairwise_consumed_arrive(pairwise_consumed_mbar) -> None:
+ """Release a pairwise SMEM stage after tcgen05 consumption."""
+
+ if prims.elect_sync():
+ prims.tcgen05_commit(pairwise_consumed_mbar, group=prims.CTAGroup.CTA_1)
+
+
+@cute.jit
+def pairwise_consumed_wait(pairwise_consumed_mbar, pairwise_consumed_phase):
+ """Wait until one pairwise stage can be overwritten."""
+
+ while not prims.mbarrier_wait_parity(
+ pairwise_consumed_mbar,
+ pairwise_consumed_phase,
+ prims.MBarrierWait.TRY,
+ ):
+ pass
+ return pairwise_consumed_phase ^ cutlass.Int32(1)
+
+
+@cute.jit
+def output_ready_arrive(output_ready_mbar) -> None:
+ """Signal that this CG1 warp has finished staging output SMEM."""
+
+ if prims.elect_sync():
+ prims.mbarrier_arrive(output_ready_mbar)
+
+
+@cute.jit
+def output_ready_wait(output_ready_mbar, output_ready_phase):
+ """Wait until all CG1 warps have staged the output tile."""
+
+ while not prims.mbarrier_wait_parity(
+ output_ready_mbar,
+ output_ready_phase,
+ prims.MBarrierWait.TRY,
+ ):
+ pass
+ return output_ready_phase ^ cutlass.Int32(1)
+
+
+@cute.jit
+def q_k_restore_ready_arrive(q_k_restore_ready_mbar) -> None:
+ """Signal that this CG0 warp has staged q_decay and k_restore."""
+
+ prims.fence_proxy(
+ prims.Proxy.ASYNC_SHARED,
+ space=prims.SharedSpace.shared_cta,
+ )
+ if prims.elect_sync():
+ prims.mbarrier_arrive(q_k_restore_ready_mbar)
+
+
+@cute.jit
+def q_k_restore_ready_wait(q_k_restore_ready_mbar, q_k_restore_ready_phase):
+ """Wait until CG0 has staged q_decay and k_restore."""
+
+ while not prims.mbarrier_wait_parity(
+ q_k_restore_ready_mbar,
+ q_k_restore_ready_phase,
+ prims.MBarrierWait.TRY,
+ ):
+ pass
+ return q_k_restore_ready_phase ^ cutlass.Int32(1)
+
+
+@cute.jit
+def cg0_k_ready_arrive(cg0_k_ready_mbar) -> None:
+ """Publish k_decay and k_inv before CG0 finishes Q operands."""
+
+ prims.fence_proxy(
+ prims.Proxy.ASYNC_SHARED,
+ space=prims.SharedSpace.shared_cta,
+ )
+ if prims.elect_sync():
+ prims.mbarrier_arrive(cg0_k_ready_mbar)
+
+
+@cute.jit
+def cg0_k_ready_wait(cg0_k_ready_mbar, cg0_k_ready_phase):
+ """Wait until CG0 has staged k_decay and k_inv."""
+
+ while not prims.mbarrier_wait_parity(
+ cg0_k_ready_mbar,
+ cg0_k_ready_phase,
+ prims.MBarrierWait.TRY,
+ ):
+ pass
+ return cg0_k_ready_phase ^ cutlass.Int32(1)
+
+
+@cute.jit
+def diag_ready_arrive(diag_ready_mbar) -> None:
+ """Publish the FP32 final-prefix diagonal before K/Q materialization."""
+
+ prims.fence_proxy(
+ prims.Proxy.ASYNC_SHARED,
+ space=prims.SharedSpace.shared_cta,
+ )
+ if prims.elect_sync():
+ prims.mbarrier_arrive(diag_ready_mbar)
+
+
+@cute.jit
+def diag_ready_wait(diag_ready_mbar, diag_ready_phase):
+ """Wait until CG0 has staged the FP32 final-prefix diagonal."""
+
+ while not prims.mbarrier_wait_parity(
+ diag_ready_mbar,
+ diag_ready_phase,
+ prims.MBarrierWait.TRY,
+ ):
+ pass
+ return diag_ready_phase ^ cutlass.Int32(1)
+
+
+@cute.jit
+def raw_ready_arrive(raw_ready_mbar) -> None:
+ """Signal that raw SMEM inputs are ready (k2-chain relay only)."""
+
+ if prims.elect_sync():
+ prims.mbarrier_arrive(raw_ready_mbar)
+
+
+@cute.jit
+def raw_ready_wait(raw_ready_mbar, raw_ready_phase):
+ """Generic mbarrier parity spin-wait.
+
+ The engine-class kernels observe raw readiness directly on the
+ chunk's tma_mbar ring slot (consumer-direct wait); the k2 chain
+ kernels still wait their raw_ready relay ring, and the ws_stored /
+ The deltas_issued relay ring uses this helper too.
+ """
+
+ while not prims.mbarrier_wait_parity(
+ raw_ready_mbar,
+ raw_ready_phase,
+ prims.MBarrierWait.TRY,
+ ):
+ pass
+ return raw_ready_phase ^ cutlass.Int32(1)
+
+
+@cute.jit
+def checkpoint_read_done_arrive(checkpoint_read_done_mbar) -> None:
+ """Signal that one producer warp finished reading checkpoint state."""
+
+ if prims.elect_sync():
+ prims.mbarrier_arrive(checkpoint_read_done_mbar)
+
+
+@cute.jit
+def checkpoint_read_done_wait(checkpoint_read_done_mbar, phase) -> None:
+ """Protect a checkpoint read from the next in-place state update."""
+
+ while not prims.mbarrier_wait_parity(
+ checkpoint_read_done_mbar,
+ phase,
+ prims.MBarrierWait.TRY,
+ ):
+ pass
+
+
+@cute.jit
+def raw_consumed_arrive(raw_consumed_mbar) -> None:
+ """Signal that one participating warp has finished raw-slot reads."""
+
+ if prims.elect_sync():
+ prims.mbarrier_arrive(raw_consumed_mbar)
+
+
+@cute.jit
+def raw_consumed_wait(raw_consumed_mbar, raw_consumed_phase):
+ """Wait until the previous chunk no longer reads raw input SMEM."""
+
+ while not prims.mbarrier_wait_parity(
+ raw_consumed_mbar,
+ raw_consumed_phase,
+ prims.MBarrierWait.TRY,
+ ):
+ pass
+ return raw_consumed_phase ^ cutlass.Int32(1)
+
+
+@cute.jit
+def state_input_ready_arrive(state_input_ready_mbar) -> None:
+ """Signal that this CG1 warp has packed state_as_input TMEM."""
+
+ if prims.elect_sync():
+ prims.mbarrier_arrive(state_input_ready_mbar)
+
+
+@cute.jit
+def state_input_ready_wait(state_input_ready_mbar, state_input_ready_phase):
+ """Wait until all CG1 warps have packed state_as_input TMEM."""
+
+ while not prims.mbarrier_wait_parity(
+ state_input_ready_mbar,
+ state_input_ready_phase,
+ prims.MBarrierWait.TRY,
+ ):
+ pass
+ return state_input_ready_phase ^ cutlass.Int32(1)
+
+
+@cute.jit
+def operand_smem_consumed_arrive(operand_smem_consumed_mbar) -> None:
+ """Signal that this warp no longer needs the current operand SMEM."""
+
+ if prims.elect_sync():
+ prims.mbarrier_arrive(operand_smem_consumed_mbar)
+
+
+@cute.jit
+def operand_smem_consumed_wait(
+ operand_smem_consumed_mbar,
+ operand_smem_consumed_phase,
+):
+ """Wait until previous chunk consumers released operand SMEM."""
+
+ while not prims.mbarrier_wait_parity(
+ operand_smem_consumed_mbar,
+ operand_smem_consumed_phase,
+ prims.MBarrierWait.TRY,
+ ):
+ pass
+ return operand_smem_consumed_phase ^ cutlass.Int32(1)
+
+
+@cute.jit
+def rhs_ready_arrive(rhs_ready_mbar) -> None:
+ """Signal that this CG1 warp has staged RHS for update MMA."""
+
+ if prims.elect_sync():
+ prims.mbarrier_arrive(rhs_ready_mbar)
+
+
+@cute.jit
+def rhs_ready_wait(rhs_ready_mbar, rhs_ready_phase):
+ """Wait until all CG1 warps have staged RHS for update MMA."""
+
+ while not prims.mbarrier_wait_parity(
+ rhs_ready_mbar,
+ rhs_ready_phase,
+ prims.MBarrierWait.TRY,
+ ):
+ pass
+ return rhs_ready_phase ^ cutlass.Int32(1)
+
+
+@cute.jit
+def update_ready_arrive(update_ready_mbar) -> None:
+ """Signal that this CG1 warp has staged update for qkv MMA."""
+
+ if prims.elect_sync():
+ prims.mbarrier_arrive(update_ready_mbar)
+
+
+@cute.jit
+def update_ready_wait(update_ready_mbar, update_ready_phase):
+ """Wait until all CG1 warps have staged update for qkv MMA."""
+
+ while not prims.mbarrier_wait_parity(
+ update_ready_mbar,
+ update_ready_phase,
+ prims.MBarrierWait.TRY,
+ ):
+ pass
+ return update_ready_phase ^ cutlass.Int32(1)
+
+
+@cute.jit
+def output_consumed_arrive(output_consumed_mbar) -> None:
+ """Signal that the epilogue warp has drained current output SMEM."""
+
+ if prims.elect_sync():
+ prims.mbarrier_arrive(output_consumed_mbar)
+
+
+@cute.jit
+def output_consumed_wait(output_consumed_mbar, output_consumed_phase):
+ """Wait until previous output SMEM contents have been stored."""
+
+ while not prims.mbarrier_wait_parity(
+ output_consumed_mbar,
+ output_consumed_phase,
+ prims.MBarrierWait.TRY,
+ ):
+ pass
+ return output_consumed_phase ^ cutlass.Int32(1)
+
+
+@cute.jit
+def final_state_stored_arrive(final_state_stored_mbar) -> None:
+ """Signal that this CG1 warp has finished draining final_state TMEM."""
+
+ if prims.elect_sync():
+ prims.mbarrier_arrive(final_state_stored_mbar)
+
+
+@cute.jit
+def final_state_stored_wait(final_state_stored_mbar, final_state_stored_phase):
+ """Wait until CG1 has drained final_state before TMEM deallocation."""
+
+ while not prims.mbarrier_wait_parity(
+ final_state_stored_mbar,
+ final_state_stored_phase,
+ prims.MBarrierWait.TRY,
+ ):
+ pass
+ return final_state_stored_phase ^ cutlass.Int32(1)
+
+
+@cute.jit
+def pack_output_b16x2_to_i32(
+ value0: cutlass.Float32,
+ value1: cutlass.Float32,
+ output_dtype: cutlass.Constexpr,
+):
+ """Pack two FP32 output values through the compile-time 16-bit output dtype."""
+
+ return (
+ cutlass.Vector.from_elements(
+ (value0, value1),
+ cutlass.Float32,
+ )
+ .to(output_dtype)
+ .bitcast(cutlass.Int32)[0]
+ )
+
+
+@cute.jit
+def tcgen05_qstate_acc_tmem_col_offset(qstate_acc_stage):
+ """Return the runtime TMEM column offset for one qstate acc stage."""
+
+ return (
+ KDA_TMEM_QSTATE_ACC_COL_OFFSET
+ + qstate_acc_stage * KDA_TMEM_QSTATE_ACC_STAGE_STRIDE_COLS
+ )
+
+
+@cute.jit
+def tcgen05_shared_acc_tmem_col_offset(shared_acc_stage):
+ """Return the runtime TMEM column offset for one shared_acc stage."""
+
+ return KDA_TMEM_SHARED_ACC_COL_OFFSET + shared_acc_stage * KDA_TMEM_N16_ACC_COLS
+
+
+@cute.jit
+def tcgen05_shared_acc_stage_from_event(shared_acc_event_id):
+ """Map one shared-acc event to its ring-buffer stage."""
+
+ return shared_acc_event_id % KDA_TMEM_SHARED_ACC_STAGE_COUNT
+
+
+@cute.jit
+def tcgen05_shared_acc_phase_from_event(shared_acc_event_id):
+ """Map one shared-acc event to the selected stage's barrier phase."""
+
+ return (shared_acc_event_id // KDA_TMEM_SHARED_ACC_STAGE_COUNT) % 2
+
+
+@cute.jit
+def tcgen05_shared_input_tmem_col_offset(shared_input_stage):
+ """Return the runtime TMEM column offset for one shared input stage."""
+
+ return (
+ KDA_TMEM_SHARED_INPUT_COL_OFFSET
+ + shared_input_stage * KDA_TMEM_SHARED_INPUT_COLS
+ )
+
+
+@cute.jit
+def advance_ring_stage(
+ stage,
+ step: cutlass.Constexpr,
+ stage_count: cutlass.Constexpr,
+):
+ """Advance a runtime ring index without division or a wrap branch."""
+
+ next_stage = stage + cutlass.Int32(step)
+ wrapped = cutlass.Int32(next_stage >= cutlass.Int32(stage_count))
+ return next_stage - wrapped * cutlass.Int32(stage_count), wrapped
+
+
+@cute.jit
+def tcgen05_store_initial_state_tmem(
+ tmem_raw_addr,
+ initial_state: cute.Tensor | None,
+ state_ckpt: cute.Tensor | None,
+ ckpt_slot,
+ state_slot,
+ bidy,
+ dv_half,
+ warp_idx,
+ lane,
+ HALF: cutlass.Constexpr,
+) -> None:
+ """Initialize recurrent TMEM from the optional external VK state.
+
+ HALF=True is the DV2 chain form (Layout F): this warp's quadrant lanes
+ 0..15 hold state rows dv_half*64 + 16q .. +15; the alignment-16 lanes are
+ zero-filled junk (kept written so the first pack reads defined data).
+ """
+
+ base_col_id = tmem_raw_addr & 0xFFFF
+ base_row_id = tmem_raw_addr >> 16
+ tmem_sp = warp_idx % TCGEN05_STATE_K_TMEM_ROW_BLOCKS
+
+ row_id = base_row_id + tmem_sp * THREADS_PER_WARP
+ if cutlass.const_expr(HALF):
+ value_dim = (
+ dv_half * cutlass.Int32(DV_HALF)
+ + tmem_sp * ROWS_PER_WARP
+ + (lane % ROWS_PER_WARP)
+ )
+ valid_lane = lane < ROWS_PER_WARP
+ else:
+ value_dim = tmem_sp * THREADS_PER_WARP + lane
+ for key_block_start in cutlass.range_constexpr(
+ 0,
+ DK,
+ TCGEN05_FINAL_STATE_TMEM_LOAD_COLS,
+ ):
+ state_block = cutlass.Array(
+ cutlass.Float32,
+ TCGEN05_FINAL_STATE_TMEM_LOAD_COLS,
+ alignment=16,
+ )
+ for col in cutlass.range_constexpr(TCGEN05_FINAL_STATE_TMEM_LOAD_COLS):
+ key_dim = key_block_start + col
+ state_value = cutlass.Float32(0.0)
+ if cutlass.const_expr(initial_state is not None):
+ if cutlass.const_expr(HALF):
+ if valid_lane:
+ state_value = initial_state[
+ state_slot, bidy, value_dim, key_dim
+ ].to(cutlass.Float32)
+ else:
+ state_value = initial_state[
+ state_slot, bidy, value_dim, key_dim
+ ].to(cutlass.Float32)
+ state_block[col] = state_value
+ if cutlass.const_expr(state_ckpt is not None):
+ if (ckpt_slot >= cutlass.Int32(0)) & (
+ ckpt_slot < cutlass.Int32(state_ckpt.shape[0])
+ ):
+ if cutlass.const_expr(HALF):
+ if valid_lane:
+ state_ckpt[ckpt_slot, bidy, value_dim, key_dim] = (
+ state_value.to(state_ckpt.element_type)
+ )
+ else:
+ state_ckpt[ckpt_slot, bidy, value_dim, key_dim] = (
+ state_value.to(state_ckpt.element_type)
+ )
+
+ projection_col_id = base_col_id + KDA_TMEM_STATE_COL_OFFSET + key_block_start
+ block_addr = (row_id << 16) | projection_col_id
+ block_ptr = cutlass.inttoptr(block_addr, 6, cutlass.Float32)
+ prims.tcgen05_st(
+ "32x32b",
+ block_ptr,
+ state_block[0:TCGEN05_FINAL_STATE_TMEM_LOAD_COLS],
+ )
+
+ prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE)
+
+
+@cute.jit
+def pack_state_input_x16(state_block, input_dtype: cutlass.Constexpr):
+ """Pack one 16-column FP32 state fragment into an 8-column A fragment."""
+
+ packed_state = cutlass.Array(
+ cutlass.Int32,
+ TCGEN05_STATE_INPUT_PACKED_COLS,
+ alignment=16,
+ )
+ for packed_col in cutlass.range_constexpr(TCGEN05_STATE_INPUT_PACKED_COLS):
+ source_pair = packed_col ^ TCGEN05_F16_A_TMEM_PAIR_XOR
+ key_dim0 = source_pair * 2
+ key_dim1 = key_dim0 + 1
+ packed_state[packed_col] = pack_input_b16x2_to_i32(
+ state_block[key_dim0],
+ state_block[key_dim1],
+ input_dtype,
+ )
+ return packed_state
+
+
+@cute.jit
+def tcgen05_stage_state_input_tmem(
+ tmem_raw_addr,
+ warp_idx,
+ output_consumed_mbar,
+ output_consumed_phase,
+ input_dtype: cutlass.Constexpr,
+ WAIT_OUTPUT_CONSUMED: cutlass.Constexpr,
+ HALF: cutlass.Constexpr,
+ DEFER_STORE_WAIT: cutlass.Constexpr,
+):
+ """Pack one 64-column half of the recurrent state as a tcgen05 A operand.
+
+ ``HALF`` selects state columns ``[HALF*64, HALF*64+64)``. Splitting the
+ pack lets the left half start as soon as the left half of the previous
+ chunk's final-state delta MMA has committed, overlapping with the right
+ delta half. Steady-state chunks overlap the prior output-stage reuse
+ wait with the asynchronous state loads.
+ """
+
+ base_col_id = tmem_raw_addr & 0xFFFF
+ base_row_id = tmem_raw_addr >> 16
+ tmem_sp = warp_idx % TCGEN05_STATE_K_TMEM_ROW_BLOCKS
+
+ row_addr = (base_row_id + tmem_sp * THREADS_PER_WARP) << 16
+ state_col_id = (
+ base_col_id
+ + KDA_TMEM_STATE_COL_OFFSET
+ + HALF * 4 * TCGEN05_STATE_INPUT_LOAD_COLS
+ )
+ state_ptr0 = cutlass.inttoptr(row_addr | state_col_id, 6, cutlass.Float32)
+ state_ptr1 = cutlass.inttoptr(
+ row_addr | (state_col_id + TCGEN05_STATE_INPUT_LOAD_COLS),
+ 6,
+ cutlass.Float32,
+ )
+ state_ptr2 = cutlass.inttoptr(
+ row_addr | (state_col_id + 2 * TCGEN05_STATE_INPUT_LOAD_COLS),
+ 6,
+ cutlass.Float32,
+ )
+ state_ptr3 = cutlass.inttoptr(
+ row_addr | (state_col_id + 3 * TCGEN05_STATE_INPUT_LOAD_COLS),
+ 6,
+ cutlass.Float32,
+ )
+
+ state0 = prims.tcgen05_ld("32x32b", state_ptr0, num=TCGEN05_STATE_INPUT_LOAD_COLS)
+ state1 = prims.tcgen05_ld("32x32b", state_ptr1, num=TCGEN05_STATE_INPUT_LOAD_COLS)
+ state2 = prims.tcgen05_ld("32x32b", state_ptr2, num=TCGEN05_STATE_INPUT_LOAD_COLS)
+ state3 = prims.tcgen05_ld("32x32b", state_ptr3, num=TCGEN05_STATE_INPUT_LOAD_COLS)
+ if cutlass.const_expr(WAIT_OUTPUT_CONSUMED):
+ output_consumed_wait(
+ output_consumed_mbar,
+ output_consumed_phase,
+ )
+ prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD)
+
+ packed_col_id = (
+ base_col_id
+ + KDA_TMEM_STATE_AS_INPUT_COL_OFFSET
+ + HALF * 4 * TCGEN05_STATE_INPUT_PACKED_COLS
+ )
+ packed_row_addr = base_row_id << 16
+ packed_ptr0 = prims.make_tmem_ptr(packed_row_addr | packed_col_id, cutlass.Int8)
+ packed_ptr1 = prims.make_tmem_ptr(
+ packed_row_addr + packed_col_id + TCGEN05_STATE_INPUT_PACKED_COLS,
+ cutlass.Int8,
+ )
+ packed_ptr2 = prims.make_tmem_ptr(
+ packed_row_addr + packed_col_id + 2 * TCGEN05_STATE_INPUT_PACKED_COLS,
+ cutlass.Int8,
+ )
+ packed_ptr3 = prims.make_tmem_ptr(
+ packed_row_addr + packed_col_id + 3 * TCGEN05_STATE_INPUT_PACKED_COLS,
+ cutlass.Int8,
+ )
+ packed0 = pack_state_input_x16(state0, input_dtype)
+ prims.tcgen05_st("32x32b", packed_ptr0, packed0[0:TCGEN05_STATE_INPUT_PACKED_COLS])
+ packed1 = pack_state_input_x16(state1, input_dtype)
+ prims.tcgen05_st("32x32b", packed_ptr1, packed1[0:TCGEN05_STATE_INPUT_PACKED_COLS])
+ packed2 = pack_state_input_x16(state2, input_dtype)
+ prims.tcgen05_st("32x32b", packed_ptr2, packed2[0:TCGEN05_STATE_INPUT_PACKED_COLS])
+ packed3 = pack_state_input_x16(state3, input_dtype)
+ prims.tcgen05_st("32x32b", packed_ptr3, packed3[0:TCGEN05_STATE_INPUT_PACKED_COLS])
+
+ if cutlass.const_expr(not DEFER_STORE_WAIT):
+ prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE)
+ prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC)
+ return state0, state1, state2, state3
+
+
+@cute.jit
+def tcgen05_scale_state_x16_regs(
+ state_block,
+ state_scale_f32_ptr,
+ key_block_start: cutlass.Constexpr,
+ EXCHANGE_LAYOUT: cutlass.Constexpr,
+):
+ """Build one true-FP32 scaled state fragment without issuing its store."""
+
+ scaled_state = cutlass.Array(
+ cutlass.Float32,
+ TCGEN05_STATE_INPUT_LOAD_COLS,
+ alignment=16,
+ )
+ for vec_idx in cutlass.range_constexpr(TCGEN05_STATE_INPUT_LOAD_COLS // 4):
+ reg_base = vec_idx * 4
+ scale_dim: cutlass.Constexpr = key_block_start + reg_base
+ scale_idx = cutlass.Int32(scale_dim)
+ if cutlass.const_expr(EXCHANGE_LAYOUT):
+ scale_idx = raw_f32_exchange_smem_index(BT - 1, scale_dim)
+ scale = (state_scale_f32_ptr + scale_idx).load(count=4, alignment=16)
+ scaled_state[reg_base], scaled_state[reg_base + 1] = fmul2(
+ (state_block[reg_base], state_block[reg_base + 1]),
+ (scale[0], scale[1]),
+ )
+ scaled_state[reg_base + 2], scaled_state[reg_base + 3] = fmul2(
+ (state_block[reg_base + 2], state_block[reg_base + 3]),
+ (scale[2], scale[3]),
+ )
+ return scaled_state
+
+
+@cute.jit
+def tcgen05_rescale_state_x32(
+ state_block,
+ state_block_ptr,
+ state_scale_f32_ptr,
+ key_block_start: cutlass.Constexpr,
+ EXCHANGE_LAYOUT: cutlass.Constexpr,
+) -> None:
+ """Apply one 32-element FP32 decay fragment and store the state block."""
+
+ block_cols: cutlass.Constexpr = 2 * TCGEN05_STATE_INPUT_LOAD_COLS
+ scaled_state = cutlass.Array(cutlass.Float32, block_cols, alignment=16)
+ for vec_idx in cutlass.range_constexpr(block_cols // 4):
+ reg_base = vec_idx * 4
+ scale_dim: cutlass.Constexpr = key_block_start + reg_base
+ scale_idx = cutlass.Int32(scale_dim)
+ if cutlass.const_expr(EXCHANGE_LAYOUT):
+ scale_idx = raw_f32_exchange_smem_index(BT - 1, scale_dim)
+ scale = (state_scale_f32_ptr + scale_idx).load(count=4, alignment=16)
+ scaled_state[reg_base], scaled_state[reg_base + 1] = fmul2(
+ (state_block[reg_base], state_block[reg_base + 1]),
+ (scale[0], scale[1]),
+ )
+ scaled_state[reg_base + 2], scaled_state[reg_base + 3] = fmul2(
+ (state_block[reg_base + 2], state_block[reg_base + 3]),
+ (scale[2], scale[3]),
+ )
+ prims.tcgen05_st("32x32b", state_block_ptr, scaled_state[0:block_cols])
+
+
+@cute.jit
+def tcgen05_publish_projection_then_rescale_state_regs(
+ tmem_raw_addr,
+ state_scale_f32_smem,
+ state_input_ready_mbar,
+ warp_idx,
+ state0,
+ state1,
+ state2,
+ state3,
+ HALF: cutlass.Constexpr,
+ EXCHANGE_LAYOUT: cutlass.Constexpr,
+) -> None:
+ """Hide a projection-input store behind true-FP32 state scaling.
+
+ The packed projection columns and live recurrent-state columns do not
+ alias. Build the scaled fragments while the preceding packed stores are
+ outstanding, then publish those packed columns before issuing the scaled
+ live-state stores. The final wait/fence remains a separate update join.
+ """
+
+ base_col_id = tmem_raw_addr & 0xFFFF
+ base_row_id = tmem_raw_addr >> 16
+ tmem_sp = warp_idx % TCGEN05_STATE_K_TMEM_ROW_BLOCKS
+ row_addr = (base_row_id + tmem_sp * THREADS_PER_WARP) << 16
+ state_col_id = (
+ base_col_id
+ + KDA_TMEM_STATE_COL_OFFSET
+ + HALF * 4 * TCGEN05_STATE_INPUT_LOAD_COLS
+ )
+ state_ptr0 = cutlass.inttoptr(row_addr | state_col_id, 6, cutlass.Float32)
+ state_ptr1 = cutlass.inttoptr(
+ row_addr | (state_col_id + TCGEN05_STATE_INPUT_LOAD_COLS),
+ 6,
+ cutlass.Float32,
+ )
+ state_ptr2 = cutlass.inttoptr(
+ row_addr | (state_col_id + 2 * TCGEN05_STATE_INPUT_LOAD_COLS),
+ 6,
+ cutlass.Float32,
+ )
+ state_ptr3 = cutlass.inttoptr(
+ row_addr | (state_col_id + 3 * TCGEN05_STATE_INPUT_LOAD_COLS),
+ 6,
+ cutlass.Float32,
+ )
+ state_scale_f32_ptr = state_scale_f32_smem.data_ptr()
+ key_half_start: cutlass.Constexpr = HALF * 4 * TCGEN05_STATE_INPUT_LOAD_COLS
+ scaled0 = tcgen05_scale_state_x16_regs(
+ state0,
+ state_scale_f32_ptr,
+ key_half_start,
+ EXCHANGE_LAYOUT,
+ )
+ scaled1 = tcgen05_scale_state_x16_regs(
+ state1,
+ state_scale_f32_ptr,
+ key_half_start + TCGEN05_STATE_INPUT_LOAD_COLS,
+ EXCHANGE_LAYOUT,
+ )
+ scaled2 = tcgen05_scale_state_x16_regs(
+ state2,
+ state_scale_f32_ptr,
+ key_half_start + 2 * TCGEN05_STATE_INPUT_LOAD_COLS,
+ EXCHANGE_LAYOUT,
+ )
+ scaled3 = tcgen05_scale_state_x16_regs(
+ state3,
+ state_scale_f32_ptr,
+ key_half_start + 3 * TCGEN05_STATE_INPUT_LOAD_COLS,
+ EXCHANGE_LAYOUT,
+ )
+
+ # This wait retires only the already-issued packed projection stores.
+ prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE)
+ prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC)
+ state_input_ready_arrive(state_input_ready_mbar)
+
+ prims.tcgen05_st(
+ "32x32b",
+ state_ptr0,
+ scaled0[0:TCGEN05_STATE_INPUT_LOAD_COLS],
+ )
+ prims.tcgen05_st(
+ "32x32b",
+ state_ptr1,
+ scaled1[0:TCGEN05_STATE_INPUT_LOAD_COLS],
+ )
+ prims.tcgen05_st(
+ "32x32b",
+ state_ptr2,
+ scaled2[0:TCGEN05_STATE_INPUT_LOAD_COLS],
+ )
+ prims.tcgen05_st(
+ "32x32b",
+ state_ptr3,
+ scaled3[0:TCGEN05_STATE_INPUT_LOAD_COLS],
+ )
+ prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE)
+ prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC)
+
+
+@cute.jit
+def tcgen05_pack_rescale_state_half_tmem(
+ tmem_raw_addr,
+ state_scale_f32_smem,
+ state_input_ready_mbar,
+ warp_idx,
+ input_dtype: cutlass.Constexpr,
+ HALF: cutlass.Constexpr,
+ EXCHANGE_LAYOUT: cutlass.Constexpr,
+) -> None:
+ """Pack and FP32-decay one state half from a single pair of TMEM loads.
+
+ The owning CG0 group publishes the packed left projection operand and the
+ decayed live state together. This removes both CG1's duplicate left-half
+ load/pack and CG0's former scale-only reload.
+ """
+
+ base_col_id = tmem_raw_addr & 0xFFFF
+ base_row_id = tmem_raw_addr >> 16
+ tmem_sp = warp_idx % TCGEN05_STATE_K_TMEM_ROW_BLOCKS
+ row_addr = (base_row_id + tmem_sp * THREADS_PER_WARP) << 16
+ block_cols: cutlass.Constexpr = 2 * TCGEN05_STATE_INPUT_LOAD_COLS
+ state_col_id = (
+ base_col_id
+ + KDA_TMEM_STATE_COL_OFFSET
+ + HALF * 4 * TCGEN05_STATE_INPUT_LOAD_COLS
+ )
+ state_ptr0 = cutlass.inttoptr(row_addr | state_col_id, 6, cutlass.Float32)
+ state_ptr1 = cutlass.inttoptr(
+ row_addr | (state_col_id + block_cols), 6, cutlass.Float32
+ )
+
+ packed_col_id = (
+ base_col_id
+ + KDA_TMEM_STATE_AS_INPUT_COL_OFFSET
+ + HALF * 4 * TCGEN05_STATE_INPUT_PACKED_COLS
+ )
+ packed_row_addr = base_row_id << 16
+ packed_ptr0 = prims.make_tmem_ptr(packed_row_addr | packed_col_id, cutlass.Int8)
+ packed_ptr1 = prims.make_tmem_ptr(
+ packed_row_addr | (packed_col_id + TCGEN05_STATE_INPUT_PACKED_COLS),
+ cutlass.Int8,
+ )
+ packed_ptr2 = prims.make_tmem_ptr(
+ packed_row_addr | (packed_col_id + 2 * TCGEN05_STATE_INPUT_PACKED_COLS),
+ cutlass.Int8,
+ )
+ packed_ptr3 = prims.make_tmem_ptr(
+ packed_row_addr | (packed_col_id + 3 * TCGEN05_STATE_INPUT_PACKED_COLS),
+ cutlass.Int8,
+ )
+
+ state0 = prims.tcgen05_ld("32x32b", state_ptr0, num=block_cols)
+ state1 = prims.tcgen05_ld("32x32b", state_ptr1, num=block_cols)
+ prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD)
+ packed0 = pack_state_input_x16(state0[0:TCGEN05_STATE_INPUT_LOAD_COLS], input_dtype)
+ packed1 = pack_state_input_x16(
+ state0[TCGEN05_STATE_INPUT_LOAD_COLS:block_cols], input_dtype
+ )
+ prims.tcgen05_st("32x32b", packed_ptr0, packed0[0:TCGEN05_STATE_INPUT_PACKED_COLS])
+ prims.tcgen05_st("32x32b", packed_ptr1, packed1[0:TCGEN05_STATE_INPUT_PACKED_COLS])
+ packed2 = pack_state_input_x16(state1[0:TCGEN05_STATE_INPUT_LOAD_COLS], input_dtype)
+ packed3 = pack_state_input_x16(
+ state1[TCGEN05_STATE_INPUT_LOAD_COLS:block_cols], input_dtype
+ )
+ prims.tcgen05_st("32x32b", packed_ptr2, packed2[0:TCGEN05_STATE_INPUT_PACKED_COLS])
+ prims.tcgen05_st("32x32b", packed_ptr3, packed3[0:TCGEN05_STATE_INPUT_PACKED_COLS])
+ prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE)
+ prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC)
+ state_input_ready_arrive(state_input_ready_mbar)
+
+ state_scale_f32_ptr = state_scale_f32_smem.data_ptr()
+ key_half_start: cutlass.Constexpr = HALF * 4 * TCGEN05_STATE_INPUT_LOAD_COLS
+ tcgen05_rescale_state_x32(
+ state0,
+ state_ptr0,
+ state_scale_f32_ptr,
+ key_half_start,
+ EXCHANGE_LAYOUT,
+ )
+ tcgen05_rescale_state_x32(
+ state1,
+ state_ptr1,
+ state_scale_f32_ptr,
+ key_half_start + block_cols,
+ EXCHANGE_LAYOUT,
+ )
+ prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE)
+ prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC)
+
+
+@cute.jit
+def tcgen05_stage_state_input_dv2_half_tmem(
+ tmem_raw_addr,
+ warp_idx,
+ input_dtype: cutlass.Constexpr,
+ HALF: cutlass.Constexpr,
+) -> cutlass.Array:
+ """Load one FP32 state half and asynchronously pack its MMA-A image."""
+
+ base_col_id = tmem_raw_addr & 0xFFFF
+ base_row_id = tmem_raw_addr >> 16
+ tmem_sp = warp_idx % TCGEN05_STATE_K_TMEM_ROW_BLOCKS
+ row_addr = (base_row_id + tmem_sp * THREADS_PER_WARP) << 16
+ state_half_col = (
+ base_col_id
+ + KDA_TMEM_STATE_COL_OFFSET
+ + HALF * 4 * TCGEN05_STATE_INPUT_LOAD_COLS
+ )
+ state_ptr = cutlass.inttoptr(
+ row_addr | state_half_col,
+ 6,
+ cutlass.Float32,
+ )
+ state = prims.tcgen05_ld("16x256b", state_ptr, num=8)
+ prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD)
+
+ packed_state = cutlass.Array(cutlass.Int32, 16, alignment=16)
+ for dst_repeat in cutlass.range_constexpr(8):
+ src_base = (dst_repeat ^ 1) * 4
+ dst_base = dst_repeat * 2
+ packed_state[dst_base] = pack_input_b16x2_to_i32(
+ state[src_base],
+ state[src_base + 1],
+ input_dtype,
+ )
+ packed_state[dst_base + 1] = pack_input_b16x2_to_i32(
+ state[src_base + 2],
+ state[src_base + 3],
+ input_dtype,
+ )
+
+ packed_col = (
+ base_col_id
+ + KDA_TMEM_STATE_AS_INPUT_COL_OFFSET
+ + HALF * 4 * TCGEN05_STATE_INPUT_PACKED_COLS
+ )
+ packed_ptr = prims.make_tmem_ptr(
+ (base_row_id << 16) | packed_col,
+ cutlass.Int8,
+ )
+ prims.tcgen05_st("16x128b", packed_ptr, packed_state[0:16])
+ return state
+
+
+@cute.jit
+def tcgen05_rescale_state_dv2_half_regs(
+ tmem_raw_addr,
+ state_scale_f32_smem,
+ warp_idx,
+ state,
+ state_input_ready_mbar,
+ HALF: cutlass.Constexpr,
+) -> None:
+ """Overlap true-FP32 decay with the outstanding packed-state store."""
+
+ base_col_id = tmem_raw_addr & 0xFFFF
+ base_row_id = tmem_raw_addr >> 16
+ tmem_sp = warp_idx % TCGEN05_STATE_K_TMEM_ROW_BLOCKS
+ row_addr = (base_row_id + tmem_sp * THREADS_PER_WARP) << 16
+ state_half_col = (
+ base_col_id
+ + KDA_TMEM_STATE_COL_OFFSET
+ + HALF * 4 * TCGEN05_STATE_INPUT_LOAD_COLS
+ )
+ state_ptr = cutlass.inttoptr(
+ row_addr | state_half_col,
+ 6,
+ cutlass.Float32,
+ )
+ state_scale_f32_ptr = state_scale_f32_smem.data_ptr()
+ key_half_start: cutlass.Constexpr = HALF * 4 * TCGEN05_STATE_INPUT_LOAD_COLS
+ scaled_state = cutlass.Array(cutlass.Float32, 32, alignment=16)
+ lane_col_pair = (cute.arch.lane_idx() % 4) * 2
+ for repeat in cutlass.range_constexpr(8):
+ reg_base = repeat * 4
+ scale_dim = key_half_start + repeat * 8 + lane_col_pair
+ scale = (state_scale_f32_ptr + scale_dim).load(count=2, alignment=8)
+ scaled_state[reg_base], scaled_state[reg_base + 1] = fmul2(
+ (state[reg_base], state[reg_base + 1]),
+ (scale[0], scale[1]),
+ )
+ scaled_state[reg_base + 2], scaled_state[reg_base + 3] = fmul2(
+ (state[reg_base + 2], state[reg_base + 3]),
+ (scale[0], scale[1]),
+ )
+
+ prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE)
+ prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC)
+ state_input_ready_arrive(state_input_ready_mbar)
+ prims.tcgen05_st("16x256b", state_ptr, scaled_state[0:32])
+ prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE)
+ prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC)
+
+
+@cute.jit
+def tcgen05_issue_state_projection_mma(
+ tcgen05_decay_smem,
+ tmem_raw_addr,
+ acc_ready_mbar,
+ tmem_col_offset,
+ input_dtype: cutlass.Constexpr,
+ K_BLOCK_BEGIN: cutlass.Constexpr,
+ K_BLOCK_END: cutlass.Constexpr,
+ INITIAL_SCALE_D: cutlass.Constexpr,
+ COMMIT: cutlass.Constexpr,
+ M_DIM: cutlass.Constexpr,
+) -> None:
+ """Issue state*decay K-slices through tcgen05, optionally committing."""
+
+ tmem_ptr = cutlass.inttoptr(
+ tmem_raw_addr + tmem_col_offset,
+ 6,
+ cutlass.Float32,
+ )
+ idesc = prims.Tcgen05InstrDesc.build(
+ c_dtype=cutlass.Float32,
+ a_dtype=input_dtype,
+ b_dtype=input_dtype,
+ n_dim=BT,
+ m_dim=M_DIM,
+ b_major=0,
+ )
+ desc_k_decay = prims.Tcgen05SmemDesc.build(
+ tcgen05_decay_smem.subview(0),
+ leading_byte_offset=TCGEN05_STATE_K_B_LEADING_BYTES,
+ stride_byte_offset=TCGEN05_STATE_K_B_STRIDE_BYTES,
+ layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B,
+ )
+
+ for k_block in cutlass.range_constexpr(K_BLOCK_BEGIN, K_BLOCK_END):
+ scale_d = INITIAL_SCALE_D or k_block != K_BLOCK_BEGIN
+ k_decay_offset = (
+ k_block % TCGEN05_SW128_K_PHASES_PER_SLICE
+ ) * TCGEN05_STATE_K_B_K_STEP_BYTES + (
+ k_block // TCGEN05_SW128_K_PHASES_PER_SLICE
+ ) * BT * TCGEN05_SW128_BYTES
+ state_a_tmem = prims.make_tmem_ptr(tmem_raw_addr, cutlass.Int8).subview(
+ KDA_TMEM_STATE_AS_INPUT_COL_OFFSET + k_block * (TCGEN05_F16_K_ATOM // 2)
+ )
+ if prims.elect_sync():
+ prims.tcgen05_mma(
+ prims.Tcgen05MMAKind.F16,
+ prims.CTAGroup.CTA_1,
+ tmem_ptr,
+ state_a_tmem,
+ desc_k_decay.advance_start_address(k_decay_offset),
+ idesc,
+ scale_d,
+ )
+
+ if cutlass.const_expr(COMMIT):
+ if prims.elect_sync():
+ prims.tcgen05_commit(acc_ready_mbar, group=prims.CTAGroup.CTA_1)
+
+
+@cute.jit
+def tcgen05_issue_state_k_mma(
+ tcgen05_k_decay_smem,
+ tmem_raw_addr,
+ acc_ready_mbar,
+ shared_acc_stage,
+ input_dtype: cutlass.Constexpr,
+ K_BLOCK_BEGIN: cutlass.Constexpr,
+ K_BLOCK_END: cutlass.Constexpr,
+ INITIAL_SCALE_D: cutlass.Constexpr,
+ COMMIT: cutlass.Constexpr,
+ M_DIM: cutlass.Constexpr,
+) -> None:
+ """Issue a K-slice range of state*k into a scheduled shared_acc stage."""
+
+ tcgen05_issue_state_projection_mma(
+ tcgen05_k_decay_smem,
+ tmem_raw_addr,
+ acc_ready_mbar,
+ tcgen05_shared_acc_tmem_col_offset(shared_acc_stage),
+ input_dtype,
+ K_BLOCK_BEGIN,
+ K_BLOCK_END,
+ INITIAL_SCALE_D,
+ COMMIT,
+ M_DIM,
+ )
+
+
+@cute.jit
+def tcgen05_issue_state_q_mma(
+ tcgen05_q_decay_smem,
+ tmem_raw_addr,
+ operand_smem_consumed_mbar,
+ qstate_acc_stage,
+ input_dtype: cutlass.Constexpr,
+) -> None:
+ """Issue state*q and release q_decay when the tensor core consumes it."""
+
+ tcgen05_issue_state_projection_mma(
+ tcgen05_q_decay_smem,
+ tmem_raw_addr,
+ operand_smem_consumed_mbar,
+ tcgen05_qstate_acc_tmem_col_offset(qstate_acc_stage),
+ input_dtype,
+ 0,
+ DK // TCGEN05_F16_K_ATOM,
+ False,
+ True,
+ DV,
+ )
+
+
+@cute.jit
+def tcgen05_wait_acc_buffer_ready(
+ acc_ready_mbar,
+ acc_ready_phase,
+):
+ """Wait until a producer commit has filled the corresponding TMEM acc tile."""
+
+ while not prims.mbarrier_wait_parity(
+ acc_ready_mbar,
+ acc_ready_phase,
+ prims.MBarrierWait.TRY,
+ ):
+ pass
+ return acc_ready_phase ^ cutlass.Int32(1)
+
+
+@cute.jit
+def tcgen05_rhs_token_pair_from_16x256b_fragment(
+ fragment,
+ reg_idx: cutlass.Constexpr,
+):
+ """Select the lane-local state*k pair needed by the RHS TMEM store."""
+
+ if cutlass.const_expr(reg_idx == 0):
+ return fragment[4], fragment[5]
+ if cutlass.const_expr(reg_idx == 1):
+ return fragment[6], fragment[7]
+ if cutlass.const_expr(reg_idx == 2):
+ return fragment[0], fragment[1]
+ return fragment[2], fragment[3]
+
+
+@cute.jit
+def tcgen05_stage_rhs_input_tmem(
+ tmem_raw_addr,
+ raw_v_smem,
+ raw_beta_smem,
+ warp_idx,
+ lane,
+ shared_acc_stage,
+ shared_input_stage,
+ input_dtype: cutlass.Constexpr,
+) -> None:
+ """Produce RHS = beta * (v - state*k) into shared_input TMEM by 16-row tiles."""
+
+ base_col_id = tmem_raw_addr & 0xFFFF
+ base_row_id = tmem_raw_addr >> 16
+ tmem_sp = warp_idx % TCGEN05_STATE_K_TMEM_ROW_BLOCKS
+
+ projection_col_id = base_col_id + tcgen05_shared_acc_tmem_col_offset(
+ shared_acc_stage
+ )
+ input_col_id = base_col_id + tcgen05_shared_input_tmem_col_offset(
+ shared_input_stage
+ )
+ value_dim_base = tmem_sp * THREADS_PER_WARP
+
+ row_id0 = base_row_id + value_dim_base
+ block_addr0 = (row_id0 << 16) | projection_col_id
+ block_ptr0 = cutlass.inttoptr(block_addr0, 6, cutlass.Float32)
+ state_k0 = prims.tcgen05_ld(
+ "16x256b",
+ block_ptr0,
+ num=2,
+ )
+
+ row_id1 = row_id0 + 16
+ block_addr1 = (row_id1 << 16) | projection_col_id
+ block_ptr1 = cutlass.inttoptr(block_addr1, 6, cutlass.Float32)
+ state_k1 = prims.tcgen05_ld(
+ "16x256b",
+ block_ptr1,
+ num=2,
+ )
+
+ raw_v_regs0 = prims.ldmatrix(
+ raw_v_ldmatrix_trans_ptr(
+ raw_v_smem,
+ value_dim_base,
+ lane,
+ ),
+ 4,
+ prims.MMALayout.COL,
+ )
+ raw_v_regs1 = prims.ldmatrix(
+ raw_v_ldmatrix_trans_ptr(
+ raw_v_smem,
+ value_dim_base + 16,
+ lane,
+ ),
+ 4,
+ prims.MMALayout.COL,
+ )
+ prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD)
+
+ packed_rhs0 = cutlass.Array(
+ cutlass.Int32,
+ 4,
+ space=cutlass.AddressSpace.rmem,
+ )
+ for reg_idx in cutlass.range_constexpr(4):
+ packed_col = (reg_idx // 2) * 4 + (lane & 3)
+ source_pair = packed_col ^ TCGEN05_F16_A_TMEM_PAIR_XOR
+ token0 = source_pair * 2
+ token1 = token0 + 1
+ beta0 = raw_beta_smem[token0].to(cutlass.Float32)
+ beta1 = raw_beta_smem[token1].to(cutlass.Float32)
+ raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1)
+ state_k_val0, state_k_val1 = tcgen05_rhs_token_pair_from_16x256b_fragment(
+ state_k0,
+ reg_idx,
+ )
+ beta_pair = pack_input_b16x2_to_i32(beta0, beta1, input_dtype)
+ state_k_pair = pack_input_b16x2_to_i32(
+ state_k_val0,
+ state_k_val1,
+ input_dtype,
+ )
+ diff_pair = sub_b16x2_input_dtype(
+ raw_v_regs0[raw_matrix],
+ state_k_pair,
+ input_dtype,
+ )
+ packed_rhs0[reg_idx] = mul_b16x2_input_dtype(
+ beta_pair,
+ diff_pair,
+ input_dtype,
+ )
+
+ packed_rhs1 = cutlass.Array(
+ cutlass.Int32,
+ 4,
+ space=cutlass.AddressSpace.rmem,
+ )
+ for reg_idx in cutlass.range_constexpr(4):
+ packed_col = (reg_idx // 2) * 4 + (lane & 3)
+ source_pair = packed_col ^ TCGEN05_F16_A_TMEM_PAIR_XOR
+ token0 = source_pair * 2
+ token1 = token0 + 1
+ beta0 = raw_beta_smem[token0].to(cutlass.Float32)
+ beta1 = raw_beta_smem[token1].to(cutlass.Float32)
+ raw_matrix = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1)
+ state_k_val0, state_k_val1 = tcgen05_rhs_token_pair_from_16x256b_fragment(
+ state_k1,
+ reg_idx,
+ )
+ beta_pair = pack_input_b16x2_to_i32(beta0, beta1, input_dtype)
+ state_k_pair = pack_input_b16x2_to_i32(
+ state_k_val0,
+ state_k_val1,
+ input_dtype,
+ )
+ diff_pair = sub_b16x2_input_dtype(
+ raw_v_regs1[raw_matrix],
+ state_k_pair,
+ input_dtype,
+ )
+ packed_rhs1[reg_idx] = mul_b16x2_input_dtype(
+ beta_pair,
+ diff_pair,
+ input_dtype,
+ )
+
+ input_block_addr0 = (base_row_id << 16) | input_col_id
+ input_block_ptr0 = prims.make_tmem_ptr(input_block_addr0, cutlass.Int8)
+ prims.tcgen05_st("16x128b", input_block_ptr0, packed_rhs0[0:4])
+
+ input_block_addr1 = ((base_row_id + 16) << 16) | input_col_id
+ input_block_ptr1 = prims.make_tmem_ptr(input_block_addr1, cutlass.Int8)
+ prims.tcgen05_st("16x128b", input_block_ptr1, packed_rhs1[0:4])
+
+ prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE)
+ prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC)
+
+
+@cute.jit
+def tcgen05_stage_update_input_tmem(
+ tmem_raw_addr,
+ warp_idx,
+ shared_acc_stage,
+ shared_input_stage,
+ input_dtype: cutlass.Constexpr,
+) -> None:
+ """Move update from shared_acc TMEM into packed shared_input TMEM."""
+
+ base_col_id = tmem_raw_addr & 0xFFFF
+ base_row_id = tmem_raw_addr >> 16
+ tmem_sp = warp_idx % TCGEN05_STATE_K_TMEM_ROW_BLOCKS
+
+ row_id = base_row_id + tmem_sp * THREADS_PER_WARP
+ projection_col_id = base_col_id + tcgen05_shared_acc_tmem_col_offset(
+ shared_acc_stage
+ )
+ block_addr = (row_id << 16) | projection_col_id
+ block_ptr = cutlass.inttoptr(block_addr, 6, cutlass.Float32)
+ update = prims.tcgen05_ld(
+ "32x32b",
+ block_ptr,
+ num=TCGEN05_TMEM_LOAD_COLS,
+ )
+ prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD)
+
+ packed_update = cutlass.Array(
+ cutlass.Int32,
+ KDA_TMEM_SHARED_INPUT_COLS,
+ alignment=16,
+ )
+ for packed_col in cutlass.range_constexpr(KDA_TMEM_SHARED_INPUT_COLS):
+ source_pair = packed_col ^ TCGEN05_F16_A_TMEM_PAIR_XOR
+ token0 = source_pair * 2
+ token1 = token0 + 1
+ packed_update[packed_col] = pack_input_b16x2_to_i32(
+ update[token0],
+ update[token1],
+ input_dtype,
+ )
+
+ col_id = base_col_id + tcgen05_shared_input_tmem_col_offset(shared_input_stage)
+ input_block_addr = (base_row_id << 16) | col_id
+ input_block_ptr = prims.make_tmem_ptr(input_block_addr, cutlass.Int8)
+ prims.tcgen05_st(
+ "32x32b",
+ input_block_ptr,
+ packed_update[0:KDA_TMEM_SHARED_INPUT_COLS],
+ )
+ prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE)
+ prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC)
+
+
+@cute.jit
+def tcgen05_issue_value_pairwise_mma(
+ pairwise_stage_smem,
+ pairwise_tile_offset: cutlass.Constexpr,
+ tmem_raw_addr,
+ acc_ready_mbar,
+ shared_input_stage,
+ tmem_col_offset,
+ scale_d: cutlass.Constexpr,
+ input_dtype: cutlass.Constexpr,
+) -> None:
+ """Issue a `[DV,BT] @ [BT,BT]` value-side tcgen05 MMA.
+
+ The A operand is a staged value tile (`rhs` or `update`) in the scheduled
+ shared_input TMEM slot. The B operand is a pairwise tile staged as
+ `[N=token_i, K=token_j]` for tcgen05.
+ """
+
+ tmem_ptr = cutlass.inttoptr(
+ tmem_raw_addr + tmem_col_offset,
+ 6,
+ cutlass.Float32,
+ )
+ idesc = prims.Tcgen05InstrDesc.build(
+ c_dtype=cutlass.Float32,
+ a_dtype=input_dtype,
+ b_dtype=input_dtype,
+ n_dim=BT,
+ m_dim=DV,
+ b_major=0,
+ )
+ lhs_tmem = prims.make_tmem_ptr(tmem_raw_addr, cutlass.Int8).subview(
+ tcgen05_shared_input_tmem_col_offset(shared_input_stage)
+ )
+ desc_pairwise = prims.Tcgen05SmemDesc.build(
+ pairwise_stage_smem.subview(pairwise_tile_offset),
+ leading_byte_offset=TCGEN05_VALUE_PAIRWISE_B_LEADING_BYTES,
+ stride_byte_offset=TCGEN05_VALUE_PAIRWISE_B_STRIDE_BYTES,
+ layout=prims.Tcgen05SmemSwizzle.SWIZZLE_32B,
+ )
+
+ if prims.elect_sync():
+ prims.tcgen05_mma(
+ prims.Tcgen05MMAKind.F16,
+ prims.CTAGroup.CTA_1,
+ tmem_ptr,
+ lhs_tmem,
+ desc_pairwise,
+ idesc,
+ scale_d,
+ )
+ prims.tcgen05_commit(acc_ready_mbar, group=prims.CTAGroup.CTA_1)
+
+
+@cute.jit
+def tcgen05_issue_update_mma(
+ pairwise_stage_smem,
+ tmem_raw_addr,
+ acc_ready_mbar,
+ shared_acc_stage,
+ shared_input_stage,
+ input_dtype: cutlass.Constexpr,
+) -> None:
+ """Issue update = A_inv @ rhs into a scheduled shared_acc TMEM stage."""
+
+ tcgen05_issue_value_pairwise_mma(
+ pairwise_stage_smem,
+ PAIRWISE_SMEM_AINV_OFFSET,
+ tmem_raw_addr,
+ acc_ready_mbar,
+ shared_input_stage,
+ tcgen05_shared_acc_tmem_col_offset(shared_acc_stage),
+ False,
+ input_dtype,
+ )
+
+
+@cute.jit
+def tcgen05_issue_qkv_mma(
+ pairwise_stage_smem,
+ tmem_raw_addr,
+ acc_ready_mbar,
+ shared_input_stage,
+ qstate_acc_stage,
+ input_dtype: cutlass.Constexpr,
+) -> None:
+ """Accumulate qkv = qk @ update into the live qstate_acc TMEM slot."""
+
+ tcgen05_issue_value_pairwise_mma(
+ pairwise_stage_smem,
+ PAIRWISE_SMEM_QK_OFFSET,
+ tmem_raw_addr,
+ acc_ready_mbar,
+ shared_input_stage,
+ tcgen05_qstate_acc_tmem_col_offset(qstate_acc_stage),
+ True,
+ input_dtype,
+ )
+
+
+@cute.jit
+def tcgen05_load_qstate_output_tmem(
+ tmem_raw_addr,
+ o_smem,
+ warp_idx,
+ lane,
+ o_stage_base,
+ qstate_acc_stage,
+ scale: cutlass.Float32,
+ output_dtype: cutlass.Constexpr,
+) -> None:
+ """Drain `state_q + qkv` from qstate_acc TMEM with STSM.T output staging."""
+
+ base_col_id = tmem_raw_addr & 0xFFFF
+ base_row_id = tmem_raw_addr >> 16
+ tmem_sp = warp_idx % TCGEN05_STATE_K_TMEM_ROW_BLOCKS
+
+ projection_col_id = base_col_id + tcgen05_qstate_acc_tmem_col_offset(
+ qstate_acc_stage
+ )
+ value_dim_base = tmem_sp * THREADS_PER_WARP
+
+ row_id0 = base_row_id + value_dim_base
+ row_id1 = row_id0 + 16
+ block_addr0 = (row_id0 << 16) | projection_col_id
+ block_addr1 = (row_id1 << 16) | projection_col_id
+ block_ptr0 = cutlass.inttoptr(block_addr0, 6, cutlass.Float32)
+ block_ptr1 = cutlass.inttoptr(block_addr1, 6, cutlass.Float32)
+ loaded0 = prims.tcgen05_ld(
+ "16x256b",
+ block_ptr0,
+ num=2,
+ )
+ loaded1 = prims.tcgen05_ld(
+ "16x256b",
+ block_ptr1,
+ num=2,
+ )
+ prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD)
+
+ stsm_regs0 = cutlass.Array(
+ cutlass.Int32,
+ 4,
+ space=cutlass.AddressSpace.rmem,
+ )
+ stsm_regs1 = cutlass.Array(
+ cutlass.Int32,
+ 4,
+ space=cutlass.AddressSpace.rmem,
+ )
+ for reg_idx in cutlass.range_constexpr(4):
+ scaled0_0, scaled0_1 = fmul2(
+ (loaded0[2 * reg_idx], loaded0[2 * reg_idx + 1]),
+ (scale, scale),
+ )
+ scaled1_0, scaled1_1 = fmul2(
+ (loaded1[2 * reg_idx], loaded1[2 * reg_idx + 1]),
+ (scale, scale),
+ )
+ stsm_regs0[reg_idx] = pack_output_b16x2_to_i32(
+ scaled0_0,
+ scaled0_1,
+ output_dtype,
+ )
+ stsm_regs1[reg_idx] = pack_output_b16x2_to_i32(
+ scaled1_0,
+ scaled1_1,
+ output_dtype,
+ )
+
+ smem_dst0 = o_smem_stmatrix_128b_ptr(
+ o_smem,
+ o_stage_base,
+ value_dim_base,
+ lane,
+ )
+ smem_dst1 = o_smem_stmatrix_128b_ptr(
+ o_smem,
+ o_stage_base,
+ value_dim_base + 16,
+ lane,
+ )
+ prims.stmatrix(
+ smem_dst0,
+ stsm_regs0.data_ptr().load(count=4, alignment=4),
+ prims.MMALayout.COL,
+ shape=prims.StoreShape.M8N8,
+ )
+ prims.stmatrix(
+ smem_dst1,
+ stsm_regs1.data_ptr().load(count=4, alignment=4),
+ prims.MMALayout.COL,
+ shape=prims.StoreShape.M8N8,
+ )
+ cute.arch.fence_view_async_shared()
+
+
+@cute.jit
+def epilogue_stage_store(
+ tma_desc_o: cutlass.GridConstant[cuda.TensorMap],
+ o_smem,
+ sequence_start,
+ head_idx,
+ chunk_start,
+ o_stage_base,
+) -> None:
+ """Store the staged `[BT, DV]` output tile to global memory with TMA."""
+
+ global_chunk_start = sequence_start + chunk_start
+ if prims.elect_sync():
+ for value_segment in cutlass.range_constexpr(O_TMA_SEGMENTS):
+ segment_base = (
+ o_stage_base + O_OUT_OFFSET + value_segment * BT * O_TMA_SWIZZLE_ELEMS
+ )
+ o_coord = (
+ cutlass.Int32(value_segment * O_TMA_SWIZZLE_ELEMS),
+ global_chunk_start,
+ head_idx,
+ cutlass.Int32(0),
+ )
+ prims.cp_async_bulk_tensor_global_shared_cta(
+ tma_desc_o.get_ptr(),
+ o_smem.subview(segment_base),
+ o_coord,
+ )
+ prims.cp_async_bulk_commit_group()
+ prims.cp_async_bulk_wait_group(0, read=True)
+ prims.bar_warp_sync(cute.arch.FULL_MASK)
+
+
+@cute.jit
+def epilogue_tail_store(
+ out,
+ o_smem,
+ sequence_start,
+ head_idx,
+ chunk_start,
+ seqlen,
+ o_stage_base,
+ lane,
+) -> None:
+ """Store a partial packed-sequence tail without crossing its boundary."""
+
+ valid_tokens = seqlen - chunk_start
+ for elem_iter in cutlass.range_constexpr((BT * DV) // THREADS_PER_WARP):
+ linear_idx = elem_iter * THREADS_PER_WARP + lane
+ token_coord = linear_idx // DV
+ value_dim = linear_idx - token_coord * DV
+ if token_coord < valid_tokens:
+ smem_idx = o_smem_swizzle_128b_elem_index(
+ o_stage_base,
+ value_dim,
+ token_coord,
+ )
+ out[0, sequence_start + chunk_start + token_coord, head_idx, value_dim] = (
+ o_smem[smem_idx]
+ )
+ prims.bar_warp_sync(cute.arch.FULL_MASK)
+
+
+@cute.jit
+def epilogue_wait_and_store_full_output(
+ tma_desc_o: cutlass.GridConstant[cuda.TensorMap],
+ o_smem,
+ output_ready_mbar,
+ output_consumed_mbar,
+ sequence_start,
+ head_idx,
+ output_chunk,
+ O_STAGES: cutlass.Constexpr,
+):
+ """Drain one full staged output chunk from SMEM with TMA."""
+
+ output_chunk_start = output_chunk * BT
+ o_stage = output_chunk % O_STAGES
+ o_stage_base = o_stage * O_SMEM_STAGE_SIZE
+ output_ready_wait(
+ output_ready_mbar.subview(o_stage),
+ (output_chunk // O_STAGES) % 2,
+ )
+ epilogue_stage_store(
+ tma_desc_o,
+ o_smem,
+ sequence_start,
+ head_idx,
+ output_chunk_start,
+ o_stage_base,
+ )
+ output_consumed_arrive(output_consumed_mbar.subview(o_stage))
+
+
+@cute.jit
+def epilogue_wait_and_store_final_output(
+ tma_desc_o: cutlass.GridConstant[cuda.TensorMap],
+ out,
+ o_smem,
+ output_ready_mbar,
+ output_consumed_mbar,
+ sequence_start,
+ head_idx,
+ seqlen,
+ output_chunk,
+ lane,
+ O_STAGES: cutlass.Constexpr,
+):
+ """Drain the final output chunk, guarding a partial packed tail."""
+
+ if seqlen % BT == 0:
+ epilogue_wait_and_store_full_output(
+ tma_desc_o,
+ o_smem,
+ output_ready_mbar,
+ output_consumed_mbar,
+ sequence_start,
+ head_idx,
+ output_chunk,
+ O_STAGES,
+ )
+ else:
+ output_chunk_start = output_chunk * BT
+ o_stage = output_chunk % O_STAGES
+ o_stage_base = o_stage * O_SMEM_STAGE_SIZE
+ output_ready_wait(
+ output_ready_mbar.subview(o_stage),
+ (output_chunk // O_STAGES) % 2,
+ )
+ epilogue_tail_store(
+ out,
+ o_smem,
+ sequence_start,
+ head_idx,
+ output_chunk_start,
+ seqlen,
+ o_stage_base,
+ lane,
+ )
+ output_consumed_arrive(output_consumed_mbar.subview(o_stage))
+
+
+@cute.jit
+def tcgen05_issue_final_state_delta_mma(
+ tcgen05_k_restore_smem,
+ tmem_raw_addr,
+ k_restore_consumed_l_mbar,
+ k_restore_consumed_mbar,
+ shared_input_stage,
+ input_dtype: cutlass.Constexpr,
+) -> None:
+ """Accumulate final_state += update @ k_restore in two N=64 halves.
+
+ The left half commits to its own mbarrier so the next chunk's left
+ state pack can begin while the right half is still in flight; the
+ right commit (which orders after every prior MMA) keeps the original
+ full-completion contract for CG0 and the right pack.
+ """
+
+ half_n = DK // 2
+ tmem_ptr_l = cutlass.inttoptr(
+ tmem_raw_addr + KDA_TMEM_FINAL_STATE_ACC_COL_OFFSET,
+ 6,
+ cutlass.Float32,
+ )
+ tmem_ptr_r = cutlass.inttoptr(
+ tmem_raw_addr + KDA_TMEM_FINAL_STATE_ACC_COL_OFFSET + half_n,
+ 6,
+ cutlass.Float32,
+ )
+ idesc = prims.Tcgen05InstrDesc.build(
+ c_dtype=cutlass.Float32,
+ a_dtype=input_dtype,
+ b_dtype=input_dtype,
+ n_dim=half_n,
+ m_dim=DV,
+ b_major=1,
+ )
+ update_tmem = prims.make_tmem_ptr(tmem_raw_addr, cutlass.Int8).subview(
+ tcgen05_shared_input_tmem_col_offset(shared_input_stage)
+ )
+ desc_k_restore = prims.Tcgen05SmemDesc.build(
+ tcgen05_k_restore_smem.subview(0),
+ leading_byte_offset=TCGEN05_FINAL_STATE_B_LEADING_BYTES,
+ stride_byte_offset=TCGEN05_FINAL_STATE_B_STRIDE_BYTES,
+ layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B,
+ )
+
+ if prims.elect_sync():
+ prims.tcgen05_mma(
+ prims.Tcgen05MMAKind.F16,
+ prims.CTAGroup.CTA_1,
+ tmem_ptr_l,
+ update_tmem,
+ desc_k_restore,
+ idesc,
+ True,
+ )
+ prims.tcgen05_commit(
+ k_restore_consumed_l_mbar,
+ group=prims.CTAGroup.CTA_1,
+ )
+ prims.tcgen05_mma(
+ prims.Tcgen05MMAKind.F16,
+ prims.CTAGroup.CTA_1,
+ tmem_ptr_r,
+ update_tmem,
+ desc_k_restore.advance_start_address(TCGEN05_FINAL_STATE_B_LEADING_BYTES),
+ idesc,
+ True,
+ )
+ prims.tcgen05_commit(
+ k_restore_consumed_mbar,
+ group=prims.CTAGroup.CTA_1,
+ )
+
+
+@cute.jit
+def tcgen05_store_final_state_tmem(
+ tmem_raw_addr,
+ state_col_offset,
+ final_state: cute.Tensor,
+ bidx,
+ bidy,
+ dv_half,
+ warp_idx,
+ lane,
+ HALF: cutlass.Constexpr,
+) -> None:
+ """Store the live recurrent TMEM state to the VK final_state tensor.
+
+ HALF=True is the DV2 chain form: each M=64 half stores its own rows
+ (Layout F -- only the quadrant lanes 0..15 carry state).
+ """
+
+ base_col_id = tmem_raw_addr & 0xFFFF
+ base_row_id = tmem_raw_addr >> 16
+ tmem_sp = warp_idx % TCGEN05_STATE_K_TMEM_ROW_BLOCKS
+
+ row_id = base_row_id + tmem_sp * THREADS_PER_WARP
+ if cutlass.const_expr(HALF):
+ value_dim = (
+ dv_half * cutlass.Int32(DV_HALF)
+ + tmem_sp * ROWS_PER_WARP
+ + (lane % ROWS_PER_WARP)
+ )
+ valid_lane = lane < ROWS_PER_WARP
+ else:
+ value_dim = tmem_sp * THREADS_PER_WARP + lane
+ for key_block_start in cutlass.range_constexpr(
+ 0,
+ DK,
+ TCGEN05_FINAL_STATE_TMEM_LOAD_COLS,
+ ):
+ projection_col_id = base_col_id + state_col_offset + key_block_start
+ block_addr = (row_id << 16) | projection_col_id
+ block_ptr = cutlass.inttoptr(block_addr, 6, cutlass.Float32)
+ loaded = prims.tcgen05_ld(
+ "32x32b",
+ block_ptr,
+ num=TCGEN05_FINAL_STATE_TMEM_LOAD_COLS,
+ )
+ prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD)
+
+ for col in cutlass.range_constexpr(TCGEN05_FINAL_STATE_TMEM_LOAD_COLS):
+ key_dim = key_block_start + col
+ if cutlass.const_expr(HALF):
+ if valid_lane:
+ final_state[bidx, bidy, value_dim, key_dim] = loaded[col].to(
+ final_state.element_type
+ )
+ else:
+ final_state[bidx, bidy, value_dim, key_dim] = loaded[col].to(
+ final_state.element_type
+ )
+
+
+PACKED_METADATA_THREADS: int = 256
+
+
+@cute.kernel
+def packed_metadata_kernel(
+ cu_seqlens: cute.Tensor,
+ seq_order: cute.Tensor,
+ cu_chunks: cute.Tensor,
+ generate_seq_order: cutlass.Int32,
+ generate_cu_chunks: cutlass.Int32,
+) -> None:
+ """Build packed scheduling metadata from device-resident prefix offsets."""
+
+ tidx, _, _ = cute.arch.thread_idx()
+ n_seq = cutlass.Int32(seq_order.shape[0])
+ if generate_seq_order != cutlass.Int32(0):
+ for index in cutlass.range(tidx, n_seq, PACKED_METADATA_THREADS):
+ seq_order[index] = cutlass.Int32(index)
+ cute.arch.sync_threads()
+
+ # Stable descending-length order. The issue workload has at most 64
+ # sequences; odd-even transposition keeps the implementation compact
+ # and supports a runtime sequence count without temporary storage.
+ for phase in cutlass.range(n_seq):
+ parity = phase & cutlass.Int32(1)
+ pair_count = (n_seq - parity) // cutlass.Int32(2)
+ for pair in cutlass.range(tidx, pair_count, PACKED_METADATA_THREADS):
+ left = parity + pair * cutlass.Int32(2)
+ right = left + cutlass.Int32(1)
+ left_seq = cutlass.Int32(seq_order[left])
+ right_seq = cutlass.Int32(seq_order[right])
+ left_len = cu_seqlens[left_seq + 1] - cu_seqlens[left_seq]
+ right_len = cu_seqlens[right_seq + 1] - cu_seqlens[right_seq]
+ swap = (right_len > left_len) | (
+ (right_len == left_len) & (right_seq < left_seq)
+ )
+ if swap:
+ seq_order[left] = right_seq
+ seq_order[right] = left_seq
+ cute.arch.sync_threads()
+
+ if (generate_cu_chunks != cutlass.Int32(0)) & (tidx == cutlass.Int32(0)):
+ cu_chunks[0] = cutlass.Int32(0)
+ total_chunks = cutlass.Int32(0)
+ for index in cutlass.range(n_seq):
+ seq_len = cutlass.Int32(cu_seqlens[index + 1] - cu_seqlens[index])
+ total_chunks += (seq_len + cutlass.Int32(BT - 1)) // cutlass.Int32(BT)
+ cu_chunks[index + 1] = total_chunks
+
+
+@cute.kernel
+def kernel(
+ tma_desc_q: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_k: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_v: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_gate: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_beta: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_o: cutlass.GridConstant[cuda.TensorMap],
+ q: cute.Tensor,
+ k: cute.Tensor,
+ v: cute.Tensor,
+ raw_gate: cute.Tensor,
+ a_log: cute.Tensor,
+ dt_bias: cute.Tensor,
+ beta: cute.Tensor,
+ cu_seqlens: cute.Tensor,
+ seq_order: cute.Tensor,
+ state_indices: cute.Tensor | None,
+ initial_state: cute.Tensor | None,
+ out: cute.Tensor,
+ final_state: cute.Tensor | None,
+ SCALE: cutlass.Float32,
+ state_ckpt: cute.Tensor | None,
+ cu_ckpts: cute.Tensor | None,
+ checkpoint_state_indices: cute.Tensor | None,
+ checkpoint_stride_chunks: cutlass.Int32,
+ SAFE_GATE: cutlass.Constexpr,
+ GATE_SCALE_LOG2: cutlass.Constexpr,
+ gate_dtype: cutlass.Constexpr,
+) -> None:
+ """BT=16 KDA forward kernel.
+
+ Grid: `(heads, num_sequences, 1)`. Each CTA owns one packed sequence/head
+ and iterates over `ceil(sequence_length / 16)` chunks in order.
+ """
+
+ tidx, _, _ = cute.arch.thread_idx()
+ # Put heads in grid-x so every sequence's head CTAs are contiguous in the
+ # linear launch order. sequence_slot follows the automatic longest-first
+ # permutation; bidx/bidy remain the original sequence/head indices below.
+ bidy, sequence_slot, _ = cute.arch.block_idx()
+ bidx = cutlass.Int32(seq_order[sequence_slot])
+ warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+ lane = tidx % THREADS_PER_WARP
+
+ sequence_start = cutlass.Int32(cu_seqlens[bidx])
+ # beta TMA transport shape: heads from grid-x; g = 8/gcd(heads, 8) via the
+ # lowest set bit (g==1 -> 8-head-group box, g>1 -> pair-packed rows).
+ heads32 = cutlass.Int32(cute.arch.grid_dim()[0])
+ beta_lsb = heads32 & (-heads32)
+ beta_g = cutlass.Int32(8) // cutlass.min(beta_lsb, cutlass.Int32(8))
+ sequence_end = cutlass.Int32(cu_seqlens[bidx + 1])
+ seqlen = sequence_end - sequence_start
+ num_chunks = cute.ceil_div(seqlen, BT)
+ input_dtype = q.element_type
+ if cutlass.const_expr(
+ k.element_type != input_dtype or v.element_type != input_dtype
+ ):
+ raise TypeError(
+ "KDA CUTLASS primitives kernel expects q/k/v to use the same 16-bit dtype"
+ )
+ if cutlass.const_expr(raw_gate.element_type != gate_dtype):
+ raise TypeError(
+ "KDA CUTLASS primitives kernel expects raw_gate to match gate_dtype"
+ )
+ if cutlass.const_expr(beta.element_type != cutlass.BFloat16):
+ raise TypeError("KDA CUTLASS primitives kernel expects beta logits to use BF16")
+ if cutlass.const_expr(cu_seqlens.element_type != cutlass.Int64):
+ raise TypeError("KDA CUTLASS primitives kernel expects cu_seqlens to use int64")
+ if cutlass.const_expr(initial_state is not None):
+ if cutlass.const_expr(
+ initial_state.element_type not in (cutlass.BFloat16, cutlass.Float32)
+ ):
+ raise TypeError(
+ "KDA CUTLASS primitives kernel input state dtype must be BF16 or FP32"
+ )
+ if cutlass.const_expr(final_state is not None):
+ if cutlass.const_expr(
+ final_state.element_type not in (cutlass.BFloat16, cutlass.Float32)
+ ):
+ raise TypeError(
+ "KDA CUTLASS primitives kernel output state dtype must be BF16 or FP32"
+ )
+ if cutlass.const_expr(initial_state is not None and final_state is not None):
+ if cutlass.const_expr(initial_state.element_type != final_state.element_type):
+ raise TypeError(
+ "KDA CUTLASS primitives kernel expects matching state input/output dtypes"
+ )
+ # Buffers are declaration-ordered and intentionally non-aliased.
+ tma_mbar = cutlass.Array(
+ cutlass.Int64,
+ TMA_MBAR_STAGE_COUNT,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ qstate_acc_ready_mbar = cutlass.Array(
+ cutlass.Int64,
+ 1,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ output_ready_mbar = cutlass.Array(
+ cutlass.Int64,
+ O_STAGE_COUNT,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ q_k_restore_ready_mbar = cutlass.Array(
+ cutlass.Int64,
+ Q_K_RESTORE_READY_STAGE_COUNT,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ cg0_k_ready_mbar = cutlass.Array(
+ cutlass.Int64,
+ DECAY_STAGE_COUNT,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ cg0_k_half_ready_mbar = cutlass.Array(
+ cutlass.Int64,
+ DECAY_STAGE_COUNT,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ diag_ready_mbar = cutlass.Array(
+ cutlass.Int64,
+ RAW_STAGE_COUNT,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ raw_consumed_mbar = cutlass.Array(
+ cutlass.Int64,
+ RAW_STAGE_COUNT,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ state_input_ready_mbar = cutlass.Array(
+ cutlass.Int64,
+ 1,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ state_input_ready_l_mbar = cutlass.Array(
+ cutlass.Int64,
+ 1,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ initial_state_ready_mbar = cutlass.Array(
+ cutlass.Int64,
+ 1,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ checkpoint_read_done_mbar = cutlass.Array(
+ cutlass.Int64,
+ 1,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ operand_smem_consumed_mbar = cutlass.Array(
+ cutlass.Int64,
+ DECAY_STAGE_COUNT,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ rhs_ready_mbar = cutlass.Array(
+ cutlass.Int64,
+ 1,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ update_ready_mbar = cutlass.Array(
+ cutlass.Int64,
+ 1,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ output_consumed_mbar = cutlass.Array(
+ cutlass.Int64,
+ O_STAGE_COUNT,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ final_state_stored_mbar = cutlass.Array(
+ cutlass.Int64,
+ 1,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ pairwise_ready_mbar = cutlass.Array(
+ cutlass.Int64,
+ PAIRWISE_STAGE_COUNT,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ qk_ready_mbar = cutlass.Array(
+ cutlass.Int64,
+ PAIRWISE_STAGE_COUNT,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ pairwise_consumed_mbar = cutlass.Array(
+ cutlass.Int64,
+ PAIRWISE_STAGE_COUNT,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ k_restore_consumed_mbar = cutlass.Array(
+ cutlass.Int64,
+ DECAY_STAGE_COUNT,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ k_restore_consumed_l_mbar = cutlass.Array(
+ cutlass.Int64,
+ DECAY_STAGE_COUNT,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ shared_acc_ready_mbar = cutlass.Array(
+ cutlass.Int64,
+ KDA_TMEM_SHARED_ACC_STAGE_COUNT,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ # The hand-written K-box-major SW128 mapping is normalized to phase 0,
+ # so both tcgen05 and ldmatrix can share 1KB-aligned operand buffers.
+ tmem_ptr_i32 = cutlass.Array(
+ cutlass.Int32,
+ 1,
+ space=cutlass.AddressSpace.smem,
+ alignment=4,
+ )
+ tcgen05_k_decay_smem = cutlass.Array(
+ input_dtype,
+ TCGEN05_K_DECAY_SMEM_TILE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=1024,
+ )
+ tcgen05_q_decay_smem = cutlass.Array(
+ input_dtype,
+ TCGEN05_Q_DECAY_SMEM_TILE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=1024,
+ )
+ tcgen05_k_restore_smem = cutlass.Array(
+ input_dtype,
+ TCGEN05_K_RESTORE_SMEM_TILE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=RAW_F16_TMA_SWIZZLE_ALIGNMENT_BYTES,
+ )
+ pairwise_smem = cutlass.Array(
+ input_dtype,
+ PAIRWISE_SMEM_TILE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=1024,
+ )
+ raw_q_smem = cutlass.Array(
+ q.element_type,
+ RAW_Q_SMEM_TILE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=RAW_F16_TMA_SWIZZLE_ALIGNMENT_BYTES,
+ )
+ raw_k_smem = cutlass.Array(
+ k.element_type,
+ RAW_K_SMEM_TILE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=RAW_F16_TMA_SWIZZLE_ALIGNMENT_BYTES,
+ )
+ raw_v_smem = cutlass.Array(
+ v.element_type,
+ RAW_V_SMEM_TILE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=RAW_F16_TMA_SWIZZLE_ALIGNMENT_BYTES,
+ )
+ # Gate SMEM. Sizes above are ELEMENT counts, so the ring's byte size
+ # follows gate_dtype (FP32: 8 x 8 KB = 64 KB; BF16: 8 x 4 KB = 32 KB).
+ # Under FP32 the exp2(g_prefix) exchange keeps aliasing this ring exactly
+ # as it always did and NOTHING extra is allocated; a 16-bit gate cannot
+ # host an FP32 exchange, so it gets a dedicated 4-deep FP32 ring (16 KB).
+ if cutlass.const_expr(gate_dtype_is_f32(gate_dtype)):
+ raw_gate_smem = cutlass.Array(
+ cutlass.Float32,
+ RAW_GATE_SMEM_TILE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=RAW_F32_TMA_SWIZZLE_ALIGNMENT_BYTES,
+ )
+ gate_exchange_smem = None
+ else:
+ raw_gate_smem = cutlass.Array(
+ gate_dtype,
+ RAW_GATE_SMEM_TILE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=RAW_F16_TMA_SWIZZLE_ALIGNMENT_BYTES,
+ )
+ gate_exchange_smem = cutlass.Array(
+ cutlass.Float32,
+ GATE_EXCHANGE_SMEM_TILE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=RAW_F32_TMA_SWIZZLE_ALIGNMENT_BYTES,
+ )
+ k_inv_smem = cutlass.Array(
+ input_dtype,
+ K_INV_SMEM_TILE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=RAW_F16_TMA_SWIZZLE_ALIGNMENT_BYTES,
+ )
+ o_smem = cutlass.Array(
+ out.element_type,
+ O_SMEM_TILE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ # The scalar CG1 store computes W128 offsets relative to this buffer.
+ # Align to the full s128b period so absolute SMEM address bits do not
+ # add a hidden phase to the TMA store-side swizzle.
+ alignment=O_TMA_SWIZZLE_ALIGNMENT_BYTES,
+ )
+ raw_dt_bias_smem = cutlass.Array(
+ cutlass.Float32,
+ RAW_DT_BIAS_SMEM_TILE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=1024,
+ )
+ raw_beta_smem = cutlass.Array(
+ cutlass.Float32,
+ RAW_BETA_SMEM_TILE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=1024,
+ )
+ beta_tile_smem = cutlass.Array(
+ cutlass.BFloat16,
+ BETA_TILE_STAGE_COUNT * BETA_TILE_STAGE_ELEMS,
+ space=cutlass.AddressSpace.smem,
+ alignment=1024,
+ )
+ beta_tile_mbar = cutlass.Array(
+ cutlass.Int64,
+ BETA_TILE_STAGE_COUNT,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ # Actual bytes one beta tile transfers (the ring stage is padded to a
+ # fixed stride; expect_tx must be the real box size).
+ beta_tile_tx = cutlass.Int32(8 * BT * 2)
+ if beta_g > cutlass.Int32(1):
+ beta_tile_tx = heads32 * (cutlass.Int32(BT) + beta_g) * cutlass.Int32(2)
+ tma_tx_bytes = cutlass.const_expr(
+ DK * BT * q.element_type.width // 8
+ + DK * BT * k.element_type.width // 8
+ + DV * BT * v.element_type.width // 8
+ + DK * BT * raw_gate.element_type.width // 8
+ )
+
+ if warp_idx == ROLES.tma_load:
+ if prims.elect_sync():
+ for stage in cutlass.range_constexpr(TMA_MBAR_STAGE_COUNT):
+ prims.mbarrier_init(tma_mbar.subview(stage), 1)
+ for stage in cutlass.range_constexpr(RAW_STAGE_COUNT):
+ # Four owner-CG0 warps consume q/k/gate, four CG1 warps
+ # consume v/beta, and the standalone auxiliary-MMA warp consumes
+ # beta while constructing the pairwise tile.
+ prims.mbarrier_init(raw_consumed_mbar.subview(stage), 9)
+ for stage in cutlass.range_constexpr(BETA_TILE_STAGE_COUNT):
+ prims.mbarrier_init(beta_tile_mbar.subview(stage), 1)
+ elif warp_idx == ROLES.tcgen05_mma:
+ if prims.elect_sync():
+ prims.mbarrier_init(qstate_acc_ready_mbar, 1)
+ for stage in cutlass.range_constexpr(KDA_TMEM_SHARED_ACC_STAGE_COUNT):
+ prims.mbarrier_init(shared_acc_ready_mbar.subview(stage), 1)
+ prims.mbarrier_init(state_input_ready_mbar, 4)
+ prims.mbarrier_init(state_input_ready_l_mbar, 4)
+ prims.mbarrier_init(initial_state_ready_mbar, 4)
+ if cutlass.const_expr(state_ckpt is not None):
+ prims.mbarrier_init(checkpoint_read_done_mbar, 4)
+ for stage in cutlass.range_constexpr(DECAY_STAGE_COUNT):
+ prims.mbarrier_init(operand_smem_consumed_mbar.subview(stage), 3)
+ prims.mbarrier_init(k_restore_consumed_mbar.subview(stage), 1)
+ prims.mbarrier_init(k_restore_consumed_l_mbar.subview(stage), 1)
+ prims.mbarrier_init(rhs_ready_mbar, 4)
+ prims.mbarrier_init(update_ready_mbar, 8)
+ prims.mbarrier_init(final_state_stored_mbar, 4)
+ elif warp_idx == ROLES.super_mma:
+ if prims.elect_sync():
+ for stage in cutlass.range_constexpr(PAIRWISE_STAGE_COUNT):
+ prims.mbarrier_init(pairwise_ready_mbar.subview(stage), 1)
+ prims.mbarrier_init(qk_ready_mbar.subview(stage), 1)
+ prims.mbarrier_init(pairwise_consumed_mbar.subview(stage), 1)
+ for stage in cutlass.range_constexpr(Q_K_RESTORE_READY_STAGE_COUNT):
+ prims.mbarrier_init(q_k_restore_ready_mbar.subview(stage), 4)
+ for stage in cutlass.range_constexpr(DECAY_STAGE_COUNT):
+ prims.mbarrier_init(cg0_k_ready_mbar.subview(stage), 4)
+ prims.mbarrier_init(cg0_k_half_ready_mbar.subview(stage), 4)
+ for stage in cutlass.range_constexpr(RAW_STAGE_COUNT):
+ prims.mbarrier_init(diag_ready_mbar.subview(stage), 4)
+ elif warp_idx == ROLES.epilogue:
+ if prims.elect_sync():
+ for stage in cutlass.range_constexpr(O_STAGE_COUNT):
+ prims.mbarrier_init(output_ready_mbar.subview(stage), 4)
+ prims.mbarrier_init(output_consumed_mbar.subview(stage), 1)
+ prims.fence_mbarrier_init()
+ cta_sync()
+ if is_tmem_user_warp(warp_idx):
+ if warp_idx == ROLES.tcgen05_mma:
+ prims.tcgen05_alloc(tmem_ptr_i32, KDA_TMEM_ALLOC_COLS, group="cta_1")
+ tmem_user_sync()
+ if warp_idx == ROLES.tcgen05_mma:
+ prims.tcgen05_relinquish_alloc_permit(group="cta_1")
+ tmem_user_sync()
+ qstate_acc_ready_phase = cutlass.Int32(0)
+ state_input_ready_phase = cutlass.Int32(0)
+ state_input_ready_l_phase = cutlass.Int32(0)
+ rhs_ready_phase = cutlass.Int32(0)
+ update_ready_phase = cutlass.Int32(0)
+ final_state_stored_phase = cutlass.Int32(0)
+
+ # Actual SMEM/TMEM buffers for the BT=16 schedule:
+ # q/k/v : 16 x 128 each
+ # gate_log2 : 16 x 128
+ # beta : 16
+ # q/k inverse norm : 16 each, staged once per decay stage
+ # exp_g_last : 128, CG0-local and staged once per decay stage
+ # state decay : row 15 of the fp32 gate-prefix exchange tile
+ # q_decay/k_decay : tcgen05 SW128 operands shared with auxiliary MMA
+ # k_restore : tcgen05 SW128 N-major final-state operand
+ # k_inv : 16 x 128 token-major for auxiliary-MMA RHS
+ # A inverse/QK : 16 x 16 each, plus transposed tcgen05 operands
+ # state : external/kernel ABI is VK `[DV, DK]`; reference
+ # math can view it as KV `[DK, DV]` by transposing.
+ # The TS A-staging path keeps VK in TMEM so state*k
+ # is `[DV, DK] @ [DK, BT] -> [DV, BT]` with M=128.
+
+ if is_service_warpgroup(warp_idx):
+ prims.setmaxregister(KDA_SERVICE_REGS, prims.SetMaxRegisterAction.DECREASE)
+ if warp_idx == ROLES.tma_load:
+ # Gate constants are needed by BOTH the safe sigmoid and the
+ # FLA non-safe softplus activations (a_log_exp == exp(A_log),
+ # dt_bias per key dim). Materialize unconditionally; for
+ # SAFE_GATE=True this traces exactly as before (byte-identical
+ # generated code), the non-safe path now reads the same staged constants.
+ if lane == 0:
+ raw_dt_bias_smem[RAW_DT_BIAS_A_LOG_EXP_OFFSET] = cute.math.exp2(
+ a_log[bidy].to(cutlass.Float32) * LOG2_E,
+ fastmath=True,
+ )
+ for dim_group in cutlass.range_constexpr(DK // THREADS_PER_WARP):
+ dim = dim_group * THREADS_PER_WARP + lane
+ raw_dt_bias_smem[dim] = dt_bias[bidy, dim].to(cutlass.Float32)
+
+ # Prime the beta tile ring: LOOKAHEAD tiles in flight before the
+ # first consumption so the wait below is never exposed. Group
+ # coords walk (head-group, token); pair coords walk packed rows
+ # from the g-aligned floor of the sequence start.
+ beta_row0 = sequence_start // beta_g
+ beta_rows_per_chunk = cutlass.Int32(BT) // beta_g
+ beta_head_group = (bidy // cutlass.Int32(8)) * cutlass.Int32(8)
+ for pre in cutlass.range_constexpr(BETA_TMA_LOOKAHEAD):
+ if cutlass.Int32(pre) < num_chunks:
+ if prims.elect_sync():
+ prims.mbarrier_arrive_expect_tx(
+ beta_tile_mbar.subview(pre % BETA_TILE_STAGE_COUNT),
+ beta_tile_tx,
+ )
+ beta_c0 = beta_head_group
+ beta_c1 = sequence_start + cutlass.Int32(pre * BT)
+ if beta_g > cutlass.Int32(1):
+ beta_c0 = cutlass.Int32(0)
+ beta_c1 = (
+ beta_row0 + cutlass.Int32(pre) * beta_rows_per_chunk
+ )
+ prims.cp_async_bulk_tensor_shared_cta_global(
+ beta_tile_smem.subview(
+ (pre % BETA_TILE_STAGE_COUNT) * BETA_TILE_STAGE_ELEMS
+ ),
+ tma_desc_beta.get_ptr(),
+ (beta_c0, beta_c1),
+ beta_tile_mbar.subview(pre % BETA_TILE_STAGE_COUNT),
+ )
+ raw_stage = cutlass.Int32(0)
+ raw_consumed_phase = cutlass.Int32(1)
+ for chunk in cutlass.range(num_chunks, unroll=1):
+ chunk_start = chunk * BT
+ beta_tile_slot = chunk % BETA_TILE_STAGE_COUNT
+ tma_transfer_wait(
+ beta_tile_mbar.subview(beta_tile_slot),
+ (chunk // BETA_TILE_STAGE_COUNT) % 2,
+ )
+ raw_q_stage = raw_q_smem.subview(raw_stage * RAW_Q_STAGE_SIZE)
+ raw_k_stage = raw_k_smem.subview(raw_stage * RAW_K_STAGE_SIZE)
+ raw_v_stage = raw_v_smem.subview(raw_stage * RAW_V_STAGE_SIZE)
+ raw_gate_stage = raw_gate_smem.subview(raw_stage * RAW_GATE_STAGE_SIZE)
+ raw_beta_stage = raw_beta_smem.subview(raw_stage * RAW_BETA_STAGE_SIZE)
+ raw_consumed_wait(
+ raw_consumed_mbar.subview(raw_stage),
+ raw_consumed_phase,
+ )
+ # TMA issues q/k/v/gate into typed SMEM against the chunk's
+ # tma_mbar ring slot (consumers wait that slot's parity
+ # directly; this warp never waits completions -- raw_consumed
+ # alone throttles issue to <= RAW_STAGE_COUNT in flight).
+ tma_stage_load_inputs(
+ tma_desc_q,
+ tma_desc_k,
+ tma_desc_v,
+ tma_desc_gate,
+ beta_tile_smem.subview(beta_tile_slot * BETA_TILE_STAGE_ELEMS),
+ beta_g,
+ heads32,
+ raw_q_stage,
+ raw_k_stage,
+ raw_v_stage,
+ raw_gate_stage,
+ raw_beta_stage,
+ sequence_start,
+ bidy,
+ lane,
+ chunk_start,
+ seqlen,
+ tma_mbar.subview(raw_stage),
+ tma_tx_bytes,
+ gate_dtype,
+ )
+ beta_next = chunk + BETA_TMA_LOOKAHEAD
+ if beta_next < num_chunks:
+ beta_next_slot = beta_next % BETA_TILE_STAGE_COUNT
+ if prims.elect_sync():
+ prims.mbarrier_arrive_expect_tx(
+ beta_tile_mbar.subview(beta_next_slot),
+ beta_tile_tx,
+ )
+ beta_c0 = beta_head_group
+ beta_c1 = sequence_start + beta_next * BT
+ if beta_g > cutlass.Int32(1):
+ beta_c0 = cutlass.Int32(0)
+ beta_c1 = beta_row0 + beta_next * beta_rows_per_chunk
+ prims.cp_async_bulk_tensor_shared_cta_global(
+ beta_tile_smem.subview(
+ beta_next_slot * BETA_TILE_STAGE_ELEMS
+ ),
+ tma_desc_beta.get_ptr(),
+ (beta_c0, beta_c1),
+ beta_tile_mbar.subview(beta_next_slot),
+ )
+ raw_stage, raw_wrapped = advance_ring_stage(
+ raw_stage,
+ 1,
+ RAW_STAGE_COUNT,
+ )
+ raw_consumed_phase = raw_consumed_phase ^ raw_wrapped
+
+ elif warp_idx == ROLES.super_mma:
+ raw_stage = cutlass.Int32(0)
+ for chunk in cutlass.range(num_chunks, unroll=1):
+ decay_stage = chunk % DECAY_STAGE_COUNT
+ pairwise_stage = chunk % PAIRWISE_STAGE_COUNT
+ raw_beta_stage = raw_beta_smem.subview(raw_stage * RAW_BETA_STAGE_SIZE)
+ k_inv_stage = k_inv_smem.subview(decay_stage * K_INV_STAGE_SIZE)
+ tcgen05_k_decay_stage = tcgen05_k_decay_smem.subview(
+ decay_stage * TCGEN05_K_DECAY_STAGE_SIZE
+ )
+ pairwise_stage_smem = pairwise_smem.subview(
+ pairwise_stage * PAIRWISE_SMEM_STAGE_SIZE
+ )
+
+ pairwise_consumed_wait(
+ pairwise_consumed_mbar.subview(pairwise_stage),
+ ((chunk // PAIRWISE_STAGE_COUNT) + 1) % 2,
+ )
+ cg0_k_ready_wait(
+ cg0_k_half_ready_mbar.subview(decay_stage),
+ (chunk // DECAY_STAGE_COUNT) % 2,
+ )
+ # KDA schedule owner: standalone auxiliary-MMA warp. The first
+ # KK half runs on the half-DK arrival; the pipeline waits the
+ # full cg0_k_ready on its own before the second half.
+ super_mma_stage_pairwise_pipeline(
+ tcgen05_k_decay_stage,
+ k_inv_stage,
+ pairwise_stage_smem,
+ raw_beta_stage,
+ cg0_k_ready_mbar.subview(decay_stage),
+ (chunk // DECAY_STAGE_COUNT) % 2,
+ lane,
+ input_dtype,
+ )
+ pairwise_ready_arrive(pairwise_ready_mbar.subview(pairwise_stage))
+ operand_smem_consumed_arrive(
+ operand_smem_consumed_mbar.subview(decay_stage)
+ )
+ raw_consumed_arrive(raw_consumed_mbar.subview(raw_stage))
+ raw_stage, _ = advance_ring_stage(raw_stage, 1, RAW_STAGE_COUNT)
+
+ elif warp_idx == ROLES.tcgen05_mma:
+ tmem_raw_addr = tmem_ptr_i32.load()
+ shared_acc_event_id = cutlass.Int32(0)
+ late_operand_stage = cutlass.Int32(0)
+ q_k_restore_ready_phase = cutlass.Int32(0)
+ for chunk in cutlass.range(num_chunks, unroll=1):
+ o_stage = chunk % O_STAGE_COUNT
+ qstate_acc_stage = chunk % KDA_TMEM_QSTATE_ACC_STAGE_COUNT
+ decay_stage = chunk % DECAY_STAGE_COUNT
+ q_k_restore_ready_stage = late_operand_stage
+ pairwise_stage = chunk % PAIRWISE_STAGE_COUNT
+ shared_input_stage = o_stage
+ tcgen05_k_decay_stage = tcgen05_k_decay_smem.subview(
+ decay_stage * TCGEN05_K_DECAY_STAGE_SIZE
+ )
+ tcgen05_q_decay_stage = tcgen05_q_decay_smem.subview(
+ decay_stage * TCGEN05_Q_DECAY_STAGE_SIZE
+ )
+ tcgen05_k_restore_stage = tcgen05_k_restore_smem.subview(
+ decay_stage * TCGEN05_K_RESTORE_STAGE_SIZE
+ )
+ pairwise_stage_smem = pairwise_smem.subview(
+ pairwise_stage * PAIRWISE_SMEM_STAGE_SIZE
+ )
+
+ cg0_k_ready_wait(
+ cg0_k_ready_mbar.subview(decay_stage),
+ (chunk // DECAY_STAGE_COUNT) % 2,
+ )
+
+ state_input_ready_l_phase = state_input_ready_wait(
+ state_input_ready_l_mbar,
+ state_input_ready_l_phase,
+ )
+ prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC)
+ state_k_acc_stage = tcgen05_shared_acc_stage_from_event(
+ shared_acc_event_id
+ )
+ tcgen05_issue_state_k_mma(
+ tcgen05_k_decay_stage,
+ tmem_raw_addr,
+ shared_acc_ready_mbar.subview(state_k_acc_stage),
+ state_k_acc_stage,
+ input_dtype,
+ 0,
+ (DK // TCGEN05_F16_K_ATOM) // 2,
+ False,
+ False,
+ DV,
+ )
+ state_input_ready_phase = state_input_ready_wait(
+ state_input_ready_mbar,
+ state_input_ready_phase,
+ )
+ prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC)
+ tcgen05_issue_state_k_mma(
+ tcgen05_k_decay_stage,
+ tmem_raw_addr,
+ shared_acc_ready_mbar.subview(state_k_acc_stage),
+ state_k_acc_stage,
+ input_dtype,
+ (DK // TCGEN05_F16_K_ATOM) // 2,
+ DK // TCGEN05_F16_K_ATOM,
+ True,
+ True,
+ DV,
+ )
+ shared_acc_event_id += cutlass.Int32(1)
+
+ q_k_restore_ready_wait(
+ q_k_restore_ready_mbar.subview(q_k_restore_ready_stage),
+ q_k_restore_ready_phase,
+ )
+
+ qstate_acc_reuse_phase = (
+ chunk // KDA_TMEM_QSTATE_ACC_STAGE_COUNT + cutlass.Int32(1)
+ ) % cutlass.Int32(2)
+ output_ready_wait(
+ output_ready_mbar.subview(qstate_acc_stage),
+ qstate_acc_reuse_phase,
+ )
+
+ # State*Q tile producer. This uses the dedicated qstate_acc slot
+ # because the tile stays live until qkv is fused into output.
+ tcgen05_issue_state_q_mma(
+ tcgen05_q_decay_stage,
+ tmem_raw_addr,
+ operand_smem_consumed_mbar.subview(decay_stage),
+ qstate_acc_stage,
+ input_dtype,
+ )
+ pairwise_ready_wait(
+ pairwise_ready_mbar.subview(pairwise_stage),
+ (chunk // PAIRWISE_STAGE_COUNT) % 2,
+ )
+ rhs_ready_phase = rhs_ready_wait(
+ rhs_ready_mbar,
+ rhs_ready_phase,
+ )
+ # KDA schedule owner: tcgen05-MMA warp.
+ #
+ # Update tile producer:
+ # update = A_inv @ rhs
+ update_acc_stage = tcgen05_shared_acc_stage_from_event(
+ shared_acc_event_id
+ )
+ tcgen05_issue_update_mma(
+ pairwise_stage_smem,
+ tmem_raw_addr,
+ shared_acc_ready_mbar.subview(update_acc_stage),
+ update_acc_stage,
+ shared_input_stage,
+ input_dtype,
+ )
+ shared_acc_event_id += cutlass.Int32(1)
+
+ update_ready_phase = update_ready_wait(
+ update_ready_mbar,
+ update_ready_phase,
+ )
+ prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC)
+ # Final-state producer:
+ # final_state_acc += update @ k_restore
+ tcgen05_issue_final_state_delta_mma(
+ tcgen05_k_restore_stage,
+ tmem_raw_addr,
+ k_restore_consumed_l_mbar.subview(decay_stage),
+ k_restore_consumed_mbar.subview(decay_stage),
+ shared_input_stage,
+ input_dtype,
+ )
+
+ pairwise_ready_wait(
+ qk_ready_mbar.subview(pairwise_stage),
+ (chunk // PAIRWISE_STAGE_COUNT) % 2,
+ )
+ # Output tile producer:
+ # staged_o = (state_q + qk @ update) * scale
+ tcgen05_issue_qkv_mma(
+ pairwise_stage_smem,
+ tmem_raw_addr,
+ qstate_acc_ready_mbar,
+ shared_input_stage,
+ qstate_acc_stage,
+ input_dtype,
+ )
+ pairwise_consumed_arrive(pairwise_consumed_mbar.subview(pairwise_stage))
+ late_operand_stage, late_operand_wrapped = advance_ring_stage(
+ late_operand_stage,
+ 1,
+ Q_K_RESTORE_READY_STAGE_COUNT,
+ )
+ q_k_restore_ready_phase = q_k_restore_ready_phase ^ late_operand_wrapped
+
+ final_state_stored_phase = final_state_stored_wait(
+ final_state_stored_mbar,
+ final_state_stored_phase,
+ )
+ tmem_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Float32)
+ prims.tcgen05_dealloc(tmem_ptr, KDA_TMEM_ALLOC_COLS, group="cta_1")
+
+ elif warp_idx == ROLES.epilogue:
+ q_k_restore_ready_stage = cutlass.Int32(0)
+ q_k_restore_ready_phase = cutlass.Int32(0)
+ for chunk in cutlass.range(num_chunks, unroll=1):
+ decay_stage = chunk % DECAY_STAGE_COUNT
+ pairwise_stage = chunk % PAIRWISE_STAGE_COUNT
+ k_inv_stage = k_inv_smem.subview(decay_stage * K_INV_STAGE_SIZE)
+ tcgen05_q_decay_stage = tcgen05_q_decay_smem.subview(
+ decay_stage * TCGEN05_Q_DECAY_STAGE_SIZE
+ )
+ pairwise_stage_smem = pairwise_smem.subview(
+ pairwise_stage * PAIRWISE_SMEM_STAGE_SIZE
+ )
+
+ pairwise_consumed_wait(
+ pairwise_consumed_mbar.subview(pairwise_stage),
+ ((chunk // PAIRWISE_STAGE_COUNT) + 1) % 2,
+ )
+ q_k_restore_ready_wait(
+ q_k_restore_ready_mbar.subview(q_k_restore_ready_stage),
+ q_k_restore_ready_phase,
+ )
+ super_mma_stage_qk(
+ tcgen05_q_decay_stage,
+ k_inv_stage,
+ pairwise_stage_smem,
+ lane,
+ input_dtype,
+ )
+ pairwise_ready_arrive(qk_ready_mbar.subview(pairwise_stage))
+ operand_smem_consumed_arrive(
+ operand_smem_consumed_mbar.subview(decay_stage)
+ )
+ q_k_restore_ready_stage, q_k_restore_wrapped = advance_ring_stage(
+ q_k_restore_ready_stage,
+ 1,
+ Q_K_RESTORE_READY_STAGE_COUNT,
+ )
+ q_k_restore_ready_phase = q_k_restore_ready_phase ^ q_k_restore_wrapped
+
+ if chunk > 0:
+ output_chunk = chunk - cutlass.Int32(1)
+ epilogue_wait_and_store_full_output(
+ tma_desc_o,
+ o_smem,
+ output_ready_mbar,
+ output_consumed_mbar,
+ sequence_start,
+ bidy,
+ output_chunk,
+ O_STAGE_COUNT,
+ )
+ if num_chunks > 0:
+ output_chunk = num_chunks - cutlass.Int32(1)
+ epilogue_wait_and_store_final_output(
+ tma_desc_o,
+ out,
+ o_smem,
+ output_ready_mbar,
+ output_consumed_mbar,
+ sequence_start,
+ bidy,
+ seqlen,
+ output_chunk,
+ lane,
+ O_STAGE_COUNT,
+ )
+ elif is_compute_group0_warp(warp_idx):
+ prims.setmaxregister(KDA_CG0_REGS, prims.SetMaxRegisterAction.INCREASE)
+ cg0_warp = warp_idx - ROLES.compute_group0_first
+ cg0_group_id = cg0_warp // CG0_WARPS_PER_GROUP
+ cg0_local_warp = cg0_warp % CG0_WARPS_PER_GROUP
+ cg0_a_log_exp = cutlass.Float32(1.0)
+ cg0_dt_bias_value = cutlass.Float32(0.0)
+ tmem_raw_addr = cutlass.Int32(0)
+ if num_chunks > 0:
+ raw_ready_wait(tma_mbar.subview(0), 0)
+ prefix_dim = cg0_local_warp * THREADS_PER_WARP + lane
+ cg0_a_log_exp = raw_dt_bias_smem[RAW_DT_BIAS_A_LOG_EXP_OFFSET]
+ cg0_dt_bias_value = raw_dt_bias_smem[prefix_dim]
+ state_input_ready_wait(initial_state_ready_mbar, cutlass.Int32(0))
+ prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC)
+ # The allocator publishes one immutable TMEM base for the CTA.
+ # Cache it after the initial-state handoff instead of reloading
+ # the SMEM pointer in every steady and peeled CG0 iteration.
+ tmem_raw_addr = tmem_ptr_i32.load()
+ raw_stage = cg0_group_id
+ raw_ready_phase = cutlass.Int32(0)
+ late_operand_stage = cg0_group_id
+ cg0_ckpt_stride = cutlass.Int32(0)
+ if cutlass.const_expr(state_ckpt is not None):
+ cg0_ckpt_stride = checkpoint_stride_chunks
+ # v27_peel: STRUCTURAL in-kernel tail-peel of the CG0 producer loop.
+ # The interior mainloop runs the GUARD-FREE body over chunks
+ # [cg0_group_id, num_chunks-1): NO per-chunk `chunk_start+BT>seqlen`
+ # compare and NO tail zeroing (materialize called with the FULL branch).
+ # The single genuinely-partial tail chunk (num_chunks-1) is peeled out
+ # once below and run with the masked tail fixup. This makes ONE kernel
+ # correct for both aligned and ragged inputs (the peeled tail is a
+ # runtime no-op when the sequence is BT-aligned). Ring-state locals
+ # (raw_stage / phase, late_operand_stage / phase) are loop-carried and
+ # left positioned for the peeled chunk by the mainloop.
+ for chunk in cutlass.range(
+ cg0_group_id,
+ num_chunks - cutlass.Int32(1),
+ CG0_GROUP_COUNT,
+ unroll=1,
+ ):
+ decay_stage = chunk % DECAY_STAGE_COUNT
+ q_k_restore_ready_stage = late_operand_stage
+ raw_q_stage = raw_q_smem.subview(raw_stage * RAW_Q_STAGE_SIZE)
+ raw_k_stage = raw_k_smem.subview(raw_stage * RAW_K_STAGE_SIZE)
+ raw_gate_stage = raw_gate_smem.subview(raw_stage * RAW_GATE_STAGE_SIZE)
+ # FP32 exp2(g_prefix) exchange tile. With an FP32 gate it IS the
+ # raw-gate stage (historical aliasing, zero cost); with a 16-bit gate
+ # it comes from the dedicated 4-deep gate_exchange ring.
+ if cutlass.const_expr(gate_dtype_is_f32(gate_dtype)):
+ gate_exchange_stage = raw_gate_stage
+ else:
+ gate_exchange_stage = gate_exchange_smem.subview(
+ (chunk % GATE_EXCHANGE_STAGE_COUNT) * GATE_EXCHANGE_STAGE_SIZE
+ )
+ k_inv_stage = k_inv_smem.subview(decay_stage * K_INV_STAGE_SIZE)
+ tcgen05_k_decay_stage = tcgen05_k_decay_smem.subview(
+ decay_stage * TCGEN05_K_DECAY_STAGE_SIZE
+ )
+ tcgen05_q_decay_stage = tcgen05_q_decay_smem.subview(
+ decay_stage * TCGEN05_Q_DECAY_STAGE_SIZE
+ )
+ tcgen05_k_restore_stage = tcgen05_k_restore_smem.subview(
+ decay_stage * TCGEN05_K_RESTORE_STAGE_SIZE
+ )
+ raw_ready_wait(
+ tma_mbar.subview(raw_stage),
+ raw_ready_phase,
+ )
+
+ # KDA schedule owner: compute warp group 0.
+ #
+ # The decay materialization stages full gate prefixes once, then
+ # issues the two gate MUFU families for the full KDA operand set:
+ # exp2(g_prefix) : 16 * 128
+ # exp2(-g_prefix) : 16 * 128
+ cg0_materialize_decay_operands(
+ raw_q_stage,
+ raw_k_stage,
+ raw_gate_stage,
+ gate_exchange_stage,
+ cg0_a_log_exp,
+ cg0_dt_bias_value,
+ k_inv_stage,
+ tcgen05_k_decay_stage,
+ tcgen05_q_decay_stage,
+ tcgen05_k_restore_stage,
+ cg0_k_ready_mbar.subview(decay_stage),
+ cg0_k_half_ready_mbar.subview(decay_stage),
+ diag_ready_mbar.subview(raw_stage),
+ operand_smem_consumed_mbar.subview(decay_stage),
+ k_restore_consumed_mbar.subview(decay_stage),
+ chunk,
+ seqlen,
+ input_dtype,
+ gate_dtype,
+ SAFE_GATE,
+ GATE_SCALE_LOG2,
+ True,
+ cg0_group_id,
+ cg0_local_warp,
+ lane,
+ )
+ q_k_restore_ready_arrive(
+ q_k_restore_ready_mbar.subview(q_k_restore_ready_stage)
+ )
+ if chunk > 0:
+ prev_state_chunk = chunk - cutlass.Int32(1)
+ tcgen05_wait_acc_buffer_ready(
+ k_restore_consumed_l_mbar.subview(
+ prev_state_chunk % DECAY_STAGE_COUNT
+ ),
+ (prev_state_chunk // DECAY_STAGE_COUNT) % 2,
+ )
+ if cutlass.const_expr(state_ckpt is not None):
+ if (chunk > cutlass.Int32(0)) & (chunk % cg0_ckpt_stride == 0):
+ checkpoint_read_done_wait(
+ checkpoint_read_done_mbar,
+ ((chunk // cg0_ckpt_stride) + cutlass.Int32(1)) % 2,
+ )
+ prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC)
+ tcgen05_pack_rescale_state_half_tmem(
+ tmem_raw_addr,
+ gate_exchange_stage,
+ state_input_ready_l_mbar,
+ warp_idx,
+ input_dtype,
+ 0,
+ True,
+ )
+ raw_consumed_arrive(raw_consumed_mbar.subview(raw_stage))
+ update_ready_arrive(update_ready_mbar)
+ raw_stage, raw_wrapped = advance_ring_stage(
+ raw_stage,
+ CG0_GROUP_COUNT,
+ RAW_STAGE_COUNT,
+ )
+ raw_ready_phase = raw_ready_phase ^ raw_wrapped
+ late_operand_stage, late_operand_wrapped = advance_ring_stage(
+ late_operand_stage,
+ CG0_GROUP_COUNT,
+ Q_K_RESTORE_READY_STAGE_COUNT,
+ )
+
+ # Peeled final (tail) chunk num_chunks-1, executed once by the CG0 group
+ # that owns it. num_chunks==1 -> the single chunk IS the tail (group 0);
+ # num_chunks==0 -> no work. materialize runs the tail branch (its
+ # built-in `tail_valid cutlass.Int32(0):
+ if (num_chunks - cutlass.Int32(1)) % CG0_GROUP_COUNT == cg0_group_id:
+ chunk = num_chunks - cutlass.Int32(1)
+ chunk_start = chunk * BT
+ decay_stage = chunk % DECAY_STAGE_COUNT
+ q_k_restore_ready_stage = late_operand_stage
+ raw_q_stage = raw_q_smem.subview(raw_stage * RAW_Q_STAGE_SIZE)
+ raw_k_stage = raw_k_smem.subview(raw_stage * RAW_K_STAGE_SIZE)
+ raw_v_stage = raw_v_smem.subview(raw_stage * RAW_V_STAGE_SIZE)
+ raw_gate_stage = raw_gate_smem.subview(raw_stage * RAW_GATE_STAGE_SIZE)
+ # FP32 exp2(g_prefix) exchange tile. With an FP32 gate it IS the
+ # raw-gate stage (historical aliasing, zero cost); with a 16-bit gate
+ # it comes from the dedicated 4-deep gate_exchange ring.
+ if cutlass.const_expr(gate_dtype_is_f32(gate_dtype)):
+ gate_exchange_stage = raw_gate_stage
+ else:
+ gate_exchange_stage = gate_exchange_smem.subview(
+ (chunk % GATE_EXCHANGE_STAGE_COUNT) * GATE_EXCHANGE_STAGE_SIZE
+ )
+ k_inv_stage = k_inv_smem.subview(decay_stage * K_INV_STAGE_SIZE)
+ tcgen05_k_decay_stage = tcgen05_k_decay_smem.subview(
+ decay_stage * TCGEN05_K_DECAY_STAGE_SIZE
+ )
+ tcgen05_q_decay_stage = tcgen05_q_decay_smem.subview(
+ decay_stage * TCGEN05_Q_DECAY_STAGE_SIZE
+ )
+ tcgen05_k_restore_stage = tcgen05_k_restore_smem.subview(
+ decay_stage * TCGEN05_K_RESTORE_STAGE_SIZE
+ )
+ raw_ready_wait(
+ tma_mbar.subview(raw_stage),
+ raw_ready_phase,
+ )
+
+ if chunk_start + cutlass.Int32(BT) > seqlen:
+ if cg0_local_warp == 0:
+ cg0_zero_tail_raw_operands(
+ raw_q_stage,
+ raw_k_stage,
+ raw_v_stage,
+ raw_gate_stage,
+ lane,
+ chunk_start,
+ seqlen,
+ q.element_type,
+ gate_dtype,
+ )
+ cg0_sync(cg0_group_id)
+
+ cg0_materialize_decay_operands(
+ raw_q_stage,
+ raw_k_stage,
+ raw_gate_stage,
+ gate_exchange_stage,
+ cg0_a_log_exp,
+ cg0_dt_bias_value,
+ k_inv_stage,
+ tcgen05_k_decay_stage,
+ tcgen05_q_decay_stage,
+ tcgen05_k_restore_stage,
+ cg0_k_ready_mbar.subview(decay_stage),
+ cg0_k_half_ready_mbar.subview(decay_stage),
+ diag_ready_mbar.subview(raw_stage),
+ operand_smem_consumed_mbar.subview(decay_stage),
+ k_restore_consumed_mbar.subview(decay_stage),
+ chunk,
+ seqlen,
+ input_dtype,
+ gate_dtype,
+ SAFE_GATE,
+ GATE_SCALE_LOG2,
+ False,
+ cg0_group_id,
+ cg0_local_warp,
+ lane,
+ )
+ q_k_restore_ready_arrive(
+ q_k_restore_ready_mbar.subview(q_k_restore_ready_stage)
+ )
+ if chunk > 0:
+ prev_state_chunk = chunk - cutlass.Int32(1)
+ tcgen05_wait_acc_buffer_ready(
+ k_restore_consumed_l_mbar.subview(
+ prev_state_chunk % DECAY_STAGE_COUNT
+ ),
+ (prev_state_chunk // DECAY_STAGE_COUNT) % 2,
+ )
+ if cutlass.const_expr(state_ckpt is not None):
+ if (chunk > cutlass.Int32(0)) & (chunk % cg0_ckpt_stride == 0):
+ checkpoint_read_done_wait(
+ checkpoint_read_done_mbar,
+ ((chunk // cg0_ckpt_stride) + cutlass.Int32(1)) % 2,
+ )
+ prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC)
+ tcgen05_pack_rescale_state_half_tmem(
+ tmem_raw_addr,
+ gate_exchange_stage,
+ state_input_ready_l_mbar,
+ warp_idx,
+ input_dtype,
+ 0,
+ True,
+ )
+ raw_consumed_arrive(raw_consumed_mbar.subview(raw_stage))
+ update_ready_arrive(update_ready_mbar)
+
+ elif is_compute_group1_warp(warp_idx):
+ prims.setmaxregister(KDA_CG1_REGS, prims.SetMaxRegisterAction.INCREASE)
+ tmem_raw_addr = tmem_ptr_i32.load()
+ # Only CG1 touches recurrent state. Keep the pool lookup outside the
+ # chunk loop and out of producer/MMAs warps.
+ state_slot = bidx
+ if cutlass.const_expr(state_indices is not None):
+ state_slot = cutlass.Int32(state_indices[bidx])
+ ckpt_slot = cutlass.Int32(0)
+ if cutlass.const_expr(state_ckpt is not None):
+ ckpt_stride = checkpoint_stride_chunks
+ ckpt_next = ckpt_stride
+ ckpt_slot = cutlass.Int32(cu_ckpts[bidx])
+ if cutlass.const_expr(checkpoint_state_indices is not None):
+ tcgen05_store_initial_state_tmem(
+ tmem_raw_addr,
+ initial_state,
+ None,
+ cutlass.Int32(0),
+ state_slot,
+ bidy,
+ cutlass.Int32(0),
+ warp_idx,
+ lane,
+ HALF=False,
+ )
+ else:
+ tcgen05_store_initial_state_tmem(
+ tmem_raw_addr,
+ initial_state,
+ state_ckpt,
+ ckpt_slot,
+ state_slot,
+ bidy,
+ cutlass.Int32(0),
+ warp_idx,
+ lane,
+ HALF=False,
+ )
+ prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC)
+ state_input_ready_arrive(initial_state_ready_mbar)
+ shared_acc_event_id = cutlass.Int32(0)
+ if cutlass.const_expr(
+ state_ckpt is not None and checkpoint_state_indices is None
+ ):
+ # cu_ckpts supplies the per-sequence base slot offsets. A
+ # loop-carried counter replaces a per-chunk div/mod.
+ ckpt_slot += cutlass.Int32(1)
+
+ if num_chunks > 0:
+ shared_input_stage = cutlass.Int32(0)
+ raw_v_stage = raw_v_smem.subview(0)
+ raw_beta_stage = raw_beta_smem.subview(0)
+ raw_gate_stage = raw_gate_smem.subview(0)
+ # FP32 exp2(g_prefix) exchange tile for the CG1 prologue (chunk 0).
+ # With an FP32 gate it IS the raw-gate stage (historical aliasing,
+ # zero cost); with a 16-bit gate it is stage 0 of the dedicated
+ # 4-deep gate_exchange ring (chunk 0 -> stage 0).
+ if cutlass.const_expr(gate_dtype_is_f32(gate_dtype)):
+ gate_exchange_stage = raw_gate_stage
+ else:
+ gate_exchange_stage = gate_exchange_smem.subview(0)
+
+ state0, state1, state2, state3 = tcgen05_stage_state_input_tmem(
+ tmem_raw_addr,
+ warp_idx,
+ output_consumed_mbar.subview(0),
+ cutlass.Int32(0),
+ input_dtype,
+ False,
+ 1,
+ True,
+ )
+ # Decay the retained right half as soon as CG0 publishes the FP32
+ # diagonal; K/Q materialization continues independently.
+ diag_ready_wait(diag_ready_mbar.subview(0), 0)
+ tcgen05_publish_projection_then_rescale_state_regs(
+ tmem_raw_addr,
+ gate_exchange_stage,
+ state_input_ready_mbar,
+ warp_idx,
+ state0,
+ state1,
+ state2,
+ state3,
+ 1,
+ True,
+ )
+
+ state_k_acc_stage = tcgen05_shared_acc_stage_from_event(shared_acc_event_id)
+ state_k_acc_phase = tcgen05_shared_acc_phase_from_event(shared_acc_event_id)
+ tcgen05_wait_acc_buffer_ready(
+ shared_acc_ready_mbar.subview(state_k_acc_stage),
+ state_k_acc_phase,
+ )
+ rhs_lane = cute.arch.lane_idx()
+ tcgen05_stage_rhs_input_tmem(
+ tmem_raw_addr,
+ raw_v_stage,
+ raw_beta_stage,
+ warp_idx,
+ rhs_lane,
+ state_k_acc_stage,
+ shared_input_stage,
+ input_dtype,
+ )
+ shared_acc_event_id += cutlass.Int32(1)
+ rhs_ready_arrive(rhs_ready_mbar)
+
+ raw_consumed_arrive(raw_consumed_mbar.subview(0))
+ post_scale_lane = cute.arch.lane_idx()
+
+ update_acc_stage = tcgen05_shared_acc_stage_from_event(shared_acc_event_id)
+ update_acc_phase = tcgen05_shared_acc_phase_from_event(shared_acc_event_id)
+ tcgen05_wait_acc_buffer_ready(
+ shared_acc_ready_mbar.subview(update_acc_stage),
+ update_acc_phase,
+ )
+ tcgen05_stage_update_input_tmem(
+ tmem_raw_addr,
+ warp_idx,
+ update_acc_stage,
+ shared_input_stage,
+ input_dtype,
+ )
+ shared_acc_event_id += cutlass.Int32(1)
+ update_ready_arrive(update_ready_mbar)
+
+ if cutlass.const_expr(state_ckpt is not None):
+ # Peeled chunk 0's checkpoint (fires only for stride 1, i.e. a
+ # checkpoint every BT tokens).
+ checkpoint_due = ckpt_next == cutlass.Int32(1)
+ if cutlass.const_expr(checkpoint_state_indices is not None):
+ checkpoint_due = checkpoint_due & (cutlass.Int32(BT) <= seqlen)
+ else:
+ checkpoint_due = checkpoint_due & (cutlass.Int32(1) < num_chunks)
+ if checkpoint_due:
+ tcgen05_wait_acc_buffer_ready(k_restore_consumed_mbar.subview(0), 0)
+ checkpoint_output_slot = ckpt_slot
+ if cutlass.const_expr(checkpoint_state_indices is not None):
+ checkpoint_output_slot = cutlass.Int32(
+ checkpoint_state_indices[ckpt_slot]
+ )
+ if (checkpoint_output_slot >= cutlass.Int32(0)) & (
+ checkpoint_output_slot < cutlass.Int32(state_ckpt.shape[0])
+ ):
+ tcgen05_store_final_state_tmem(
+ tmem_raw_addr,
+ KDA_TMEM_FINAL_STATE_ACC_COL_OFFSET,
+ state_ckpt,
+ checkpoint_output_slot,
+ bidy,
+ cutlass.Int32(0),
+ warp_idx,
+ post_scale_lane,
+ HALF=False,
+ )
+ prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC)
+ checkpoint_read_done_arrive(checkpoint_read_done_mbar)
+ ckpt_slot += cutlass.Int32(1)
+ ckpt_next += ckpt_stride
+
+ # Peel chunk 0 so the steady-state loop always drains a prior output.
+ # Each actual shared-acc commit advances the event ring independently
+ # of which GEMM produced it.
+ raw_stage = cutlass.Int32(1)
+ for chunk in cutlass.range(1, num_chunks, 1, unroll=1):
+ o_stage = chunk % O_STAGE_COUNT
+ decay_stage = chunk % DECAY_STAGE_COUNT
+ shared_input_stage = o_stage
+ raw_v_stage = raw_v_smem.subview(raw_stage * RAW_V_STAGE_SIZE)
+ raw_beta_stage = raw_beta_smem.subview(raw_stage * RAW_BETA_STAGE_SIZE)
+ raw_gate_stage = raw_gate_smem.subview(raw_stage * RAW_GATE_STAGE_SIZE)
+ # FP32 exp2(g_prefix) exchange tile. With an FP32 gate it IS the
+ # raw-gate stage (historical aliasing, zero cost); with a 16-bit gate
+ # it comes from the dedicated 4-deep gate_exchange ring.
+ if cutlass.const_expr(gate_dtype_is_f32(gate_dtype)):
+ gate_exchange_stage = raw_gate_stage
+ else:
+ gate_exchange_stage = gate_exchange_smem.subview(
+ (chunk % GATE_EXCHANGE_STAGE_COUNT) * GATE_EXCHANGE_STAGE_SIZE
+ )
+
+ prev_output_chunk = chunk - cutlass.Int32(1)
+ prev_o_stage = prev_output_chunk % O_STAGE_COUNT
+ prev_qstate_acc_stage = prev_output_chunk % KDA_TMEM_QSTATE_ACC_STAGE_COUNT
+ prev_o_stage_base = prev_o_stage * O_SMEM_STAGE_SIZE
+ prev_decay_stage = prev_output_chunk % DECAY_STAGE_COUNT
+ prev_k_restore_phase = (prev_output_chunk // DECAY_STAGE_COUNT) % 2
+ tcgen05_wait_acc_buffer_ready(
+ k_restore_consumed_mbar.subview(prev_decay_stage),
+ prev_k_restore_phase,
+ )
+ state0, state1, state2, state3 = tcgen05_stage_state_input_tmem(
+ tmem_raw_addr,
+ warp_idx,
+ output_consumed_mbar.subview(prev_o_stage),
+ ((prev_output_chunk // O_STAGE_COUNT) + 1) % 2,
+ input_dtype,
+ True,
+ 1,
+ True,
+ )
+ # The diagonal is complete before k_decay/q_decay/k_restore; use
+ # that narrower dependency to hide right-half FP32 scaling.
+ diag_ready_wait(
+ diag_ready_mbar.subview(raw_stage),
+ (chunk // RAW_STAGE_COUNT) % 2,
+ )
+ tcgen05_publish_projection_then_rescale_state_regs(
+ tmem_raw_addr,
+ gate_exchange_stage,
+ state_input_ready_mbar,
+ warp_idx,
+ state0,
+ state1,
+ state2,
+ state3,
+ 1,
+ True,
+ )
+ pre_scale_lane = cute.arch.lane_idx()
+ qstate_acc_ready_phase = tcgen05_wait_acc_buffer_ready(
+ qstate_acc_ready_mbar,
+ qstate_acc_ready_phase,
+ )
+ tcgen05_load_qstate_output_tmem(
+ tmem_raw_addr,
+ o_smem,
+ warp_idx,
+ pre_scale_lane,
+ prev_o_stage_base,
+ prev_qstate_acc_stage,
+ SCALE,
+ out.element_type,
+ )
+ output_ready_arrive(output_ready_mbar.subview(prev_o_stage))
+
+ state_k_acc_stage = tcgen05_shared_acc_stage_from_event(shared_acc_event_id)
+ state_k_acc_phase = tcgen05_shared_acc_phase_from_event(shared_acc_event_id)
+ tcgen05_wait_acc_buffer_ready(
+ shared_acc_ready_mbar.subview(state_k_acc_stage),
+ state_k_acc_phase,
+ )
+ # KDA schedule owner: compute warp group 1.
+ #
+ # Epilogue for tcgen05 state*k, staged directly as the next
+ # tcgen05 A operand:
+ # shared_input = beta_i * (v_i - state*k_i)
+ tcgen05_stage_rhs_input_tmem(
+ tmem_raw_addr,
+ raw_v_stage,
+ raw_beta_stage,
+ warp_idx,
+ pre_scale_lane,
+ state_k_acc_stage,
+ shared_input_stage,
+ input_dtype,
+ )
+ shared_acc_event_id += cutlass.Int32(1)
+ rhs_ready_arrive(rhs_ready_mbar)
+
+ raw_consumed_arrive(raw_consumed_mbar.subview(raw_stage))
+ post_scale_lane = cute.arch.lane_idx()
+
+ update_acc_stage = tcgen05_shared_acc_stage_from_event(shared_acc_event_id)
+ update_acc_phase = tcgen05_shared_acc_phase_from_event(shared_acc_event_id)
+ tcgen05_wait_acc_buffer_ready(
+ shared_acc_ready_mbar.subview(update_acc_stage),
+ update_acc_phase,
+ )
+ tcgen05_stage_update_input_tmem(
+ tmem_raw_addr,
+ warp_idx,
+ update_acc_stage,
+ shared_input_stage,
+ input_dtype,
+ )
+ shared_acc_event_id += cutlass.Int32(1)
+ update_ready_arrive(update_ready_mbar)
+
+ if cutlass.const_expr(state_ckpt is not None):
+ # State checkpoint: at every ckpt_stride-th chunk boundary the
+ # TMEM accumulator holds exactly the state a run truncated here
+ # would emit as final_state. Wait the SAME parity slot the
+ # post-loop final store waits (the chunk's state-update MMA has
+ # committed), then reuse the final-state store routine with the
+ # flat checkpoint slot standing in for the sequence index.
+ checkpoint_due = chunk + cutlass.Int32(1) == ckpt_next
+ if cutlass.const_expr(checkpoint_state_indices is not None):
+ checkpoint_due = checkpoint_due & (
+ (chunk + cutlass.Int32(1)) * cutlass.Int32(BT) <= seqlen
+ )
+ else:
+ checkpoint_due = checkpoint_due & (
+ chunk + cutlass.Int32(1) < num_chunks
+ )
+ if checkpoint_due:
+ tcgen05_wait_acc_buffer_ready(
+ k_restore_consumed_mbar.subview(chunk % DECAY_STAGE_COUNT),
+ (chunk // DECAY_STAGE_COUNT) % 2,
+ )
+ checkpoint_output_slot = ckpt_slot
+ if cutlass.const_expr(checkpoint_state_indices is not None):
+ checkpoint_output_slot = cutlass.Int32(
+ checkpoint_state_indices[ckpt_slot]
+ )
+ if (checkpoint_output_slot >= cutlass.Int32(0)) & (
+ checkpoint_output_slot < cutlass.Int32(state_ckpt.shape[0])
+ ):
+ tcgen05_store_final_state_tmem(
+ tmem_raw_addr,
+ KDA_TMEM_FINAL_STATE_ACC_COL_OFFSET,
+ state_ckpt,
+ checkpoint_output_slot,
+ bidy,
+ cutlass.Int32(0),
+ warp_idx,
+ post_scale_lane,
+ HALF=False,
+ )
+ prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC)
+ checkpoint_read_done_arrive(checkpoint_read_done_mbar)
+ ckpt_slot += cutlass.Int32(1)
+ ckpt_next += ckpt_stride
+
+ raw_stage, _ = advance_ring_stage(raw_stage, 1, RAW_STAGE_COUNT)
+
+ if num_chunks > 0:
+ final_lane = cute.arch.lane_idx()
+ output_chunk = num_chunks - cutlass.Int32(1)
+ last_decay_stage = output_chunk % DECAY_STAGE_COUNT
+ tcgen05_wait_acc_buffer_ready(
+ k_restore_consumed_mbar.subview(last_decay_stage),
+ (output_chunk // DECAY_STAGE_COUNT) % 2,
+ )
+ final_o_stage = output_chunk % O_STAGE_COUNT
+ final_qstate_acc_stage = output_chunk % KDA_TMEM_QSTATE_ACC_STAGE_COUNT
+ final_o_stage_base = final_o_stage * O_SMEM_STAGE_SIZE
+ output_consumed_wait(
+ output_consumed_mbar.subview(final_o_stage),
+ ((output_chunk // O_STAGE_COUNT) + 1) % 2,
+ )
+
+ qstate_acc_ready_phase = tcgen05_wait_acc_buffer_ready(
+ qstate_acc_ready_mbar,
+ qstate_acc_ready_phase,
+ )
+ tcgen05_load_qstate_output_tmem(
+ tmem_raw_addr,
+ o_smem,
+ warp_idx,
+ final_lane,
+ final_o_stage_base,
+ final_qstate_acc_stage,
+ SCALE,
+ out.element_type,
+ )
+ output_ready_arrive(output_ready_mbar.subview(final_o_stage))
+
+ if cutlass.const_expr(final_state is not None):
+ final_lane = cute.arch.lane_idx()
+ tcgen05_store_final_state_tmem(
+ tmem_raw_addr,
+ KDA_TMEM_FINAL_STATE_ACC_COL_OFFSET,
+ final_state,
+ state_slot,
+ bidy,
+ cutlass.Int32(0),
+ warp_idx,
+ final_lane,
+ HALF=False,
+ )
+ final_state_stored_arrive(final_state_stored_mbar)
+
+
+@cute.jit
+def host(
+ q: cute.Tensor,
+ k: cute.Tensor,
+ v: cute.Tensor,
+ raw_gate: cute.Tensor,
+ a_log: cute.Tensor,
+ dt_bias: cute.Tensor,
+ beta: cute.Tensor,
+ cu_seqlens: cute.Tensor,
+ seq_order: cute.Tensor,
+ state_indices: cute.Tensor | None,
+ initial_state: cute.Tensor | None,
+ out: cute.Tensor,
+ final_state: cute.Tensor | None,
+ stream,
+ SCALE: cutlass.Float32,
+ state_ckpt: cute.Tensor | None,
+ cu_ckpts: cute.Tensor | None,
+ checkpoint_state_indices: cute.Tensor | None,
+ checkpoint_stride_chunks: cutlass.Int32,
+ SAFE_GATE: cutlass.Constexpr,
+ GATE_SCALE_LOG2: cutlass.Constexpr,
+ THREADS: cutlass.Constexpr,
+ gate_dtype: cutlass.Constexpr,
+) -> None:
+ packed_batch = q.shape[0]
+ num_sequences = cu_seqlens.shape[0] - 1
+ seqlen = q.shape[1]
+ heads = q.shape[2]
+ # Token-major activations: memory order [T, H, D] (token stride = D*heads),
+ # so packed [1, T_total, H, D] / batched [B, T, H, D] callers are read
+ # directly with no transpose.
+ # Keep the unused outer stride representable by CuTe IR for packed inputs.
+ # CuTe layouts currently lower strides through signed Int32, so the natural
+ # D*T*H batch stride overflows once a singleton packed token buffer exceeds
+ # that range even though the batch coordinate is always zero.
+ qk_batch_stride = DK
+ v_batch_stride = DV
+ if packed_batch != cutlass.Int32(1):
+ qk_batch_stride = DK * seqlen * heads
+ v_batch_stride = DV * seqlen * heads
+ qk_layout = cute.make_layout(
+ (DK, seqlen, heads, packed_batch),
+ stride=(1, DK * heads, DK, qk_batch_stride),
+ )
+ v_layout = cute.make_layout(
+ (DV, seqlen, heads, packed_batch),
+ stride=(1, DV * heads, DV, v_batch_stride),
+ )
+ q_tma = cute.make_tensor(q.iterator, qk_layout)
+ k_tma = cute.make_tensor(k.iterator, qk_layout)
+ v_tma = cute.make_tensor(v.iterator, v_layout)
+ gate_tma = cute.make_tensor(raw_gate.iterator, qk_layout)
+ out_tma = cute.make_tensor(out.iterator, v_layout)
+ raw_f16_tma_box = (RAW_F16_TMA_SWIZZLE_ELEMS, BT, 1, 1)
+ raw_f32_tma_box = (RAW_F32_TMA_SWIZZLE_ELEMS, BT, 1, 1)
+ o_tma_box = (O_TMA_SWIZZLE_ELEMS, BT, 1, 1)
+ tma_desc_q = cuda.create_tensor_map_tiled_from_view(
+ q_tma,
+ box_dims=raw_f16_tma_box,
+ stride_order=(0, 1, 2, 3),
+ swizzle=cuda.TensorMapSwizzle.s128b,
+ )
+ tma_desc_k = cuda.create_tensor_map_tiled_from_view(
+ k_tma,
+ box_dims=raw_f16_tma_box,
+ stride_order=(0, 1, 2, 3),
+ swizzle=cuda.TensorMapSwizzle.s128b,
+ )
+ tma_desc_v = cuda.create_tensor_map_tiled_from_view(
+ v_tma,
+ box_dims=raw_f16_tma_box,
+ stride_order=(0, 1, 2, 3),
+ swizzle=cuda.TensorMapSwizzle.s128b,
+ )
+ # The gate box/swizzle family follows `gate_dtype`: FP32 keeps the wide
+ # 4 x 128 B box, a 16-bit gate uses the same family as q/k/v.
+ gate_tma_box = raw_f32_tma_box if gate_dtype_is_f32(gate_dtype) else raw_f16_tma_box
+ tma_desc_gate = cuda.create_tensor_map_tiled_from_view(
+ gate_tma,
+ box_dims=gate_tma_box,
+ stride_order=(0, 1, 2, 3),
+ swizzle=cuda.TensorMapSwizzle.s128b,
+ )
+ # beta transport: one descriptor family over the SAME contiguous [T, H]
+ # memory, runtime-shaped by g = 8/gcd(heads, 8) (see the constants block).
+ # g == 1: view (heads, T) box (8, BT) -- per-head-group 16B rows.
+ # g > 1 : packed view (g*heads, ceil(T/g)) box (g*heads, BT/g + 1).
+ beta_heads = cutlass.Int32(beta.shape[2])
+ beta_lsb = beta_heads & (-beta_heads)
+ beta_g = cutlass.Int32(8) // cutlass.min(beta_lsb, cutlass.Int32(8))
+ beta_is_pair = cutlass.Int32(1)
+ if beta_g == cutlass.Int32(1):
+ beta_is_pair = cutlass.Int32(0)
+ beta_rows = (cutlass.Int32(beta.shape[1]) + beta_g - cutlass.Int32(1)) // beta_g
+ beta_inner = beta_g * beta_heads
+ beta_box_inner = (
+ cutlass.Int32(8) * (cutlass.Int32(1) - beta_is_pair) + beta_inner * beta_is_pair
+ )
+ beta_box_outer = (
+ cutlass.Int32(BT) * (cutlass.Int32(1) - beta_is_pair)
+ + (cutlass.Int32(BT) // beta_g + cutlass.Int32(1)) * beta_is_pair
+ )
+ beta_tma_view = cute.make_tensor(
+ beta.iterator,
+ cute.make_layout((beta_inner, beta_rows), stride=(1, beta_inner)),
+ )
+ tma_desc_beta = cuda.create_tensor_map_tiled_from_view(
+ beta_tma_view,
+ box_dims=(beta_box_inner, beta_box_outer),
+ stride_order=(0, 1),
+ swizzle=cuda.TensorMapSwizzle.none,
+ )
+ tma_desc_o = cuda.create_tensor_map_tiled_from_view(
+ out_tma,
+ box_dims=o_tma_box,
+ stride_order=(0, 1, 2, 3),
+ swizzle=cuda.TensorMapSwizzle.s128b,
+ )
+ kernel(
+ tma_desc_q,
+ tma_desc_k,
+ tma_desc_v,
+ tma_desc_gate,
+ tma_desc_beta,
+ tma_desc_o,
+ q,
+ k,
+ v,
+ raw_gate,
+ a_log,
+ dt_bias,
+ beta,
+ cu_seqlens,
+ seq_order,
+ state_indices,
+ initial_state,
+ out,
+ final_state,
+ SCALE,
+ state_ckpt,
+ cu_ckpts,
+ checkpoint_state_indices,
+ checkpoint_stride_chunks,
+ SAFE_GATE,
+ GATE_SCALE_LOG2,
+ gate_dtype,
+ ).launch(
+ grid=(heads, num_sequences, 1),
+ block=(THREADS, 1, 1),
+ stream=stream,
+ min_blocks_per_mp=1,
+ )
+
+
+# ---------------------------------------------------------------------------
+# Host-side plan/decision cache keys: every cu_seqlens-keyed cache
+# keys on the tensor CONTENTS, obtained sync-free through the memo below.
+# ---------------------------------------------------------------------------
+
+# data_ptr -> (weakref(tensor), version, contents). See _cu_seqlens_contents.
+_CU_CONTENTS_MEMO: dict = {}
+
+
+def _cu_seqlens_contents(cu_seqlens: torch.Tensor) -> tuple:
+ """cu_seqlens contents as a host tuple, sync-free on the steady path.
+
+ Every plan/decision cache keys on the CONTENTS of cu_seqlens: keying on
+ buffer identity (data_ptr) is unsafe because torch's caching allocator
+ reuses freed blocks, so a recycled same-shape buffer with different
+ contents would hit a stale plan (wrong partition boundaries => wrong
+ output rows). But reading the contents is a D2H copy that synchronizes
+ the current stream — done per call it serializes the host into every
+ call and costs 5-14% wall time on the cached routes (measured, session
+ 29). This memo removes the read when the SAME live tensor object is
+ passed again (the steady case): a hit requires the memoized weakref to
+ resolve to the argument tensor ITSELF (`wr() is cu_seqlens`), so a
+ recycled allocator block can never produce a false hit — its previous
+ owner is a different (dead or dying) object; and an unchanged
+ `_version` (catches in-place mutation, which shares the version
+ counter across views). Tensors without version counters (inference
+ mode) memoize with version None and still get the identity protection.
+ """
+
+ key = cu_seqlens.data_ptr()
+ try:
+ ver = cu_seqlens._version
+ except Exception: # inference tensors track no version counter
+ ver = None
+ ent = _CU_CONTENTS_MEMO.get(key)
+ if ent is not None:
+ wr, memo_ver, contents = ent
+ if wr() is cu_seqlens and memo_ver == ver:
+ return contents
+ contents = tuple(int(x) for x in cu_seqlens.cpu().tolist())
+ _CU_CONTENTS_MEMO[key] = (weakref.ref(cu_seqlens), ver, contents)
+ return contents
+
+
+# ---------------------------------------------------------------------------
+# Single host entry: ONE callable that routes the plain engine vs
+# the two-kernel decomposition (DV2 chain) by a one-line occupancy rule
+# (_occupancy_pick_route). The engine is compiled eagerly, the decomp route
+# lazily on first use. The decomp route takes a user-allocated opaque
+# workspace (query bytes with workspace_size(); engine needs none).
+# ---------------------------------------------------------------------------
+
+_SM_COUNT_CACHE: dict = {}
+
+
+def _device_sm_count(device) -> int:
+ key = str(device)
+ count = _SM_COUNT_CACHE.get(key)
+ if count is None:
+ count = torch.cuda.get_device_properties(device).multi_processor_count
+ _SM_COUNT_CACHE[key] = count
+ return count
+
+
+# The route is a one-line occupancy rule: run the decomp DV2 split (2 CTAs
+# per sequence-head walk) iff the doubled k2 grid still fits one CTA-per-SM
+# wave, else the plain engine. DV2 is the only decomp variant, so decomp
+# always dispatches at dv=2.
+def _occupancy_pick_route(
+ cu_list: tuple[int, ...], heads: int, sm_count: int
+) -> tuple[str, int]:
+ """Pick ("decomp", 2) or ("engine", 0) by a simple occupancy rule.
+
+ n_seq * heads is the number of sequence-head walks (B*H for uniform
+ seqlen; n_seq = len(cu_list) - 1). DV2 doubles the k2 grid to 2 CTAs
+ per walk; take it iff that grid still fits one wave (<= sm_count), else
+ the plain engine.
+ """
+
+ n_seq = len(cu_list) - 1
+ if n_seq * heads * 2 <= sm_count:
+ return ("decomp", 2)
+ return ("engine", 0)
+
+
+# =============================================================================
+# Two-kernel decomposition of the same math.
+#
+# Kernel 1 (`kernel_prep`) is chunk-parallel prep: it reuses the engine's CG0
+# (gate prefix / L2 norm / decay operand staging) and warp-12 (KK / L / A_inv
+# A_inv) code paths, then composes and stores the rank-16 factors.
+# Kernel 2 (`kernel_chain_dv2`) is the lean serial recursion.
+#
+# Factor math:
+# Ainvb = A_inv (.) beta-col [16, 16]
+# W = k_restore^T @ Ainvb [DK, 16]
+# QK' = QK @ Ainvb [16, 16]
+# S_next = D (.) S + W @ (v - X), X = k_decay @ round_bf16(S)
+# o = (q_decay @ round_bf16(S) + QK' @ (v - X)) * scale
+# The value-side u_v / o_v tiles are algebraically folded into W and QK'
+# (k_restore^T @ u_v == W @ v and QK @ u_v == QK' @ v exactly), so kernel 2
+# reads the ORIGINAL v tensor as the value-side A operand and issues one extra
+# 16x16 qkv MMA instead of moving 8 KB/chunk of extra factors.
+#
+# Workspace pre-permutations (so kernel 2's raw s128 TMA loads land the exact
+# SMEM images the engine's tcgen05 descriptors expect):
+# ws_kd/ws_qd[.., c*16 + t, k ^ 8 ^ ((t & 2) * 16)] = decay[t, k]
+# (the K-box-major SW128 storage key of tcgen05_decay_b_key_storage_dim_
+# runtime; the TMA s128b row pattern then reproduces
+# tcgen05_swizzle_128b_elem_index staging exactly)
+# ws_w[.., c*16 + (j ^ 8), k] = W^T[j, k]
+# (the engine's k_restore staging is raw_f16_s128 with token rows ^8)
+# ws_qk[.., c, pairwise_sw32_smem_index(0, i, j)] = QK'[i, j]
+# ws_diag[.., c, k] = 2^{g_last[k]} (fp32; kernel 2's warp 15 scatters it
+# into the engine's SW32 block-diagonal ring off-chain)
+# =============================================================================
+
+K2_RAW_STAGE_COUNT: int = 8
+# TMA completion-mbar ring depth (< raw ring depth: at depth 8 the raw slot
+# recycles before w14 republishes, serializing slot reuse on consumer latency).
+K2_TMA_MBAR_STAGE_COUNT: int = 6
+if K2_TMA_MBAR_STAGE_COUNT > K2_RAW_STAGE_COUNT:
+ raise ValueError("TMA mbar ring cannot be deeper than the raw ring")
+TILE_ELEMS: int = BT * DK
+DIAG_REC_ELEMS: int = DK # per-chunk fp32 diag record
+QK_REC_ELEMS: int = BT * BT # per-chunk SW32-permuted QK' record
+# kernel-2 TMEM: engine layout (cols 0..271) + 2-slot packed V A-operand ring.
+TMEM_V_INPUT_COL: int = 272
+TMEM_ALLOC_COLS: int = 512
+# kernel-2 qstate/output acc ring depth. Stages keep the engine's 48-col
+# stride (208 / 256 / ...); with the drain on its own warpgroup (below) a
+# shallow ring keeps the issuer from running ahead of the drain.
+K2_QSTATE_STAGE_COUNT: int = 2
+# The kernel-2 qstate output drain (TMEM read -> o_smem stage) runs on the
+# otherwise-idle warps 0-3 (a lone warp cannot drain a 128-lane accumulator
+# because tcgen05 lane access is restricted to the warp's `warp_id % 4`
+# quadrant).
+# --- kernel-2 DV split -----------------------------------------------------
+# mode="decomp" routes through kernel_chain_dv2: one CTA per (sequence, head,
+# DV-half), M=64 MMAs throughout (PTX Layout F). DV2 is the only decomp split.
+DV_HALF: int = DV // 2
+ROWS_PER_WARP: int = DV_HALF // 4 # Layout F: 16 rows per quadrant
+V_TILE_ELEMS: int = BT * DV_HALF # one 64-elem s128 TMA segment
+K2_O_SMEM_STAGE_SIZE: int = BT * DV_HALF
+K2_O_SMEM_TILE_SIZE: int = K2_QSTATE_STAGE_COUNT * K2_O_SMEM_STAGE_SIZE
+K2_TX_BYTES: int = (
+ 3 * TILE_ELEMS * 2 # kd + w + qd (bf16, full DK)
+ + V_TILE_ELEMS * 2 # v (bf16, DV half)
+ + DIAG_REC_ELEMS * 4 # diag (fp32)
+ + QK_REC_ELEMS * 2 # qk' (bf16)
+)
+_K2_QSTATE_LAST_COL = (
+ KDA_TMEM_QSTATE_ACC_COL_OFFSET
+ + (K2_QSTATE_STAGE_COUNT - 1) * KDA_TMEM_QSTATE_ACC_STAGE_STRIDE_COLS
+ + KDA_TMEM_N16_ACC_COLS
+)
+if _K2_QSTATE_LAST_COL > TMEM_ALLOC_COLS:
+ raise ValueError("qstate acc ring exceeds the TMEM allocation")
+if (
+ K2_QSTATE_STAGE_COUNT > 2
+ and KDA_TMEM_QSTATE_ACC_COL_OFFSET + 2 * KDA_TMEM_QSTATE_ACC_STAGE_STRIDE_COLS
+ < TMEM_V_INPUT_COL + 2 * KDA_TMEM_SHARED_INPUT_COLS
+):
+ raise ValueError("qstate acc stage 2 overlaps the packed V ring")
+
+
+@cute.jit
+def super_mma_beta_col_scale_a_frags(
+ a_frags,
+ raw_beta_smem,
+ lane,
+ ab_frag,
+ input_dtype: cutlass.Constexpr,
+) -> None:
+ """Scale a packed 16x16 A-layout fragment by beta per COLUMN.
+
+ A-layout m16n8k16 fragments hold columns 2*(lane%4)/+1 in regs 0/1 and
+ columns 8+2*(lane%4)/+1 in regs 2/3, so one packed beta pair per column
+ half scales the whole fragment (Ainvb = A_inv (.) beta-col). The scaled
+ registers land in the caller-provided `ab_frag` array.
+ """
+
+ col_lo = 2 * (lane % 4)
+ beta_lo = pack_input_b16x2_to_i32(
+ raw_beta_smem[col_lo].to(cutlass.Float32),
+ raw_beta_smem[col_lo + 1].to(cutlass.Float32),
+ input_dtype,
+ )
+ beta_hi = pack_input_b16x2_to_i32(
+ raw_beta_smem[col_lo + 8].to(cutlass.Float32),
+ raw_beta_smem[col_lo + 9].to(cutlass.Float32),
+ input_dtype,
+ )
+ ab_frag[0] = mul_b16x2_input_dtype(a_frags[0], beta_lo, input_dtype)
+ ab_frag[1] = mul_b16x2_input_dtype(a_frags[1], beta_lo, input_dtype)
+ ab_frag[2] = mul_b16x2_input_dtype(a_frags[2], beta_hi, input_dtype)
+ ab_frag[3] = mul_b16x2_input_dtype(a_frags[3], beta_hi, input_dtype)
+
+
+@cute.jit
+def super_mma_load_ainv_a_frags(pairwise_stage_smem, lane):
+ """Reload the staged A_inv pairwise tile as an m16n8k16 A-layout fragment."""
+
+ return prims.ldmatrix(
+ pairwise_stmatrix_m8n8x4_ptr(
+ pairwise_stage_smem,
+ PAIRWISE_SMEM_AINV_OFFSET,
+ lane,
+ ),
+ 4,
+ prims.MMALayout.ROW,
+ )
+
+
+@cute.jit
+def super_mma_load_k_restore_a_fragment(
+ tcgen05_k_restore_smem,
+ lane,
+ n_group: cutlass.Constexpr,
+):
+ """Load an A-layout fragment of the k_restore tile (tokens x 16 key cols).
+
+ The engine stages k_restore in the raw s128 layout with token rows XORed
+ by TCGEN05_SW32_BT_HALF_XOR (see cg0_materialize_decay_operands), so the
+ per-lane row coordinate applies the same xor before the s128 index.
+ """
+
+ lane_div8 = lane // 8
+ lane_mod8 = lane % 8
+ row_offset = cutlass.Int32(8) if (lane_div8 % 2) else cutlass.Int32(0)
+ col_offset = cutlass.Int32(8) if (lane_div8 // 2) else cutlass.Int32(0)
+ row_coord = (lane_mod8 + row_offset) ^ cutlass.Int32(TCGEN05_SW32_BT_HALF_XOR)
+ col_coord = n_group * 16 + col_offset
+ ptr = tcgen05_k_restore_smem.subview(raw_f16_s128_smem_index(row_coord, col_coord))
+ return prims.ldmatrix(ptr.data_ptr(), 4, prims.MMALayout.ROW)
+
+
+@cute.jit
+def w_out_stmatrix_ptr(
+ w_out_smem,
+ n_group: cutlass.Constexpr,
+ lane,
+):
+ """Per-lane row-segment pointer for one 16x16 STSM tile of the W image.
+
+ The target SMEM image is the engine's N-major k_restore staging:
+ raw_f16_s128 with token rows XORed by 8; a TMA bulk store of this tile
+ then lands the pre-permuted ws_w global rows directly.
+ """
+
+ matrix_id = lane // 8
+ row_coord = lane & 7
+ col_coord = cutlass.Int32(n_group * 16)
+ if matrix_id & 1:
+ row_coord = row_coord + cutlass.Int32(8)
+ if matrix_id >= 2:
+ col_coord = col_coord + cutlass.Int32(8)
+ storage_row = row_coord ^ cutlass.Int32(TCGEN05_SW32_BT_HALF_XOR)
+ return w_out_smem.subview(
+ raw_f16_s128_smem_index(storage_row, col_coord)
+ ).data_ptr()
+
+
+@cute.jit
+def super_mma_stage_qk_prime(
+ tcgen05_q_decay_smem,
+ k_inv_smem,
+ pairwise_stage_smem,
+ raw_beta_smem,
+ ws_qk: cute.Tensor,
+ bidy,
+ ws_chunk,
+ lane,
+ input_dtype: cutlass.Constexpr,
+) -> None:
+ """Produce QK' = QK @ Ainvb and store it in SW32 pairwise order to ws_qk.
+
+ The QK accumulation matches super_mma_stage_qk (16 m16n8k16 K-blocks +
+ causal mask) but keeps the tile in registers instead of staging it; the
+ packed causal QK fragment is then contracted with the beta-scaled A_inv
+ fragment and the f32 product is stored bf16-rounded, element-addressed
+ with pairwise_sw32_smem_index so kernel 2's linear TMA copy of the
+ 512-byte record IS the engine's pairwise SMEM tile.
+ """
+
+ qk_n0_acc = cutlass.Array(
+ cutlass.Float32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ qk_n1_acc = cutlass.Array(
+ cutlass.Float32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ for accum_idx in cutlass.range_constexpr(SUPER_MMA_ACCUMULATORS_PER_LANE):
+ qk_n0_acc[accum_idx] = cutlass.Float32(0.0)
+ qk_n1_acc[accum_idx] = cutlass.Float32(0.0)
+
+ for k_block in cutlass.range_constexpr(SUPER_MMA_K_BLOCKS):
+ rhs_vec = super_mma_load_k_inv_rhs_fragment(
+ k_inv_smem,
+ lane,
+ k_block,
+ )
+ qk_lhs_vec = super_mma_load_decay_lhs_fragment(
+ tcgen05_q_decay_smem,
+ lane,
+ k_block,
+ )
+ qk_n0_d0, qk_n0_d1, qk_n0_d2, qk_n0_d3 = ptx_mma_m16n8k16_b16_f32(
+ qk_lhs_vec[0],
+ qk_lhs_vec[1],
+ qk_lhs_vec[2],
+ qk_lhs_vec[3],
+ rhs_vec[0],
+ rhs_vec[1],
+ qk_n0_acc[0],
+ qk_n0_acc[1],
+ qk_n0_acc[2],
+ qk_n0_acc[3],
+ input_dtype,
+ )
+ qk_n0_acc[0] = qk_n0_d0
+ qk_n0_acc[1] = qk_n0_d1
+ qk_n0_acc[2] = qk_n0_d2
+ qk_n0_acc[3] = qk_n0_d3
+ qk_n1_d0, qk_n1_d1, qk_n1_d2, qk_n1_d3 = ptx_mma_m16n8k16_b16_f32(
+ qk_lhs_vec[0],
+ qk_lhs_vec[1],
+ qk_lhs_vec[2],
+ qk_lhs_vec[3],
+ rhs_vec[2],
+ rhs_vec[3],
+ qk_n1_acc[0],
+ qk_n1_acc[1],
+ qk_n1_acc[2],
+ qk_n1_acc[3],
+ input_dtype,
+ )
+ qk_n1_acc[0] = qk_n1_d0
+ qk_n1_acc[1] = qk_n1_d1
+ qk_n1_acc[2] = qk_n1_d2
+ qk_n1_acc[3] = qk_n1_d3
+
+ qk_n0_acc[0] = super_mma_qk_causal_value(qk_n0_acc[0], lane, 0, 0)
+ qk_n0_acc[1] = super_mma_qk_causal_value(qk_n0_acc[1], lane, 0, 1)
+ qk_n0_acc[2] = super_mma_qk_causal_value(qk_n0_acc[2], lane, 0, 2)
+ qk_n0_acc[3] = super_mma_qk_causal_value(qk_n0_acc[3], lane, 0, 3)
+ qk_n1_acc[0] = super_mma_qk_causal_value(qk_n1_acc[0], lane, 1, 0)
+ qk_n1_acc[1] = super_mma_qk_causal_value(qk_n1_acc[1], lane, 1, 1)
+ qk_n1_acc[2] = super_mma_qk_causal_value(qk_n1_acc[2], lane, 1, 2)
+ qk_n1_acc[3] = super_mma_qk_causal_value(qk_n1_acc[3], lane, 1, 3)
+
+ qk_frag = cutlass.Array(
+ cutlass.Int32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ super_mma_pack_pairwise_accumulator(
+ qk_n0_acc,
+ qk_n1_acc,
+ qk_frag,
+ input_dtype,
+ )
+ ainv = super_mma_load_ainv_a_frags(pairwise_stage_smem, lane)
+ ab_frag = cutlass.Array(
+ cutlass.Int32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ super_mma_beta_col_scale_a_frags(
+ ainv,
+ raw_beta_smem,
+ lane,
+ ab_frag,
+ input_dtype,
+ )
+ p0_acc = cutlass.Array(
+ cutlass.Float32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ p1_acc = cutlass.Array(
+ cutlass.Float32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ super_mma_pairwise_product_from_ab_regs(
+ qk_frag,
+ ab_frag,
+ p0_acc,
+ p1_acc,
+ input_dtype,
+ )
+
+ row_lo = super_mma_accumulator_row(lane, 0)
+ row_hi = super_mma_accumulator_row(lane, 2)
+ n0_col0 = super_mma_accumulator_col(lane, 0, 0)
+ n0_col1 = super_mma_accumulator_col(lane, 0, 1)
+ n1_col0 = super_mma_accumulator_col(lane, 1, 0)
+ n1_col1 = super_mma_accumulator_col(lane, 1, 1)
+ ws_qk[0, bidy, ws_chunk, pairwise_sw32_smem_index(0, row_lo, n0_col0)] = p0_acc[
+ 0
+ ].to(input_dtype)
+ ws_qk[0, bidy, ws_chunk, pairwise_sw32_smem_index(0, row_lo, n0_col1)] = p0_acc[
+ 1
+ ].to(input_dtype)
+ ws_qk[0, bidy, ws_chunk, pairwise_sw32_smem_index(0, row_hi, n0_col0)] = p0_acc[
+ 2
+ ].to(input_dtype)
+ ws_qk[0, bidy, ws_chunk, pairwise_sw32_smem_index(0, row_hi, n0_col1)] = p0_acc[
+ 3
+ ].to(input_dtype)
+ ws_qk[0, bidy, ws_chunk, pairwise_sw32_smem_index(0, row_lo, n1_col0)] = p1_acc[
+ 0
+ ].to(input_dtype)
+ ws_qk[0, bidy, ws_chunk, pairwise_sw32_smem_index(0, row_lo, n1_col1)] = p1_acc[
+ 1
+ ].to(input_dtype)
+ ws_qk[0, bidy, ws_chunk, pairwise_sw32_smem_index(0, row_hi, n1_col0)] = p1_acc[
+ 2
+ ].to(input_dtype)
+ ws_qk[0, bidy, ws_chunk, pairwise_sw32_smem_index(0, row_hi, n1_col1)] = p1_acc[
+ 3
+ ].to(input_dtype)
+
+
+@cute.jit
+def chunk_coords(cu_seqlens, cu_chunks, gchunk):
+ """Return coordinates for a global chunk using the cumulative chunk prefix."""
+
+ # Binary-search the compact [num_sequences + 1] prefix instead of carrying
+ # a redundant dense chunk_to_seq[total_chunks] inverse map.
+ lo = cutlass.Int32(0)
+ hi = cutlass.Int32(cu_chunks.shape[0] - 1)
+ while lo + cutlass.Int32(1) < hi:
+ mid = (lo + hi) // cutlass.Int32(2)
+ if cutlass.Int32(cu_chunks[mid]) <= gchunk:
+ lo = mid
+ else:
+ hi = mid
+ seq = lo
+ sequence_start = cutlass.Int32(cu_seqlens[seq])
+ seqlen = cutlass.Int32(cu_seqlens[seq + 1]) - sequence_start
+ chunk_start_tok = (gchunk - cutlass.Int32(cu_chunks[seq])) * cutlass.Int32(BT)
+ return sequence_start, seqlen, chunk_start_tok
+
+
+@cute.jit
+def tma_prep_stage_load_inputs(
+ tma_desc_q: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_k: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_gate: cutlass.GridConstant[cuda.TensorMap],
+ beta_logits: cute.Tensor,
+ raw_q_smem,
+ raw_k_smem,
+ raw_gate_smem,
+ raw_beta_smem,
+ sequence_start,
+ head_idx,
+ lane,
+ chunk_start_tok,
+ seqlen,
+ tma_mbar,
+ tma_tx_bytes: cutlass.Constexpr,
+ gate_dtype: cutlass.Constexpr,
+) -> None:
+ """Kernel-1 chunk input transaction: q/k/gate TMA + beta sigmoid (no v)."""
+
+ global_chunk_start = sequence_start + chunk_start_tok
+ if prims.elect_sync():
+ prims.mbarrier_arrive_expect_tx(tma_mbar, tma_tx_bytes)
+ if prims.elect_sync():
+ for segment in cutlass.range_constexpr(RAW_F16_TMA_SEGMENTS):
+ tma_coord = (
+ cutlass.Int32(segment * RAW_F16_TMA_SWIZZLE_ELEMS),
+ global_chunk_start,
+ head_idx,
+ cutlass.Int32(0),
+ )
+ smem_offset = segment * RAW_F16_TMA_SEGMENT_ELEMS
+ prims.cp_async_bulk_tensor_shared_cta_global(
+ raw_q_smem.subview(smem_offset),
+ tma_desc_q.get_ptr(),
+ tma_coord,
+ tma_mbar,
+ )
+ prims.cp_async_bulk_tensor_shared_cta_global(
+ raw_k_smem.subview(smem_offset),
+ tma_desc_k.get_ptr(),
+ tma_coord,
+ tma_mbar,
+ )
+ if cutlass.const_expr(not gate_dtype_is_f32(gate_dtype)):
+ # 16-bit gate: shares the q/k 16-bit box/swizzle family.
+ prims.cp_async_bulk_tensor_shared_cta_global(
+ raw_gate_smem.subview(smem_offset),
+ tma_desc_gate.get_ptr(),
+ tma_coord,
+ tma_mbar,
+ )
+ if cutlass.const_expr(gate_dtype_is_f32(gate_dtype)):
+ for segment in cutlass.range_constexpr(RAW_F32_TMA_SEGMENTS):
+ tma_coord = (
+ cutlass.Int32(segment * RAW_F32_TMA_SWIZZLE_ELEMS),
+ global_chunk_start,
+ head_idx,
+ cutlass.Int32(0),
+ )
+ smem_offset = segment * RAW_F32_TMA_SEGMENT_ELEMS
+ prims.cp_async_bulk_tensor_shared_cta_global(
+ raw_gate_smem.subview(smem_offset),
+ tma_desc_gate.get_ptr(),
+ tma_coord,
+ tma_mbar,
+ )
+
+ if lane < BT:
+ token_idx = chunk_start_tok + lane
+ beta_value = cutlass.Float32(0.0)
+ if token_idx < seqlen:
+ beta_logit = beta_logits[0, sequence_start + token_idx, head_idx].to(
+ cutlass.Float32
+ )
+ half = cutlass.Float32(0.5)
+ beta_value = cute.math.tanh(beta_logit * half, approx=True) * half + half
+ raw_beta_smem[lane] = beta_value
+
+
+@cute.jit
+def host_prep(
+ q: cute.Tensor,
+ k: cute.Tensor,
+ raw_gate: cute.Tensor,
+ a_log: cute.Tensor,
+ dt_bias: cute.Tensor,
+ beta: cute.Tensor,
+ cu_seqlens: cute.Tensor,
+ cu_chunks: cute.Tensor,
+ ws_kd: cute.Tensor,
+ ws_qd: cute.Tensor,
+ ws_w: cute.Tensor,
+ ws_qk: cute.Tensor,
+ ws_diag: cute.Tensor,
+ stream,
+ num_ctas: cutlass.Int32,
+ chunks_per_cta: cutlass.Int32,
+ head_base: cutlass.Int32,
+ launch_heads: cutlass.Int32,
+ SAFE_GATE: cutlass.Constexpr,
+ GATE_SCALE_LOG2: cutlass.Constexpr,
+ THREADS: cutlass.Constexpr,
+ gate_dtype: cutlass.Constexpr,
+) -> None:
+ seqlen = q.shape[1]
+ heads = q.shape[2]
+ # Token-major user activations (see host()).
+ qk_layout = cute.make_layout(
+ (DK, seqlen, heads, 1),
+ stride=(1, DK * heads, DK, DK * seqlen * heads),
+ )
+ raw_f16_tma_box = (RAW_F16_TMA_SWIZZLE_ELEMS, BT, 1, 1)
+ raw_f32_tma_box = (RAW_F32_TMA_SWIZZLE_ELEMS, BT, 1, 1)
+ tma_desc_q = cuda.create_tensor_map_tiled_from_view(
+ cute.make_tensor(q.iterator, qk_layout),
+ box_dims=raw_f16_tma_box,
+ stride_order=(0, 1, 2, 3),
+ swizzle=cuda.TensorMapSwizzle.s128b,
+ )
+ tma_desc_k = cuda.create_tensor_map_tiled_from_view(
+ cute.make_tensor(k.iterator, qk_layout),
+ box_dims=raw_f16_tma_box,
+ stride_order=(0, 1, 2, 3),
+ swizzle=cuda.TensorMapSwizzle.s128b,
+ )
+ gate_tma_box = raw_f32_tma_box if gate_dtype_is_f32(gate_dtype) else raw_f16_tma_box
+ tma_desc_gate = cuda.create_tensor_map_tiled_from_view(
+ cute.make_tensor(raw_gate.iterator, qk_layout),
+ box_dims=gate_tma_box,
+ stride_order=(0, 1, 2, 3),
+ swizzle=cuda.TensorMapSwizzle.s128b,
+ )
+ ws_rows = ws_kd.shape[2]
+ # Canonicalize the unused head stride for H=1. Leaving it proportional to
+ # ws_rows reaches the grid-constant TensorMap launch boundary at 1M tokens,
+ # even though a singleton mode never contributes to an address.
+ ws_head_stride = DK
+ if heads != cutlass.Int32(1):
+ ws_head_stride = DK * ws_rows
+ ws_tile_layout = cute.make_layout(
+ (DK, ws_rows, heads, 1),
+ stride=(1, DK, ws_head_stride, DK),
+ )
+ tma_desc_ws_kd = cuda.create_tensor_map_tiled_from_view(
+ cute.make_tensor(ws_kd.iterator, ws_tile_layout),
+ box_dims=raw_f16_tma_box,
+ stride_order=(0, 1, 2, 3),
+ swizzle=cuda.TensorMapSwizzle.s128b,
+ )
+ tma_desc_ws_qd = cuda.create_tensor_map_tiled_from_view(
+ cute.make_tensor(ws_qd.iterator, ws_tile_layout),
+ box_dims=raw_f16_tma_box,
+ stride_order=(0, 1, 2, 3),
+ swizzle=cuda.TensorMapSwizzle.s128b,
+ )
+ tma_desc_ws_w = cuda.create_tensor_map_tiled_from_view(
+ cute.make_tensor(ws_w.iterator, ws_tile_layout),
+ box_dims=raw_f16_tma_box,
+ stride_order=(0, 1, 2, 3),
+ swizzle=cuda.TensorMapSwizzle.s128b,
+ )
+ # per-chunk-tile kernel 1 (4-warp CTAs; THREADS is the prep ABI
+ # placeholder and is ignored — the launch uses K1_THREADS with a
+ # min-CTAs/SM occupancy floor so several CTAs co-reside per SM).
+ kernel_prep(
+ tma_desc_q,
+ tma_desc_k,
+ tma_desc_gate,
+ tma_desc_ws_kd,
+ tma_desc_ws_qd,
+ tma_desc_ws_w,
+ q,
+ k,
+ raw_gate,
+ a_log,
+ dt_bias,
+ beta,
+ cu_seqlens,
+ cu_chunks,
+ ws_kd,
+ ws_qd,
+ ws_w,
+ ws_qk,
+ ws_diag,
+ chunks_per_cta,
+ head_base,
+ SAFE_GATE,
+ GATE_SCALE_LOG2,
+ gate_dtype,
+ ).launch(
+ grid=(num_ctas, launch_heads, 1),
+ block=(K1_THREADS, 1, 1),
+ stream=stream,
+ min_blocks_per_mp=K1_MIN_BLOCKS,
+ # Explicit carveout: (a) several 38-KB CTAs per SM need the max
+ # SMEM split anyway, (b) leaving it None makes min_blocks_per_mp>1
+ # query the device at COMPILE time, which breaks the login-node
+ # CUTE_DSL_DRYRUN flow.
+ preferred_smem_carveout=100,
+ )
+
+
+_PLAN_CACHE: dict = {}
+
+
+def _plan(cu_seqlens: torch.Tensor) -> dict:
+ """Cached cumulative chunk-prefix tensor and host metadata.
+
+ Keyed by the cu_seqlens CONTENTS (data_ptr identity is unsafe under
+ allocator block reuse — the s27 stale-plan hazard; the read is
+ memoized sync-free by `_cu_seqlens_contents`).
+ """
+
+ cu_list = list(_cu_seqlens_contents(cu_seqlens))
+ key = (tuple(cu_list), str(cu_seqlens.device))
+ plan = _PLAN_CACHE.get(key)
+ if plan is None:
+ chunks = [
+ (cu_list[i + 1] - cu_list[i] + BT - 1) // BT
+ for i in range(len(cu_list) - 1)
+ ]
+ cu_chunks = [0]
+ for c in chunks:
+ cu_chunks.append(cu_chunks[-1] + c)
+ total_chunks = cu_chunks[-1]
+ plan = {
+ "total_chunks": total_chunks,
+ "cu_chunks": torch.tensor(
+ cu_chunks, dtype=torch.int32, device=cu_seqlens.device
+ ),
+ "cu_list": cu_list,
+ }
+ _PLAN_CACHE[key] = plan
+ return plan
+
+
+# ---------------------------------------------------------------------------
+# CUTLASS-style user-allocated decomp workspace (single opaque uint8 buffer).
+#
+# The decomp route needs the kernel-1 -> kernel-2 factor tiles plus a tiny
+# flag array. Instead of a hidden per-shape cache, the caller queries
+# `workspace_size(cu_seqlens, heads[, mode])` and allocates one raw uint8
+# buffer; `host` partitions it itself and does ALL init (it zeroes the
+# flag region each call). The buffer is kernel-managed: pass a raw/uninitialized
+# buffer, reuse it freely across calls, never touch it otherwise.
+#
+# Layout (each region rounded up to WS_REGION_ALIGN so every view is
+# >=16B aligned for TMA):
+# [ ws_kd | ws_qd | ws_w | ws_qk | ws_diag ]
+# ws_kd/ws_qd/ws_w : bf16 (1, heads, total_chunks*BT, DK)
+# ws_qk : bf16 (1, heads, total_chunks, QK_REC_ELEMS)
+# ws_diag : f32 (1, heads, total_chunks, DIAG_REC_ELEMS)
+# ---------------------------------------------------------------------------
+
+WS_REGION_ALIGN: int = 256
+
+
+def _align_up(n: int, a: int) -> int:
+ return -(-int(n) // a) * a
+
+
+def _decomp_ws_regions(heads: int, total_chunks: int) -> list[tuple]:
+ """Region byte sizes and shapes for the decomp workspace, in buffer order.
+
+ Each entry is (nbytes, torch_dtype, shape).
+ """
+
+ rows = total_chunks * BT
+ tile = (heads * rows * DK * 2, torch.bfloat16, (1, heads, rows, DK))
+ return [
+ tile,
+ tile,
+ tile,
+ (
+ heads * total_chunks * QK_REC_ELEMS * 2,
+ torch.bfloat16,
+ (1, heads, total_chunks, QK_REC_ELEMS),
+ ),
+ (
+ heads * total_chunks * DIAG_REC_ELEMS * 4,
+ torch.float32,
+ (1, heads, total_chunks, DIAG_REC_ELEMS),
+ ),
+ ]
+
+
+def _decomp_ws_bytes(heads: int, total_chunks: int) -> int:
+ return sum(
+ _align_up(nbytes, WS_REGION_ALIGN)
+ for nbytes, _, _ in _decomp_ws_regions(heads, total_chunks)
+ )
+
+
+def _partition_workspace(workspace, heads: int, total_chunks: int) -> dict:
+ """View the user uint8 buffer as the six factor/flag tensors (no copy)."""
+
+ if workspace.dtype != torch.uint8 or workspace.dim() != 1:
+ raise ValueError("workspace must be a 1-D uint8 tensor")
+ need = _decomp_ws_bytes(heads, total_chunks)
+ if workspace.numel() < need:
+ raise ValueError(
+ f"workspace too small: {workspace.numel()} bytes < {need} required; "
+ "call workspace_size() to size it"
+ )
+ keys = ("kd", "qd", "w", "qk", "diag")
+ out: dict = {}
+ off = 0
+ for key, (nbytes, dtype, shape) in zip(
+ keys, _decomp_ws_regions(heads, total_chunks), strict=True
+ ):
+ seg = workspace.narrow(0, off, nbytes)
+ out[key] = seg.view(dtype).view(*shape)
+ off += _align_up(nbytes, WS_REGION_ALIGN)
+ return out
+
+
+def _route_for_workspace(num_sequences: int, heads: int, device, mode: str) -> str:
+ if mode == "engine":
+ return "engine"
+ if mode == "decomp":
+ return "decomp"
+ if mode != "auto":
+ raise ValueError(f"Unknown mode: {mode}")
+ sm_count = _device_sm_count(device)
+ return "decomp" if num_sequences * heads * 2 <= sm_count else "engine"
+
+
+def workspace_size(cu_seqlens, heads: int, mode: str = "auto") -> int:
+ """Bytes for the opaque decomp workspace of this problem shape.
+
+ Returns 0 for the engine route (engine needs no workspace), else the total
+ partitioned + 256B-region-aligned byte count. Uses the same CTA-count
+ occupancy rule as `host`, so the query and the launch agree.
+ ``cu_seqlens`` is the int64 cumulative-length tensor (or a host sequence).
+ """
+
+ heads = int(heads)
+ if isinstance(cu_seqlens, torch.Tensor):
+ device = cu_seqlens.device
+ num_sequences = cu_seqlens.numel() - 1
+ else:
+ cu_list = tuple(int(x) for x in cu_seqlens)
+ device = torch.cuda.current_device()
+ num_sequences = len(cu_list) - 1
+ if _route_for_workspace(num_sequences, heads, device, mode) == "engine":
+ return 0
+ if isinstance(cu_seqlens, torch.Tensor):
+ cu_list = _cu_seqlens_contents(cu_seqlens)
+ total_chunks = sum(
+ (cu_list[i + 1] - cu_list[i] + BT - 1) // BT for i in range(len(cu_list) - 1)
+ )
+ return _decomp_ws_bytes(heads, total_chunks)
+
+
+def _k1_grid(total_chunks: int, heads: int) -> tuple[int, int]:
+ """Return (num_ctas, chunks_per_cta) for the per-chunk-tile grid.
+
+ One CTA per K1_CPC contiguous chunks (the measured ladder puts the
+ optimum at 4); `heads` is the grid's y dimension and does not shape x.
+ """
+
+ if total_chunks <= 0:
+ return 1, 1
+ chunks_per_cta = max(1, K1_CPC)
+ num_ctas = -(-total_chunks // chunks_per_cta)
+ return num_ctas, chunks_per_cta
+
+
+# =============================================================================
+# kernel 1: FLA-style per-chunk-tile factor prep.
+#
+# One small CTA per (chunk-group, head) tile; 4 warps, no warp
+# specialization, no cross-chunk rings. All factor math and workspace
+# pre-permutations are the ring-based donors (byte-identical output); only the
+# CTA structure changes. Warp roles inside the CTA:
+# warp 0: TMA issue + beta, CG0 rows 0-3, KK -> L -> A_inv,
+# overlap-flag publish
+# warp 1: CG0 rows 4-7, W fold columns [0, 64) + W TMA store (seg 0)
+# warp 2: CG0 rows 8-11, QK' fold (plain global stores)
+# warp 3: CG0 rows 12-15, kd/qd TMA stores, W fold columns [64, 128)
+# + W TMA store (seg 1)
+# =============================================================================
+
+K1_WARPS: int = 4
+K1_THREADS: int = K1_WARPS * THREADS_PER_WARP
+# Chunks walked per CTA (grid x = ceil(total_chunks / CPC)). 1 = pure
+# per-chunk-tile grid; >1 amortizes launch/preamble cost per CTA and
+# enables the mid-chunk prefetch of the next chunk's raw tiles (measured
+# ladder: CPC=4 optimal).
+K1_CPC: int = 4
+# Launch occupancy floor (min CTAs per SM): caps regs/thread so several
+# small CTAs co-reside per SM — the entire point of this kernel shape.
+# Measured: 5 (~102 regs/thread) beats 4 (128 regs) and 8 (64 regs).
+K1_MIN_BLOCKS: int = 5
+K1_W_N_GROUPS: int = DK // 16
+K1_W_N_SPLIT: int = K1_W_N_GROUPS // 2
+# Which warp issues the CPC>1 mid-chunk prefetch of the next chunk's raw
+# tiles (must be a warp that waits q_k_restore_ready: 1 = W-half warp
+# (lighter tail), 2 = QK' warp).
+K1_PREFETCH_WARP: int = 1
+
+
+@cute.jit
+def prep_cta_sync() -> None:
+ """Synchronize the 4-warp kernel-1 CTA."""
+
+ prims.barrier_cta_sync(0, thread_count=K1_THREADS)
+
+
+@cute.jit
+def prefetch_next_chunk(
+ tma_desc_q: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_k: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_gate: cutlass.GridConstant[cuda.TensorMap],
+ beta: cute.Tensor,
+ raw_q_smem,
+ raw_k_smem,
+ raw_gate_smem,
+ raw_beta_smem,
+ cu_seqlens: cute.Tensor,
+ cu_chunks: cute.Tensor,
+ local_chunk,
+ my_chunks,
+ gchunk,
+ gchunk_stride,
+ bidy,
+ lane,
+ tma_mbar,
+ tma_tx_bytes: cutlass.Constexpr,
+ gate_dtype: cutlass.Constexpr,
+) -> None:
+ """CPC>1 prefetch: overwrite the raw tiles with chunk c+1's transaction.
+
+ Caller contract: issued only AFTER this chunk's q_k_restore_ready
+ completed (all 4 warps arrived), i.e. every raw q/k/gate read of the
+ chunk in flight is done. The FP32 gate-prefix exchange is produced AND
+ consumed inside materialize, either aliased into raw_gate_smem (FP32 gate)
+ or in its own gate_exchange_smem tile (16-bit gate); this prefetch only
+ rewrites the raw tiles, so it cannot disturb it either way. Beta lands in the
+ other half of the double-buffered record; the transaction completes
+ tma_mbar[(c+1) % 2].
+ """
+
+ if local_chunk + cutlass.Int32(1) < my_chunks:
+ next_gchunk = gchunk + gchunk_stride
+ (
+ next_sequence_start,
+ next_seqlen,
+ next_chunk_start_tok,
+ ) = chunk_coords(
+ cu_seqlens,
+ cu_chunks,
+ next_gchunk,
+ )
+ tma_prep_stage_load_inputs(
+ tma_desc_q,
+ tma_desc_k,
+ tma_desc_gate,
+ beta,
+ raw_q_smem,
+ raw_k_smem,
+ raw_gate_smem,
+ raw_beta_smem.subview(((local_chunk + 1) % 2) * RAW_BETA_STAGE_SIZE),
+ next_sequence_start,
+ bidy,
+ lane,
+ next_chunk_start_tok,
+ next_seqlen,
+ tma_mbar.subview((local_chunk + 1) % 2),
+ tma_tx_bytes,
+ gate_dtype,
+ )
+
+
+@cute.jit
+def cg0_materialize_prep_operands(
+ raw_q_smem,
+ raw_k_smem,
+ raw_gate_smem,
+ gate_exchange_smem,
+ a_log_exp,
+ dt_bias_value,
+ k_inv_smem,
+ tcgen05_k_decay_smem,
+ tcgen05_q_decay_smem,
+ tcgen05_k_restore_smem,
+ cg0_k_ready_mbar,
+ cg0_k_half_ready_mbar,
+ ws_diag: cute.Tensor,
+ bidy,
+ ws_chunk,
+ chunk_start_tok,
+ seqlen,
+ input_dtype: cutlass.Constexpr,
+ gate_dtype: cutlass.Constexpr,
+ SAFE_GATE: cutlass.Constexpr,
+ GATE_SCALE_LOG2: cutlass.Constexpr,
+ cg0_local_warp,
+ lane,
+) -> None:
+ """One-shot ring-free CG0 materializer.
+
+ All 4 CTA warps form a single CG0 group (cg0_local_warp = warp index);
+ the ring-stage consumed waits of the ring-based version are gone (single
+ stage SMEM, chunk serialization by the caller's CTA sync). The
+ cg0_k_half/full ready arrives are kept so warp 0's KK K-blocks 0..3
+ can start while the second half-DK is still being staged.
+ """
+
+ row_group_start = cg0_local_warp * CG0_TOKEN_ROWS_PER_WARP
+ lane_row_group = lane // RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ lane_in_row_group = lane - lane_row_group * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ decay_row = row_group_start + lane_row_group
+
+ raw_q_ptr = raw_q_smem.data_ptr()
+ raw_k_ptr = raw_k_smem.data_ptr()
+ # The FP32 gate-prefix exchange buffer. With an FP32 gate the caller
+ # passes the raw-gate stage itself (the historical aliasing); with a 16-bit
+ # gate it is the dedicated `gate_exchange_smem` tile.
+ g_prefix_ptr = gate_exchange_smem.data_ptr()
+ k_inv_ptr = k_inv_smem.data_ptr()
+ tcgen05_k_decay_ptr = tcgen05_k_decay_smem.data_ptr()
+ tcgen05_q_decay_ptr = tcgen05_q_decay_smem.data_ptr()
+ tcgen05_k_restore_ptr = tcgen05_k_restore_smem.data_ptr()
+
+ prefix_dim = cg0_local_warp * THREADS_PER_WARP + lane
+ g_prefix_regs = cutlass.Array(
+ cutlass.Float32,
+ BT,
+ alignment=16,
+ )
+ if cutlass.const_expr(SAFE_GATE):
+ # Fold tanh's exact 0.5 scale into the chunk-uniform coefficient.
+ a_log_exp_half = a_log_exp * cutlass.Float32(0.5)
+ for row_pair in cutlass.range_constexpr(BT // 2):
+ row0 = row_pair * 2
+ row1 = row0 + 1
+ # Only the MEMORY FORMAT of the raw gate depends on gate_dtype:
+ # a 16-bit gate lands in the q/k s128 geometry and is widened to
+ # FP32 right here -- ALL gate arithmetic below stays FP32.
+ if cutlass.const_expr(gate_dtype_is_f32(gate_dtype)):
+ prefix_idx0 = raw_f32_s128_smem_index(row0, prefix_dim)
+ prefix_idx1 = raw_f32_s128_smem_index(row1, prefix_dim)
+ gate0 = raw_gate_smem[prefix_idx0]
+ gate1 = raw_gate_smem[prefix_idx1]
+ else:
+ prefix_idx0 = raw_f16_s128_smem_index(row0, prefix_dim)
+ prefix_idx1 = raw_f16_s128_smem_index(row1, prefix_dim)
+ gate0 = raw_gate_smem[prefix_idx0].to(cutlass.Float32)
+ gate1 = raw_gate_smem[prefix_idx1].to(cutlass.Float32)
+ gate0 = a_log_exp_half * (gate0 + dt_bias_value)
+ gate1 = a_log_exp_half * (gate1 + dt_bias_value)
+ gate0 = safe_gate_log2_increment_prehalved(
+ gate0,
+ SAFE_GATE,
+ GATE_SCALE_LOG2,
+ )
+ gate1 = safe_gate_log2_increment_prehalved(
+ gate1,
+ SAFE_GATE,
+ GATE_SCALE_LOG2,
+ )
+ gate_pair = cutlass.Vector.from_elements((gate0, gate1), cutlass.Float32)
+ g_prefix_regs[row0] = gate_pair[0]
+ g_prefix_regs[row1] = gate_pair[1]
+ else:
+ # FLA-compatible non-safe gate: accept raw gate logits + A_log +
+ # dt_bias and compute the log2-domain decay increment in-kernel as
+ # ``-exp(A_log) * softplus(raw_gate + dt_bias) * log2(e)``. a_log_exp
+ # already carries ``exp2(A_log * log2(e)) == exp(A_log)`` and the
+ # softplus is evaluated in base-2 units (softplus_log2_f32).
+ for row in cutlass.range_constexpr(BT):
+ if cutlass.const_expr(gate_dtype_is_f32(gate_dtype)):
+ prefix_idx = raw_f32_s128_smem_index(row, prefix_dim)
+ gate = raw_gate_smem[prefix_idx]
+ else:
+ prefix_idx = raw_f16_s128_smem_index(row, prefix_dim)
+ gate = raw_gate_smem[prefix_idx].to(cutlass.Float32)
+ gate = -a_log_exp * softplus_log2_f32(gate + dt_bias_value)
+ g_prefix_regs[row] = gate
+
+ # Masked-zeroing tail fixup: a warp-uniform runtime compare (one per chunk)
+ # applies only to a sequence's genuinely partial tail chunk; full chunks
+ # skip it via the `tail_valid_rows < BT` test.
+ tail_valid_rows = seqlen - chunk_start_tok
+ if tail_valid_rows < cutlass.Int32(BT):
+ tail_mask_pt = cutlass.vector.create_mask([BT], [tail_valid_rows])
+ for row_pair_pt in cutlass.range_constexpr(BT // 2):
+ row0_pt = row_pair_pt * 2
+ row1_pt = row0_pt + 1
+ gate_pair_pt = cutlass.Vector.from_elements(
+ (g_prefix_regs[row0_pt], g_prefix_regs[row1_pt]),
+ cutlass.Float32,
+ )
+ gate_pair_pt = cutlass.vector.where(
+ tail_mask_pt[row0_pt : row1_pt + 1], gate_pair_pt, 0.0
+ )
+ g_prefix_regs[row0_pt] = gate_pair_pt[0]
+ g_prefix_regs[row1_pt] = gate_pair_pt[1]
+
+ prefix_acc = cutlass.Float32(0.0)
+ for row_pair in cutlass.range_constexpr(BT // 2):
+ row0 = row_pair * 2
+ row1 = row0 + 1
+ # The scalar scan has the same dependency depth as packed fadd2 but
+ # avoids the register copy needed to construct its first input pair.
+ prefix0 = prefix_acc + g_prefix_regs[row0]
+ prefix1 = prefix0 + g_prefix_regs[row1]
+ g_prefix_regs[row0] = prefix0
+ g_prefix_regs[row1] = prefix1
+ prefix_acc = prefix1
+
+ for row in cutlass.range_constexpr(BT):
+ g_prefix_regs[row] = cute.math.exp2(g_prefix_regs[row], fastmath=True)
+
+ exp_g_last = g_prefix_regs[BT - 1]
+ for row in cutlass.range_constexpr(BT):
+ prefix_idx = raw_f32_exchange_smem_index(row, prefix_dim)
+ gate_exchange_smem[prefix_idx] = g_prefix_regs[row]
+
+ ws_diag[0, bidy, ws_chunk, prefix_dim] = exp_g_last
+
+ cg0_sync(cutlass.Int32(0))
+
+ k_inv_regs = cutlass.Array(
+ input_dtype,
+ 2 * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=16,
+ )
+ k_restore_all_regs = cutlass.Array(
+ input_dtype,
+ 2 * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=16,
+ )
+ raw_q_regs = cutlass.Array(
+ cutlass.Float32,
+ 2 * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=16,
+ )
+ raw_k_regs = cutlass.Array(
+ cutlass.Float32,
+ 2 * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=16,
+ )
+ q_sum_sq = cutlass.Float32(0.0)
+ k_sum_sq = cutlass.Float32(0.0)
+ for dim_half in cutlass.range_constexpr(2):
+ dim_base = (
+ dim_half * (DK // 2) + lane_in_row_group * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ )
+ reg_base = dim_half * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ raw_f16_idx = raw_f16_s128_smem_index(decay_row, dim_base)
+ raw_q_vec = (raw_q_ptr + raw_f16_idx).load(
+ count=RAW_F16_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=RAW_F16_TMA_SWIZZLE_GROUP_BYTES,
+ )
+ raw_k_vec = (raw_k_ptr + raw_f16_idx).load(
+ count=RAW_F16_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=RAW_F16_TMA_SWIZZLE_GROUP_BYTES,
+ )
+ raw_q_vec_f32 = raw_q_vec.to(cutlass.Float32)
+ raw_k_vec_f32 = raw_k_vec.to(cutlass.Float32)
+ for dim_offset in cutlass.range_constexpr(RAW_F16_TMA_SWIZZLE_GROUP_ELEMS):
+ q_val = raw_q_vec_f32[dim_offset]
+ k_val = raw_k_vec_f32[dim_offset]
+ raw_q_regs[reg_base + dim_offset] = q_val
+ raw_k_regs[reg_base + dim_offset] = k_val
+ q_sum_sq = q_sum_sq + q_val * q_val
+ k_sum_sq = k_sum_sq + k_val * k_val
+
+ q_sum_sq = warp_group_sum_8(q_sum_sq)
+ k_sum_sq = warp_group_sum_8(k_sum_sq)
+ norm_floor_sq = cutlass.Float32(L2_NORM_EPS * L2_NORM_EPS)
+ q_inv_norm = cute.math.rsqrt(
+ cute.math.max(q_sum_sq, norm_floor_sq, ftz=True),
+ fastmath=True,
+ )
+ k_inv_norm = cute.math.rsqrt(
+ cute.math.max(k_sum_sq, norm_floor_sq, ftz=True),
+ fastmath=True,
+ )
+
+ exp_g_regs = cutlass.Array(
+ cutlass.Float32,
+ 2 * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=16,
+ )
+ exp_g_last_regs = cutlass.Array(
+ cutlass.Float32,
+ 2 * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=16,
+ )
+ for dim_half in cutlass.range_constexpr(2):
+ dim_base = (
+ dim_half * (DK // 2) + lane_in_row_group * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ )
+ reg_base = dim_half * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ exp_neg_g_regs = cutlass.Array(
+ cutlass.Float32,
+ RAW_F16_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=16,
+ )
+ for f32_group in cutlass.range_constexpr(
+ RAW_F16_TMA_SWIZZLE_GROUP_ELEMS // RAW_F32_TMA_SWIZZLE_GROUP_ELEMS
+ ):
+ f32_dim_base = dim_base + f32_group * RAW_F32_TMA_SWIZZLE_GROUP_ELEMS
+ g_prefix_idx = raw_f32_exchange_smem_index(decay_row, f32_dim_base)
+ exp_g_vec = (g_prefix_ptr + g_prefix_idx).load(
+ count=RAW_F32_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=RAW_F32_TMA_SWIZZLE_GROUP_BYTES,
+ )
+ exp_g_last_idx = raw_f32_exchange_smem_index(BT - 1, f32_dim_base)
+ exp_g_last_vec = (g_prefix_ptr + exp_g_last_idx).load(
+ count=RAW_F32_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=RAW_F32_TMA_SWIZZLE_GROUP_BYTES,
+ )
+ half_reg_base = f32_group * RAW_F32_TMA_SWIZZLE_GROUP_ELEMS
+ f32_reg_base = reg_base + half_reg_base
+ exp_g_regs[f32_reg_base] = exp_g_vec[0]
+ exp_g_regs[f32_reg_base + 1] = exp_g_vec[1]
+ exp_g_regs[f32_reg_base + 2] = exp_g_vec[2]
+ exp_g_regs[f32_reg_base + 3] = exp_g_vec[3]
+ exp_neg_g_regs[half_reg_base] = cute.math.rcp(
+ exp_g_vec[0], approx=True, ftz=True
+ )
+ exp_neg_g_regs[half_reg_base + 1] = cute.math.rcp(
+ exp_g_vec[1], approx=True, ftz=True
+ )
+ exp_neg_g_regs[half_reg_base + 2] = cute.math.rcp(
+ exp_g_vec[2], approx=True, ftz=True
+ )
+ exp_neg_g_regs[half_reg_base + 3] = cute.math.rcp(
+ exp_g_vec[3], approx=True, ftz=True
+ )
+ exp_g_last_regs[f32_reg_base] = exp_g_last_vec[0]
+ exp_g_last_regs[f32_reg_base + 1] = exp_g_last_vec[1]
+ exp_g_last_regs[f32_reg_base + 2] = exp_g_last_vec[2]
+ exp_g_last_regs[f32_reg_base + 3] = exp_g_last_vec[3]
+
+ k_decay_vec_regs = cutlass.Array(
+ input_dtype,
+ RAW_F16_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=RAW_F16_TMA_SWIZZLE_GROUP_BYTES,
+ )
+ for pair_idx in cutlass.range_constexpr(RAW_F16_TMA_SWIZZLE_GROUP_ELEMS // 2):
+ dim0 = pair_idx * 2
+ dim1 = dim0 + 1
+ raw_reg_idx0 = reg_base + dim0
+ raw_reg_idx1 = reg_base + dim1
+ k_value0, k_value1 = fmul2(
+ (raw_k_regs[raw_reg_idx0], raw_k_regs[raw_reg_idx1]),
+ (k_inv_norm, k_inv_norm),
+ )
+ k_decay0, k_decay1 = fmul2(
+ (k_value0, k_value1),
+ (exp_g_regs[raw_reg_idx0], exp_g_regs[raw_reg_idx1]),
+ )
+ k_inv0, k_inv1 = fmul2(
+ (k_value0, k_value1),
+ (exp_neg_g_regs[dim0], exp_neg_g_regs[dim1]),
+ )
+ k_inv_regs[reg_base + dim0] = k_inv0.to(input_dtype)
+ k_inv_regs[reg_base + dim1] = k_inv1.to(input_dtype)
+ k_restore0, k_restore1 = fmul2(
+ (k_inv0, k_inv1),
+ (
+ exp_g_last_regs[reg_base + dim0],
+ exp_g_last_regs[reg_base + dim1],
+ ),
+ )
+ k_restore_all_regs[reg_base + dim0] = k_restore0.to(input_dtype)
+ k_restore_all_regs[reg_base + dim1] = k_restore1.to(input_dtype)
+ k_decay_vec_regs[dim0] = k_decay0.to(input_dtype)
+ k_decay_vec_regs[dim1] = k_decay1.to(input_dtype)
+
+ k_inv_vec = cutlass.Vector.from_elements(
+ (
+ k_inv_regs[reg_base],
+ k_inv_regs[reg_base + 1],
+ k_inv_regs[reg_base + 2],
+ k_inv_regs[reg_base + 3],
+ k_inv_regs[reg_base + 4],
+ k_inv_regs[reg_base + 5],
+ k_inv_regs[reg_base + 6],
+ k_inv_regs[reg_base + 7],
+ ),
+ input_dtype,
+ )
+ k_decay_vec = cutlass.Vector.from_elements(
+ (
+ k_decay_vec_regs[0],
+ k_decay_vec_regs[1],
+ k_decay_vec_regs[2],
+ k_decay_vec_regs[3],
+ k_decay_vec_regs[4],
+ k_decay_vec_regs[5],
+ k_decay_vec_regs[6],
+ k_decay_vec_regs[7],
+ ),
+ input_dtype,
+ )
+ k_inv_swizzled_idx = k_inv_s128_smem_index(decay_row, dim_base)
+ (k_inv_ptr + k_inv_swizzled_idx).store(
+ k_inv_vec,
+ alignment=RAW_F16_TMA_SWIZZLE_GROUP_BYTES,
+ )
+ decay_storage_dim_base = tcgen05_decay_b_key_storage_dim_runtime(
+ decay_row,
+ dim_base,
+ )
+ decay_linear_idx_base = decay_row * DK + decay_storage_dim_base
+ decay_swizzled_idx_base = tcgen05_swizzle_128b_elem_index(
+ decay_linear_idx_base,
+ TCGEN05_F16_ELEM_BYTES,
+ BT,
+ )
+ (tcgen05_k_decay_ptr + decay_swizzled_idx_base).store(
+ k_decay_vec,
+ alignment=RAW_F16_TMA_SWIZZLE_GROUP_BYTES,
+ )
+ if cutlass.const_expr(dim_half == 0):
+ # Publish the first half-DK of k_inv/k_decay early: warp 0's
+ # KK K-blocks 0..3 only read key dims [0, 64).
+ cg0_k_ready_arrive(cg0_k_half_ready_mbar)
+ cg0_k_ready_arrive(cg0_k_ready_mbar)
+
+ for dim_half in cutlass.range_constexpr(2):
+ dim_base = (
+ dim_half * (DK // 2) + lane_in_row_group * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ )
+ reg_base = dim_half * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ q_decay_vec_regs = cutlass.Array(
+ input_dtype,
+ RAW_F16_TMA_SWIZZLE_GROUP_ELEMS,
+ alignment=RAW_F16_TMA_SWIZZLE_GROUP_BYTES,
+ )
+ for pair_idx in cutlass.range_constexpr(RAW_F16_TMA_SWIZZLE_GROUP_ELEMS // 2):
+ dim0 = pair_idx * 2
+ dim1 = dim0 + 1
+ raw_reg_idx0 = reg_base + dim0
+ raw_reg_idx1 = reg_base + dim1
+ q_value0, q_value1 = fmul2(
+ (raw_q_regs[raw_reg_idx0], raw_q_regs[raw_reg_idx1]),
+ (q_inv_norm, q_inv_norm),
+ )
+ q_decay0, q_decay1 = fmul2(
+ (q_value0, q_value1),
+ (exp_g_regs[raw_reg_idx0], exp_g_regs[raw_reg_idx1]),
+ )
+ q_decay_vec_regs[dim0] = q_decay0.to(input_dtype)
+ q_decay_vec_regs[dim1] = q_decay1.to(input_dtype)
+
+ q_decay_vec = cutlass.Vector.from_elements(
+ (
+ q_decay_vec_regs[0],
+ q_decay_vec_regs[1],
+ q_decay_vec_regs[2],
+ q_decay_vec_regs[3],
+ q_decay_vec_regs[4],
+ q_decay_vec_regs[5],
+ q_decay_vec_regs[6],
+ q_decay_vec_regs[7],
+ ),
+ input_dtype,
+ )
+ decay_storage_dim_base = tcgen05_decay_b_key_storage_dim_runtime(
+ decay_row,
+ dim_base,
+ )
+ decay_linear_idx_base = decay_row * DK + decay_storage_dim_base
+ decay_swizzled_idx_base = tcgen05_swizzle_128b_elem_index(
+ decay_linear_idx_base,
+ TCGEN05_F16_ELEM_BYTES,
+ BT,
+ )
+ (tcgen05_q_decay_ptr + decay_swizzled_idx_base).store(
+ q_decay_vec,
+ alignment=RAW_F16_TMA_SWIZZLE_GROUP_BYTES,
+ )
+
+ for dim_half in cutlass.range_constexpr(2):
+ dim_base = (
+ dim_half * (DK // 2) + lane_in_row_group * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ )
+ reg_base = dim_half * RAW_F16_TMA_SWIZZLE_GROUP_ELEMS
+ storage_row = decay_row ^ TCGEN05_SW32_BT_HALF_XOR
+ k_restore_idx = raw_f16_s128_smem_index(storage_row, dim_base)
+ k_restore_vec = cutlass.Vector.from_elements(
+ (
+ k_restore_all_regs[reg_base],
+ k_restore_all_regs[reg_base + 1],
+ k_restore_all_regs[reg_base + 2],
+ k_restore_all_regs[reg_base + 3],
+ k_restore_all_regs[reg_base + 4],
+ k_restore_all_regs[reg_base + 5],
+ k_restore_all_regs[reg_base + 6],
+ k_restore_all_regs[reg_base + 7],
+ ),
+ input_dtype,
+ )
+ (tcgen05_k_restore_ptr + k_restore_idx).store(
+ k_restore_vec,
+ alignment=RAW_F16_TMA_SWIZZLE_GROUP_BYTES,
+ )
+
+
+@cute.jit
+def super_mma_stage_w_half(
+ tcgen05_k_restore_smem,
+ pairwise_stage_smem,
+ raw_beta_smem,
+ w_out_smem,
+ lane,
+ input_dtype: cutlass.Constexpr,
+ n_group_lo: cutlass.Constexpr,
+ n_group_hi: cutlass.Constexpr,
+) -> None:
+ """W^T = Ainvb^T @ k_restore for key-column groups [n_lo, n_hi).
+
+ Column-split variant of super_mma_stage_w: groups [0, 4) cover key
+ dims [0, 64) = W TMA segment 0, groups [4, 8) segment 1, so two warps
+ stage and bulk-store their halves fully independently. The Ainvb^T
+ fragment build is duplicated per warp (cheap: one ldmatrix + 4
+ movmatrix + the beta packs).
+ """
+
+ ainv = super_mma_load_ainv_a_frags(pairwise_stage_smem, lane)
+ ab_frag = cutlass.Array(
+ cutlass.Int32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ super_mma_beta_col_scale_a_frags(
+ ainv,
+ raw_beta_smem,
+ lane,
+ ab_frag,
+ input_dtype,
+ )
+ # A-layout fragment of Ainvb^T: transpose quadrants, swap off-diagonals.
+ t0 = movmatrix_b16(ab_frag[0])
+ t1 = movmatrix_b16(ab_frag[2])
+ t2 = movmatrix_b16(ab_frag[1])
+ t3 = movmatrix_b16(ab_frag[3])
+
+ # Literal FP32 operands into the inline-PTX MMA are rejected downstream; route the
+ # zero accumulators through an array exactly like the donor code does.
+ zero_acc = cutlass.Array(
+ cutlass.Float32,
+ SUPER_MMA_ACCUMULATORS_PER_LANE,
+ alignment=16,
+ )
+ for accum_idx in cutlass.range_constexpr(SUPER_MMA_ACCUMULATORS_PER_LANE):
+ zero_acc[accum_idx] = cutlass.Float32(0.0)
+
+ for n_group_off in cutlass.range_constexpr(n_group_hi - n_group_lo):
+ n_group = n_group_lo + n_group_off
+ kr = super_mma_load_k_restore_a_fragment(
+ tcgen05_k_restore_smem,
+ lane,
+ n_group,
+ )
+ b0 = movmatrix_b16(kr[0])
+ b1 = movmatrix_b16(kr[1])
+ b2 = movmatrix_b16(kr[2])
+ b3 = movmatrix_b16(kr[3])
+ n0_d0, n0_d1, n0_d2, n0_d3 = ptx_mma_m16n8k16_b16_f32(
+ t0,
+ t1,
+ t2,
+ t3,
+ b0,
+ b1,
+ zero_acc[0],
+ zero_acc[1],
+ zero_acc[2],
+ zero_acc[3],
+ input_dtype,
+ )
+ n1_d0, n1_d1, n1_d2, n1_d3 = ptx_mma_m16n8k16_b16_f32(
+ t0,
+ t1,
+ t2,
+ t3,
+ b2,
+ b3,
+ zero_acc[0],
+ zero_acc[1],
+ zero_acc[2],
+ zero_acc[3],
+ input_dtype,
+ )
+ prims.stmatrix(
+ w_out_stmatrix_ptr(w_out_smem, n_group, lane),
+ [
+ pack_input_b16x2_to_i32(n0_d0, n0_d1, input_dtype),
+ pack_input_b16x2_to_i32(n0_d2, n0_d3, input_dtype),
+ pack_input_b16x2_to_i32(n1_d0, n1_d1, input_dtype),
+ pack_input_b16x2_to_i32(n1_d2, n1_d3, input_dtype),
+ ],
+ prims.MMALayout.ROW,
+ shape=prims.StoreShape.M8N8,
+ )
+
+
+@cute.kernel
+def kernel_prep(
+ tma_desc_q: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_k: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_gate: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_ws_kd: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_ws_qd: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_ws_w: cutlass.GridConstant[cuda.TensorMap],
+ q: cute.Tensor,
+ k: cute.Tensor,
+ raw_gate: cute.Tensor,
+ a_log: cute.Tensor,
+ dt_bias: cute.Tensor,
+ beta: cute.Tensor,
+ cu_seqlens: cute.Tensor,
+ cu_chunks: cute.Tensor,
+ ws_kd: cute.Tensor,
+ ws_qd: cute.Tensor,
+ ws_w: cute.Tensor,
+ ws_qk: cute.Tensor,
+ ws_diag: cute.Tensor,
+ chunks_per_cta: cutlass.Int32,
+ head_base: cutlass.Int32,
+ SAFE_GATE: cutlass.Constexpr,
+ GATE_SCALE_LOG2: cutlass.Constexpr,
+ gate_dtype: cutlass.Constexpr,
+) -> None:
+ """kernel 1: per-chunk-tile factor prep (4-warp CTA, no rings).
+
+ Grid `(ceil(total_chunks / chunks_per_cta), heads, 1)`; each CTA walks a
+ CONTIGUOUS range of chunks_per_cta chunks (K1_CPC), which keeps its TMA
+ loads and workspace stores cache-local.
+ """
+
+ tidx, _, _ = cute.arch.thread_idx()
+ bidx, bidy, _ = cute.arch.block_idx()
+ bidy = bidy + head_base
+ warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+ lane = tidx % THREADS_PER_WARP
+
+ # The workspace extent is a graph-static capacity. The device-generated
+ # prefix carries the actual count for the current replay.
+ total_chunks = cutlass.Int32(cu_chunks[cu_chunks.shape[0] - 1])
+ gchunk_stride = cutlass.Int32(1)
+ chunk_lo = bidx * chunks_per_cta
+ my_chunks = total_chunks - chunk_lo
+ if my_chunks > chunks_per_cta:
+ my_chunks = chunks_per_cta
+ if my_chunks < 0:
+ my_chunks = cutlass.Int32(0)
+ input_dtype = q.element_type
+
+ # 2 slots: chunk c completes tma_mbar[c % 2] (parity (c//2) % 2), so the
+ # CPC>1 prefetch of chunk c+1 (issued mid-chunk by warp 2, see below)
+ # never aliases the wait slot of the chunk in flight.
+ tma_mbar = cutlass.Array(
+ cutlass.Int64,
+ 2,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ cg0_k_ready_mbar = cutlass.Array(
+ cutlass.Int64,
+ 1,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ cg0_k_half_ready_mbar = cutlass.Array(
+ cutlass.Int64,
+ 1,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ q_k_restore_ready_mbar = cutlass.Array(
+ cutlass.Int64,
+ 1,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ pairwise_ready_mbar = cutlass.Array(
+ cutlass.Int64,
+ 1,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+
+ tcgen05_k_decay_smem = cutlass.Array(
+ input_dtype,
+ TCGEN05_K_DECAY_STAGE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=1024,
+ )
+ tcgen05_q_decay_smem = cutlass.Array(
+ input_dtype,
+ TCGEN05_Q_DECAY_STAGE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=1024,
+ )
+ tcgen05_k_restore_smem = cutlass.Array(
+ input_dtype,
+ TCGEN05_K_RESTORE_STAGE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=RAW_F16_TMA_SWIZZLE_ALIGNMENT_BYTES,
+ )
+ pairwise_smem = cutlass.Array(
+ input_dtype,
+ PAIRWISE_SMEM_STAGE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=1024,
+ )
+ raw_q_smem = cutlass.Array(
+ q.element_type,
+ RAW_Q_STAGE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=RAW_F16_TMA_SWIZZLE_ALIGNMENT_BYTES,
+ )
+ raw_k_smem = cutlass.Array(
+ k.element_type,
+ RAW_K_STAGE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=RAW_F16_TMA_SWIZZLE_ALIGNMENT_BYTES,
+ )
+ # Same policy as the engine, ring-free: k1's exchange is produced AND
+ # consumed inside cg0_materialize_prep_operands by the same 4-warp CTA,
+ # and the CPC>1 prefetch (which only overwrites raw_gate_smem) is issued
+ # after q_k_restore_ready, so one stage suffices on the 16-bit path.
+ if cutlass.const_expr(gate_dtype_is_f32(gate_dtype)):
+ raw_gate_smem = cutlass.Array(
+ cutlass.Float32,
+ RAW_GATE_STAGE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=RAW_F32_TMA_SWIZZLE_ALIGNMENT_BYTES,
+ )
+ gate_exchange_smem = raw_gate_smem
+ else:
+ raw_gate_smem = cutlass.Array(
+ gate_dtype,
+ RAW_GATE_STAGE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=RAW_F16_TMA_SWIZZLE_ALIGNMENT_BYTES,
+ )
+ gate_exchange_smem = cutlass.Array(
+ cutlass.Float32,
+ GATE_EXCHANGE_STAGE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=RAW_F32_TMA_SWIZZLE_ALIGNMENT_BYTES,
+ )
+ k_inv_smem = cutlass.Array(
+ input_dtype,
+ K_INV_STAGE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=RAW_F16_TMA_SWIZZLE_ALIGNMENT_BYTES,
+ )
+ # Beta is double-buffered (2 x 64 B): the CPC>1 prefetch writes chunk
+ # c+1's beta while chunk c's consumers still read theirs.
+ raw_beta_smem = cutlass.Array(
+ cutlass.Float32,
+ 2 * RAW_BETA_STAGE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=128,
+ )
+ # W image staging (warps 1/3 stmatrix -> per-segment TMA bulk store).
+ w_out_smem = cutlass.Array(
+ input_dtype,
+ TILE_ELEMS,
+ space=cutlass.AddressSpace.smem,
+ alignment=RAW_F16_TMA_SWIZZLE_ALIGNMENT_BYTES,
+ )
+ tma_tx_bytes = cutlass.const_expr(
+ DK * BT * q.element_type.width // 8
+ + DK * BT * k.element_type.width // 8
+ + DK * BT * raw_gate.element_type.width // 8
+ )
+
+ if warp_idx == 0:
+ if prims.elect_sync():
+ prims.mbarrier_init(tma_mbar.subview(0), 1)
+ prims.mbarrier_init(tma_mbar.subview(1), 1)
+ prims.mbarrier_init(cg0_k_ready_mbar.subview(0), K1_WARPS)
+ prims.mbarrier_init(cg0_k_half_ready_mbar.subview(0), K1_WARPS)
+ prims.mbarrier_init(q_k_restore_ready_mbar.subview(0), K1_WARPS)
+ prims.mbarrier_init(pairwise_ready_mbar.subview(0), 1)
+ prims.fence_mbarrier_init()
+ prep_cta_sync()
+
+ # Per-CTA gate constants: direct global loads (one element per
+ # (warp, lane) prefix column; no SMEM staging hop needed at this CTA
+ # size).
+ # Gate constants for BOTH the safe sigmoid and the FLA non-safe softplus
+ # activation. Loaded unconditionally; SAFE_GATE=True traces identically
+ # (init-then-overwrite, byte-identical generated code) and the non-safe path reuses
+ # the same constants.
+ a_log_exp = cutlass.Float32(1.0)
+ dt_bias_value = cutlass.Float32(0.0)
+ prefix_dim = warp_idx * THREADS_PER_WARP + lane
+ a_log_exp = cute.math.exp2(
+ a_log[bidy].to(cutlass.Float32) * LOG2_E,
+ fastmath=True,
+ )
+ dt_bias_value = dt_bias[bidy, prefix_dim].to(cutlass.Float32)
+
+ for local_chunk in cutlass.range(my_chunks, unroll=1):
+ gchunk = chunk_lo + local_chunk * gchunk_stride
+ sequence_start, seqlen, chunk_start_tok = chunk_coords(
+ cu_seqlens,
+ cu_chunks,
+ gchunk,
+ )
+ phase = local_chunk % 2
+ beta_stage = local_chunk % 2
+ tma_slot = local_chunk % 2
+ tma_phase = (local_chunk // 2) % 2
+ ws_row_start = gchunk * cutlass.Int32(BT)
+
+ if warp_idx == 0:
+ if local_chunk == 0:
+ # First chunk load; chunks c+1 are PREFETCHED mid-chunk by
+ # warp 2 (below), overlapping the DRAM round trip with the
+ # A_inv/W/QK' phase of the chunk in flight.
+ tma_prep_stage_load_inputs(
+ tma_desc_q,
+ tma_desc_k,
+ tma_desc_gate,
+ beta,
+ raw_q_smem,
+ raw_k_smem,
+ raw_gate_smem,
+ raw_beta_smem.subview(beta_stage * RAW_BETA_STAGE_SIZE),
+ sequence_start,
+ bidy,
+ lane,
+ chunk_start_tok,
+ seqlen,
+ tma_mbar.subview(tma_slot),
+ tma_tx_bytes,
+ gate_dtype,
+ )
+ tma_transfer_wait(tma_mbar.subview(tma_slot), tma_phase)
+
+ if chunk_start_tok + cutlass.Int32(BT) > seqlen:
+ if warp_idx == 0:
+ # No raw V in kernel 1: pass raw_q twice (idempotent
+ # double zero of the q tail rows).
+ cg0_zero_tail_raw_operands(
+ raw_q_smem,
+ raw_k_smem,
+ raw_q_smem,
+ raw_gate_smem,
+ lane,
+ chunk_start_tok,
+ seqlen,
+ q.element_type,
+ gate_dtype,
+ )
+ prep_cta_sync()
+
+ cg0_materialize_prep_operands(
+ raw_q_smem,
+ raw_k_smem,
+ raw_gate_smem,
+ gate_exchange_smem,
+ a_log_exp,
+ dt_bias_value,
+ k_inv_smem,
+ tcgen05_k_decay_smem,
+ tcgen05_q_decay_smem,
+ tcgen05_k_restore_smem,
+ cg0_k_ready_mbar.subview(0),
+ cg0_k_half_ready_mbar.subview(0),
+ ws_diag,
+ bidy,
+ gchunk,
+ chunk_start_tok,
+ seqlen,
+ input_dtype,
+ gate_dtype,
+ SAFE_GATE,
+ GATE_SCALE_LOG2,
+ warp_idx,
+ lane,
+ )
+ q_k_restore_ready_arrive(q_k_restore_ready_mbar.subview(0))
+
+ if warp_idx == 0:
+ # KK -> L -> blockwise A_inv (starts on the half-DK arrival).
+ cg0_k_ready_wait(cg0_k_half_ready_mbar.subview(0), phase)
+ super_mma_stage_pairwise_pipeline(
+ tcgen05_k_decay_smem,
+ k_inv_smem,
+ pairwise_smem,
+ raw_beta_smem.subview(beta_stage * RAW_BETA_STAGE_SIZE),
+ cg0_k_ready_mbar.subview(0),
+ phase,
+ lane,
+ input_dtype,
+ )
+ pairwise_ready_arrive(pairwise_ready_mbar.subview(0))
+ elif warp_idx == 1:
+ q_k_restore_ready_wait(q_k_restore_ready_mbar.subview(0), phase)
+ if cutlass.const_expr(K1_PREFETCH_WARP == 1):
+ prefetch_next_chunk(
+ tma_desc_q,
+ tma_desc_k,
+ tma_desc_gate,
+ beta,
+ raw_q_smem,
+ raw_k_smem,
+ raw_gate_smem,
+ raw_beta_smem,
+ cu_seqlens,
+ cu_chunks,
+ local_chunk,
+ my_chunks,
+ gchunk,
+ gchunk_stride,
+ bidy,
+ lane,
+ tma_mbar,
+ tma_tx_bytes,
+ gate_dtype,
+ )
+ pairwise_ready_wait(pairwise_ready_mbar.subview(0), phase)
+ super_mma_stage_w_half(
+ tcgen05_k_restore_smem,
+ pairwise_smem,
+ raw_beta_smem.subview(beta_stage * RAW_BETA_STAGE_SIZE),
+ w_out_smem,
+ lane,
+ input_dtype,
+ 0,
+ K1_W_N_SPLIT,
+ )
+ prims.fence_proxy(
+ prims.Proxy.ASYNC_SHARED,
+ space=prims.SharedSpace.shared_cta,
+ )
+ prims.bar_warp_sync(cute.arch.FULL_MASK)
+ if prims.elect_sync():
+ w_coord = (
+ cutlass.Int32(0),
+ ws_row_start,
+ bidy,
+ cutlass.Int32(0),
+ )
+ prims.cp_async_bulk_tensor_global_shared_cta(
+ tma_desc_ws_w.get_ptr(),
+ w_out_smem.subview(0),
+ w_coord,
+ )
+ prims.cp_async_bulk_commit_group()
+ elif warp_idx == 2:
+ q_k_restore_ready_wait(q_k_restore_ready_mbar.subview(0), phase)
+ if cutlass.const_expr(K1_PREFETCH_WARP == 2):
+ prefetch_next_chunk(
+ tma_desc_q,
+ tma_desc_k,
+ tma_desc_gate,
+ beta,
+ raw_q_smem,
+ raw_k_smem,
+ raw_gate_smem,
+ raw_beta_smem,
+ cu_seqlens,
+ cu_chunks,
+ local_chunk,
+ my_chunks,
+ gchunk,
+ gchunk_stride,
+ bidy,
+ lane,
+ tma_mbar,
+ tma_tx_bytes,
+ gate_dtype,
+ )
+ pairwise_ready_wait(pairwise_ready_mbar.subview(0), phase)
+ super_mma_stage_qk_prime(
+ tcgen05_q_decay_smem,
+ k_inv_smem,
+ pairwise_smem,
+ raw_beta_smem.subview(beta_stage * RAW_BETA_STAGE_SIZE),
+ ws_qk,
+ bidy,
+ gchunk,
+ lane,
+ input_dtype,
+ )
+ else:
+ # kd/qd workspace TMA stores as soon as both tiles are staged.
+ cg0_k_ready_wait(cg0_k_ready_mbar.subview(0), phase)
+ q_k_restore_ready_wait(q_k_restore_ready_mbar.subview(0), phase)
+ if prims.elect_sync():
+ for segment in cutlass.range_constexpr(RAW_F16_TMA_SEGMENTS):
+ ws_coord = (
+ cutlass.Int32(segment * RAW_F16_TMA_SWIZZLE_ELEMS),
+ ws_row_start,
+ bidy,
+ cutlass.Int32(0),
+ )
+ prims.cp_async_bulk_tensor_global_shared_cta(
+ tma_desc_ws_kd.get_ptr(),
+ tcgen05_k_decay_smem.subview(
+ segment * RAW_F16_TMA_SEGMENT_ELEMS
+ ),
+ ws_coord,
+ )
+ prims.cp_async_bulk_tensor_global_shared_cta(
+ tma_desc_ws_qd.get_ptr(),
+ tcgen05_q_decay_smem.subview(
+ segment * RAW_F16_TMA_SEGMENT_ELEMS
+ ),
+ ws_coord,
+ )
+ prims.cp_async_bulk_commit_group()
+ pairwise_ready_wait(pairwise_ready_mbar.subview(0), phase)
+ super_mma_stage_w_half(
+ tcgen05_k_restore_smem,
+ pairwise_smem,
+ raw_beta_smem.subview(beta_stage * RAW_BETA_STAGE_SIZE),
+ w_out_smem,
+ lane,
+ input_dtype,
+ K1_W_N_SPLIT,
+ K1_W_N_GROUPS,
+ )
+ prims.fence_proxy(
+ prims.Proxy.ASYNC_SHARED,
+ space=prims.SharedSpace.shared_cta,
+ )
+ prims.bar_warp_sync(cute.arch.FULL_MASK)
+ if prims.elect_sync():
+ w_coord = (
+ cutlass.Int32(K1_W_N_SPLIT * 16),
+ ws_row_start,
+ bidy,
+ cutlass.Int32(0),
+ )
+ prims.cp_async_bulk_tensor_global_shared_cta(
+ tma_desc_ws_w.get_ptr(),
+ w_out_smem.subview(RAW_F16_TMA_SEGMENT_ELEMS),
+ w_coord,
+ )
+ prims.cp_async_bulk_commit_group()
+
+ if (warp_idx == 1) | (warp_idx == 3):
+ if prims.elect_sync():
+ prims.cp_async_bulk_wait_group(0, read=True)
+ prep_cta_sync()
+
+
+# ---------------------------------------------------------------------------
+# kernel-2 device helpers (value-correct forms taken from the engine path).
+# ---------------------------------------------------------------------------
+
+
+@cute.jit
+def tma_chain_load_tile(
+ tma_desc: cutlass.GridConstant[cuda.TensorMap],
+ tile_smem,
+ row_start,
+ head_idx,
+ tma_mbar,
+) -> None:
+ """Issue the two s128 segments of one [16, DK] bf16 workspace tile."""
+
+ for segment in cutlass.range_constexpr(RAW_F16_TMA_SEGMENTS):
+ tma_coord = (
+ cutlass.Int32(segment * RAW_F16_TMA_SWIZZLE_ELEMS),
+ row_start,
+ head_idx,
+ cutlass.Int32(0),
+ )
+ prims.cp_async_bulk_tensor_shared_cta_global(
+ tile_smem.subview(segment * RAW_F16_TMA_SEGMENT_ELEMS),
+ tma_desc.get_ptr(),
+ tma_coord,
+ tma_mbar,
+ )
+
+
+@cute.jit
+def tcgen05_commit(mbar) -> None:
+ """Commit the tensor pipe's progress to an mbarrier (single elected lane)."""
+
+ if prims.elect_sync():
+ prims.tcgen05_commit(mbar, group=prims.CTAGroup.CTA_1)
+
+
+# --- chain state-input helpers --------------------------------------------
+
+
+# =============================================================================
+# kernel-2 DV split (M = 64, PTX Layout F).
+#
+# Everything column-indexed (TMEM layout, K blocks, B-operand SMEM
+# descriptors) is IDENTICAL to kernel 2; only M changes. Layout F places
+# an M=64 tile's rows 16q..16q+15 at lanes 32q..32q+15 (q = warp quadrant,
+# lane alignment 0), so:
+# - the state pack stays the shared `tcgen05_stage_state_input_tmem`
+# verbatim: it is lane-local (each lane converts its own row), valid
+# lanes map to valid lanes and the alignment-16 junk lanes feed junk
+# lanes no M=64 MMA reads;
+# - per-warp ROW-OWNING helpers (V ldmatrix stage, -X repack, qstate
+# drain, state ABI) drop from 32 to 16 rows per warp (one 16-lane
+# tcgen05 access instead of two);
+# - all mbar arrival counts are unchanged (all four warps of every
+# group still participate).
+# =============================================================================
+
+
+@cute.jit
+def tcgen05_chain_issue_delta_half_mma(
+ b_smem,
+ tmem_raw_addr,
+ a_input_col,
+ input_dtype: cutlass.Constexpr,
+ HALF: cutlass.Constexpr,
+) -> None:
+ """M=64 delta-half MMA issue."""
+
+ half_n = DK // 2
+ tmem_ptr = cutlass.inttoptr(
+ tmem_raw_addr + KDA_TMEM_STATE_COL_OFFSET + HALF * half_n,
+ 6,
+ cutlass.Float32,
+ )
+ desc_b = prims.Tcgen05SmemDesc.build(
+ b_smem.subview(0),
+ leading_byte_offset=TCGEN05_FINAL_STATE_B_LEADING_BYTES,
+ stride_byte_offset=TCGEN05_FINAL_STATE_B_STRIDE_BYTES,
+ layout=prims.Tcgen05SmemSwizzle.SWIZZLE_128B,
+ )
+ idesc = prims.Tcgen05InstrDesc.build(
+ c_dtype=cutlass.Float32,
+ a_dtype=input_dtype,
+ b_dtype=input_dtype,
+ n_dim=half_n,
+ m_dim=DV_HALF,
+ b_major=1,
+ )
+ half_byte_offset = HALF * TCGEN05_FINAL_STATE_B_LEADING_BYTES
+ a_tmem = prims.make_tmem_ptr(tmem_raw_addr, cutlass.Int8).subview(a_input_col)
+ if prims.elect_sync():
+ prims.tcgen05_mma(
+ prims.Tcgen05MMAKind.F16,
+ prims.CTAGroup.CTA_1,
+ tmem_ptr,
+ a_tmem,
+ desc_b.advance_start_address(half_byte_offset),
+ idesc,
+ True,
+ )
+
+
+@cute.jit
+def tcgen05_chain_issue_qkv_mma(
+ qk_stage_smem,
+ tmem_raw_addr,
+ a_input_col,
+ qstate_acc_stage,
+ acc_ready_mbar,
+ input_dtype: cutlass.Constexpr,
+ COMMIT: cutlass.Constexpr,
+) -> None:
+ """M=64 qkv MMA issue."""
+
+ tmem_ptr = cutlass.inttoptr(
+ tmem_raw_addr + tcgen05_qstate_acc_tmem_col_offset(qstate_acc_stage),
+ 6,
+ cutlass.Float32,
+ )
+ idesc = prims.Tcgen05InstrDesc.build(
+ c_dtype=cutlass.Float32,
+ a_dtype=input_dtype,
+ b_dtype=input_dtype,
+ n_dim=BT,
+ m_dim=DV_HALF,
+ b_major=0,
+ )
+ lhs_tmem = prims.make_tmem_ptr(tmem_raw_addr, cutlass.Int8).subview(a_input_col)
+ desc_pairwise = prims.Tcgen05SmemDesc.build(
+ qk_stage_smem.subview(0),
+ leading_byte_offset=TCGEN05_VALUE_PAIRWISE_B_LEADING_BYTES,
+ stride_byte_offset=TCGEN05_VALUE_PAIRWISE_B_STRIDE_BYTES,
+ layout=prims.Tcgen05SmemSwizzle.SWIZZLE_32B,
+ )
+ if prims.elect_sync():
+ prims.tcgen05_mma(
+ prims.Tcgen05MMAKind.F16,
+ prims.CTAGroup.CTA_1,
+ tmem_ptr,
+ lhs_tmem,
+ desc_pairwise,
+ idesc,
+ True,
+ )
+ if cutlass.const_expr(COMMIT):
+ prims.tcgen05_commit(acc_ready_mbar, group=prims.CTAGroup.CTA_1)
+
+
+@cute.jit
+def tcgen05_chain_stage_vmx_input_tmem(
+ tmem_raw_addr,
+ raw_v_smem,
+ warp_idx,
+ lane,
+ shared_acc_stage,
+ shared_input_stage,
+ input_dtype: cutlass.Constexpr,
+) -> None:
+ """M=64 (v - X) fused repack staging.
+
+ ONE v ldmatrix ([16,64] half tile) + ONE 16x256b X
+ fragment read + sub.b16x2 + ONE 16x128b store.
+ """
+
+ base_col_id = tmem_raw_addr & 0xFFFF
+ base_row_id = tmem_raw_addr >> 16
+ tmem_sp = warp_idx % TCGEN05_STATE_K_TMEM_ROW_BLOCKS
+
+ projection_col_id = base_col_id + tcgen05_shared_acc_tmem_col_offset(
+ shared_acc_stage
+ )
+ input_col_id = base_col_id + tcgen05_shared_input_tmem_col_offset(
+ shared_input_stage
+ )
+ value_dim_base = tmem_sp * ROWS_PER_WARP
+
+ row_id0 = base_row_id + tmem_sp * THREADS_PER_WARP
+ block_ptr0 = cutlass.inttoptr(
+ (row_id0 << 16) | projection_col_id, 6, cutlass.Float32
+ )
+ state_k0 = prims.tcgen05_ld("16x256b", block_ptr0, num=2)
+
+ raw_v_regs0 = prims.ldmatrix(
+ raw_v_ldmatrix_trans_ptr(raw_v_smem, value_dim_base, lane),
+ 4,
+ prims.MMALayout.COL,
+ )
+ prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD)
+
+ packed0 = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.rmem)
+ for reg_idx in cutlass.range_constexpr(4):
+ raw_matrix: cutlass.Constexpr[int] = (1 - (reg_idx // 2)) * 2 + (reg_idx & 1)
+ val0, val1 = tcgen05_rhs_token_pair_from_16x256b_fragment(
+ state_k0,
+ reg_idx,
+ )
+ packed0[reg_idx] = sub_b16x2_input_dtype(
+ raw_v_regs0[raw_matrix],
+ pack_input_b16x2_to_i32(val0, val1, input_dtype),
+ input_dtype,
+ )
+
+ input_block_addr0 = (base_row_id << 16) | input_col_id
+ input_block_ptr0 = prims.make_tmem_ptr(input_block_addr0, cutlass.Int8)
+ prims.tcgen05_st("16x128b", input_block_ptr0, packed0[0:4])
+ prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE)
+ prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC)
+
+
+@cute.jit
+def tcgen05_chain_load_qstate_output_tmem(
+ tmem_raw_addr,
+ o_smem,
+ warp_idx,
+ lane,
+ o_stage_base,
+ qstate_acc_stage,
+ scale: cutlass.Float32,
+ output_dtype: cutlass.Constexpr,
+) -> None:
+ """M=64 variant of tcgen05_load_qstate_output_tmem (16 rows per warp)."""
+
+ base_col_id = tmem_raw_addr & 0xFFFF
+ base_row_id = tmem_raw_addr >> 16
+ tmem_sp = warp_idx % TCGEN05_STATE_K_TMEM_ROW_BLOCKS
+
+ projection_col_id = base_col_id + tcgen05_qstate_acc_tmem_col_offset(
+ qstate_acc_stage
+ )
+ value_dim_base = tmem_sp * ROWS_PER_WARP
+
+ row_id0 = base_row_id + tmem_sp * THREADS_PER_WARP
+ block_addr0 = (row_id0 << 16) | projection_col_id
+ block_ptr0 = cutlass.inttoptr(block_addr0, 6, cutlass.Float32)
+ loaded0 = prims.tcgen05_ld(
+ "16x256b",
+ block_ptr0,
+ num=2,
+ )
+ prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD)
+
+ stsm_regs0 = cutlass.Array(
+ cutlass.Int32,
+ 4,
+ space=cutlass.AddressSpace.rmem,
+ )
+ for reg_idx in cutlass.range_constexpr(4):
+ scaled0_0, scaled0_1 = fmul2(
+ (loaded0[2 * reg_idx], loaded0[2 * reg_idx + 1]),
+ (scale, scale),
+ )
+ stsm_regs0[reg_idx] = pack_output_b16x2_to_i32(
+ scaled0_0,
+ scaled0_1,
+ output_dtype,
+ )
+
+ smem_dst0 = o_smem_stmatrix_128b_ptr(
+ o_smem,
+ o_stage_base,
+ value_dim_base,
+ lane,
+ )
+ prims.stmatrix(
+ smem_dst0,
+ stsm_regs0.data_ptr().load(count=4, alignment=4),
+ prims.MMALayout.COL,
+ shape=prims.StoreShape.M8N8,
+ )
+ cute.arch.fence_view_async_shared()
+
+
+@cute.jit
+def tma_chain_stage_load_inputs(
+ tma_desc_kd: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_w: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_qd: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_v: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_diag: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_qk: cutlass.GridConstant[cuda.TensorMap],
+ kd_stage,
+ w_stage,
+ qd_stage,
+ v_stage,
+ diag_stage,
+ qk_stage,
+ head_idx,
+ dv_half,
+ ws_row_start,
+ ws_chunk,
+ v_row_start,
+ tma_mbar,
+ tma_tx_bytes: cutlass.Constexpr,
+) -> None:
+ """DV2 chunk operand transaction: full kd/w/qd, HALF v (one segment)."""
+
+ if prims.elect_sync():
+ prims.mbarrier_arrive_expect_tx(tma_mbar, tma_tx_bytes)
+ if prims.elect_sync():
+ tma_chain_load_tile(tma_desc_kd, kd_stage, ws_row_start, head_idx, tma_mbar)
+ tma_chain_load_tile(tma_desc_w, w_stage, ws_row_start, head_idx, tma_mbar)
+ tma_chain_load_tile(tma_desc_qd, qd_stage, ws_row_start, head_idx, tma_mbar)
+ v_coord = (
+ dv_half * cutlass.Int32(DV_HALF),
+ v_row_start,
+ head_idx,
+ cutlass.Int32(0),
+ )
+ prims.cp_async_bulk_tensor_shared_cta_global(
+ v_stage.subview(0),
+ tma_desc_v.get_ptr(),
+ v_coord,
+ tma_mbar,
+ )
+ diag_coord = (
+ cutlass.Int32(0),
+ ws_chunk,
+ head_idx,
+ cutlass.Int32(0),
+ )
+ prims.cp_async_bulk_tensor_shared_cta_global(
+ diag_stage.subview(0),
+ tma_desc_diag.get_ptr(),
+ diag_coord,
+ tma_mbar,
+ )
+ prims.cp_async_bulk_tensor_shared_cta_global(
+ qk_stage.subview(0),
+ tma_desc_qk.get_ptr(),
+ diag_coord,
+ tma_mbar,
+ )
+
+
+@cute.jit
+def epilogue_chain_stage_store(
+ tma_desc_o: cutlass.GridConstant[cuda.TensorMap],
+ o_smem,
+ sequence_start,
+ head_idx,
+ dv_half,
+ chunk_start,
+ o_stage_base,
+) -> None:
+ """Store the staged `[BT, 64]` half-output tile (one s128 segment)."""
+
+ global_chunk_start = sequence_start + chunk_start
+ if prims.elect_sync():
+ o_coord = (
+ dv_half * cutlass.Int32(DV_HALF),
+ global_chunk_start,
+ head_idx,
+ cutlass.Int32(0),
+ )
+ prims.cp_async_bulk_tensor_global_shared_cta(
+ tma_desc_o.get_ptr(),
+ o_smem.subview(o_stage_base + O_OUT_OFFSET),
+ o_coord,
+ )
+ prims.cp_async_bulk_commit_group()
+ prims.cp_async_bulk_wait_group(0, read=True)
+ prims.bar_warp_sync(cute.arch.FULL_MASK)
+
+
+@cute.jit
+def epilogue_chain_tail_store(
+ out,
+ o_smem,
+ sequence_start,
+ head_idx,
+ dv_half,
+ chunk_start,
+ seqlen,
+ o_stage_base,
+ lane,
+) -> None:
+ """Store a partial packed-sequence half-tail without crossing its boundary."""
+
+ valid_tokens = seqlen - chunk_start
+ value_base = dv_half * cutlass.Int32(DV_HALF)
+ for elem_iter in cutlass.range_constexpr((BT * DV_HALF) // THREADS_PER_WARP):
+ linear_idx = elem_iter * THREADS_PER_WARP + lane
+ token_coord = linear_idx // DV_HALF
+ value_dim = linear_idx - token_coord * DV_HALF
+ if token_coord < valid_tokens:
+ smem_idx = o_smem_swizzle_128b_elem_index(
+ o_stage_base,
+ value_dim,
+ token_coord,
+ )
+ out[
+ 0,
+ sequence_start + chunk_start + token_coord,
+ head_idx,
+ value_base + value_dim,
+ ] = o_smem[smem_idx]
+ prims.bar_warp_sync(cute.arch.FULL_MASK)
+
+
+@cute.jit
+def epilogue_chain_wait_and_store_full_output(
+ tma_desc_o: cutlass.GridConstant[cuda.TensorMap],
+ o_smem,
+ output_ready_mbar,
+ output_consumed_mbar,
+ sequence_start,
+ head_idx,
+ dv_half,
+ output_chunk,
+ O_STAGES: cutlass.Constexpr,
+):
+ """Drain one full staged half-output chunk from SMEM with TMA."""
+
+ output_chunk_start = output_chunk * BT
+ o_stage = output_chunk % O_STAGES
+ o_stage_base = o_stage * K2_O_SMEM_STAGE_SIZE
+ output_ready_wait(
+ output_ready_mbar.subview(o_stage),
+ (output_chunk // O_STAGES) % 2,
+ )
+ epilogue_chain_stage_store(
+ tma_desc_o,
+ o_smem,
+ sequence_start,
+ head_idx,
+ dv_half,
+ output_chunk_start,
+ o_stage_base,
+ )
+ output_consumed_arrive(output_consumed_mbar.subview(o_stage))
+
+
+@cute.jit
+def epilogue_chain_wait_and_store_final_output(
+ tma_desc_o: cutlass.GridConstant[cuda.TensorMap],
+ out,
+ o_smem,
+ output_ready_mbar,
+ output_consumed_mbar,
+ sequence_start,
+ head_idx,
+ dv_half,
+ seqlen,
+ output_chunk,
+ lane,
+ O_STAGES: cutlass.Constexpr,
+):
+ """Drain the final half-output chunk, guarding a partial packed tail."""
+
+ if seqlen % BT == 0:
+ epilogue_chain_wait_and_store_full_output(
+ tma_desc_o,
+ o_smem,
+ output_ready_mbar,
+ output_consumed_mbar,
+ sequence_start,
+ head_idx,
+ dv_half,
+ output_chunk,
+ O_STAGES,
+ )
+ else:
+ output_chunk_start = output_chunk * BT
+ o_stage = output_chunk % O_STAGES
+ o_stage_base = o_stage * K2_O_SMEM_STAGE_SIZE
+ output_ready_wait(
+ output_ready_mbar.subview(o_stage),
+ (output_chunk // O_STAGES) % 2,
+ )
+ epilogue_chain_tail_store(
+ out,
+ o_smem,
+ sequence_start,
+ head_idx,
+ dv_half,
+ output_chunk_start,
+ seqlen,
+ o_stage_base,
+ lane,
+ )
+ output_consumed_arrive(output_consumed_mbar.subview(o_stage))
+
+
+@cute.kernel
+def kernel_chain_dv2(
+ tma_desc_kd: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_w: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_qd: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_v: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_diag: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_qk: cutlass.GridConstant[cuda.TensorMap],
+ tma_desc_o: cutlass.GridConstant[cuda.TensorMap],
+ v: cute.Tensor,
+ cu_seqlens: cute.Tensor,
+ cu_chunks: cute.Tensor,
+ state_indices: cute.Tensor | None,
+ initial_state: cute.Tensor | None,
+ out: cute.Tensor,
+ final_state: cute.Tensor | None,
+ head_base: cutlass.Int32,
+ SCALE: cutlass.Float32,
+ state_ckpt: cute.Tensor | None,
+ cu_ckpts: cute.Tensor | None,
+ checkpoint_state_indices: cute.Tensor | None,
+ checkpoint_stride_chunks: cutlass.Int32,
+) -> None:
+ """kernel 2, DV-split: each CTA owns half the hidden dimension.
+
+ Grid `(num_sequences, launch_heads * 2, 1)`; bidy = head * 2 + half.
+ Identical schedule and mbar topology to the base chain with M=64
+ MMAs (PTX Layout F: 16 rows per warp quadrant, lane alignment 0);
+ kd/W/QK'/diag are read identically by both halves, v/o/state are
+ DV-split (one 64-elem s128 TMA segment at value offset half*64).
+ """
+
+ tidx, _, _ = cute.arch.thread_idx()
+ bidx, bidy, _ = cute.arch.block_idx()
+ dv_half = bidy % 2
+ bidy = head_base + bidy // 2
+ warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
+ lane = tidx % THREADS_PER_WARP
+
+ sequence_start = cutlass.Int32(cu_seqlens[bidx])
+ sequence_end = cutlass.Int32(cu_seqlens[bidx + 1])
+ seqlen = sequence_end - sequence_start
+ num_chunks = cute.ceil_div(seqlen, BT)
+ chunk_base = cutlass.Int32(cu_chunks[bidx])
+ input_dtype = v.element_type
+
+ # --- mbarriers (identical topology to the base chain) ----------------
+ tma_mbar = cutlass.Array(
+ cutlass.Int64,
+ K2_TMA_MBAR_STAGE_COUNT,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ raw_ready_mbar = cutlass.Array(
+ cutlass.Int64,
+ K2_RAW_STAGE_COUNT,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ raw_consumed_mbar = cutlass.Array(
+ cutlass.Int64,
+ K2_RAW_STAGE_COUNT,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ state_input_ready_l_mbar = cutlass.Array(
+ cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8
+ )
+ state_input_ready_mbar = cutlass.Array(
+ cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8
+ )
+ u_input_ready_mbar = cutlass.Array(
+ cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8
+ )
+ update_ready_mbar = cutlass.Array(
+ cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8
+ )
+ shared_acc_ready_mbar = cutlass.Array(
+ cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8
+ )
+ k_restore_consumed_l_mbar = cutlass.Array(
+ cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8
+ )
+ k_restore_consumed_mbar = cutlass.Array(
+ cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8
+ )
+ qstate_acc_ready_mbar = cutlass.Array(
+ cutlass.Int64,
+ K2_QSTATE_STAGE_COUNT,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ stateq_done_mbar = cutlass.Array(
+ cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8
+ )
+ output_ready_mbar = cutlass.Array(
+ cutlass.Int64,
+ K2_QSTATE_STAGE_COUNT,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ output_consumed_mbar = cutlass.Array(
+ cutlass.Int64,
+ K2_QSTATE_STAGE_COUNT,
+ space=cutlass.AddressSpace.smem,
+ alignment=8,
+ )
+ final_state_stored_mbar = cutlass.Array(
+ cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8
+ )
+ checkpoint_read_done_mbar = cutlass.Array(
+ cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8
+ )
+ tmem_ptr_i32 = cutlass.Array(
+ cutlass.Int32, 1, space=cutlass.AddressSpace.smem, alignment=4
+ )
+
+ # --- SMEM rings (v halved to one segment; o halved) --------------------
+ kd_smem = cutlass.Array(
+ input_dtype,
+ K2_RAW_STAGE_COUNT * TILE_ELEMS,
+ space=cutlass.AddressSpace.smem,
+ alignment=RAW_F16_TMA_SWIZZLE_ALIGNMENT_BYTES,
+ )
+ w_smem = cutlass.Array(
+ input_dtype,
+ K2_RAW_STAGE_COUNT * TILE_ELEMS,
+ space=cutlass.AddressSpace.smem,
+ alignment=RAW_F16_TMA_SWIZZLE_ALIGNMENT_BYTES,
+ )
+ qd_smem = cutlass.Array(
+ input_dtype,
+ K2_RAW_STAGE_COUNT * TILE_ELEMS,
+ space=cutlass.AddressSpace.smem,
+ alignment=RAW_F16_TMA_SWIZZLE_ALIGNMENT_BYTES,
+ )
+ v_raw_smem = cutlass.Array(
+ input_dtype,
+ K2_RAW_STAGE_COUNT * V_TILE_ELEMS,
+ space=cutlass.AddressSpace.smem,
+ alignment=RAW_F16_TMA_SWIZZLE_ALIGNMENT_BYTES,
+ )
+ diag_raw_smem = cutlass.Array(
+ cutlass.Float32,
+ K2_RAW_STAGE_COUNT * DIAG_REC_ELEMS,
+ space=cutlass.AddressSpace.smem,
+ alignment=1024,
+ )
+ qk_smem = cutlass.Array(
+ input_dtype,
+ K2_RAW_STAGE_COUNT * QK_REC_ELEMS,
+ space=cutlass.AddressSpace.smem,
+ alignment=1024,
+ )
+ o_smem = cutlass.Array(
+ out.element_type,
+ K2_O_SMEM_TILE_SIZE,
+ space=cutlass.AddressSpace.smem,
+ alignment=O_TMA_SWIZZLE_ALIGNMENT_BYTES,
+ )
+
+ # --- init (identical to the base chain) ------------------------------
+ if warp_idx == ROLES.tma_load:
+ if prims.elect_sync():
+ for stage in cutlass.range_constexpr(K2_TMA_MBAR_STAGE_COUNT):
+ prims.mbarrier_init(tma_mbar.subview(stage), 1)
+ for stage in cutlass.range_constexpr(K2_RAW_STAGE_COUNT):
+ prims.mbarrier_init(raw_ready_mbar.subview(stage), 1)
+ prims.mbarrier_init(
+ raw_consumed_mbar.subview(stage),
+ 5,
+ )
+ elif warp_idx == ROLES.tcgen05_mma:
+ if prims.elect_sync():
+ prims.mbarrier_init(state_input_ready_l_mbar, 4)
+ prims.mbarrier_init(state_input_ready_mbar, 4)
+ for stage in cutlass.range_constexpr(2):
+ prims.mbarrier_init(u_input_ready_mbar.subview(stage), 4)
+ prims.mbarrier_init(shared_acc_ready_mbar.subview(stage), 1)
+ prims.mbarrier_init(k_restore_consumed_l_mbar.subview(stage), 1)
+ prims.mbarrier_init(k_restore_consumed_mbar.subview(stage), 1)
+ prims.mbarrier_init(stateq_done_mbar.subview(stage), 1)
+ # CG1 fuses the right-half pack/scale while idle CG0 warps 4-7
+ # fuse the left half. Both groups must finish before delta.
+ prims.mbarrier_init(update_ready_mbar, 8)
+ for stage in cutlass.range_constexpr(K2_QSTATE_STAGE_COUNT):
+ prims.mbarrier_init(qstate_acc_ready_mbar.subview(stage), 1)
+ prims.mbarrier_init(final_state_stored_mbar, 8)
+ if cutlass.const_expr(state_ckpt is not None):
+ prims.mbarrier_init(checkpoint_read_done_mbar, 4)
+ elif warp_idx == ROLES.epilogue:
+ if prims.elect_sync():
+ for stage in cutlass.range_constexpr(K2_QSTATE_STAGE_COUNT):
+ prims.mbarrier_init(output_ready_mbar.subview(stage), 4)
+ prims.mbarrier_init(output_consumed_mbar.subview(stage), 1)
+ prims.fence_mbarrier_init()
+ cta_sync()
+
+ if is_tmem_user_warp(warp_idx):
+ if warp_idx == ROLES.tcgen05_mma:
+ prims.tcgen05_alloc(tmem_ptr_i32, TMEM_ALLOC_COLS, group="cta_1")
+ tmem_user_sync()
+ if warp_idx == ROLES.tcgen05_mma:
+ prims.tcgen05_relinquish_alloc_permit(group="cta_1")
+ tmem_user_sync()
+ cta_sync()
+
+ if is_service_warpgroup(warp_idx):
+ prims.setmaxregister(KDA_SERVICE_REGS, prims.SetMaxRegisterAction.DECREASE)
+
+ # ======================= warp 14: TMA ring =========================
+ if warp_idx == ROLES.tma_load:
+ raw_stage = cutlass.Int32(0)
+ raw_consumed_phase = cutlass.Int32(1)
+ issue_mbar_slot = cutlass.Int32(0)
+ wait_mbar_slot = cutlass.Int32(0)
+ ready_stage = cutlass.Int32(0)
+ tma_phase = cutlass.Int32(0)
+ for chunk in cutlass.range(num_chunks, unroll=1):
+ ws_chunk = chunk_base + chunk
+ ws_row_start = ws_chunk * cutlass.Int32(BT)
+ v_row_start = sequence_start + chunk * cutlass.Int32(BT)
+ raw_consumed_wait(
+ raw_consumed_mbar.subview(raw_stage),
+ raw_consumed_phase,
+ )
+ tma_chain_stage_load_inputs(
+ tma_desc_kd,
+ tma_desc_w,
+ tma_desc_qd,
+ tma_desc_v,
+ tma_desc_diag,
+ tma_desc_qk,
+ kd_smem.subview(raw_stage * TILE_ELEMS),
+ w_smem.subview(raw_stage * TILE_ELEMS),
+ qd_smem.subview(raw_stage * TILE_ELEMS),
+ v_raw_smem.subview(raw_stage * V_TILE_ELEMS),
+ diag_raw_smem.subview(raw_stage * DIAG_REC_ELEMS),
+ qk_smem.subview(raw_stage * QK_REC_ELEMS),
+ bidy,
+ dv_half,
+ ws_row_start,
+ ws_chunk,
+ v_row_start,
+ tma_mbar.subview(issue_mbar_slot),
+ K2_TX_BYTES,
+ )
+ issue_mbar_slot, _ = advance_ring_stage(
+ issue_mbar_slot, 1, K2_TMA_MBAR_STAGE_COUNT
+ )
+ if chunk >= cutlass.Int32(K2_TMA_MBAR_STAGE_COUNT - 1):
+ tma_transfer_wait(tma_mbar.subview(wait_mbar_slot), tma_phase)
+ wait_mbar_slot, wait_wrapped = advance_ring_stage(
+ wait_mbar_slot, 1, K2_TMA_MBAR_STAGE_COUNT
+ )
+ tma_phase = tma_phase ^ wait_wrapped
+ raw_ready_arrive(raw_ready_mbar.subview(ready_stage))
+ ready_stage, _ = advance_ring_stage(ready_stage, 1, K2_RAW_STAGE_COUNT)
+ raw_stage, raw_wrapped = advance_ring_stage(
+ raw_stage, 1, K2_RAW_STAGE_COUNT
+ )
+ raw_consumed_phase = raw_consumed_phase ^ raw_wrapped
+ tma_full = cutlass.Int32(
+ num_chunks >= cutlass.Int32(K2_TMA_MBAR_STAGE_COUNT - 1)
+ )
+ tma_tail = (
+ tma_full * cutlass.Int32(K2_TMA_MBAR_STAGE_COUNT - 1)
+ + (cutlass.Int32(1) - tma_full) * num_chunks
+ )
+ for _tail in cutlass.range(tma_tail, unroll=1):
+ tma_transfer_wait(tma_mbar.subview(wait_mbar_slot), tma_phase)
+ wait_mbar_slot, wait_wrapped = advance_ring_stage(
+ wait_mbar_slot, 1, K2_TMA_MBAR_STAGE_COUNT
+ )
+ tma_phase = tma_phase ^ wait_wrapped
+ raw_ready_arrive(raw_ready_mbar.subview(ready_stage))
+ ready_stage, _ = advance_ring_stage(ready_stage, 1, K2_RAW_STAGE_COUNT)
+
+ # ====== warp 12: output-group issuer (OUT_ISSUER12) / flag relay ======
+ elif warp_idx == ROLES.super_mma:
+ tmem_raw_addr = tmem_ptr_i32.load()
+ si_phase12 = cutlass.Int32(0)
+ upd_phase12 = cutlass.Int32(0)
+ for chunk in cutlass.range(num_chunks, unroll=1):
+ raw_stage = chunk % K2_RAW_STAGE_COUNT
+ qstate_stage = chunk % K2_QSTATE_STAGE_COUNT
+ kr_stage = chunk % 2
+ acc_stage = chunk % 2
+ if chunk > 0:
+ prev12 = chunk - cutlass.Int32(1)
+ tcgen05_wait_acc_buffer_ready(
+ qstate_acc_ready_mbar.subview(prev12 % K2_QSTATE_STAGE_COUNT),
+ (prev12 // K2_QSTATE_STAGE_COUNT) % 2,
+ )
+ raw_consumed_arrive(
+ raw_consumed_mbar.subview(prev12 % K2_RAW_STAGE_COUNT)
+ )
+ raw_ready_wait(
+ raw_ready_mbar.subview(raw_stage),
+ (chunk // K2_RAW_STAGE_COUNT) % 2,
+ )
+ si_phase12 = state_input_ready_wait(state_input_ready_mbar, si_phase12)
+ prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC)
+ output_ready_wait(
+ output_ready_mbar.subview(qstate_stage),
+ (chunk // K2_QSTATE_STAGE_COUNT + cutlass.Int32(1)) % 2,
+ )
+ tcgen05_issue_state_projection_mma(
+ qd_smem.subview(raw_stage * TILE_ELEMS),
+ tmem_raw_addr,
+ stateq_done_mbar.subview(kr_stage),
+ tcgen05_qstate_acc_tmem_col_offset(qstate_stage),
+ input_dtype,
+ 0,
+ DK // TCGEN05_F16_K_ATOM,
+ False,
+ True,
+ DV_HALF,
+ )
+ upd_phase12 = update_ready_wait(update_ready_mbar, upd_phase12)
+ prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC)
+ xpack_col12 = tcgen05_shared_input_tmem_col_offset(acc_stage)
+ tcgen05_chain_issue_qkv_mma(
+ qk_smem.subview(raw_stage * QK_REC_ELEMS),
+ tmem_raw_addr,
+ xpack_col12,
+ qstate_stage,
+ qstate_acc_ready_mbar.subview(qstate_stage),
+ input_dtype,
+ True,
+ )
+
+ # =============== warp 13: tcgen05 issuer (the chain owner) =============
+ elif warp_idx == ROLES.tcgen05_mma:
+ tmem_raw_addr = tmem_ptr_i32.load()
+ si_l_phase = cutlass.Int32(0)
+ si_phase = cutlass.Int32(0)
+ upd_phase = cutlass.Int32(0)
+ for chunk in cutlass.range(num_chunks, unroll=1):
+ raw_stage = chunk % K2_RAW_STAGE_COUNT
+ raw_phase = (chunk // K2_RAW_STAGE_COUNT) % 2
+ acc_stage = chunk % 2
+ qstate_stage = chunk % K2_QSTATE_STAGE_COUNT
+ kr_stage = chunk % 2
+ kd_stage_smem = kd_smem.subview(raw_stage * TILE_ELEMS)
+ w_stage_smem = w_smem.subview(raw_stage * TILE_ELEMS)
+
+ raw_ready_wait(raw_ready_mbar.subview(raw_stage), raw_phase)
+
+ si_l_phase = state_input_ready_wait(state_input_ready_l_mbar, si_l_phase)
+ prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC)
+ tcgen05_issue_state_k_mma(
+ kd_stage_smem,
+ tmem_raw_addr,
+ shared_acc_ready_mbar.subview(acc_stage),
+ acc_stage,
+ input_dtype,
+ 0,
+ (DK // TCGEN05_F16_K_ATOM) // 2,
+ False,
+ False,
+ DV_HALF,
+ )
+ si_phase = state_input_ready_wait(state_input_ready_mbar, si_phase)
+ prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC)
+ tcgen05_issue_state_k_mma(
+ kd_stage_smem,
+ tmem_raw_addr,
+ shared_acc_ready_mbar.subview(acc_stage),
+ acc_stage,
+ input_dtype,
+ (DK // TCGEN05_F16_K_ATOM) // 2,
+ DK // TCGEN05_F16_K_ATOM,
+ True,
+ True,
+ DV_HALF,
+ )
+
+ xpack_col = tcgen05_shared_input_tmem_col_offset(acc_stage)
+ upd_phase = update_ready_wait(update_ready_mbar, upd_phase)
+ prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC)
+ tcgen05_chain_issue_delta_half_mma(
+ w_stage_smem, tmem_raw_addr, xpack_col, input_dtype, 0
+ )
+ tcgen05_commit(k_restore_consumed_l_mbar.subview(kr_stage))
+ tcgen05_chain_issue_delta_half_mma(
+ w_stage_smem, tmem_raw_addr, xpack_col, input_dtype, 1
+ )
+ tcgen05_commit(k_restore_consumed_mbar.subview(kr_stage))
+
+ final_state_stored_wait(final_state_stored_mbar, cutlass.Int32(0))
+ tmem_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Float32)
+ prims.tcgen05_dealloc(tmem_ptr, TMEM_ALLOC_COLS, group="cta_1")
+
+ # ======================= warp 15: output store =========================
+ elif warp_idx == ROLES.epilogue:
+ for chunk in cutlass.range(num_chunks, unroll=1):
+ if chunk > 0:
+ epilogue_chain_wait_and_store_full_output(
+ tma_desc_o,
+ o_smem,
+ output_ready_mbar,
+ output_consumed_mbar,
+ sequence_start,
+ bidy,
+ dv_half,
+ chunk - cutlass.Int32(1),
+ K2_QSTATE_STAGE_COUNT,
+ )
+ if num_chunks > 0:
+ epilogue_chain_wait_and_store_final_output(
+ tma_desc_o,
+ out,
+ o_smem,
+ output_ready_mbar,
+ output_consumed_mbar,
+ sequence_start,
+ bidy,
+ dv_half,
+ seqlen,
+ num_chunks - cutlass.Int32(1),
+ lane,
+ K2_QSTATE_STAGE_COUNT,
+ )
+
+ # ================= CG1 (warps 8-11): TMEM epilogues ====================
+ elif is_compute_group1_warp(warp_idx):
+ prims.setmaxregister(KDA_CG1_REGS, prims.SetMaxRegisterAction.INCREASE)
+ tmem_raw_addr = tmem_ptr_i32.load()
+ # Only CG1 touches recurrent state. Keep the pool lookup outside the
+ # chunk loop and out of producer/MMAs warps.
+ state_slot = bidx
+ if cutlass.const_expr(state_indices is not None):
+ state_slot = cutlass.Int32(state_indices[bidx])
+ ckpt_slot = cutlass.Int32(0)
+ if cutlass.const_expr(state_ckpt is not None):
+ # cu_ckpts supplies the per-sequence base slot offsets. A
+ # loop-carried counter replaces a per-chunk div/mod.
+ ckpt_stride = checkpoint_stride_chunks
+ ckpt_next = ckpt_stride
+ ckpt_slot = cutlass.Int32(cu_ckpts[bidx])
+ if cutlass.const_expr(checkpoint_state_indices is not None):
+ tcgen05_store_initial_state_tmem(
+ tmem_raw_addr,
+ initial_state,
+ None,
+ cutlass.Int32(0),
+ state_slot,
+ bidy,
+ dv_half,
+ warp_idx,
+ lane,
+ HALF=True,
+ )
+ else:
+ tcgen05_store_initial_state_tmem(
+ tmem_raw_addr,
+ initial_state,
+ state_ckpt,
+ ckpt_slot,
+ state_slot,
+ bidy,
+ dv_half,
+ warp_idx,
+ lane,
+ HALF=True,
+ )
+ if cutlass.const_expr(
+ state_ckpt is not None and checkpoint_state_indices is None
+ ):
+ ckpt_slot += cutlass.Int32(1)
+ if num_chunks > 0:
+ diag_raw_stage = diag_raw_smem.subview(0)
+ state_left = tcgen05_stage_state_input_dv2_half_tmem(
+ tmem_raw_addr,
+ warp_idx,
+ input_dtype,
+ 0,
+ )
+ raw_ready_wait(raw_ready_mbar.subview(0), 0)
+ tcgen05_rescale_state_dv2_half_regs(
+ tmem_raw_addr,
+ diag_raw_stage,
+ warp_idx,
+ state_left,
+ state_input_ready_l_mbar,
+ 0,
+ )
+ state_right = tcgen05_stage_state_input_dv2_half_tmem(
+ tmem_raw_addr,
+ warp_idx,
+ input_dtype,
+ 1,
+ )
+ tcgen05_rescale_state_dv2_half_regs(
+ tmem_raw_addr,
+ diag_raw_stage,
+ warp_idx,
+ state_right,
+ state_input_ready_mbar,
+ 1,
+ )
+ tcgen05_wait_acc_buffer_ready(shared_acc_ready_mbar.subview(0), 0)
+ vmx_lane = cute.arch.lane_idx()
+ tcgen05_chain_stage_vmx_input_tmem(
+ tmem_raw_addr,
+ v_raw_smem.subview(0),
+ warp_idx,
+ vmx_lane,
+ 0,
+ 0,
+ input_dtype,
+ )
+ post_scale_lane = cute.arch.lane_idx()
+ update_ready_arrive(update_ready_mbar)
+ if cutlass.const_expr(state_ckpt is not None):
+ # Peeled chunk 0's checkpoint (stride 1 only).
+ checkpoint_due = ckpt_next == cutlass.Int32(1)
+ if cutlass.const_expr(checkpoint_state_indices is not None):
+ checkpoint_due = checkpoint_due & (cutlass.Int32(BT) <= seqlen)
+ else:
+ checkpoint_due = checkpoint_due & (cutlass.Int32(1) < num_chunks)
+ if checkpoint_due:
+ tcgen05_wait_acc_buffer_ready(
+ k_restore_consumed_l_mbar.subview(0), 0
+ )
+ tcgen05_wait_acc_buffer_ready(k_restore_consumed_mbar.subview(0), 0)
+ checkpoint_output_slot = ckpt_slot
+ if cutlass.const_expr(checkpoint_state_indices is not None):
+ checkpoint_output_slot = cutlass.Int32(
+ checkpoint_state_indices[ckpt_slot]
+ )
+ if (checkpoint_output_slot >= cutlass.Int32(0)) & (
+ checkpoint_output_slot < cutlass.Int32(state_ckpt.shape[0])
+ ):
+ tcgen05_store_final_state_tmem(
+ tmem_raw_addr,
+ KDA_TMEM_STATE_COL_OFFSET,
+ state_ckpt,
+ checkpoint_output_slot,
+ bidy,
+ dv_half,
+ warp_idx,
+ post_scale_lane,
+ HALF=True,
+ )
+ checkpoint_read_done_arrive(checkpoint_read_done_mbar)
+ ckpt_slot += cutlass.Int32(1)
+ ckpt_next += ckpt_stride
+ for chunk in cutlass.range(1, num_chunks, 1, unroll=1):
+ prev = chunk - cutlass.Int32(1)
+ raw_stage = chunk % K2_RAW_STAGE_COUNT
+ diag_raw_stage = diag_raw_smem.subview(raw_stage * DIAG_REC_ELEMS)
+ acc_stage = chunk % 2
+ prev_kr_phase = (prev // 2) % 2
+ tcgen05_wait_acc_buffer_ready(
+ stateq_done_mbar.subview(prev % 2), prev_kr_phase
+ )
+ tcgen05_wait_acc_buffer_ready(
+ k_restore_consumed_mbar.subview(prev % 2), prev_kr_phase
+ )
+ state_right = tcgen05_stage_state_input_dv2_half_tmem(
+ tmem_raw_addr,
+ warp_idx,
+ input_dtype,
+ 1,
+ )
+ raw_consumed_arrive(raw_consumed_mbar.subview(prev % K2_RAW_STAGE_COUNT))
+ raw_ready_wait(
+ raw_ready_mbar.subview(raw_stage),
+ (chunk // K2_RAW_STAGE_COUNT) % 2,
+ )
+ tcgen05_rescale_state_dv2_half_regs(
+ tmem_raw_addr,
+ diag_raw_stage,
+ warp_idx,
+ state_right,
+ state_input_ready_mbar,
+ 1,
+ )
+ tcgen05_wait_acc_buffer_ready(
+ shared_acc_ready_mbar.subview(acc_stage), (chunk // 2) % 2
+ )
+ vmx_lane = cute.arch.lane_idx()
+ tcgen05_chain_stage_vmx_input_tmem(
+ tmem_raw_addr,
+ v_raw_smem.subview(raw_stage * V_TILE_ELEMS),
+ warp_idx,
+ vmx_lane,
+ acc_stage,
+ acc_stage,
+ input_dtype,
+ )
+ post_scale_lane = cute.arch.lane_idx()
+ update_ready_arrive(update_ready_mbar)
+ if cutlass.const_expr(state_ckpt is not None):
+ # Same contract as the engine-side checkpoint; both M=64 TMEM
+ # halves wait their own k_restore parity slot (the same wait
+ # the post-loop final store performs) before draining.
+ checkpoint_due = chunk + cutlass.Int32(1) == ckpt_next
+ if cutlass.const_expr(checkpoint_state_indices is not None):
+ checkpoint_due = checkpoint_due & (
+ (chunk + cutlass.Int32(1)) * cutlass.Int32(BT) <= seqlen
+ )
+ else:
+ checkpoint_due = checkpoint_due & (
+ chunk + cutlass.Int32(1) < num_chunks
+ )
+ if checkpoint_due:
+ tcgen05_wait_acc_buffer_ready(
+ k_restore_consumed_l_mbar.subview(chunk % 2),
+ (chunk // 2) % 2,
+ )
+ tcgen05_wait_acc_buffer_ready(
+ k_restore_consumed_mbar.subview(chunk % 2),
+ (chunk // 2) % 2,
+ )
+ checkpoint_output_slot = ckpt_slot
+ if cutlass.const_expr(checkpoint_state_indices is not None):
+ checkpoint_output_slot = cutlass.Int32(
+ checkpoint_state_indices[ckpt_slot]
+ )
+ if (checkpoint_output_slot >= cutlass.Int32(0)) & (
+ checkpoint_output_slot < cutlass.Int32(state_ckpt.shape[0])
+ ):
+ tcgen05_store_final_state_tmem(
+ tmem_raw_addr,
+ KDA_TMEM_STATE_COL_OFFSET,
+ state_ckpt,
+ checkpoint_output_slot,
+ bidy,
+ dv_half,
+ warp_idx,
+ post_scale_lane,
+ HALF=True,
+ )
+ checkpoint_read_done_arrive(checkpoint_read_done_mbar)
+ ckpt_slot += cutlass.Int32(1)
+ ckpt_next += ckpt_stride
+ if num_chunks > 0:
+ last = num_chunks - cutlass.Int32(1)
+ last_kr_phase = (last // 2) % 2
+ tcgen05_wait_acc_buffer_ready(
+ k_restore_consumed_l_mbar.subview(last % 2), last_kr_phase
+ )
+ tcgen05_wait_acc_buffer_ready(
+ k_restore_consumed_mbar.subview(last % 2), last_kr_phase
+ )
+ if cutlass.const_expr(final_state is not None):
+ final_lane = cute.arch.lane_idx()
+ tcgen05_store_final_state_tmem(
+ tmem_raw_addr,
+ KDA_TMEM_STATE_COL_OFFSET,
+ final_state,
+ state_slot,
+ bidy,
+ dv_half,
+ warp_idx,
+ final_lane,
+ HALF=True,
+ )
+ final_state_stored_arrive(final_state_stored_mbar)
+ elif is_compute_group0_warp(warp_idx):
+ if warp_idx < 4:
+ prims.setmaxregister(KDA_CG1_REGS, prims.SetMaxRegisterAction.INCREASE)
+ tmem_raw_addr = tmem_ptr_i32.load()
+ for chunk in cutlass.range(num_chunks, unroll=1):
+ o_stage = chunk % K2_QSTATE_STAGE_COUNT
+ o_phase = (chunk // K2_QSTATE_STAGE_COUNT) % 2
+ output_consumed_wait(
+ output_consumed_mbar.subview(o_stage),
+ (o_phase + 1) % 2,
+ )
+ tcgen05_wait_acc_buffer_ready(
+ qstate_acc_ready_mbar.subview(o_stage), o_phase
+ )
+ tcgen05_chain_load_qstate_output_tmem(
+ tmem_raw_addr,
+ o_smem,
+ warp_idx,
+ lane,
+ o_stage * K2_O_SMEM_STAGE_SIZE,
+ o_stage,
+ SCALE,
+ out.element_type,
+ )
+ output_ready_arrive(output_ready_mbar.subview(o_stage))
+ final_state_stored_arrive(final_state_stored_mbar)
+ else:
+ # Chunk 0 is initialized by CG1. Afterwards, warps 4-7 own the
+ # fused left-half state pack and FP32 decay while CG1 owns right.
+ prims.setmaxregister(KDA_CG1_REGS, prims.SetMaxRegisterAction.INCREASE)
+ tmem_raw_addr = tmem_ptr_i32.load()
+ if cutlass.const_expr(state_ckpt is not None):
+ cg0_ckpt_stride = checkpoint_stride_chunks
+ if num_chunks > 0:
+ state_input_ready_wait(
+ state_input_ready_l_mbar,
+ cutlass.Int32(0),
+ )
+ prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC)
+ update_ready_arrive(update_ready_mbar)
+ for chunk in cutlass.range(1, num_chunks, 1, unroll=1):
+ prev = chunk - cutlass.Int32(1)
+ raw_stage = chunk % K2_RAW_STAGE_COUNT
+ prev_kr_phase = (prev // 2) % 2
+ tcgen05_wait_acc_buffer_ready(
+ stateq_done_mbar.subview(prev % 2),
+ prev_kr_phase,
+ )
+ tcgen05_wait_acc_buffer_ready(
+ k_restore_consumed_l_mbar.subview(prev % 2),
+ prev_kr_phase,
+ )
+ if cutlass.const_expr(state_ckpt is not None):
+ if chunk % cg0_ckpt_stride == 0:
+ checkpoint_read_done_wait(
+ checkpoint_read_done_mbar,
+ ((chunk // cg0_ckpt_stride) + cutlass.Int32(1)) % 2,
+ )
+ prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC)
+ state_left = tcgen05_stage_state_input_dv2_half_tmem(
+ tmem_raw_addr,
+ warp_idx,
+ input_dtype,
+ 0,
+ )
+ raw_ready_wait(
+ raw_ready_mbar.subview(raw_stage),
+ (chunk // K2_RAW_STAGE_COUNT) % 2,
+ )
+ diag_raw_stage = diag_raw_smem.subview(raw_stage * DIAG_REC_ELEMS)
+ tcgen05_rescale_state_dv2_half_regs(
+ tmem_raw_addr,
+ diag_raw_stage,
+ warp_idx,
+ state_left,
+ state_input_ready_l_mbar,
+ 0,
+ )
+ update_ready_arrive(update_ready_mbar)
+
+
+@cute.jit
+def host_chain_dv2(
+ v: cute.Tensor,
+ cu_seqlens: cute.Tensor,
+ cu_chunks: cute.Tensor,
+ ws_kd: cute.Tensor,
+ ws_qd: cute.Tensor,
+ ws_w: cute.Tensor,
+ ws_qk: cute.Tensor,
+ ws_diag: cute.Tensor,
+ state_indices: cute.Tensor | None,
+ initial_state: cute.Tensor | None,
+ out: cute.Tensor,
+ final_state: cute.Tensor | None,
+ stream,
+ head_base: cutlass.Int32,
+ launch_heads: cutlass.Int32,
+ SCALE: cutlass.Float32,
+ state_ckpt: cute.Tensor | None,
+ cu_ckpts: cute.Tensor | None,
+ checkpoint_state_indices: cute.Tensor | None,
+ checkpoint_stride_chunks: cutlass.Int32,
+ THREADS: cutlass.Constexpr,
+) -> None:
+ """DV2 host: identical tensor maps to the base chain host, doubled grid-y."""
+
+ num_sequences = cu_seqlens.shape[0] - 1
+ seqlen = v.shape[1]
+ heads = v.shape[2]
+ ws_rows = ws_kd.shape[2]
+ num_chunks_total = ws_qk.shape[2]
+ # Keep singleton TensorMap modes canonical for the same reason as K1.
+ tile_head_stride = DK
+ diag_head_stride = DIAG_REC_ELEMS
+ qk_head_stride = QK_REC_ELEMS
+ if heads != cutlass.Int32(1):
+ tile_head_stride = DK * ws_rows
+ diag_head_stride = DIAG_REC_ELEMS * num_chunks_total
+ qk_head_stride = QK_REC_ELEMS * num_chunks_total
+ tile_layout = cute.make_layout(
+ (DK, ws_rows, heads, 1),
+ stride=(1, DK, tile_head_stride, DK),
+ )
+ v_layout = cute.make_layout(
+ (DV, seqlen, heads, 1),
+ stride=(1, DV * heads, DV, DV * seqlen * heads),
+ )
+ diag_layout = cute.make_layout(
+ (DIAG_REC_ELEMS, num_chunks_total, heads, 1),
+ stride=(
+ 1,
+ DIAG_REC_ELEMS,
+ diag_head_stride,
+ DIAG_REC_ELEMS,
+ ),
+ )
+ qk_layout = cute.make_layout(
+ (QK_REC_ELEMS, num_chunks_total, heads, 1),
+ stride=(
+ 1,
+ QK_REC_ELEMS,
+ qk_head_stride,
+ QK_REC_ELEMS,
+ ),
+ )
+ f16_box = (RAW_F16_TMA_SWIZZLE_ELEMS, BT, 1, 1)
+ tma_desc_kd = cuda.create_tensor_map_tiled_from_view(
+ cute.make_tensor(ws_kd.iterator, tile_layout),
+ box_dims=f16_box,
+ stride_order=(0, 1, 2, 3),
+ swizzle=cuda.TensorMapSwizzle.s128b,
+ )
+ tma_desc_w = cuda.create_tensor_map_tiled_from_view(
+ cute.make_tensor(ws_w.iterator, tile_layout),
+ box_dims=f16_box,
+ stride_order=(0, 1, 2, 3),
+ swizzle=cuda.TensorMapSwizzle.s128b,
+ )
+ tma_desc_qd = cuda.create_tensor_map_tiled_from_view(
+ cute.make_tensor(ws_qd.iterator, tile_layout),
+ box_dims=f16_box,
+ stride_order=(0, 1, 2, 3),
+ swizzle=cuda.TensorMapSwizzle.s128b,
+ )
+ tma_desc_v = cuda.create_tensor_map_tiled_from_view(
+ cute.make_tensor(v.iterator, v_layout),
+ box_dims=f16_box,
+ stride_order=(0, 1, 2, 3),
+ swizzle=cuda.TensorMapSwizzle.s128b,
+ )
+ tma_desc_diag = cuda.create_tensor_map_tiled_from_view(
+ cute.make_tensor(ws_diag.iterator, diag_layout),
+ box_dims=(DIAG_REC_ELEMS, 1, 1, 1),
+ stride_order=(0, 1, 2, 3),
+ swizzle=cuda.TensorMapSwizzle.none,
+ )
+ tma_desc_qk = cuda.create_tensor_map_tiled_from_view(
+ cute.make_tensor(ws_qk.iterator, qk_layout),
+ box_dims=(QK_REC_ELEMS, 1, 1, 1),
+ stride_order=(0, 1, 2, 3),
+ swizzle=cuda.TensorMapSwizzle.none,
+ )
+ tma_desc_o = cuda.create_tensor_map_tiled_from_view(
+ cute.make_tensor(out.iterator, v_layout),
+ box_dims=(O_TMA_SWIZZLE_ELEMS, BT, 1, 1),
+ stride_order=(0, 1, 2, 3),
+ swizzle=cuda.TensorMapSwizzle.s128b,
+ )
+ kernel_chain_dv2(
+ tma_desc_kd,
+ tma_desc_w,
+ tma_desc_qd,
+ tma_desc_v,
+ tma_desc_diag,
+ tma_desc_qk,
+ tma_desc_o,
+ v,
+ cu_seqlens,
+ cu_chunks,
+ state_indices,
+ initial_state,
+ out,
+ final_state,
+ head_base,
+ SCALE,
+ state_ckpt,
+ cu_ckpts,
+ checkpoint_state_indices,
+ checkpoint_stride_chunks,
+ ).launch(
+ grid=(num_sequences, launch_heads * 2, 1),
+ block=(THREADS, 1, 1),
+ stream=stream,
+ min_blocks_per_mp=1,
+ preferred_smem_carveout=100,
+ )
+
+
+# =============================================================================
+# Unified single-`cute.compile`-target host.
+#
+# ONE @cute.jit that runtime-branches the plain engine vs the two-kernel
+# decomposition (kernel-1 factor prep + kernel-2 DV2 chain), calling `host` /
+# `host_prep` / `host_chain_dv2` as NESTED jit sub-calls. The occupancy rule
+# lives inside the host (`n_seq * heads * 2 <= sm_count`, sm_count passed as a
+# runtime arg). Decomp always runs prep-first on one stream: k1 completes its
+# full grid, then k2 reads a workspace k1 has fully written -- the kernel
+# boundary is the only handoff, so neither kernel carries any cross-kernel
+# synchronization.
+#
+# `cute.compile(host_unified, ...)` is the single target; `export_to_c` emits it
+# to one `.o` with every device cubin embedded under one host entry.
+# =============================================================================
+@cute.jit
+def host_unified(
+ q: cute.Tensor,
+ k: cute.Tensor,
+ v: cute.Tensor,
+ raw_gate: cute.Tensor,
+ a_log: cute.Tensor,
+ dt_bias: cute.Tensor,
+ beta: cute.Tensor,
+ cu_seqlens: cute.Tensor,
+ seq_order: cute.Tensor,
+ generate_seq_order: cutlass.Int32,
+ state_indices: cute.Tensor | None,
+ cu_chunks: cute.Tensor,
+ generate_cu_chunks: cutlass.Int32,
+ ws_kd: cute.Tensor,
+ ws_qd: cute.Tensor,
+ ws_w: cute.Tensor,
+ ws_qk: cute.Tensor,
+ ws_diag: cute.Tensor,
+ initial_state: cute.Tensor | None,
+ out: cute.Tensor,
+ final_state: cute.Tensor | None,
+ stream_a: cuda_driver.CUstream,
+ seq_route: cutlass.Int32,
+ sm_count: cutlass.Int32,
+ checkpoint_stride_chunks: cutlass.Int32,
+ prep_num_ctas: cutlass.Int32,
+ prep_cpc: cutlass.Int32,
+ scale: cutlass.Float32,
+ state_ckpt: cute.Tensor | None,
+ cu_ckpts: cute.Tensor | None,
+ checkpoint_state_indices: cute.Tensor | None,
+ SAFE_GATE: cutlass.Constexpr,
+ GATE_SCALE_LOG2: cutlass.Constexpr,
+ THREADS: cutlass.Constexpr,
+ MODE: cutlass.Constexpr,
+ gate_dtype: cutlass.Constexpr,
+) -> None:
+ n_seq = cu_seqlens.shape[0] - 1
+ heads = q.shape[2]
+ h32 = cutlass.Int32(heads)
+ z = cutlass.Int32(0)
+ if (generate_seq_order != z) | (generate_cu_chunks != z):
+ # A separate launch supplies the grid-wide publication boundary needed
+ # before either the persistent engine or decomposed kernels consume the
+ # generated metadata. The launch is captured as part of a CUDA Graph.
+ packed_metadata_kernel(
+ cu_seqlens,
+ seq_order,
+ cu_chunks,
+ generate_seq_order,
+ generate_cu_chunks,
+ ).launch(
+ grid=(1, 1, 1),
+ block=(PACKED_METADATA_THREADS, 1, 1),
+ stream=stream_a,
+ )
+ # Route selection. MODE is a COMPILE-TIME constexpr:
+ # MODE is None -> RUNTIME routing (both routes emitted, one .o handles
+ # all shapes) — the default build.
+ # MODE == "engine"-> engine only (decomp branch const-folds away).
+ # MODE == "decomp"-> decomp only (engine branch const-folds away).
+ # For MODE None the route is the single occupancy test: run the decomposition
+ # iff its doubled k2 grid fits one wave (n_seq*heads*2 <= sm_count). The
+ # wrapper re-derives the SAME decision host-side.
+ if cutlass.const_expr(MODE == "engine"):
+ route_decomp = False
+ elif cutlass.const_expr(MODE == "decomp"):
+ route_decomp = True
+ else:
+ free_sms = sm_count - n_seq * heads * 2
+ route_decomp = free_sms >= cutlass.Int32(0)
+ if route_decomp:
+ # ------------------------------ DECOMP ------------------------------
+ # PREP-FIRST on a single stream, always. k1 runs its full grid and
+ # completes; k2 then reads a workspace that is already whole, so no
+ # flag-ring, no per-chunk fences and no cross-stream ordering are needed
+ # -- the kernel boundary supplies the visibility the flags used to. The
+ # co-resident (chain-first, two-stream) variant was removed: it required
+ # the flag machinery on both kernels, and that cost more than the
+ # concurrency returned.
+ host_prep(
+ q,
+ k,
+ raw_gate,
+ a_log,
+ dt_bias,
+ beta,
+ cu_seqlens,
+ cu_chunks,
+ ws_kd,
+ ws_qd,
+ ws_w,
+ ws_qk,
+ ws_diag,
+ stream_a,
+ prep_num_ctas,
+ prep_cpc,
+ z,
+ h32,
+ SAFE_GATE,
+ GATE_SCALE_LOG2,
+ THREADS,
+ gate_dtype,
+ )
+ host_chain_dv2(
+ v,
+ cu_seqlens,
+ cu_chunks,
+ ws_kd,
+ ws_qd,
+ ws_w,
+ ws_qk,
+ ws_diag,
+ state_indices,
+ initial_state,
+ out,
+ final_state,
+ stream_a,
+ z,
+ h32,
+ scale,
+ state_ckpt,
+ cu_ckpts,
+ checkpoint_state_indices,
+ checkpoint_stride_chunks,
+ THREADS,
+ )
+ else:
+ # ------------------------------ ENGINE ------------------------------
+ host(
+ q,
+ k,
+ v,
+ raw_gate,
+ a_log,
+ dt_bias,
+ beta,
+ cu_seqlens,
+ seq_order,
+ state_indices,
+ initial_state,
+ out,
+ final_state,
+ stream_a,
+ scale,
+ state_ckpt,
+ cu_ckpts,
+ checkpoint_state_indices,
+ checkpoint_stride_chunks,
+ SAFE_GATE,
+ GATE_SCALE_LOG2,
+ THREADS,
+ gate_dtype,
+ )
+
+
+def _unified_fakes(
+ dtype: type,
+ state_dtype: type,
+ has_state_in: bool,
+ has_state_out: bool,
+ has_state_ckpt: bool,
+ has_state_indices: bool,
+ has_checkpoint_state_indices: bool,
+ gate_dtype: type = cutlass.BFloat16,
+) -> tuple:
+ """Union fake-tensor set for the single `host_unified` compile.
+
+ Independent symbols where lengths are unrelated: cu_chunks (`scu2`) is NOT
+ tied to cu_seqlens (`scu`), so the engine route can pass a length-1 dummy;
+ ws_qk / ws_diag share `sc` (chunk count), while ws_kd/qd/w use the
+ 16-divisible `sc16` (chunk*BT rows).
+ """
+ F = make_fake_compact_tensor
+ scu = cute.sym_int64(divisibility=1)
+ scu2 = cute.sym_int64(divisibility=1)
+ sn = cute.sym_int64(divisibility=1)
+ sp = cute.sym_int64(divisibility=1)
+ sh = cute.sym_int64(divisibility=1)
+ ss = cute.sym_int64(divisibility=1)
+ sc = cute.sym_int64(divisibility=1)
+ sc16 = cute.sym_int64(divisibility=16)
+ fq = F(dtype, (1, ss, sh, DK), stride_order=(3, 2, 1, 0), assumed_align=16)
+ fk = F(dtype, (1, ss, sh, DK), stride_order=(3, 2, 1, 0), assumed_align=16)
+ fv = F(dtype, (1, ss, sh, DV), stride_order=(3, 2, 1, 0), assumed_align=16)
+ fgate = F(gate_dtype, (1, ss, sh, DK), stride_order=(3, 2, 1, 0), assumed_align=16)
+ fa_log = F(cutlass.Float32, (sh,), stride_order=(0,), assumed_align=16)
+ fdt = F(cutlass.Float32, (sh, DK), stride_order=(1, 0), assumed_align=16)
+ fbeta = F(cutlass.BFloat16, (1, ss, sh), stride_order=(2, 1, 0), assumed_align=16)
+ fcu = F(cutlass.Int64, (scu,), stride_order=(0,), assumed_align=8)
+ forder = F(cutlass.Int32, (sn,), stride_order=(0,), assumed_align=8)
+ fstate_indices = (
+ F(cutlass.Int32, (sn,), stride_order=(0,), assumed_align=4)
+ if has_state_indices
+ else None
+ )
+ fcuc = F(cutlass.Int32, (scu2,), stride_order=(0,), assumed_align=8)
+
+ def ws_tile():
+ return F(dtype, (1, sh, sc16, DK), stride_order=(3, 2, 1, 0), assumed_align=16)
+
+ fkd, fqd, fw = ws_tile(), ws_tile(), ws_tile()
+ fqk = F(
+ dtype, (1, sh, sc, QK_REC_ELEMS), stride_order=(3, 2, 1, 0), assumed_align=16
+ )
+ fdiag = F(
+ cutlass.Float32,
+ (1, sh, sc, DIAG_REC_ELEMS),
+ stride_order=(3, 2, 1, 0),
+ assumed_align=16,
+ )
+
+ def state_pool_fake():
+ return cute.runtime.make_fake_tensor(
+ state_dtype,
+ shape=(sp, sh, DV, DK),
+ stride=(cute.sym_int64(divisibility=8), DV * DK, DK, 1),
+ assumed_align=16,
+ )
+
+ fstate = state_pool_fake() if has_state_in else None
+ fout = F(dtype, (1, ss, sh, DV), stride_order=(3, 2, 1, 0), assumed_align=16)
+ ffinal = state_pool_fake() if has_state_out else None
+ # State checkpoints: a flat [total_ckpts, H, DV, DK] tensor (same VK layout
+ # as final_state) plus per-sequence base offsets; own symbols so capacity is
+ # unrelated to n_seq.
+ # The symbols are only minted for checkpoint builds: sym_int64() registers
+ # into the shared symbol table, and a dangling entry shifts every kernel's
+ # constant-bank layout of the non-checkpoint build.
+ if has_state_ckpt:
+ sk = cute.sym_int64(divisibility=1)
+ scuk = cute.sym_int64(divisibility=1)
+ fckpt = F(
+ state_dtype, (sk, sh, DV, DK), stride_order=(3, 2, 1, 0), assumed_align=16
+ )
+ fcuk = F(cutlass.Int64, (scuk,), stride_order=(0,), assumed_align=8)
+ fcheckpoint_state_indices = (
+ F(
+ cutlass.Int32,
+ (cute.sym_int64(divisibility=1),),
+ stride_order=(0,),
+ assumed_align=4,
+ )
+ if has_checkpoint_state_indices
+ else None
+ )
+ else:
+ fckpt = None
+ fcuk = None
+ fcheckpoint_state_indices = None
+ return (
+ (
+ fq,
+ fk,
+ fv,
+ fgate,
+ fa_log,
+ fdt,
+ fbeta,
+ fcu,
+ forder,
+ 0, # generate_seq_order
+ fstate_indices,
+ fcuc,
+ 0, # generate_cu_chunks
+ fkd,
+ fqd,
+ fw,
+ fqk,
+ fdiag,
+ fstate,
+ fout,
+ ffinal,
+ ),
+ fckpt,
+ fcuk,
+ fcheckpoint_state_indices,
+ )
+
+
+_ENGINE_DUMMIES: dict = {}
+
+
+def _engine_dummies(heads: int, device, dtype: type):
+ """Minimal length-1 workspace views for the engine route.
+
+ The unified ABI takes the workspace tiles unconditionally; the engine branch
+ never reads them, so a cached length-1 set (satisfying the compiled symbolic
+ shape constraints) is passed instead of allocating a real workspace.
+ """
+ torch_dtype = DTYPE_MAP[dtype]
+ key = (heads, str(device), torch_dtype)
+ d = _ENGINE_DUMMIES.get(key)
+ if d is None:
+ kd = torch.zeros(1, heads, 16, DK, dtype=torch_dtype, device=device)
+ qk = torch.zeros(1, heads, 1, QK_REC_ELEMS, dtype=torch_dtype, device=device)
+ diag = torch.zeros(
+ 1, heads, 1, DIAG_REC_ELEMS, dtype=torch.float32, device=device
+ )
+ cuc = torch.zeros(1, dtype=torch.int32, device=device)
+ d = (kd, kd, kd, qk, diag, cuc)
+ _ENGINE_DUMMIES[key] = d
+ return d
+
+
+_CKPT_OFFSETS: dict = {}
+
+
+_UNIFORM_CU_CACHE: dict = {}
+
+
+_IDENTITY_SEQUENCE_ORDER_CACHE: dict = {}
+_DEVICE_SEQUENCE_ORDER_CACHE: dict = {}
+
+
+def _stream_pointer(stream) -> int:
+ value = getattr(stream, "value", stream)
+ return int(value or 0)
+
+
+def _device_sequence_order_buffer(num_sequences: int, device, stream) -> torch.Tensor:
+ """Reusable per-stream output for device-side packed sequence sorting."""
+
+ key = (str(device), _stream_pointer(stream), int(num_sequences))
+ order = _DEVICE_SEQUENCE_ORDER_CACHE.get(key)
+ if order is None:
+ if torch.cuda.is_current_stream_capturing():
+ raise RuntimeError(
+ "KDA device sequence order must be warmed before CUDA graph capture"
+ )
+ order = torch.empty(num_sequences, dtype=torch.int32, device=device)
+ _DEVICE_SEQUENCE_ORDER_CACHE[key] = order
+ return order
+
+
+def _identity_sequence_order(num_sequences: int, device) -> torch.Tensor:
+ """Cached original sequence order for the low-level callable."""
+
+ key = (str(device), int(num_sequences))
+ order = _IDENTITY_SEQUENCE_ORDER_CACHE.get(key)
+ if order is None:
+ if torch.cuda.is_current_stream_capturing():
+ raise RuntimeError(
+ "KDA identity sequence order must be warmed before CUDA graph capture"
+ )
+ order = torch.arange(num_sequences, dtype=torch.int32, device=device)
+ _IDENTITY_SEQUENCE_ORDER_CACHE[key] = order
+ return order
+
+
+def _uniform_cu_seqlens(batch: int, seqlen: int, device) -> torch.Tensor:
+ """Cached uniform cu_seqlens for the batched [B, T, H, D] entry."""
+
+ key = (batch, seqlen, str(device))
+ cu = _UNIFORM_CU_CACHE.get(key)
+ if cu is None:
+ cu = torch.arange(
+ 0, (batch + 1) * seqlen, seqlen, dtype=torch.int64, device=device
+ )
+ _UNIFORM_CU_CACHE[key] = cu
+ return cu
+
+
+def _make_call(unified: Callable, spec: dict) -> CompiledKDA:
+ """Build the single user host callable over the one compiled `host_unified`.
+
+ Runtime ABI: (q, k, v, raw_gate, a_log, dt_bias, beta, cu_seqlens,
+ initial_state, out, final_state, workspace, stream_a, scale, ..., seq_order)
+ — packed engine calls with ``seq_order=None`` generate an LPT order on the
+ device; fixed and decomp calls retain the original sequence order unless a
+ wrapper requests device-generated metadata. An explicit packed CUDA int32
+ permutation overrides the implicit engine order. The remaining
+ positional ABI is
+ the reference host ABI (stream_a == the reference `stream`, scale in the same
+ position) plus the `workspace` operand. Decomp always runs prep-first on
+ overlap. `.workspace_size(cu_seqlens, heads)` is attached. There is NO
+ runtime `mode`: the route (engine/decomp) is fixed by the compile-time `mode`
+ constexpr — None (default) keeps the runtime occupancy rule, "engine"/"decomp"
+ specialize the build to one route.
+
+ Both routes launch on the caller's single stream.
+ """
+
+ decisions: dict = {}
+ compile_mode = spec["mode"]
+ dtype = spec["dtype"]
+ built_states = (
+ spec["has_state_in"],
+ spec["has_state_out"],
+ spec["has_state_ckpt"],
+ spec["has_state_indices"],
+ spec.get("has_checkpoint_state_indices", False),
+ )
+
+ def call(
+ q,
+ k,
+ v,
+ raw_gate,
+ a_log,
+ dt_bias,
+ beta,
+ cu_seqlens,
+ initial_state,
+ out,
+ final_state,
+ workspace,
+ stream_a,
+ scale=DEFAULT_SCALE,
+ state_ckpt=None,
+ checkpoint_cu_starts=None,
+ ckpt_interval=0,
+ seq_order=None,
+ planned_cu_chunks=None,
+ planned_total_chunks=None,
+ generate_planned_metadata=False,
+ state_indices=None,
+ checkpoint_state_indices=None,
+ ):
+ if checkpoint_state_indices is not None and state_ckpt is None:
+ raise ValueError("checkpoint_state_indices requires state_ckpt")
+ # The state specialization is DERIVED from which state tensors the call
+ # actually passes (initial/final/checkpoint present or None); if it
+ # differs from this build, transparently delegate to the matching one
+ # (compile() is cached -- the first such call pays one compilation).
+ want = (
+ initial_state is not None,
+ final_state is not None,
+ state_ckpt is not None,
+ state_indices is not None,
+ checkpoint_state_indices is not None,
+ )
+ if want != built_states:
+ sibling = compile(
+ dtype=spec["dtype"],
+ state_dtype=spec["state_dtype"],
+ gate_dtype=spec["gate_dtype"],
+ safe_gate=spec["safe_gate"],
+ gate_lower_bound=spec["gate_lower_bound"],
+ has_state_in=want[0],
+ has_state_out=want[1],
+ has_state_ckpt=want[2],
+ has_state_indices=want[3],
+ has_checkpoint_state_indices=want[4],
+ mode=spec["mode"],
+ )
+ return sibling(
+ q,
+ k,
+ v,
+ raw_gate,
+ a_log,
+ dt_bias,
+ beta,
+ cu_seqlens,
+ initial_state,
+ out,
+ final_state,
+ workspace,
+ stream_a,
+ scale,
+ state_indices=state_indices,
+ state_ckpt=state_ckpt,
+ checkpoint_cu_starts=checkpoint_cu_starts,
+ checkpoint_state_indices=checkpoint_state_indices,
+ ckpt_interval=ckpt_interval,
+ seq_order=seq_order,
+ planned_cu_chunks=planned_cu_chunks,
+ planned_total_chunks=planned_total_chunks,
+ generate_planned_metadata=generate_planned_metadata,
+ )
+ # One stream. There used to be an optional `stream_b` that opted into a
+ # co-resident k1/k2 overlap; it was removed because the flag-ring it
+ # required cost more than the concurrency it bought (see host_unified).
+ # Decomp is now always prep-first on `stream_a`.
+ # Token-major user convention: cu_seqlens given -> packed
+ # [1, T_total, H, D] activations (beta [1, T_total, H]); cu_seqlens
+ # None -> batched [B, T, H, D] (beta [B, T, H]), flattened zero-copy
+ # into a packed call with synthesized uniform boundaries. The tensors
+ # feed the kernels directly in token-major memory order: the hosts
+ # build every TMA descriptor from explicit (shape, stride) tuples, so
+ # no logical permute is needed anywhere.
+ packed_layout = cu_seqlens is not None
+ if not packed_layout:
+ bsz = int(q.shape[0])
+ tlen = int(q.shape[1])
+ cu_seqlens = _uniform_cu_seqlens(bsz, tlen, q.device)
+ q = q.reshape(1, bsz * tlen, q.shape[2], q.shape[3])
+ k = k.reshape(1, bsz * tlen, k.shape[2], k.shape[3])
+ v = v.reshape(1, bsz * tlen, v.shape[2], v.shape[3])
+ raw_gate = raw_gate.reshape(
+ 1, bsz * tlen, raw_gate.shape[2], raw_gate.shape[3]
+ )
+ beta = beta.reshape(1, bsz * tlen, beta.shape[2])
+ out = out.reshape(1, bsz * tlen, out.shape[2], out.shape[3])
+ heads = int(q.shape[2])
+ device = q.device
+ sm_count = _device_sm_count(device)
+ n_seq = cu_seqlens.numel() - 1
+ if state_indices is not None and (
+ initial_state is None
+ or state_indices.device != device
+ or state_indices.dtype != torch.int32
+ or state_indices.ndim != 1
+ or not state_indices.is_contiguous()
+ or state_indices.numel() != n_seq
+ ):
+ raise ValueError(
+ "state_indices requires initial_state and must be a contiguous "
+ "CUDA int32 tensor with one entry per sequence"
+ )
+ if checkpoint_state_indices is not None and (
+ state_ckpt is None
+ or checkpoint_state_indices.device != device
+ or checkpoint_state_indices.dtype != torch.int32
+ or checkpoint_state_indices.ndim != 1
+ or not checkpoint_state_indices.is_contiguous()
+ ):
+ raise ValueError(
+ "checkpoint_state_indices requires state_ckpt and must be a "
+ "contiguous CUDA int32 tensor"
+ )
+ key = (n_seq, str(device), heads, compile_mode)
+ kind = decisions.get(key)
+ if kind is None:
+ kind = _route_for_workspace(n_seq, heads, device, compile_mode or "auto")
+ decisions[key] = kind
+ generate_seq_order = bool(generate_planned_metadata)
+ if seq_order is None:
+ if packed_layout and kind == "engine":
+ generate_seq_order = True
+ seq_order = _device_sequence_order_buffer(n_seq, device, stream_a)
+ else:
+ seq_order = _identity_sequence_order(n_seq, device)
+ elif (
+ seq_order.device != device
+ or seq_order.dtype != torch.int32
+ or seq_order.ndim != 1
+ or not seq_order.is_contiguous()
+ or seq_order.numel() != n_seq
+ ):
+ raise ValueError(
+ "seq_order must be a contiguous CUDA int32 tensor with one "
+ "entry per sequence"
+ )
+ has_planned_chunks = (
+ planned_cu_chunks is not None and planned_total_chunks is not None
+ )
+ if (planned_cu_chunks is None) != (planned_total_chunks is None):
+ raise ValueError(
+ "planned_cu_chunks and planned_total_chunks must be provided together"
+ )
+ if generate_planned_metadata and not has_planned_chunks:
+ raise ValueError(
+ "device-generated chunk metadata requires planned_cu_chunks and "
+ "planned_total_chunks capacity"
+ )
+ cu_list = None
+ if (kind == "decomp" and not has_planned_chunks) or (
+ state_ckpt is not None and checkpoint_cu_starts is None
+ ):
+ cu_list = _cu_seqlens_contents(cu_seqlens)
+ checkpoint_stride_chunks = 0
+ if state_ckpt is not None:
+ if ckpt_interval <= 0 or ckpt_interval % BT:
+ raise ValueError(
+ f"ckpt_interval must be a positive multiple of {BT}, "
+ f"got {ckpt_interval}"
+ )
+ if checkpoint_cu_starts is not None and (
+ checkpoint_cu_starts.device != device
+ or checkpoint_cu_starts.dtype != torch.int64
+ or checkpoint_cu_starts.ndim != 1
+ or not checkpoint_cu_starts.is_contiguous()
+ or checkpoint_cu_starts.numel() != n_seq + 1
+ ):
+ raise ValueError(
+ "checkpoint_cu_starts must be a contiguous CUDA int64 "
+ "tensor with one entry per sequence plus one"
+ )
+ checkpoint_stride_chunks = ckpt_interval // BT
+ if checkpoint_cu_starts is None:
+ assert cu_list is not None
+ offsets, total = [0], 0
+ for i in range(n_seq):
+ seq_len = cu_list[i + 1] - cu_list[i]
+ if checkpoint_state_indices is None:
+ total += (seq_len + ckpt_interval - 1) // ckpt_interval
+ else:
+ total += seq_len // ckpt_interval
+ offsets.append(total)
+ cu_ckpts_arg = torch.tensor(offsets, dtype=torch.int64, device=device)
+ else:
+ cu_ckpts_arg = checkpoint_cu_starts
+ else:
+ cu_ckpts_arg = None
+ seq_route = 1 # single stream; kept as an arg so host_unified re-derives
+ # The route rule lives in exactly two places: `_route_for_workspace`
+ # (which `workspace_size` uses, so the caller allocates the right
+ # buffer) and `host_unified` (which re-derives it to pick the branch to
+ # execute). This wrapper does NOT own a third copy -- it only asks the
+ # workspace helper which operands to marshal.
+ if kind == "engine":
+ # beta rides a TMA descriptor built over the compact [T, H] layout:
+ # the base must be 16B-aligned (cuTensorMapEncodeTiled hard rule)
+ # and the memory contiguous. Fresh allocator tensors always pass;
+ # an odd-token-offset VIEW at heads % 8 != 0 is the one real way
+ # to violate it -- fail loudly with the one-line fix.
+ if (not beta.is_contiguous()) or beta.data_ptr() % 16:
+ raise ValueError(
+ "beta must be contiguous [1, T, H] with a 16-byte-aligned "
+ "base for the TMA beta transport; pass the unsliced tensor "
+ "or call .contiguous()"
+ )
+ if (
+ heads % 8
+ and heads
+ * (BT + 8 // (heads & -heads if heads & -heads < 8 else 8))
+ * 2
+ > 2 * BETA_TILE_STAGE_ELEMS
+ ):
+ raise ValueError(
+ f"heads={heads}: the pair-packed beta tile exceeds the SMEM "
+ "stage; supported are heads % 8 == 0 or heads <= 14"
+ )
+ kd, qd, wt, qk, diag, cuc = _engine_dummies(heads, device, dtype)
+ unified(
+ q,
+ k,
+ v,
+ raw_gate,
+ a_log,
+ dt_bias,
+ beta,
+ cu_seqlens,
+ seq_order,
+ int(generate_seq_order),
+ state_indices,
+ cuc,
+ 0, # engine does not consume a chunk prefix
+ kd,
+ qd,
+ wt,
+ qk,
+ diag,
+ initial_state,
+ out,
+ final_state,
+ stream_a,
+ seq_route,
+ sm_count,
+ checkpoint_stride_chunks,
+ 1,
+ 1,
+ scale,
+ state_ckpt,
+ cu_ckpts_arg,
+ checkpoint_state_indices,
+ )
+ return
+
+ # ---- decomp route ----
+ if workspace is None:
+ raise ValueError(
+ "workspace required for this shape; call workspace_size() and allocate"
+ )
+ if has_planned_chunks:
+ if (
+ planned_cu_chunks.device != device
+ or planned_cu_chunks.dtype != torch.int32
+ or planned_cu_chunks.ndim != 1
+ or not planned_cu_chunks.is_contiguous()
+ or planned_cu_chunks.numel() != n_seq + 1
+ or not isinstance(planned_total_chunks, int)
+ or planned_total_chunks <= 0
+ ):
+ raise ValueError(
+ "planned chunk metadata must be a contiguous CUDA int32 "
+ "cu_chunks[N+1] tensor and a positive host total_chunks"
+ )
+ plan = {
+ "total_chunks": planned_total_chunks,
+ "cu_chunks": planned_cu_chunks,
+ }
+ else:
+ plan = _plan(cu_seqlens)
+ total_chunks = plan["total_chunks"]
+ ws = _partition_workspace(workspace, heads, total_chunks)
+ # Decomp is prep-first on one stream: k1 runs its full grid, then k2
+ # consumes what k1 left behind. The co-resident overlap variant that
+ # used to live here was removed because it never paid: enabling it forced
+ # the per-chunk flag publish/wait onto BOTH kernels, and that tax exceeds
+ # the concurrency it buys. Measured on the historical dual-variant
+ # build (flag-free and flag-carrying kernels side by side), dropping
+ # the flags makes prep 30-44% cheaper (H32-T8192 130.7 -> 91.2 us,
+ # H12-2048x4 73.6 -> 41.5) and 4 of 5 decomp shapes tie or improve
+ # end-to-end (H12-2048x4 -22.6%, H32-T8192 -13.0%) against a single 1.3%
+ # loss on H12-varlen -- with verify() bit-identical between the two.
+ num_ctas, cpc = _k1_grid(total_chunks, heads)
+ unified(
+ q,
+ k,
+ v,
+ raw_gate,
+ a_log,
+ dt_bias,
+ beta,
+ cu_seqlens,
+ seq_order,
+ int(generate_seq_order),
+ state_indices,
+ plan["cu_chunks"],
+ int(generate_planned_metadata),
+ ws["kd"],
+ ws["qd"],
+ ws["w"],
+ ws["qk"],
+ ws["diag"],
+ initial_state,
+ out,
+ final_state,
+ stream_a,
+ seq_route,
+ sm_count,
+ checkpoint_stride_chunks,
+ num_ctas,
+ cpc,
+ scale,
+ state_ckpt,
+ cu_ckpts_arg,
+ checkpoint_state_indices,
+ )
+
+ def _ws_size(cu_seqlens, heads, mode=None, *, batch=None, seqlen=None):
+ if cu_seqlens is None:
+ if batch is None or seqlen is None:
+ raise ValueError(
+ "workspace_size(None, heads) needs batch= and seqlen= "
+ "for the non-varlen [B, T, H, D] entry"
+ )
+ cu_seqlens = tuple(int(seqlen) * i for i in range(int(batch) + 1))
+ return workspace_size(cu_seqlens, heads, mode or (compile_mode or "auto"))
+
+ def _ws_size_from_total_chunks(num_sequences, heads, total_chunks, device):
+ route_mode = compile_mode or "auto"
+ if _route_for_workspace(num_sequences, heads, device, route_mode) == "engine":
+ return 0
+ return _decomp_ws_bytes(int(heads), int(total_chunks))
+
+ compiled_call = cast(CompiledKDA, call)
+ compiled_call.workspace_size = _ws_size
+ compiled_call.workspace_size_from_total_chunks = _ws_size_from_total_chunks
+ return compiled_call
+
+
+@lru_cache(maxsize=None)
+def compile( # noqa: A001
+ dtype: type = cutlass.BFloat16,
+ state_dtype: type = cutlass.Float32,
+ gate_dtype: type = cutlass.BFloat16,
+ safe_gate: bool = True,
+ gate_lower_bound: float = DEFAULT_GATE_LOWER_BOUND,
+ has_state_in: bool = True,
+ has_state_out: bool = True,
+ has_state_ckpt: bool = False,
+ has_state_indices: bool = False,
+ has_checkpoint_state_indices: bool = False,
+ mode=None,
+) -> CompiledKDA:
+ """Compile KDA; the returned callable is the single host entry.
+
+ The signature matches the reference `kda_chunked_bt16.compile`; the only
+ user-visible additions are the decomp `workspace` (a single opaque uint8
+ buffer sized by the attached `.workspace_size(cu_seqlens, heads)` and passed
+ before the stream).
+
+ ONE `cute.compile(host_unified, ...)` target. `mode` is a COMPILE-TIME
+ constexpr fixing the route:
+ * None (default) -> the returned callable routes engine vs the two-kernel
+ decomposition at RUNTIME by the occupancy rule (decomp-DV2 iff
+ `n_seq*heads*2 <= sm_count`, else engine); both paths are in the one .o.
+ * "engine" -> engine-only build (decomp branch const-folds away).
+ * "decomp" -> decomp-only build (engine branch const-folds away).
+ ("auto" is accepted as an alias for None.)
+ """
+
+ if dtype not in CLI_DTYPES.values():
+ raise ValueError(f"Unsupported dtype: {dtype}")
+ if state_dtype not in (cutlass.BFloat16, cutlass.Float32):
+ raise ValueError(f"Unsupported state dtype: {state_dtype}")
+ if has_checkpoint_state_indices and not has_state_ckpt:
+ raise ValueError("checkpoint state indices require checkpoint state output")
+ validate_gate_dtype(gate_dtype)
+ if mode == "auto":
+ mode = None
+ if mode not in (None, "engine", "decomp"):
+ raise ValueError(f"Unknown mode: {mode}")
+ gate_scale_log2 = gate_lower_bound * LOG2_E
+
+ fakes, fckpt, fcuk, fcheckpoint_state_indices = _unified_fakes(
+ dtype,
+ state_dtype,
+ has_state_in,
+ has_state_out,
+ has_state_ckpt,
+ has_state_indices,
+ has_checkpoint_state_indices,
+ gate_dtype,
+ )
+ unified = cute.compile(
+ host_unified,
+ *fakes,
+ make_fake_stream(),
+ 0, # seq_route
+ 0, # sm_count
+ 0, # checkpoint_stride_chunks
+ 0, # prep_num_ctas
+ 0, # prep_cpc
+ DEFAULT_SCALE,
+ fckpt,
+ fcuk,
+ fcheckpoint_state_indices,
+ safe_gate,
+ gate_scale_log2,
+ THREADS_PER_CTA,
+ mode, # MODE constexpr (route specialization)
+ gate_dtype, # gate ABI dtype constexpr (32-bit vs 16-bit gate path)
+ options="--enable-tvm-ffi",
+ )
+ fn = _make_call(
+ unified,
+ dict(
+ dtype=dtype,
+ state_dtype=state_dtype,
+ gate_dtype=gate_dtype,
+ safe_gate=safe_gate,
+ gate_lower_bound=gate_lower_bound,
+ has_state_in=has_state_in,
+ has_state_out=has_state_out,
+ has_state_ckpt=has_state_ckpt,
+ has_state_indices=has_state_indices,
+ has_checkpoint_state_indices=has_checkpoint_state_indices,
+ mode=mode,
+ ),
+ )
+ # Eager device-module load: keep the module resident from build time with one
+ # tiny decomp launch, so no runtime call pays a lazy cuModuleLoad. Guarded so
+ # a login-node dry-run (no GPU) still returns the compiled callable.
+ if torch.cuda.is_available():
+ _eager_module_load(
+ fn,
+ dtype,
+ state_dtype,
+ has_state_in,
+ has_state_out,
+ has_state_ckpt,
+ has_state_indices,
+ has_checkpoint_state_indices,
+ gate_dtype,
+ )
+ return fn
+
+
+def _eager_module_load(
+ fn: CompiledKDA,
+ dtype: type,
+ state_dtype: type,
+ has_state_in: bool,
+ has_state_out: bool,
+ has_state_ckpt: bool = False,
+ has_state_indices: bool = False,
+ has_checkpoint_state_indices: bool = False,
+ gate_dtype: type = cutlass.BFloat16,
+) -> None:
+ """One launch to make the device module resident so no real call lazy-loads.
+
+ heads=1 keeps free_sms >= 0, so a MODE=None build routes this to DECOMP and
+ loads the prep+chain cubins (an engine-routed preload would not).
+ MODE="engine" routes it to engine (the only cubin there); MODE="decomp" to
+ decomp — either way the single module's cubins become resident."""
+ try:
+ device = torch.cuda.current_device()
+ heads = 1
+ seqlen = 64 * BT # a few chunks: enough to exercise prep + chain
+ q, k, v, raw_gate, a_log, dt_bias, beta, cu_seqlens, state = _make_inputs(
+ 1,
+ heads,
+ seqlen,
+ dtype,
+ state_dtype,
+ device,
+ random_state=False,
+ gate_dtype=gate_dtype,
+ )
+ out = torch.zeros_like(v)
+ initial_state = state if has_state_in else None
+ final_state = torch.zeros_like(state) if has_state_out else None
+ ws_bytes = fn.workspace_size(cu_seqlens, heads)
+ workspace = (
+ torch.empty(ws_bytes, dtype=torch.uint8, device=cu_seqlens.device)
+ if ws_bytes > 0
+ else None
+ )
+ s = cuda_driver.CUstream(torch.cuda.current_stream().cuda_stream)
+ # Prep-first decomp launch: loads the single module's cubins eagerly.
+ fn(
+ q,
+ k,
+ v,
+ raw_gate,
+ a_log,
+ dt_bias,
+ beta,
+ cu_seqlens,
+ initial_state,
+ out,
+ final_state,
+ workspace,
+ s,
+ DEFAULT_SCALE,
+ state_indices=(
+ torch.zeros(1, dtype=torch.int32, device="cuda")
+ if has_state_indices
+ else None
+ ),
+ state_ckpt=(
+ torch.zeros(
+ 64,
+ heads,
+ DV,
+ DK,
+ dtype=DTYPE_MAP[state_dtype],
+ device="cuda",
+ )
+ if has_state_ckpt
+ else None
+ ),
+ ckpt_interval=(BT if has_state_ckpt else 0),
+ checkpoint_state_indices=(
+ torch.arange(64, dtype=torch.int32, device="cuda")
+ if has_checkpoint_state_indices
+ else None
+ ),
+ )
+ torch.cuda.synchronize()
+ except Exception: # noqa: BLE001 — best-effort module preload; never fatal
+ pass
+
+
+def _make_inputs(
+ batch: int,
+ heads: int,
+ seqlen: int,
+ dtype: type,
+ state_dtype: type,
+ device: str,
+ random_state: bool = False,
+ safe_gate: bool = True,
+ seq_lens: tuple[int, ...] | None = None,
+ gate_dtype: type = cutlass.BFloat16,
+) -> tuple[
+ torch.Tensor,
+ torch.Tensor,
+ torch.Tensor,
+ torch.Tensor,
+ torch.Tensor,
+ torch.Tensor,
+ torch.Tensor,
+ torch.Tensor,
+ torch.Tensor,
+]:
+ if seq_lens is None:
+ seq_lens = (seqlen,) * batch
+ if len(seq_lens) != batch or any(length < 0 for length in seq_lens):
+ raise ValueError("seq_lens must contain one non-negative length per sequence")
+ total_seqlen = sum(seq_lens)
+ torch_dtype = DTYPE_MAP[dtype]
+ torch_state_dtype = DTYPE_MAP[state_dtype]
+ # The gate is materialized in FP32 (RNG parity with the historical harness,
+ # byte-for-byte identical draws) and then ROUNDED to the ABI dtype. The
+ # torch reference consumes THIS tensor, so it models the input rounding
+ # exactly while keeping every internal gate op in FP32.
+ torch_gate_dtype = DTYPE_MAP[gate_dtype]
+ q = 0.25 * torch.randn(1, total_seqlen, heads, DK, dtype=torch_dtype, device=device)
+ k = 0.25 * torch.randn(1, total_seqlen, heads, DK, dtype=torch_dtype, device=device)
+ v = torch.randn(1, total_seqlen, heads, DV, dtype=torch_dtype, device=device)
+ if safe_gate:
+ raw_gate = (
+ 0.25
+ * torch.randn(
+ 1,
+ total_seqlen,
+ heads,
+ DK,
+ dtype=torch.float32,
+ device=device,
+ )
+ ).to(torch_gate_dtype)
+ # Safe path draws are unchanged (zeros consume no RNG): reproducible
+ # byte-for-byte vs the original harness.
+ a_log = torch.zeros(heads, dtype=torch.float32, device=device)
+ dt_bias = torch.zeros(heads, DK, dtype=torch.float32, device=device)
+ else:
+ # FLA non-safe gate: RAW standard-normal gate logits plus non-trivial
+ # A_log and dt_bias so the in-kernel softplus activation is exercised.
+ raw_gate = torch.randn(
+ 1,
+ total_seqlen,
+ heads,
+ DK,
+ dtype=torch.float32,
+ device=device,
+ ).to(torch_gate_dtype)
+ a_log = torch.rand(heads, dtype=torch.float32, device=device)
+ dt_bias = torch.rand(heads, DK, dtype=torch.float32, device=device)
+ beta = torch.randn(
+ 1,
+ total_seqlen,
+ heads,
+ dtype=torch.bfloat16,
+ device=device,
+ )
+ cu_seqlens = torch.zeros(batch + 1, dtype=torch.int64, device=device)
+ if batch > 0:
+ cu_seqlens[1:] = torch.tensor(
+ seq_lens, dtype=torch.int64, device=device
+ ).cumsum(0)
+ if random_state:
+ state = 0.05 * torch.randn(
+ batch, heads, DV, DK, dtype=torch_state_dtype, device=device
+ )
+ else:
+ state = torch.zeros(
+ batch,
+ heads,
+ DV,
+ DK,
+ dtype=torch_state_dtype,
+ device=device,
+ )
+ return q, k, v, raw_gate, a_log, dt_bias, beta, cu_seqlens, state
+
+
+def run(
+ compiled_fn: CompiledKDA,
+ batch: int = DEFAULT_BATCH,
+ heads: int = DEFAULT_HEADS,
+ seqlen: int = DEFAULT_SEQLEN,
+ dtype: type = cutlass.BFloat16,
+ state_dtype: type = cutlass.Float32,
+ gate_dtype: type = cutlass.BFloat16,
+ has_state_in: bool = True,
+ has_state_out: bool = True,
+ seq_lens: tuple[int, ...] | None = None,
+ scale: float = DEFAULT_SCALE,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+ """Run packed KDA with runtime scale and optionally return the VK state.
+
+ This public runtime helper is also re-exported by the Rubin specialization.
+ ``compiled_fn`` must use the same state-presence specialization selected by
+ ``has_state_in`` and ``has_state_out``.
+ """
+
+ if not torch.cuda.is_available():
+ raise RuntimeError("CUDA GPU required")
+ q, k, v, raw_gate, a_log, dt_bias, beta, cu_seqlens, state = _make_inputs(
+ batch,
+ heads,
+ seqlen,
+ dtype,
+ state_dtype,
+ device="cuda",
+ seq_lens=seq_lens,
+ gate_dtype=gate_dtype,
+ )
+ out = torch.zeros_like(v)
+ initial_state = state if has_state_in else None
+ final_state = torch.zeros_like(state) if has_state_out else None
+
+ # User-allocated opaque workspace: query bytes for the route this host will
+ # take, then allocate a raw (uninitialized) buffer. The host initializes it.
+ ws_bytes = compiled_fn.workspace_size(cu_seqlens, heads)
+ workspace = (
+ torch.empty(ws_bytes, dtype=torch.uint8, device=cu_seqlens.device)
+ if ws_bytes > 0
+ else None
+ )
+
+ stream = torch.cuda.current_stream(cu_seqlens.device).cuda_stream
+ compiled_fn(
+ q,
+ k,
+ v,
+ raw_gate,
+ a_log,
+ dt_bias,
+ beta,
+ cu_seqlens,
+ initial_state,
+ out,
+ final_state,
+ workspace,
+ stream,
+ scale,
+ )
+ torch.cuda.synchronize()
+ return out, final_state
+
+
+def verify(
+ batch: int = DEFAULT_BATCH,
+ heads: int = DEFAULT_HEADS,
+ seqlen: int = BT,
+ dtype: type = cutlass.BFloat16,
+ state_dtype: type = cutlass.Float32,
+ gate_dtype: type = cutlass.BFloat16,
+ compile_only: bool = True,
+ safe_gate: bool = True,
+ has_state_in: bool = True,
+ has_state_out: bool = True,
+ seq_lens: tuple[int, ...] | None = None,
+ scale: float = DEFAULT_SCALE,
+ mode: str = "engine",
+) -> None:
+ """Compile the kernel and optionally compare a device run to the MMA reference."""
+
+ compiled_fn = compile(
+ dtype=dtype,
+ state_dtype=state_dtype,
+ gate_dtype=gate_dtype,
+ safe_gate=safe_gate,
+ has_state_in=has_state_in,
+ has_state_out=has_state_out,
+ mode=mode,
+ )
+ print(
+ f"Compile KDA kernel (BT={BT}, runtime seqlen, dtype={dtype}, "
+ f"gate_dtype={gate_dtype}) OK"
+ )
+
+ if compile_only:
+ return
+
+ if not torch.cuda.is_available():
+ raise RuntimeError("CUDA GPU required")
+
+ q, k, v, raw_gate, a_log, dt_bias, beta, cu_seqlens, state = _make_inputs(
+ batch,
+ heads,
+ seqlen,
+ dtype,
+ state_dtype,
+ device="cuda",
+ random_state=True,
+ safe_gate=safe_gate,
+ seq_lens=seq_lens,
+ gate_dtype=gate_dtype,
+ )
+ out = torch.zeros_like(v)
+ initial_state = state if has_state_in else None
+ final_state = torch.zeros_like(state) if has_state_out else None
+
+ ws_bytes = compiled_fn.workspace_size(cu_seqlens, heads)
+ workspace = (
+ torch.empty(ws_bytes, dtype=torch.uint8, device=cu_seqlens.device)
+ if ws_bytes > 0
+ else None
+ )
+
+ stream_a = torch.cuda.current_stream(cu_seqlens.device).cuda_stream
+ compiled_fn(
+ q,
+ k,
+ v,
+ raw_gate,
+ a_log,
+ dt_bias,
+ beta,
+ cu_seqlens,
+ initial_state,
+ out,
+ final_state,
+ workspace,
+ stream_a,
+ scale,
+ )
+ torch.cuda.synchronize()
+
+ state_kv = _state_vk_to_kv(state) if has_state_in else None
+ # Reference keeps the head-major logical convention; permute the
+ # token-major user tensors to zero-copy [1, H, T, D] / [1, H, T] views.
+ out_ref, state_ref_kv = kda_chunked_mma_reference(
+ q.permute(0, 2, 1, 3),
+ k.permute(0, 2, 1, 3),
+ v.permute(0, 2, 1, 3),
+ raw_gate.permute(0, 2, 1, 3),
+ beta.permute(0, 2, 1),
+ state_kv,
+ mma_dtype=DTYPE_MAP[dtype],
+ a_log=a_log,
+ dt_bias=dt_bias,
+ safe_gate=safe_gate,
+ cu_seqlens=cu_seqlens,
+ scale=scale,
+ )
+ state_ref_vk = _state_kv_to_vk(state_ref_kv)
+ verify_atol = 1.0e-2 if dtype is cutlass.BFloat16 else VERIFY_ATOL
+ out_hm = out.permute(0, 2, 1, 3)
+ torch.testing.assert_close(
+ out_hm.float(), out_ref.float(), rtol=VERIFY_RTOL, atol=verify_atol
+ )
+ if has_state_out:
+ torch.testing.assert_close(
+ final_state,
+ state_ref_vk.to(final_state.dtype),
+ rtol=VERIFY_RTOL,
+ atol=verify_atol,
+ )
+ out_diff = (out_hm.float() - out_ref.float()).abs()
+ state_diff_text = "disabled"
+ if has_state_out:
+ state_diff = (final_state.float() - state_ref_vk.float()).abs()
+ state_diff_text = f"{state_diff.max().item():.6g}"
+ print(
+ "Run KDA staged-SMEM/TMA dtype-rounded MMA reference check: PASS "
+ f"max_out_abs={out_diff.max().item():.6g}, "
+ f"max_state_abs={state_diff_text}"
+ )
+
+
+def _parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(
+ description="BT=16 KDA CUTLASS primitives kernel",
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+ )
+ parser.add_argument("--batch", type=int, default=DEFAULT_BATCH)
+ parser.add_argument("--heads", type=int, default=DEFAULT_HEADS)
+ parser.add_argument("--seqlen", type=int, default=BT)
+ parser.add_argument("--scale", type=float, default=DEFAULT_SCALE)
+ parser.add_argument("--dtype", choices=sorted(CLI_DTYPES), default="bf16")
+ parser.add_argument(
+ "--state-dtype",
+ choices=sorted(CLI_STATE_DTYPES),
+ default="fp32",
+ )
+ parser.add_argument(
+ "--gate-dtype",
+ choices=sorted(CLI_GATE_DTYPES),
+ default="bf16",
+ help="gate input ABI dtype (bf16 = FlashKDA-aligned default)",
+ )
+ parser.add_argument(
+ "--no-initial-state",
+ action="store_true",
+ help="initialize the recurrent state to zero inside the kernel",
+ )
+ parser.add_argument(
+ "--no-final-state",
+ action="store_true",
+ help="skip the final recurrent-state global store",
+ )
+ parser.add_argument(
+ "--disable-safe-gate",
+ action="store_true",
+ help="interpret raw gate input as a precomputed log2-domain increment",
+ )
+ parser.add_argument(
+ "--run-kernel",
+ action="store_true",
+ help="run the current staged-SMEM/TMA device path after compile",
+ )
+ parser.add_argument(
+ "--mode",
+ choices=("engine", "decomp", "auto"),
+ default="engine",
+ help="engine = plain kernel; decomp = two-kernel decomposition; "
+ "auto = two-way shape router (plain engine / decomp-DV2)",
+ )
+ return parser.parse_args()
+
+
+if __name__ == "__main__":
+ args = _parse_args()
+ verify(
+ batch=args.batch,
+ heads=args.heads,
+ seqlen=args.seqlen,
+ dtype=CLI_DTYPES[args.dtype],
+ state_dtype=CLI_STATE_DTYPES[args.state_dtype],
+ gate_dtype=CLI_GATE_DTYPES[args.gate_dtype],
+ compile_only=not args.run_kernel,
+ safe_gate=not args.disable_safe_gate,
+ has_state_in=not args.no_initial_state,
+ has_state_out=not args.no_final_state,
+ scale=args.scale,
+ mode=args.mode,
+ )
diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_chunked_bt16_small_bh.py b/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_chunked_bt16_small_bh.py
new file mode 100644
index 000000000000..7478dc054a5a
--- /dev/null
+++ b/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_chunked_bt16_small_bh.py
@@ -0,0 +1,3196 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2026 by FlashInfer team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""FlashInfer Small-BH KDA, with indexed recurrent-state pool access.
+
+Vendored from flashinfer-ai/flashinfer commit
+b8e81a7dae6c2aa5c9aa19c06e81e9a729271493,
+flashinfer/kda_kernels/kda_chunked_bt16_small_bh.py (Apache-2.0).
+Local changes add state-slot indirection and padded pool-row strides, reject
+unwarmed graph specializations, and use the active device SM count for scheduling.
+The K1/K2 arithmetic and layouts are retained from upstream.
+"""
+
+import functools
+from typing import Optional
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.pipeline as pipeline
+import cutlass.utils as utils
+import torch
+from cutlass._mlir.dialects import llvm
+from cutlass.cute.core import _pack_shape
+from cutlass.cute.nvgpu import cpasync, warp
+from cutlass.cutlass_dsl import dsl_user_op
+
+
+RCP_LN2 = 1.4426950408889634
+
+
+def layout_separate(thr, src, ref):
+ lt = cute.make_layout(())
+ ge = cute.make_layout(())
+
+ for k, v in enumerate(ref):
+ if cutlass.const_expr(v < thr):
+ lt = cute.append(lt, src[k])
+ else:
+ ge = cute.append(ge, src[k])
+
+ r = None
+ if cutlass.const_expr(cute.rank(lt) == 1):
+ r = cute.append(lt, ge)
+ else:
+ r = cute.append(cute.append(cute.make_layout(()), lt), ge)
+ return r
+
+
+@cute.jit
+def layout_acc_mn(tiled_mma, acc):
+ separated = layout_separate(
+ tiled_mma.shape_mnk[0], acc[0], tiled_mma.tv_layout_C.stride[1]
+ )
+
+ V_M = separated[0]
+ V_N = separated[1]
+ V_M1 = None
+ V_N1 = None
+ if cutlass.const_expr(cute.rank(V_M) == 1):
+ V_M1 = cute.append(V_M, acc[1])
+ else:
+ V_M1 = cute.append(cute.append(cute.make_layout(()), V_M), acc[1])
+
+ if cutlass.const_expr(cute.rank(V_N) == 1):
+ V_N1 = cute.append(V_N, acc[2])
+ else:
+ V_N1 = cute.append(cute.append(cute.make_layout(()), V_N), acc[2])
+
+ r = None
+ if cutlass.const_expr(cute.rank(V_M1) == 1):
+ r = cute.append(V_M1, V_N1)
+ else:
+ r = cute.append(cute.append(cute.make_layout(()), V_M1), V_N1)
+ return r
+
+
+@cute.jit
+def get_seq_len(cu_seqlens: cute.Tensor, b_idx: cutlass.Int32) -> cutlass.Int32:
+ return cu_seqlens[b_idx + 1] - cu_seqlens[b_idx]
+
+
+@cute.jit
+def get_seq_offset(
+ cu_seqlens: cute.Tensor,
+ b_idx: cutlass.Int32,
+ *,
+ padding: cutlass.Int32 = 0,
+ index_unit: cutlass.Int32 = 1,
+) -> cutlass.Int32:
+ """Return a packed sequence offset in tokens or token-sized records.
+
+ Unpadded tensors start sequence ``b`` at ``cu_seqlens[b]``. Padded
+ intermediates reserve ``padding`` tokens per preceding sequence and start
+ it at ``cu_seqlens[b] + b * padding``. ``index_unit`` converts that token
+ offset to another token-indexed storage unit, such as BT-sized records.
+ """
+ offset = cu_seqlens[b_idx]
+ if cutlass.const_expr(padding > 0):
+ offset += b_idx * padding
+ if cutlass.const_expr(index_unit != 1):
+ offset = offset // index_unit
+ return offset
+
+
+@dsl_user_op
+def tanh(
+ x: cutlass.Float32,
+ *,
+ loc=None,
+ ip=None,
+) -> cutlass.Float32:
+ return cutlass.Float32(
+ llvm.inline_asm(
+ cutlass.Float32.mlir_type,
+ [x.ir_value(loc=loc, ip=ip)],
+ "tanh.approx.f32 $0, $1;",
+ "=f,f",
+ has_side_effects=False,
+ is_align_stack=False,
+ asm_dialect=llvm.AsmDialect.AD_ATT,
+ )
+ )
+
+
+@dsl_user_op
+def movmatrix(
+ x: cutlass.Uint32,
+ *,
+ loc=None,
+ ip=None,
+) -> cutlass.Uint32:
+ return cutlass.Uint32(
+ llvm.inline_asm(
+ cutlass.Uint32.mlir_type,
+ [x.ir_value(loc=loc, ip=ip)],
+ "movmatrix.sync.aligned.m8n8.trans.b16 $0, $1;",
+ "=r,r",
+ has_side_effects=True,
+ is_align_stack=False,
+ asm_dialect=llvm.AsmDialect.AD_ATT,
+ )
+ )
+
+
+class ChunkKdaFwdK1:
+ def __init__(
+ self,
+ D: int = 128,
+ H: int = 32,
+ BT: int = 16,
+ gate_lower_bound: bool = False,
+ max_active_clusters: int = 152,
+ ):
+ self.D = D
+ self.H = H
+ self.BT = BT
+ self.gate_lower_bound = gate_lower_bound
+ self.max_active_clusters = max_active_clusters
+
+ self.CT = 4
+ self.TT = self.CT * self.BT
+ self.load_stages = 2
+ self.gate_mma_stages = 2
+
+ self.beta_tma_heads = 8
+ # H-major BF16 beta needs a 16-byte contiguous TMA dimension
+ self.beta_use_tma = self.H % self.beta_tma_heads == 0
+ self.beta_stages = (
+ self.load_stages + self.gate_mma_stages if self.beta_use_tma else 4
+ )
+
+ self.warp_threads = 32
+ self.warpgroup_threads = 128
+ self.warps_per_warpgroup = self.warpgroup_threads // self.warp_threads
+
+ self.threads_per_cta = self.warpgroup_threads * 3
+ self.load_warp_id = 0
+ self.beta_warp_id = 1
+ self.aux_warpgroup_id = 0
+ self.gate_warpgroup_id = 1
+ self.mma_warpgroup_id = 2
+
+ self.load_register_requirement = 48
+ self.gate_register_requirement = 256
+ self.mma_register_requirement = 200
+
+ @cute.jit
+ def __call__(
+ self,
+ mQ: cute.Tensor, # (B, T, H, D)
+ mK: cute.Tensor, # (B, T, H, D)
+ mG: cute.Tensor, # (B, T, H, D)
+ mBeta: cute.Tensor, # (B, T, H)
+ mDtBias: cute.Tensor, # (H, D)
+ mAlog: cute.Tensor, # (H,)
+ mQd: cute.Tensor, # (B, T, H, D)
+ mKd: cute.Tensor, # (B, T, H, D)
+ mKr: cute.Tensor, # (B, T, H, D)
+ mMqk: cute.Tensor, # (B, NT, H, BT, BT)
+ mMkk: cute.Tensor, # (B, NT, H, BT, BT)
+ mGk: cute.Tensor, # (B, NT, H, D)
+ scale: cutlass.Float32,
+ gate_scale: cutlass.Float32,
+ mCuSeqlens: Optional[cute.Tensor] = None, # (B + 1,)
+ max_seqlen: cutlass.Int32 = 0,
+ stream=None,
+ ):
+ if cutlass.const_expr(mCuSeqlens is not None):
+ B = cute.size(mCuSeqlens) - 1
+ num_tiles = cute.ceil_div(max_seqlen, self.TT)
+ mQ = cute.make_tensor(
+ mQ.iterator, cute.select(mQ.layout, mode=[0, 2, 1])
+ ) # (T, D, H)
+ mK = cute.make_tensor(
+ mK.iterator, cute.select(mK.layout, mode=[0, 2, 1])
+ ) # (T, D, H)
+ mG = cute.make_tensor(
+ mG.iterator, cute.select(mG.layout, mode=[0, 2, 1])
+ ) # (T, D, H)
+ mBeta = cute.make_tensor(
+ mBeta.iterator, cute.select(mBeta.layout, mode=[0, 1])
+ ) # (T, H)
+ mQd = cute.make_tensor(
+ mQd.iterator, cute.select(mQd.layout, mode=[0, 2, 1])
+ ) # (Tp, D, H)
+ mKd = cute.make_tensor(
+ mKd.iterator, cute.select(mKd.layout, mode=[0, 2, 1])
+ ) # (Tp, D, H)
+ mKr = cute.make_tensor(
+ mKr.iterator, cute.select(mKr.layout, mode=[0, 2, 1])
+ ) # (Tp, D, H)
+ mMqk = cute.make_tensor(
+ mMqk.iterator, cute.select(mMqk.layout, mode=[2, 3, 0, 1])
+ ) # (BT, BT, NT, H)
+ mMkk = cute.make_tensor(
+ mMkk.iterator, cute.select(mMkk.layout, mode=[2, 3, 0, 1])
+ ) # (BT, BT, NT, H)
+ mGk = cute.make_tensor(
+ mGk.iterator, cute.select(mGk.layout, mode=[2, 0, 1])
+ ) # (D, NT, H)
+ else:
+ B, _, _, _ = mQ.shape
+ NT = cute.size(mMqk, mode=[1])
+ num_tiles = cute.ceil_div(NT, self.CT)
+ mQ = cute.make_tensor(
+ mQ.iterator, cute.select(mQ.layout, mode=[1, 3, 2, 0])
+ ) # (T, D, H, B)
+ mK = cute.make_tensor(
+ mK.iterator, cute.select(mK.layout, mode=[1, 3, 2, 0])
+ ) # (T, D, H, B)
+ mG = cute.make_tensor(
+ mG.iterator, cute.select(mG.layout, mode=[1, 3, 2, 0])
+ ) # (T, D, H, B)
+ mBeta = cute.make_tensor(
+ mBeta.iterator, cute.select(mBeta.layout, mode=[1, 2, 0])
+ ) # (T, H, B)
+ mQd = cute.make_tensor(
+ mQd.iterator, cute.select(mQd.layout, mode=[1, 3, 2, 0])
+ ) # (T, D, H, B)
+ mKd = cute.make_tensor(
+ mKd.iterator, cute.select(mKd.layout, mode=[1, 3, 2, 0])
+ ) # (T, D, H, B)
+ mKr = cute.make_tensor(
+ mKr.iterator, cute.select(mKr.layout, mode=[1, 3, 2, 0])
+ ) # (T, D, H, B)
+ mMqk = cute.make_tensor(
+ mMqk.iterator, cute.select(mMqk.layout, mode=[3, 4, 1, 2, 0])
+ ) # (BT, BT, NT, H, B)
+ mMkk = cute.make_tensor(
+ mMkk.iterator, cute.select(mMkk.layout, mode=[3, 4, 1, 2, 0])
+ ) # (BT, BT, NT, H, B)
+ mGk = cute.make_tensor(
+ mGk.iterator, cute.select(mGk.layout, mode=[3, 1, 2, 0])
+ ) # (D, NT, H, B)
+ mDtBias = cute.make_tensor(
+ mDtBias.iterator, cute.select(mDtBias.layout, mode=[1, 0])
+ ) # (H, D)
+
+ tiled_mma = cute.make_tiled_mma(
+ warp.MmaF16BF16Op(mQ.element_type, cutlass.Float32, (16, 8, 16)),
+ permutation_mnk=(self.BT, self.BT, 16),
+ )
+ copy_atom_QKd = cute.make_copy_atom(
+ cute.nvgpu.warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4),
+ mQ.element_type,
+ )
+ tiled_copy_QKd = cute.make_tiled_copy_A(copy_atom_QKd, tiled_mma)
+ tiled_copy_Ki = cute.make_tiled_copy_B(copy_atom_QKd, tiled_mma)
+
+ sQKG_layout_staged = cute.make_layout(
+ (self.TT, self.D, self.load_stages),
+ stride=(self.D, 1, self.TT * self.D),
+ )
+ sQKG_layout = cute.slice_(sQKG_layout_staged, (None, None, 0))
+
+ if cutlass.const_expr(self.beta_use_tma):
+ sBeta_layout_staged = cute.make_layout(
+ (self.TT, self.beta_tma_heads, self.beta_stages),
+ stride=(self.beta_tma_heads, 1, self.TT * self.beta_tma_heads),
+ )
+ sBeta_layout = cute.slice_(sBeta_layout_staged, (None, None, 0))
+ else:
+ sBeta_layout_staged = cute.make_layout(
+ (self.TT, self.beta_stages),
+ stride=(1, self.TT),
+ )
+ sBeta_layout = cute.slice_(sBeta_layout_staged, (None, 0))
+
+ sDtBias_layout_staged = cute.make_layout(
+ (self.D, self.load_stages),
+ stride=(1, self.D),
+ )
+ sDtBias_layout = cute.slice_(sDtBias_layout_staged, (None, 0))
+
+ sAlog_layout = cute.make_layout((self.H,))
+
+ smem_layout_atom_sw128 = cute.nvgpu.warpgroup.make_smem_layout_atom(
+ cute.nvgpu.warpgroup.SmemLayoutAtomKind.K_SW128,
+ mQd.element_type,
+ )
+ sQKd_layout_staged = cute.tile_to_shape(
+ smem_layout_atom_sw128,
+ (self.TT, self.D, self.gate_mma_stages),
+ order=(0, 1, 2),
+ )
+ sQKd_layout = cute.slice_(sQKd_layout_staged, (None, None, 0))
+
+ buffer_align_bytes = 1024
+
+ @cute.struct
+ class SharedStorage:
+ sQ: cute.struct.Align[
+ cute.struct.MemRange[mQ.element_type, cute.cosize(sQKG_layout_staged)],
+ buffer_align_bytes,
+ ]
+ sK: cute.struct.Align[
+ cute.struct.MemRange[mK.element_type, cute.cosize(sQKG_layout_staged)],
+ buffer_align_bytes,
+ ]
+ sG: cute.struct.Align[
+ cute.struct.MemRange[mG.element_type, cute.cosize(sQKG_layout_staged)],
+ buffer_align_bytes,
+ ]
+ sBeta: cute.struct.Align[
+ cute.struct.MemRange[
+ mBeta.element_type, cute.cosize(sBeta_layout_staged)
+ ],
+ buffer_align_bytes,
+ ]
+ sDtBias: cute.struct.Align[
+ cute.struct.MemRange[
+ mDtBias.element_type, cute.cosize(sDtBias_layout_staged)
+ ],
+ buffer_align_bytes,
+ ]
+ sAlog: cute.struct.Align[
+ cute.struct.MemRange[mAlog.element_type, cute.cosize(sAlog_layout)],
+ buffer_align_bytes,
+ ]
+ sQd: cute.struct.Align[
+ cute.struct.MemRange[mQ.element_type, cute.cosize(sQKd_layout_staged)],
+ buffer_align_bytes,
+ ]
+ sKd: cute.struct.Align[
+ cute.struct.MemRange[mK.element_type, cute.cosize(sQKd_layout_staged)],
+ buffer_align_bytes,
+ ]
+ sKi: cute.struct.Align[
+ cute.struct.MemRange[mK.element_type, cute.cosize(sQKd_layout_staged)],
+ buffer_align_bytes,
+ ]
+ tma_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_stages * 2]
+ gate_mma_mbar_ptr: cute.struct.MemRange[
+ cutlass.Int64, self.gate_mma_stages * 2
+ ]
+ beta_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.beta_stages * 2]
+
+ beta_tma_load_bytes = (
+ cute.size_in_bytes(mBeta.element_type, sBeta_layout)
+ if self.beta_use_tma
+ else 0
+ )
+ self.num_tma_load_bytes = (
+ cute.size_in_bytes(mQ.element_type, sQKG_layout)
+ + cute.size_in_bytes(mK.element_type, sQKG_layout)
+ + cute.size_in_bytes(mG.element_type, sQKG_layout)
+ + beta_tma_load_bytes
+ + cute.size_in_bytes(mDtBias.element_type, sDtBias_layout)
+ )
+
+ tma_atom_Q, tma_tensor_Q = cpasync.make_tiled_tma_atom(
+ cpasync.CopyBulkTensorTileG2SOp(),
+ mQ,
+ sQKG_layout,
+ (self.TT, self.D),
+ )
+ tma_atom_K, tma_tensor_K = cpasync.make_tiled_tma_atom(
+ cpasync.CopyBulkTensorTileG2SOp(),
+ mK,
+ sQKG_layout,
+ (self.TT, self.D),
+ )
+ tma_atom_G, tma_tensor_G = cpasync.make_tiled_tma_atom(
+ cpasync.CopyBulkTensorTileG2SOp(),
+ mG,
+ sQKG_layout,
+ (self.TT, self.D),
+ )
+ tma_atom_Beta, tma_tensor_Beta = None, None
+ if cutlass.const_expr(self.beta_use_tma):
+ tma_atom_Beta, tma_tensor_Beta = cpasync.make_tiled_tma_atom(
+ cpasync.CopyBulkTensorTileG2SOp(),
+ mBeta,
+ sBeta_layout,
+ (self.TT, self.beta_tma_heads),
+ )
+ tma_atom_dtBias, tma_tensor_dtBias = cpasync.make_tiled_tma_atom(
+ cpasync.CopyBulkTensorTileG2SOp(),
+ mDtBias,
+ sDtBias_layout,
+ (self.D,),
+ )
+ tma_atom_Qd, tma_tensor_Qd = cpasync.make_tiled_tma_atom(
+ cpasync.CopyBulkTensorTileS2GOp(),
+ mQd,
+ sQKd_layout,
+ (self.TT, self.D),
+ )
+ tma_atom_Kd, tma_tensor_Kd = cpasync.make_tiled_tma_atom(
+ cpasync.CopyBulkTensorTileS2GOp(),
+ mKd,
+ sQKd_layout,
+ (self.TT, self.D),
+ )
+
+ qkg_elems_per_copy = 8
+ copy_atom_QKG = cute.make_copy_atom(
+ cute.nvgpu.CopyUniversalOp(),
+ mQ.element_type,
+ num_bits_per_copy=qkg_elems_per_copy * mQ.element_type.width,
+ )
+ tv_layout_QKG = cute.make_layout(
+ (
+ (self.warp_threads // 2, 2),
+ (qkg_elems_per_copy, self.BT // 2),
+ ),
+ stride=(
+ (qkg_elems_per_copy * self.BT, self.BT // 2),
+ (self.BT, 1),
+ ),
+ )
+ tiled_copy_QKG = cute.make_tiled_copy(
+ copy_atom_QKG, tv_layout_QKG, (self.BT, self.D)
+ )
+
+ copy_atom_dtBias = cute.make_copy_atom(
+ cute.nvgpu.CopyUniversalOp(),
+ mDtBias.element_type,
+ num_bits_per_copy=qkg_elems_per_copy * mQ.element_type.width,
+ )
+ copy_atom_Gk = cute.make_copy_atom(
+ cute.nvgpu.CopyUniversalOp(),
+ mGk.element_type,
+ num_bits_per_copy=2 * mGk.element_type.width,
+ )
+ tv_layout_dtBias_Gk = cute.make_layout(
+ ((self.warp_threads // 2, 2), qkg_elems_per_copy),
+ stride=((qkg_elems_per_copy, 0), 1),
+ )
+ tiled_copy_dtBias = cute.make_tiled_copy(
+ copy_atom_dtBias, tv_layout_dtBias_Gk, (self.D,)
+ )
+ tiled_copy_Gk = cute.make_tiled_copy(
+ copy_atom_Gk, tv_layout_dtBias_Gk, (self.D,)
+ )
+
+ tile_sched_params = utils.PersistentTileSchedulerParams(
+ problem_shape_ntile_mnl=(num_tiles, self.H, B),
+ cluster_shape_mnk=(1, 1, 1),
+ )
+ grid = utils.StaticPersistentTileScheduler.get_grid_shape(
+ tile_sched_params, self.max_active_clusters
+ )
+
+ self.kernel(
+ tma_atom_Q,
+ tma_tensor_Q,
+ tma_atom_K,
+ tma_tensor_K,
+ tma_atom_G,
+ tma_tensor_G,
+ mBeta,
+ tma_atom_Beta,
+ tma_tensor_Beta,
+ tma_atom_dtBias,
+ tma_tensor_dtBias,
+ mAlog,
+ tma_atom_Qd,
+ tma_tensor_Qd,
+ tma_atom_Kd,
+ tma_tensor_Kd,
+ mKr,
+ mMqk,
+ mMkk,
+ mGk,
+ mCuSeqlens,
+ scale,
+ gate_scale,
+ sQKG_layout_staged,
+ sBeta_layout_staged,
+ sDtBias_layout_staged,
+ sAlog_layout,
+ sQKd_layout_staged,
+ tiled_mma,
+ tiled_copy_QKG,
+ tiled_copy_dtBias,
+ tiled_copy_Gk,
+ tiled_copy_QKd,
+ tiled_copy_Ki,
+ tile_sched_params,
+ SharedStorage,
+ ).launch(
+ grid=grid,
+ block=(self.threads_per_cta, 1, 1),
+ min_blocks_per_mp=1,
+ stream=stream,
+ )
+
+ @cute.kernel
+ def kernel(
+ self,
+ tma_atom_Q: cute.CopyAtom,
+ tma_tensor_Q: cute.Tensor,
+ tma_atom_K: cute.CopyAtom,
+ tma_tensor_K: cute.Tensor,
+ tma_atom_G: cute.CopyAtom,
+ tma_tensor_G: cute.Tensor,
+ mBeta: cute.Tensor,
+ tma_atom_Beta: cute.CopyAtom,
+ tma_tensor_Beta: cute.Tensor,
+ tma_atom_dtBias: cute.CopyAtom,
+ tma_tensor_dtBias: cute.Tensor,
+ mAlog: cute.Tensor,
+ tma_atom_Qd: cute.CopyAtom,
+ tma_tensor_Qd: cute.Tensor,
+ tma_atom_Kd: cute.CopyAtom,
+ tma_tensor_Kd: cute.Tensor,
+ mKr: cute.Tensor,
+ mMqk: cute.Tensor,
+ mMkk: cute.Tensor,
+ mGk: cute.Tensor,
+ mCuSeqlens: Optional[cute.Tensor],
+ scale: cutlass.Float32,
+ gate_scale: cutlass.Float32,
+ sQKG_layout_staged: cute.Layout,
+ sBeta_layout_staged: cute.Layout,
+ sDtBias_layout_staged: cute.Layout,
+ sAlog_layout: cute.Layout,
+ sQKd_layout_staged: cute.ComposedLayout,
+ tiled_mma: cute.TiledMma,
+ tiled_copy_QKG: cute.TiledCopy,
+ tiled_copy_dtBias: cute.TiledCopy,
+ tiled_copy_Gk: cute.TiledCopy,
+ tiled_copy_QKd: cute.TiledCopy,
+ tiled_copy_Ki: cute.TiledCopy,
+ tile_sched_params: utils.PersistentTileSchedulerParams,
+ SharedStorage: cutlass.Constexpr,
+ ):
+ tidx, _, _ = cute.arch.thread_idx()
+ tidx_in_wg = tidx % self.warpgroup_threads
+ lane_idx = cute.arch.lane_idx()
+ warp_idx = cute.arch.make_warp_uniform(tidx // self.warp_threads)
+ widx_in_wg = warp_idx % (self.warpgroup_threads // self.warp_threads)
+ warpgroup_idx = cute.arch.make_warp_uniform(tidx // self.warpgroup_threads)
+
+ smem = cutlass.utils.SmemAllocator()
+ storage = smem.allocate(SharedStorage)
+ sQ = storage.sQ.get_tensor(sQKG_layout_staged)
+ sK = storage.sK.get_tensor(sQKG_layout_staged)
+ sG = storage.sG.get_tensor(sQKG_layout_staged)
+ sBeta = storage.sBeta.get_tensor(sBeta_layout_staged)
+ sDtBias = storage.sDtBias.get_tensor(sDtBias_layout_staged)
+ sAlog = storage.sAlog.get_tensor(sAlog_layout)
+ sQd = storage.sQd.get_tensor(
+ sQKd_layout_staged.outer, swizzle=sQKd_layout_staged.inner
+ )
+ sKd = storage.sKd.get_tensor(
+ sQKd_layout_staged.outer, swizzle=sQKd_layout_staged.inner
+ )
+ sKi = storage.sKi.get_tensor(
+ sQKd_layout_staged.outer, swizzle=sQKd_layout_staged.inner
+ )
+
+ tma_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+ tma_consumer_group = pipeline.CooperativeGroup(
+ pipeline.Agent.Thread, self.warps_per_warpgroup
+ )
+ tma_producer, tma_consumer = pipeline.PipelineTmaAsync.create(
+ num_stages=self.load_stages,
+ producer_group=tma_producer_group,
+ consumer_group=tma_consumer_group,
+ barrier_storage=storage.tma_mbar_ptr.data_ptr(),
+ tx_count=self.num_tma_load_bytes,
+ ).make_participants()
+
+ beta_group = cutlass.pipeline.CooperativeGroup(
+ cutlass.pipeline.Agent.Thread, self.warp_threads
+ )
+ gate_mma_group = cutlass.pipeline.CooperativeGroup(
+ cutlass.pipeline.Agent.Thread, self.warpgroup_threads
+ )
+
+ beta_producer, beta_consumer = None, None
+ if cutlass.const_expr(not self.beta_use_tma):
+ beta_producer, beta_consumer = pipeline.PipelineAsync.create(
+ num_stages=self.beta_stages,
+ producer_group=beta_group,
+ consumer_group=gate_mma_group,
+ barrier_storage=storage.beta_mbar_ptr.data_ptr(),
+ ).make_participants()
+
+ gate_mma_producer, gate_mma_consumer = cutlass.pipeline.PipelineAsync.create(
+ num_stages=self.gate_mma_stages,
+ producer_group=gate_mma_group,
+ consumer_group=gate_mma_group,
+ barrier_storage=storage.gate_mma_mbar_ptr.data_ptr(),
+ ).make_participants()
+
+ store_pipeline = pipeline.PipelineTmaStore.create(
+ num_stages=1,
+ producer_group=gate_mma_group,
+ )
+
+ if warpgroup_idx == self.aux_warpgroup_id:
+ cute.arch.setmaxregister_decrease(self.load_register_requirement)
+
+ if warp_idx == self.load_warp_id:
+ beta_idx = 0
+
+ tile_sched = utils.StaticPersistentTileScheduler.create(
+ tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
+ )
+ work_tile = tile_sched.initial_work_tile_info()
+
+ while work_tile.is_valid_tile:
+ t_idx, h_idx, b_idx = work_tile.tile_idx
+
+ is_valid_tile = True
+ if cutlass.const_expr(mCuSeqlens is not None):
+ is_valid_tile = t_idx * self.TT < get_seq_len(mCuSeqlens, b_idx)
+ if is_valid_tile:
+ beta_head_tile = h_idx // self.beta_tma_heads
+ gBeta = None
+ if cutlass.const_expr(mCuSeqlens is not None):
+ seq_offset = get_seq_offset(mCuSeqlens, b_idx)
+
+ mQ_seq = cute.domain_offset(
+ (seq_offset, 0, 0), tma_tensor_Q
+ )
+ mK_seq = cute.domain_offset(
+ (seq_offset, 0, 0), tma_tensor_K
+ )
+ mG_seq = cute.domain_offset(
+ (seq_offset, 0, 0), tma_tensor_G
+ )
+
+ gQ = cute.local_tile(
+ mQ_seq[None, None, h_idx], (self.TT, self.D), (None, 0)
+ )
+ gK = cute.local_tile(
+ mK_seq[None, None, h_idx], (self.TT, self.D), (None, 0)
+ )
+ gG = cute.local_tile(
+ mG_seq[None, None, h_idx], (self.TT, self.D), (None, 0)
+ )
+ if cutlass.const_expr(self.beta_use_tma):
+ mBeta_seq = cute.domain_offset(
+ (seq_offset, 0), tma_tensor_Beta
+ )
+ gBeta = cute.local_tile(
+ mBeta_seq,
+ (self.TT, self.beta_tma_heads),
+ (None, beta_head_tile),
+ )
+ else:
+ gQ = cute.local_tile(
+ tma_tensor_Q[None, None, h_idx, b_idx],
+ (self.TT, self.D),
+ (None, 0),
+ )
+ gK = cute.local_tile(
+ tma_tensor_K[None, None, h_idx, b_idx],
+ (self.TT, self.D),
+ (None, 0),
+ )
+ gG = cute.local_tile(
+ tma_tensor_G[None, None, h_idx, b_idx],
+ (self.TT, self.D),
+ (None, 0),
+ )
+ if cutlass.const_expr(self.beta_use_tma):
+ gBeta = cute.local_tile(
+ tma_tensor_Beta[None, None, b_idx],
+ (self.TT, self.beta_tma_heads),
+ (None, beta_head_tile),
+ )
+ gDtBias = tma_tensor_dtBias[None, h_idx]
+
+ tQsQ, tQgQ = cpasync.tma_partition(
+ tma_atom_Q,
+ 0,
+ cute.make_layout(1),
+ cute.group_modes(sQ, 0, 2),
+ cute.group_modes(gQ, 0, 2),
+ )
+ tKsK, tKgK = cpasync.tma_partition(
+ tma_atom_K,
+ 0,
+ cute.make_layout(1),
+ cute.group_modes(sK, 0, 2),
+ cute.group_modes(gK, 0, 2),
+ )
+ tGsG, tGgG = cpasync.tma_partition(
+ tma_atom_G,
+ 0,
+ cute.make_layout(1),
+ cute.group_modes(sG, 0, 2),
+ cute.group_modes(gG, 0, 2),
+ )
+ if cutlass.const_expr(gBeta is not None):
+ tBeta_sBeta, tBeta_gBeta = cpasync.tma_partition(
+ tma_atom_Beta,
+ 0,
+ cute.make_layout(1),
+ cute.group_modes(sBeta, 0, 2),
+ cute.group_modes(gBeta, 0, 2),
+ )
+ tDtBias_sDtBias, tDtBias_gDtBias = cpasync.tma_partition(
+ tma_atom_dtBias,
+ 0,
+ cute.make_layout(1),
+ cute.group_modes(sDtBias, 0, 1),
+ cute.group_modes(gDtBias, 0, 1),
+ )
+
+ tma_empty = tma_producer.acquire_and_advance()
+ cute.copy(
+ tma_atom_Q,
+ tQgQ[None, t_idx],
+ tQsQ[None, tma_empty.index],
+ tma_bar_ptr=tma_empty.barrier,
+ )
+ cute.copy(
+ tma_atom_K,
+ tKgK[None, t_idx],
+ tKsK[None, tma_empty.index],
+ tma_bar_ptr=tma_empty.barrier,
+ )
+ cute.copy(
+ tma_atom_G,
+ tGgG[None, t_idx],
+ tGsG[None, tma_empty.index],
+ tma_bar_ptr=tma_empty.barrier,
+ )
+ if cutlass.const_expr(self.beta_use_tma):
+ cute.copy(
+ tma_atom_Beta,
+ tBeta_gBeta[None, t_idx],
+ tBeta_sBeta[None, beta_idx],
+ tma_bar_ptr=tma_empty.barrier,
+ )
+ cute.copy(
+ tma_atom_dtBias,
+ tDtBias_gDtBias,
+ tDtBias_sDtBias[None, tma_empty.index],
+ tma_bar_ptr=tma_empty.barrier,
+ )
+ beta_idx = (beta_idx + 1) % self.beta_stages
+
+ tile_sched.advance_to_next_work()
+ work_tile = tile_sched.get_current_work()
+
+ elif warp_idx == self.beta_warp_id:
+ if cutlass.const_expr(not self.beta_use_tma):
+ tile_sched = utils.StaticPersistentTileScheduler.create(
+ tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
+ )
+ work_tile = tile_sched.initial_work_tile_info()
+
+ while work_tile.is_valid_tile:
+ t_idx, h_idx, b_idx = work_tile.tile_idx
+
+ if cutlass.const_expr(mCuSeqlens is not None):
+ seq_len = get_seq_len(mCuSeqlens, b_idx)
+ else:
+ seq_len = cute.size(mBeta, mode=[0])
+
+ is_valid_tile = True
+ if cutlass.const_expr(mCuSeqlens is not None):
+ is_valid_tile = t_idx * self.TT < seq_len
+ if is_valid_tile:
+ if cutlass.const_expr(mCuSeqlens is not None):
+ seq_offset = get_seq_offset(mCuSeqlens, b_idx)
+ mBeta_seq = cute.domain_offset(
+ (seq_offset,), mBeta[None, h_idx]
+ )
+ gBeta = cute.local_tile(mBeta_seq, (self.TT,), (t_idx,))
+ else:
+ gBeta = cute.local_tile(
+ mBeta[None, h_idx, b_idx], (self.TT,), (t_idx,)
+ )
+
+ valid_tokens = min(self.TT, seq_len - t_idx * self.TT)
+ if valid_tokens == self.TT:
+ beta_empty = beta_producer.acquire_and_advance()
+ for t in cutlass.range_constexpr(
+ self.TT // self.warp_threads
+ ):
+ idx = t * self.warp_threads + lane_idx
+ sBeta[idx, beta_empty.index] = gBeta[idx]
+ beta_empty.commit()
+ else:
+ beta_empty = beta_producer.acquire_and_advance()
+ for t in cutlass.range_constexpr(
+ self.TT // self.warp_threads
+ ):
+ idx = t * self.warp_threads + lane_idx
+ sBeta[idx, beta_empty.index] = (
+ gBeta[idx]
+ if idx < valid_tokens
+ else sBeta.element_type(0.0)
+ )
+ beta_empty.commit()
+
+ tile_sched.advance_to_next_work()
+ work_tile = tile_sched.get_current_work()
+
+ elif warpgroup_idx == self.gate_warpgroup_id:
+ cute.arch.setmaxregister_increase(self.gate_register_requirement)
+
+ alog_scale = (
+ 0.5 if cutlass.const_expr(self.gate_lower_bound) else -gate_scale
+ )
+ for h in cutlass.range_constexpr(
+ 0,
+ self.H // self.warpgroup_threads * self.warpgroup_threads,
+ self.warpgroup_threads,
+ ):
+ sAlog[h + tidx_in_wg] = alog_scale * cute.math.exp2(
+ mAlog[h + tidx_in_wg] * RCP_LN2, fastmath=True
+ )
+ if cutlass.const_expr(self.H % self.warpgroup_threads != 0):
+ if tidx_in_wg < self.H % self.warpgroup_threads:
+ h_idx = (
+ self.H // self.warpgroup_threads * self.warpgroup_threads
+ + tidx_in_wg
+ )
+ sAlog[h_idx] = alog_scale * cute.math.exp2(
+ mAlog[h_idx] * RCP_LN2, fastmath=True
+ )
+
+ gate_scale = gate_scale * 0.5 * RCP_LN2
+
+ cutlass.pipeline.arrive_and_wait(
+ barrier_id=1, num_threads=self.warpgroup_threads
+ )
+
+ tile_sched = utils.StaticPersistentTileScheduler.create(
+ tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
+ )
+ work_tile = tile_sched.initial_work_tile_info()
+
+ while work_tile.is_valid_tile:
+ t_idx, h_idx, b_idx = work_tile.tile_idx
+
+ seq_len = 0
+ seq_offset_padded = None
+ chunk_offset = None
+ if cutlass.const_expr(mCuSeqlens is not None):
+ seq_len = get_seq_len(mCuSeqlens, b_idx)
+ seq_offset_padded = get_seq_offset(
+ mCuSeqlens, b_idx, padding=self.TT
+ )
+ chunk_offset = get_seq_offset(
+ mCuSeqlens, b_idx, padding=self.TT, index_unit=self.BT
+ )
+ else:
+ seq_len = cute.size(tma_tensor_Q, mode=[0])
+
+ is_valid_tile = True
+ if cutlass.const_expr(mCuSeqlens is not None):
+ is_valid_tile = t_idx * self.TT < seq_len
+ if is_valid_tile:
+ valid_tokens = min(self.TT, seq_len - t_idx * self.TT)
+ if valid_tokens == self.TT:
+ tma_consumer, gate_mma_producer = self.compute_gate_tile(
+ mGk,
+ mKr,
+ sQ,
+ sK,
+ sG,
+ sDtBias,
+ sAlog,
+ sQd,
+ sKd,
+ sKi,
+ tma_consumer,
+ gate_mma_producer,
+ gate_scale,
+ scale,
+ t_idx,
+ h_idx,
+ b_idx,
+ tiled_copy_QKG,
+ tiled_copy_dtBias,
+ tiled_copy_Gk,
+ lane_idx,
+ widx_in_wg,
+ seq_offset_padded,
+ chunk_offset,
+ )
+ else:
+ tma_consumer, gate_mma_producer = self.compute_gate_tile(
+ mGk,
+ mKr,
+ sQ,
+ sK,
+ sG,
+ sDtBias,
+ sAlog,
+ sQd,
+ sKd,
+ sKi,
+ tma_consumer,
+ gate_mma_producer,
+ gate_scale,
+ scale,
+ t_idx,
+ h_idx,
+ b_idx,
+ tiled_copy_QKG,
+ tiled_copy_dtBias,
+ tiled_copy_Gk,
+ lane_idx,
+ widx_in_wg,
+ seq_offset_padded,
+ chunk_offset,
+ valid_tokens,
+ )
+
+ tile_sched.advance_to_next_work()
+ work_tile = tile_sched.get_current_work()
+
+ else:
+ cute.arch.setmaxregister_increase(self.mma_register_requirement)
+
+ thr_mma = tiled_mma.get_slice(lane_idx)
+ tAcA = thr_mma.partition_C(cute.make_identity_tensor((self.BT, self.BT)))
+ a_shape_L = thr_mma.partition_shape_A((self.BT, self.BT))
+ b_shape_L = thr_mma.partition_shape_B((self.BT, self.BT))
+ c_shape_L = thr_mma.partition_shape_C((self.BT, self.BT))
+ tCrI = thr_mma.make_fragment_C(c_shape_L)
+ for e in cutlass.range_constexpr(cute.size(tAcA)):
+ m, n = tAcA[e]
+ tCrI[e] = cutlass.Float32(1.0) if m == n else cutlass.Float32(0.0)
+
+ beta_idx = 0
+
+ tile_sched = utils.StaticPersistentTileScheduler.create(
+ tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
+ )
+ work_tile = tile_sched.initial_work_tile_info()
+
+ while work_tile.is_valid_tile:
+ t_idx, h_idx, b_idx = work_tile.tile_idx
+
+ is_valid_tile = True
+ if cutlass.const_expr(mCuSeqlens is not None):
+ seq_len = get_seq_len(mCuSeqlens, b_idx)
+ is_valid_tile = t_idx * self.TT < seq_len
+ if is_valid_tile:
+ if cutlass.const_expr(mCuSeqlens is not None):
+ seq_offset_padded = get_seq_offset(
+ mCuSeqlens, b_idx, padding=self.TT
+ )
+ mQd_seq = cute.domain_offset(
+ (seq_offset_padded, 0, 0), tma_tensor_Qd
+ )
+ mKd_seq = cute.domain_offset(
+ (seq_offset_padded, 0, 0), tma_tensor_Kd
+ )
+ gQd = cute.local_tile(
+ mQd_seq[None, None, h_idx], (self.TT, self.D), (None, 0)
+ )
+ gKd = cute.local_tile(
+ mKd_seq[None, None, h_idx], (self.TT, self.D), (None, 0)
+ )
+ else:
+ gQd = cute.local_tile(
+ tma_tensor_Qd[None, None, h_idx, b_idx],
+ (self.TT, self.D),
+ (None, 0),
+ )
+ gKd = cute.local_tile(
+ tma_tensor_Kd[None, None, h_idx, b_idx],
+ (self.TT, self.D),
+ (None, 0),
+ )
+
+ tQsQd, tQgQd = cpasync.tma_partition(
+ tma_atom_Qd,
+ 0,
+ cute.make_layout(1),
+ cute.group_modes(sQd, 0, 2),
+ cute.group_modes(gQd, 0, 2),
+ )
+ tKsKd, tKgKd = cpasync.tma_partition(
+ tma_atom_Kd,
+ 0,
+ cute.make_layout(1),
+ cute.group_modes(sKd, 0, 2),
+ cute.group_modes(gKd, 0, 2),
+ )
+
+ beta_full = None
+ if cutlass.const_expr(not self.beta_use_tma):
+ beta_full = beta_consumer.wait_and_advance()
+ beta_idx = beta_full.index
+ mma_full = gate_mma_consumer.wait_and_advance()
+
+ if widx_in_wg == 0:
+ cute.copy(
+ tma_atom_Qd,
+ tQsQd[None, mma_full.index],
+ tQgQd[None, t_idx],
+ )
+ cute.copy(
+ tma_atom_Kd,
+ tKsKd[None, mma_full.index],
+ tKgKd[None, t_idx],
+ )
+ store_pipeline.producer_commit()
+
+ if cutlass.const_expr(mCuSeqlens is not None):
+ chunk_offset = get_seq_offset(
+ mCuSeqlens, b_idx, padding=self.TT, index_unit=self.BT
+ )
+ mMqk_seq = cute.domain_offset((0, 0, chunk_offset, 0), mMqk)
+ mMkk_seq = cute.domain_offset((0, 0, chunk_offset, 0), mMkk)
+ mMqk_seq = mMqk_seq[None, None, None, h_idx]
+ mMkk_seq = mMkk_seq[None, None, None, h_idx]
+ else:
+ mMqk_seq = mMqk[None, None, None, h_idx, b_idx]
+ mMkk_seq = mMkk[None, None, None, h_idx, b_idx]
+ gMqk = cute.local_tile(
+ mMqk_seq, (self.BT, self.BT, self.CT), (0, 0, t_idx)
+ )
+ gMkk = cute.local_tile(
+ mMkk_seq, (self.BT, self.BT, self.CT), (0, 0, t_idx)
+ )
+ gMqk = gMqk[None, None, widx_in_wg]
+ gMkk = gMkk[None, None, widx_in_wg]
+
+ sQd_tile = cute.local_tile(
+ sQd[None, None, mma_full.index], (self.BT, self.D), (None, 0)
+ )
+ sKd_tile = cute.local_tile(
+ sKd[None, None, mma_full.index], (self.BT, self.D), (None, 0)
+ )
+ sKi_tile = cute.local_tile(
+ sKi[None, None, mma_full.index], (self.BT, self.D), (None, 0)
+ )
+ if cutlass.const_expr(self.beta_use_tma):
+ beta_head = h_idx % self.beta_tma_heads
+ sBeta_tile = cute.local_tile(
+ sBeta[None, beta_head, beta_idx], (self.BT,), (None,)
+ )
+ else:
+ sBeta_tile = cute.local_tile(
+ sBeta[None, beta_idx], (self.BT,), (None,)
+ )
+ sQd_tile = sQd_tile[None, None, widx_in_wg]
+ sKd_tile = sKd_tile[None, None, widx_in_wg]
+ sKi_tile = sKi_tile[None, None, widx_in_wg]
+ sBeta_tile = sBeta_tile[None, widx_in_wg]
+
+ beta_stride = (
+ self.beta_tma_heads
+ if cutlass.const_expr(self.beta_use_tma)
+ else 1
+ )
+ sBeta_pre = cute.make_tensor(
+ sBeta_tile.iterator,
+ cute.make_layout((self.BT, self.BT), stride=(beta_stride, 0)),
+ )
+ sBeta_post = cute.make_tensor(
+ sBeta_tile.iterator,
+ cute.make_layout((self.BT, self.BT), stride=(0, beta_stride)),
+ )
+
+ tArQd = thr_mma.make_fragment_A(thr_mma.partition_A(sQd_tile))
+ tArKd = thr_mma.make_fragment_A(thr_mma.partition_A(sKd_tile))
+ tArKi = thr_mma.make_fragment_B(thr_mma.partition_B(sKi_tile))
+ tAgMqk = thr_mma.partition_C(gMqk)
+
+ thr_copy_QKd = tiled_copy_QKd.get_slice(lane_idx)
+ tAsQd_copy_view = thr_copy_QKd.partition_S(sQd_tile)
+ tArQd_copy_view = thr_copy_QKd.retile(tArQd)
+ tAsKd_copy_view = thr_copy_QKd.partition_S(sKd_tile)
+ tArKd_copy_view = thr_copy_QKd.retile(tArKd)
+
+ thr_copy_Ki = tiled_copy_Ki.get_slice(lane_idx)
+ tAsKi_copy_view = thr_copy_Ki.partition_S(sKi_tile)
+ tArKi_copy_view = thr_copy_Ki.retile(tArKi)
+
+ tAsBeta_pre = thr_mma.partition_C(sBeta_pre)
+ tAsBeta_pre_mn = cute.make_tensor(
+ tAsBeta_pre.iterator,
+ layout_acc_mn(tiled_mma, tAsBeta_pre.layout),
+ )
+ tAsBeta_pre_mn = cute.flatten(tAsBeta_pre_mn[None, 0])
+ tArBeta_pre_mn = cute.make_rmem_tensor_like(tAsBeta_pre_mn)
+ tArBeta_pre_mn_f32 = cute.make_rmem_tensor_like(
+ tAsBeta_pre_mn, cutlass.Float32
+ )
+
+ tMsBeta_post = thr_mma.partition_C(sBeta_post)
+ tMsBeta_post_mn = cute.make_tensor(
+ tMsBeta_post.iterator,
+ layout_acc_mn(tiled_mma, tMsBeta_post.layout),
+ )
+ tMsBeta_post_mn = tMsBeta_post_mn[0, None]
+ tMrBeta_post_mn = cute.make_rmem_tensor_like(tMsBeta_post_mn)
+ tMrBeta_post_mn_f32 = cute.make_rmem_tensor_like(
+ tMsBeta_post_mn, cutlass.Float32
+ )
+
+ tMgMkk = thr_mma.partition_C(gMkk)
+
+ cute.copy(tiled_copy_QKd, tAsQd_copy_view, tArQd_copy_view)
+ cute.copy(tiled_copy_QKd, tAsKd_copy_view, tArKd_copy_view)
+ cute.copy(tiled_copy_Ki, tAsKi_copy_view, tArKi_copy_view)
+ cute.autovec_copy(tAsBeta_pre_mn, tArBeta_pre_mn)
+ cute.autovec_copy(tMsBeta_post_mn, tMrBeta_post_mn)
+
+ tArAkk = cute.make_rmem_tensor_like(tAcA, cutlass.Float32)
+ tArAkk.fill(0.0)
+ cute.gemm(tiled_mma, tArAkk, tArKd, tArKi, tArAkk)
+ for e in cutlass.range_constexpr(cute.size(tAcA)):
+ m, n = tAcA[e]
+ tArAkk[e] = tArAkk[e] if m > n else cutlass.Float32(0.0)
+ self.sigmoid(tArBeta_pre_mn, tArBeta_pre_mn_f32)
+ tArAkk_mn = cute.make_tensor(
+ tArAkk.iterator, layout_acc_mn(tiled_mma, tArAkk.layout)
+ )
+ for m in cutlass.range_constexpr(cute.size(tArAkk_mn, mode=[0])):
+ for n in cutlass.range_constexpr(
+ cute.size(tArAkk_mn, mode=[1]) // 2
+ ):
+ tArAkk_mn[m, 2 * n], tArAkk_mn[m, 2 * n + 1] = (
+ cute.arch.mul_packed_f32x2(
+ (tArAkk_mn[m, 2 * n], tArAkk_mn[m, 2 * n + 1]),
+ (-tArBeta_pre_mn_f32[m], -tArBeta_pre_mn_f32[m]),
+ )
+ )
+
+ tArAqk = cute.make_rmem_tensor_like(tAcA, cutlass.Float32)
+ tArAqk.fill(0.0)
+ cute.gemm(tiled_mma, tArAqk, tArQd, tArKi, tArAqk)
+ for e in cutlass.range_constexpr(cute.size(tAcA)):
+ m, n = tAcA[e]
+ tArAqk[e] = tArAqk[e] if m >= n else cutlass.Float32(0.0)
+ tArAqk_f16 = cute.make_rmem_tensor_like(tArAqk, gMqk.element_type)
+ tArAqk_f16.store(tArAqk.load().to(tArAqk_f16.element_type))
+ cute.autovec_copy(tArAqk_f16, tAgMqk)
+
+ tArL = thr_mma.make_fragment_A(a_shape_L)
+ tBrL = thr_mma.make_fragment_B(_pack_shape(b_shape_L))
+ tCrL = thr_mma.make_fragment_C(c_shape_L)
+
+ tArL.store(tArAkk.load().to(tArL.element_type))
+ self.transpose(tArL, tBrL)
+ tArIL = thr_mma.make_fragment_A(a_shape_L)
+ tArIL.store(
+ self.add_identity(tArAkk, tCrI).load().to(tArIL.element_type)
+ )
+ tCrL.fill(0.0)
+ cute.gemm(tiled_mma, tCrL, tArL, tBrL, tCrL)
+
+ tArL.store(tCrL.load().to(tArL.element_type))
+ self.transpose(tArL, tBrL)
+ tArIL2 = thr_mma.make_fragment_A(a_shape_L)
+ tArIL2.store(
+ self.add_identity(tCrL, tCrI).load().to(tArIL2.element_type)
+ )
+ tBrIL2 = thr_mma.make_fragment_B(_pack_shape(b_shape_L))
+ self.transpose(tArIL2, tBrIL2)
+ tCrL.fill(0.0)
+ cute.gemm(tiled_mma, tCrL, tArL, tBrL, tCrL)
+
+ tArL.store(tCrL.load().to(tArL.element_type))
+ self.transpose(tArL, tBrL)
+ tArIL4 = thr_mma.make_fragment_A(a_shape_L)
+ tArIL4.store(
+ self.add_identity(tCrL, tCrI).load().to(tArIL4.element_type)
+ )
+ tCrIL8 = thr_mma.make_fragment_C(c_shape_L)
+ # cute.gemm(tiled_mma, tCrIL8, tArL4, tBrL, tCrI) # CuTe DSL bug not supporting different C and D
+ tCrIL8.fill(0.0)
+ cute.gemm(tiled_mma, tCrIL8, tArL, tBrL, tCrIL8)
+
+ tArIL8 = thr_mma.make_fragment_A(a_shape_L)
+ # tArIL8.store(tCrIL8.load().to(tArIL8.element_type))
+ tArIL8.store(
+ self.add_identity(tCrIL8, tCrI).load().to(tArIL8.element_type)
+ )
+ tBrIL8 = thr_mma.make_fragment_B(_pack_shape(b_shape_L))
+ self.transpose(tArIL8, tBrIL8)
+
+ tCrL.fill(0.0)
+ cute.gemm(tiled_mma, tCrL, tArIL, tBrIL2, tCrL)
+ tArL12 = thr_mma.make_fragment_A(a_shape_L)
+ tArL12.store(tCrL.load().to(tArL12.element_type))
+
+ tCrL.fill(0.0)
+ cute.gemm(tiled_mma, tCrL, tArIL4, tBrIL8, tCrL)
+ tArL48 = thr_mma.make_fragment_A(a_shape_L)
+ tArL48.store(tCrL.load().to(tArL48.element_type))
+ tBrL48 = thr_mma.make_fragment_B(_pack_shape(b_shape_L))
+ self.transpose(tArL48, tBrL48)
+
+ tCrL.fill(0.0)
+ cute.gemm(tiled_mma, tCrL, tArL12, tBrL48, tCrL)
+ self.sigmoid(tMrBeta_post_mn, tMrBeta_post_mn_f32)
+ tCrL_mn = cute.make_tensor(
+ tCrL.iterator, layout_acc_mn(tiled_mma, tCrL.layout)
+ )
+ for m in cutlass.range_constexpr(cute.size(tCrL_mn, mode=[0])):
+ for n in cutlass.range_constexpr(
+ cute.size(tCrL_mn, mode=[1]) // 2
+ ):
+ tCrL_mn[m, 2 * n], tCrL_mn[m, 2 * n + 1] = (
+ cute.arch.mul_packed_f32x2(
+ (tCrL_mn[m, 2 * n], tCrL_mn[m, 2 * n + 1]),
+ (
+ tMrBeta_post_mn_f32[2 * n],
+ tMrBeta_post_mn_f32[2 * n + 1],
+ ),
+ )
+ )
+ tMrMkk_f16 = cute.make_rmem_tensor_like(tCrL, gMkk.element_type)
+ tMrMkk_f16.store(tCrL.load().to(tMrMkk_f16.element_type))
+ cute.autovec_copy(tMrMkk_f16, tMgMkk)
+
+ if cutlass.const_expr(self.beta_use_tma):
+ beta_idx = (beta_idx + 1) % self.beta_stages
+ else:
+ beta_full.release()
+ store_pipeline.producer_acquire()
+ mma_full.release()
+
+ tile_sched.advance_to_next_work()
+ work_tile = tile_sched.get_current_work()
+
+ @cute.jit
+ def compute_gate_tile(
+ self,
+ mGk: cute.Tensor,
+ mKr: cute.Tensor,
+ sQ: cute.Tensor,
+ sK: cute.Tensor,
+ sG: cute.Tensor,
+ sDtBias: cute.Tensor,
+ sAlog: cute.Tensor,
+ sQd: cute.Tensor,
+ sKd: cute.Tensor,
+ sKi: cute.Tensor,
+ tma_consumer: pipeline.PipelineConsumer,
+ gate_mma_producer: pipeline.PipelineProducer,
+ gate_scale: cutlass.Float32,
+ scale: cutlass.Float32,
+ t_idx: cutlass.Int32,
+ h_idx: cutlass.Int32,
+ b_idx: cutlass.Int32,
+ tiled_copy_QKG: cute.TiledCopy,
+ tiled_copy_dtBias: cute.TiledCopy,
+ tiled_copy_Gk: cute.TiledCopy,
+ lane_idx: cutlass.Int32,
+ widx_in_wg: cutlass.Int32,
+ seq_offset_padded: Optional[cutlass.Int32] = None,
+ chunk_offset: Optional[cutlass.Int32] = None,
+ valid_tokens: Optional[cutlass.Int32] = None,
+ ):
+ alog = sAlog[h_idx]
+
+ tma_full = tma_consumer.wait_and_advance()
+ mma_empty = gate_mma_producer.acquire_and_advance()
+
+ if cutlass.const_expr(seq_offset_padded is not None):
+ mKr_seq = cute.domain_offset((seq_offset_padded, 0, 0), mKr)
+ mGk_seq = cute.domain_offset((0, chunk_offset, 0), mGk)
+ gGk = cute.local_tile(
+ mGk_seq[None, None, h_idx], (self.D, self.CT), (0, t_idx)
+ )
+ gKr = cute.local_tile(
+ mKr_seq[None, None, h_idx], (self.TT, self.D), (t_idx, 0)
+ )
+ else:
+ gGk = cute.local_tile(
+ mGk[None, None, h_idx, b_idx], (self.D, self.CT), (0, t_idx)
+ )
+ gKr = cute.local_tile(
+ mKr[None, None, h_idx, b_idx], (self.TT, self.D), (t_idx, 0)
+ )
+ gGk_tile = gGk[None, widx_in_wg]
+ gKr_tile = cute.local_tile(gKr, (self.BT, self.D), (widx_in_wg, 0))
+
+ sQ_tile = cute.local_tile(
+ sQ[None, None, tma_full.index], (self.BT, self.D), (widx_in_wg, 0)
+ )
+ sK_tile = cute.local_tile(
+ sK[None, None, tma_full.index], (self.BT, self.D), (widx_in_wg, 0)
+ )
+ sG_tile = cute.local_tile(
+ sG[None, None, tma_full.index], (self.BT, self.D), (widx_in_wg, 0)
+ )
+ sDtBias_tile = sDtBias[None, tma_full.index]
+ sQd_tile = cute.local_tile(
+ sQd[None, None, mma_empty.index], (self.BT, self.D), (widx_in_wg, 0)
+ )
+ sKd_tile = cute.local_tile(
+ sKd[None, None, mma_empty.index], (self.BT, self.D), (widx_in_wg, 0)
+ )
+ sKi_tile = cute.local_tile(
+ sKi[None, None, mma_empty.index], (self.BT, self.D), (widx_in_wg, 0)
+ )
+
+ thr_copy_QKG = tiled_copy_QKG.get_slice(lane_idx)
+ tQsQ = thr_copy_QKG.partition_S(sQ_tile)
+ tQrQ = cute.make_rmem_tensor_like(tQsQ)
+ tQrQ_f32 = cute.make_rmem_tensor_like(tQrQ, cutlass.Float32)
+ tKsK = thr_copy_QKG.partition_S(sK_tile)
+ tKrK = cute.make_rmem_tensor_like(tKsK)
+ tKrK_f32 = cute.make_rmem_tensor_like(tKrK, cutlass.Float32)
+ tGsG = thr_copy_QKG.partition_S(sG_tile)
+ tGrG = cute.make_rmem_tensor_like(tGsG)
+ tGrG_f32 = cute.make_rmem_tensor_like(tGrG, cutlass.Float32)
+ tGrGe_f32 = cute.make_rmem_tensor_like(tGrG, cutlass.Float32)
+ tGrGi_f32 = cute.make_rmem_tensor_like(tGrG, cutlass.Float32)
+ tQsQd = thr_copy_QKG.partition_D(sQd_tile)
+ tQrQd = cute.make_rmem_tensor_like(tQsQd)
+ tQrQd_f32 = cute.make_rmem_tensor_like(tQrQd, cutlass.Float32)
+ tKsKd = thr_copy_QKG.partition_D(sKd_tile)
+ tKrKd = cute.make_rmem_tensor_like(tKsKd)
+ tKrKd_f32 = cute.make_rmem_tensor_like(tKrKd, cutlass.Float32)
+ tKsKi = thr_copy_QKG.partition_D(sKi_tile)
+ tKrKi = cute.make_rmem_tensor_like(tKsKi)
+ tKrKi_f32 = cute.make_rmem_tensor_like(tKrKi, cutlass.Float32)
+ tKgKr = thr_copy_QKG.partition_D(gKr_tile)
+ tKrKr = cute.make_rmem_tensor_like(tKgKr)
+ tKrKr_f32 = cute.make_rmem_tensor_like(tKrKr, cutlass.Float32)
+
+ thr_copy_dtBias = tiled_copy_dtBias.get_slice(lane_idx)
+ tDtBias_sDtBias = thr_copy_dtBias.partition_S(sDtBias_tile)
+ tDtBias_rDtBias = cute.make_rmem_tensor_like(tDtBias_sDtBias)
+
+ thr_copy_Gk = tiled_copy_Gk.get_slice(lane_idx)
+ tGkgGk = thr_copy_Gk.partition_S(gGk_tile)
+ tGkrGk = cute.make_rmem_tensor_like(tGkgGk)
+
+ nr = cute.size(tQrQ, mode=[0, 1])
+ nc = cute.size(self.get_row(tQrQ, 0), mode=[1])
+ t = widx_in_wg * self.BT
+ if lane_idx >= self.warp_threads // 2:
+ t += self.BT // 2
+
+ cum = cute.make_rmem_tensor_like(self.get_row(tGrG, 0), cutlass.Float32)
+
+ precompute_qk_l2_norm = 2
+ if cutlass.const_expr(seq_offset_padded is not None):
+ if cutlass.const_expr(self.gate_lower_bound):
+ precompute_qk_l2_norm = 4
+ else:
+ precompute_qk_l2_norm = 5
+
+ cute.copy(tiled_copy_dtBias, tDtBias_sDtBias, tDtBias_rDtBias)
+ cute.copy(
+ tiled_copy_QKG, tGsG[((None, 0), None, None)], tGrG[((None, 0), None, None)]
+ )
+ if cutlass.const_expr(precompute_qk_l2_norm == nr):
+ cute.copy(
+ tiled_copy_QKG,
+ tQsQ[((None, 0), None, None)],
+ tQrQ[((None, 0), None, None)],
+ )
+ cute.copy(
+ tiled_copy_QKG,
+ tKsK[((None, 0), None, None)],
+ tKrK[((None, 0), None, None)],
+ )
+
+ for r in cutlass.range_constexpr(nr):
+ if cutlass.const_expr(r < nr - 1):
+ cute.copy(
+ tiled_copy_QKG,
+ tGsG[((None, r + 1), None, None)],
+ tGrG[((None, r + 1), None, None)],
+ )
+
+ tGrG_row = self.get_row(tGrG, r)
+ tGrG_f32_row = self.get_row(tGrG_f32, r)
+ for c in cutlass.range_constexpr(nc):
+ tGrG_f32_row[0, c], tGrG_f32_row[1, c] = cute.arch.add_packed_f32x2(
+ (
+ tGrG_row[0, c].to(cutlass.Float32),
+ tGrG_row[1, c].to(cutlass.Float32),
+ ),
+ (tDtBias_rDtBias[2 * c], tDtBias_rDtBias[2 * c + 1]),
+ )
+ if cutlass.const_expr(self.gate_lower_bound):
+ tGrG_f32_row[0, c], tGrG_f32_row[1, c] = cute.arch.mul_packed_f32x2(
+ (tGrG_f32_row[0, c], tGrG_f32_row[1, c]),
+ (alog, alog),
+ )
+ tGrG_f32_row[0, c] = tanh(tGrG_f32_row[0, c])
+ tGrG_f32_row[1, c] = tanh(tGrG_f32_row[1, c])
+ tGrG_f32_row[0, c], tGrG_f32_row[1, c] = cute.arch.fma_packed_f32x2(
+ (tGrG_f32_row[0, c], tGrG_f32_row[1, c]),
+ (gate_scale, gate_scale),
+ (gate_scale, gate_scale),
+ )
+ else:
+ tGrG_f32_row[0, c], tGrG_f32_row[1, c] = cute.arch.mul_packed_f32x2(
+ (tGrG_f32_row[0, c], tGrG_f32_row[1, c]),
+ (RCP_LN2, RCP_LN2),
+ )
+ tGrG_f32_row[0, c] = cute.math.exp2(
+ tGrG_f32_row[0, c], fastmath=True
+ )
+ tGrG_f32_row[1, c] = cute.math.exp2(
+ tGrG_f32_row[1, c], fastmath=True
+ )
+ tGrG_f32_row[0, c], tGrG_f32_row[1, c] = cute.arch.add_packed_f32x2(
+ (tGrG_f32_row[0, c], tGrG_f32_row[1, c]),
+ (cutlass.Float32(1.0), cutlass.Float32(1.0)),
+ )
+ tGrG_f32_row[0, c] = cute.math.log2(
+ tGrG_f32_row[0, c], fastmath=True
+ )
+ tGrG_f32_row[1, c] = cute.math.log2(
+ tGrG_f32_row[1, c], fastmath=True
+ )
+ tGrG_f32_row[0, c], tGrG_f32_row[1, c] = cute.arch.mul_packed_f32x2(
+ (tGrG_f32_row[0, c], tGrG_f32_row[1, c]),
+ (alog, alog),
+ )
+ if cutlass.const_expr(valid_tokens is not None):
+ tGrG_f32_row[0, c] = (
+ tGrG_f32_row[0, c]
+ if t + r < valid_tokens
+ else cutlass.Float32(0.0)
+ )
+ tGrG_f32_row[1, c] = (
+ tGrG_f32_row[1, c]
+ if t + r < valid_tokens
+ else cutlass.Float32(0.0)
+ )
+ if cutlass.const_expr(r == 0):
+ cum[0, c] = tGrG_f32_row[0, c]
+ cum[1, c] = tGrG_f32_row[1, c]
+ else:
+ cum[0, c], cum[1, c] = cute.arch.add_packed_f32x2(
+ (cum[0, c], cum[1, c]),
+ (tGrG_f32_row[0, c], tGrG_f32_row[1, c]),
+ )
+ tGrG_f32_row[0, c] = cum[0, c]
+ tGrG_f32_row[1, c] = cum[1, c]
+
+ qk_r = r - nr + precompute_qk_l2_norm
+ if cutlass.const_expr(qk_r >= -1 and r < nr - 1):
+ cute.copy(
+ tiled_copy_QKG,
+ tQsQ[((None, qk_r + 1), None, None)],
+ tQrQ[((None, qk_r + 1), None, None)],
+ )
+ cute.copy(
+ tiled_copy_QKG,
+ tKsK[((None, qk_r + 1), None, None)],
+ tKrK[((None, qk_r + 1), None, None)],
+ )
+ if cutlass.const_expr(qk_r >= 0):
+ self.qk_l2_norm_row(
+ tQrQ,
+ tQrQ_f32,
+ tKrK,
+ tKrK_f32,
+ qk_r,
+ t,
+ valid_tokens,
+ seq_offset_padded,
+ )
+
+ if cutlass.const_expr(precompute_qk_l2_norm < nr):
+ cute.copy(
+ tiled_copy_QKG,
+ tQsQ[((None, precompute_qk_l2_norm), None, None)],
+ tQrQ[((None, precompute_qk_l2_norm), None, None)],
+ )
+ cute.copy(
+ tiled_copy_QKG,
+ tKsK[((None, precompute_qk_l2_norm), None, None)],
+ tKrK[((None, precompute_qk_l2_norm), None, None)],
+ )
+
+ for c in cutlass.range_constexpr(nc):
+ cum_prefix0 = cute.arch.shuffle_sync_bfly(cum[0, c], self.warp_threads // 2)
+ cum_prefix1 = cute.arch.shuffle_sync_bfly(cum[1, c], self.warp_threads // 2)
+ if lane_idx < self.warp_threads // 2:
+ cum_prefix0 = cutlass.Float32(0.0)
+ cum_prefix1 = cutlass.Float32(0.0)
+ for r in cutlass.range_constexpr(nr):
+ tGrG_f32_row = self.get_row(tGrG_f32, r)
+ tGrG_f32_row[0, c], tGrG_f32_row[1, c] = cute.arch.add_packed_f32x2(
+ (tGrG_f32_row[0, c], tGrG_f32_row[1, c]),
+ (cum_prefix0, cum_prefix1),
+ )
+ if lane_idx >= self.warp_threads // 2:
+ tGrG_f32_row = self.get_row(tGrG_f32, nr - 1)
+ tGkrGk[2 * c] = cute.math.exp2(tGrG_f32_row[0, c], fastmath=True)
+ tGkrGk[2 * c + 1] = cute.math.exp2(tGrG_f32_row[1, c], fastmath=True)
+ Gk_shfl0 = cute.arch.shuffle_sync_bfly(tGkrGk[2 * c], 16)
+ Gk_shfl1 = cute.arch.shuffle_sync_bfly(tGkrGk[2 * c + 1], 16)
+ if lane_idx < self.warp_threads // 2:
+ tGkrGk[2 * c] = Gk_shfl0
+ tGkrGk[2 * c + 1] = Gk_shfl1
+ if lane_idx >= self.warp_threads // 2:
+ cute.copy(tiled_copy_Gk, tGkrGk, tGkgGk)
+
+ for r in cutlass.range_constexpr(nr):
+ qk_r = r + precompute_qk_l2_norm
+ if cutlass.const_expr(qk_r + 1 < nr):
+ cute.copy(
+ tiled_copy_QKG,
+ tQsQ[((None, qk_r + 1), None, None)],
+ tQrQ[((None, qk_r + 1), None, None)],
+ )
+ cute.copy(
+ tiled_copy_QKG,
+ tKsK[((None, qk_r + 1), None, None)],
+ tKrK[((None, qk_r + 1), None, None)],
+ )
+ if cutlass.const_expr(qk_r < nr):
+ self.qk_l2_norm_row(
+ tQrQ,
+ tQrQ_f32,
+ tKrK,
+ tKrK_f32,
+ qk_r,
+ t,
+ valid_tokens,
+ seq_offset_padded,
+ )
+
+ tQrQ_f32_row = self.get_row(tQrQ_f32, r)
+ tKrK_f32_row = self.get_row(tKrK_f32, r)
+ tGrG_f32_row = self.get_row(tGrG_f32, r)
+ tGrGe_f32_row = self.get_row(tGrGe_f32, r)
+ tGrGi_f32_row = self.get_row(tGrGi_f32, r)
+ tQrQd_row = self.get_row(tQrQd, r)
+ tQrQd_f32_row = self.get_row(tQrQd_f32, r)
+ tKrKd_row = self.get_row(tKrKd, r)
+ tKrKd_f32_row = self.get_row(tKrKd_f32, r)
+ tKrKi_row = self.get_row(tKrKi, r)
+ tKrKi_f32_row = self.get_row(tKrKi_f32, r)
+ tKrKr_row = self.get_row(tKrKr, r)
+ tKrKr_f32_row = self.get_row(tKrKr_f32, r)
+
+ for c in cutlass.range_constexpr(nc):
+ tGrGe_f32_row[0, c] = cute.math.exp2(tGrG_f32_row[0, c], fastmath=True)
+ tGrGe_f32_row[1, c] = cute.math.exp2(tGrG_f32_row[1, c], fastmath=True)
+
+ tQrQd_f32_row[0, c], tQrQd_f32_row[1, c] = cute.arch.mul_packed_f32x2(
+ (tGrGe_f32_row[0, c], tGrGe_f32_row[1, c]),
+ (tQrQ_f32_row[0, c], tQrQ_f32_row[1, c]),
+ )
+ tQrQd_f32_row[0, c], tQrQd_f32_row[1, c] = cute.arch.mul_packed_f32x2(
+ (tQrQd_f32_row[0, c], tQrQd_f32_row[1, c]),
+ (scale, scale),
+ )
+ tQrQd_row[None, c].store(
+ tQrQd_f32_row[None, c].load().to(tQrQd_row.element_type)
+ )
+
+ tKrKd_f32_row[0, c], tKrKd_f32_row[1, c] = cute.arch.mul_packed_f32x2(
+ (tKrK_f32_row[0, c], tKrK_f32_row[1, c]),
+ (tGrGe_f32_row[0, c], tGrGe_f32_row[1, c]),
+ )
+ tKrKd_row[None, c].store(
+ tKrKd_f32_row[None, c].load().to(tKrKd_row.element_type)
+ )
+
+ tGrGi_f32_row[0, c] = cute.math.exp2(-tGrG_f32_row[0, c], fastmath=True)
+ tGrGi_f32_row[1, c] = cute.math.exp2(-tGrG_f32_row[1, c], fastmath=True)
+
+ tKrKi_f32_row[0, c], tKrKi_f32_row[1, c] = cute.arch.mul_packed_f32x2(
+ (tKrK_f32_row[0, c], tKrK_f32_row[1, c]),
+ (tGrGi_f32_row[0, c], tGrGi_f32_row[1, c]),
+ )
+ tKrKi_row[None, c].store(
+ tKrKi_f32_row[None, c].load().to(tKrKi_row.element_type)
+ )
+
+ tKrKr_f32_row[0, c], tKrKr_f32_row[1, c] = cute.arch.mul_packed_f32x2(
+ (tKrKi_f32_row[0, c], tKrKi_f32_row[1, c]),
+ (tGkrGk[2 * c], tGkrGk[2 * c + 1]),
+ )
+ tKrKr_row[None, c].store(
+ tKrKr_f32_row[None, c].load().to(tKrKr_row.element_type)
+ )
+
+ cute.copy(
+ tiled_copy_QKG,
+ tQrQd[((None, r), None, None)],
+ tQsQd[((None, r), None, None)],
+ )
+ cute.copy(
+ tiled_copy_QKG,
+ tKrKd[((None, r), None, None)],
+ tKsKd[((None, r), None, None)],
+ )
+ cute.copy(
+ tiled_copy_QKG,
+ tKrKi[((None, r), None, None)],
+ tKsKi[((None, r), None, None)],
+ )
+ if cutlass.const_expr(
+ valid_tokens is not None and seq_offset_padded is None
+ ):
+ if t + r < valid_tokens:
+ cute.copy(
+ tiled_copy_QKG,
+ tKrKr[((None, r), None, None)],
+ tKgKr[((None, r), None, None)],
+ )
+ else:
+ cute.copy(
+ tiled_copy_QKG,
+ tKrKr[((None, r), None, None)],
+ tKgKr[((None, r), None, None)],
+ )
+
+ tma_full.release()
+ mma_empty.commit()
+ return tma_consumer, gate_mma_producer
+
+ @cute.jit
+ def qk_l2_norm_row(
+ self,
+ tQrQ: cute.Tensor,
+ tQrQ_f32: cute.Tensor,
+ tKrK: cute.Tensor,
+ tKrK_f32: cute.Tensor,
+ r: int,
+ t: cutlass.Int32,
+ valid_tokens: Optional[cutlass.Int32] = None,
+ seq_offset_padded: Optional[cutlass.Int32] = None,
+ ):
+ tQrQ_row = self.get_row(tQrQ, r)
+ tQrQ_f32_row = self.get_row(tQrQ_f32, r)
+ tKrK_row = self.get_row(tKrK, r)
+ tKrK_f32_row = self.get_row(tKrK_f32, r)
+ nc = cute.size(tQrQ_f32_row, mode=[1])
+ tQrQ_f32_row.store(tQrQ_row.load().to(tQrQ_f32_row.element_type))
+ tKrK_f32_row.store(tKrK_row.load().to(tKrK_f32_row.element_type))
+ q_sq0, q_sq1 = cute.arch.mul_packed_f32x2(
+ (tQrQ_f32_row[0, 0], tQrQ_f32_row[1, 0]),
+ (tQrQ_f32_row[0, 0], tQrQ_f32_row[1, 0]),
+ )
+ k_sq0, k_sq1 = cute.arch.mul_packed_f32x2(
+ (tKrK_f32_row[0, 0], tKrK_f32_row[1, 0]),
+ (tKrK_f32_row[0, 0], tKrK_f32_row[1, 0]),
+ )
+ for c in cutlass.range_constexpr(1, nc):
+ q_sq0, q_sq1 = cute.arch.fma_packed_f32x2(
+ (tQrQ_f32_row[0, c], tQrQ_f32_row[1, c]),
+ (tQrQ_f32_row[0, c], tQrQ_f32_row[1, c]),
+ (q_sq0, q_sq1),
+ )
+ k_sq0, k_sq1 = cute.arch.fma_packed_f32x2(
+ (tKrK_f32_row[0, c], tKrK_f32_row[1, c]),
+ (tKrK_f32_row[0, c], tKrK_f32_row[1, c]),
+ (k_sq0, k_sq1),
+ )
+ q_sq = q_sq0 + q_sq1
+ k_sq = k_sq0 + k_sq1
+ for d in cutlass.range_constexpr(3, -1, -1):
+ q_sq += cute.arch.shuffle_sync_bfly(q_sq, 1 << d)
+ k_sq += cute.arch.shuffle_sync_bfly(k_sq, 1 << d)
+ q_inv = cute.math.rsqrt(q_sq + cutlass.Float32(1e-6), fastmath=True)
+ k_inv = cute.math.rsqrt(k_sq + cutlass.Float32(1e-6), fastmath=True)
+ if cutlass.const_expr(
+ valid_tokens is not None and seq_offset_padded is not None
+ ):
+ k_inv = k_inv if t + r < valid_tokens else cutlass.Float32(0.0)
+ for c in cutlass.range_constexpr(nc):
+ tQrQ_f32_row[0, c], tQrQ_f32_row[1, c] = cute.arch.mul_packed_f32x2(
+ (tQrQ_f32_row[0, c], tQrQ_f32_row[1, c]),
+ (q_inv, q_inv),
+ )
+ tKrK_f32_row[0, c], tKrK_f32_row[1, c] = cute.arch.mul_packed_f32x2(
+ (tKrK_f32_row[0, c], tKrK_f32_row[1, c]),
+ (k_inv, k_inv),
+ )
+
+ @staticmethod
+ @cute.jit
+ def get_row(t: cute.Tensor, r: int):
+ t = t[(None, r), 0, 0]
+ return cute.make_tensor(
+ t.iterator,
+ cute.make_layout((2, cute.size(t) // 2), stride=(1, 2)),
+ )
+
+ @staticmethod
+ @cute.jit
+ def sigmoid(src_f16: cute.Tensor, dst_f32: cute.Tensor):
+ dst_f32.store(src_f16.load().to(dst_f32.element_type))
+ for i in cutlass.range_constexpr(cute.size(dst_f32) // 2):
+ dst_f32[2 * i], dst_f32[2 * i + 1] = cute.arch.mul_packed_f32x2(
+ (dst_f32[2 * i], dst_f32[2 * i + 1]),
+ (cutlass.Float32(0.5), cutlass.Float32(0.5)),
+ )
+ dst_f32[2 * i] = tanh(dst_f32[2 * i])
+ dst_f32[2 * i + 1] = tanh(dst_f32[2 * i + 1])
+ dst_f32[2 * i], dst_f32[2 * i + 1] = cute.arch.fma_packed_f32x2(
+ (dst_f32[2 * i], dst_f32[2 * i + 1]),
+ (cutlass.Float32(0.5), cutlass.Float32(0.5)),
+ (cutlass.Float32(0.5), cutlass.Float32(0.5)),
+ )
+
+ @staticmethod
+ @cute.jit
+ def add_identity(tCrC: cute.Tensor, tCrI: cute.Tensor):
+ tCrIC = cute.make_rmem_tensor_like(tCrC, tCrC.element_type)
+ for i in cutlass.range_constexpr(cute.size(tCrC) // 2):
+ tCrIC[2 * i], tCrIC[2 * i + 1] = cute.arch.add_packed_f32x2(
+ (tCrC[2 * i], tCrC[2 * i + 1]),
+ (tCrI[2 * i], tCrI[2 * i + 1]),
+ )
+ return tCrIC
+
+ @staticmethod
+ @cute.jit
+ def transpose(tArA: cute.Tensor, tBrB: cute.Tensor):
+ tArA_u32 = cute.make_tensor(
+ cute.recast_tensor(tArA, cutlass.Uint32).iterator, cute.make_layout((4,))
+ )
+ tBrB_u32 = cute.make_tensor(
+ cute.recast_tensor(tBrB, cutlass.Uint32).iterator, cute.make_layout((4,))
+ )
+ for i in cutlass.range_constexpr(4):
+ tBrB_u32[i] = movmatrix(tArA_u32[i])
+
+
+class ChunkKdaFwdK2:
+ def __init__(
+ self, D: int = 128, BT: int = 16, BV: int = 64, transpose_S: bool = False
+ ):
+ self.D = D
+ self.BT = BT
+ self.BV = BV
+ self.transpose_S = transpose_S
+
+ self.CT = 4
+ self.TT = self.CT * self.BT
+ self.load_stages = 2 if BV == 128 else 3
+ self.store_stages = 2 if BV == 128 else 3
+
+ self.warp_threads = 32
+ self.warpgroup_threads = 128
+ self.warps_per_warpgroup = self.warpgroup_threads // self.warp_threads
+
+ self.load_warp_id = 0
+ self.store_warp_id = 1
+ self.compute_warpgroup_ids = [1] if BV == 64 else [1, 2]
+ self.compute_threads = self.warpgroup_threads * len(self.compute_warpgroup_ids)
+ self.compute_warps = self.warps_per_warpgroup * len(self.compute_warpgroup_ids)
+ self.threads_per_cta = self.warpgroup_threads + self.compute_threads
+
+ @cute.jit
+ def __call__(
+ self,
+ mQ: cute.Tensor, # (B, T, H, D)
+ mK: cute.Tensor, # (B, T, H, D)
+ mKr: cute.Tensor, # (B, T, H, D)
+ mV: cute.Tensor, # (B, T, H, D)
+ mMqk: cute.Tensor, # (B, NT, H, BT, BT)
+ mMkk: cute.Tensor, # (B, NT, H, BT, BT)
+ mGk: cute.Tensor, # (B, NT, H, D)
+ mO: cute.Tensor, # (B, T, H, D)
+ mS_in: Optional[cute.Tensor] = None, # (B, H, D, D)
+ mS_out: Optional[cute.Tensor] = None, # (B, H, D, D)
+ mCuSeqlens: Optional[cute.Tensor] = None, # (B + 1,)
+ mStateIndices: Optional[cute.Tensor] = None, # (B,)
+ stream=None,
+ ):
+ if cutlass.const_expr(mCuSeqlens is not None):
+ B = cute.size(mCuSeqlens) - 1
+ H = cute.size(mQ, mode=[1])
+ mQ = cute.make_tensor(
+ mQ.iterator, cute.select(mQ.layout, mode=[0, 2, 1])
+ ) # (T, D, H)
+ mK = cute.make_tensor(
+ mK.iterator, cute.select(mK.layout, mode=[0, 2, 1])
+ ) # (T, D, H)
+ mKr = cute.make_tensor(
+ mKr.iterator, cute.select(mKr.layout, mode=[0, 2, 1])
+ ) # (T, D, H)
+ mV = cute.make_tensor(
+ mV.iterator, cute.select(mV.layout, mode=[0, 2, 1])
+ ) # (T, D, H)
+ mMqk = cute.make_tensor(
+ mMqk.iterator, cute.select(mMqk.layout, mode=[2, 3, 0, 1])
+ ) # (BT, BT, NT, H)
+ mMkk = cute.make_tensor(
+ mMkk.iterator, cute.select(mMkk.layout, mode=[2, 3, 0, 1])
+ ) # (BT, BT, NT, H)
+ mGk = cute.make_tensor(
+ mGk.iterator, cute.select(mGk.layout, mode=[2, 0, 1])
+ ) # (D, NT, H)
+ mO = cute.make_tensor(
+ mO.iterator, cute.select(mO.layout, mode=[0, 2, 1])
+ ) # (T, D, H)
+ else:
+ B, _, H, _ = mQ.shape
+ mQ = cute.make_tensor(
+ mQ.iterator, cute.select(mQ.layout, mode=[1, 3, 2, 0])
+ ) # (T, D, H, B)
+ mK = cute.make_tensor(
+ mK.iterator, cute.select(mK.layout, mode=[1, 3, 2, 0])
+ ) # (T, D, H, B)
+ mKr = cute.make_tensor(
+ mKr.iterator, cute.select(mKr.layout, mode=[1, 3, 2, 0])
+ ) # (T, D, H, B)
+ mV = cute.make_tensor(
+ mV.iterator, cute.select(mV.layout, mode=[1, 3, 2, 0])
+ ) # (T, D, H, B)
+ mMqk = cute.make_tensor(
+ mMqk.iterator, cute.select(mMqk.layout, mode=[3, 4, 1, 2, 0])
+ ) # (BT, BT, NT, H, B)
+ mMkk = cute.make_tensor(
+ mMkk.iterator, cute.select(mMkk.layout, mode=[3, 4, 1, 2, 0])
+ ) # (BT, BT, NT, H, B)
+ mGk = cute.make_tensor(
+ mGk.iterator, cute.select(mGk.layout, mode=[3, 1, 2, 0])
+ ) # (D, NT, H, B)
+ mO = cute.make_tensor(
+ mO.iterator, cute.select(mO.layout, mode=[1, 3, 2, 0])
+ ) # (T, D, H, B)
+ if cutlass.const_expr(mS_in is not None):
+ mS_in = cute.make_tensor(
+ mS_in.iterator, cute.select(mS_in.layout, mode=[2, 3, 1, 0])
+ ) # (D, D, H, B)
+ if cutlass.const_expr(mS_out is not None):
+ mS_out = cute.make_tensor(
+ mS_out.iterator, cute.select(mS_out.layout, mode=[2, 3, 1, 0])
+ ) # (D, D, H, B)
+
+ tiled_mma = cute.make_tiled_mma(
+ warp.MmaF16BF16Op(mQ.element_type, cutlass.Float32, (16, 8, 16)),
+ atom_layout_mnk=(self.BV // 16, 1, 1),
+ permutation_mnk=(self.BV, self.BT, 16),
+ )
+
+ copy_atom_n = cute.make_copy_atom(
+ cute.nvgpu.warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4),
+ mK.element_type,
+ )
+ copy_atom_t = cute.make_copy_atom(
+ cute.nvgpu.warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=4),
+ mV.element_type,
+ )
+ copy_atom_o = cute.make_copy_atom(
+ cute.nvgpu.warp.StMatrix8x8x16bOp(transpose=True, num_matrices=4),
+ mO.element_type,
+ )
+
+ tiled_copy_Q = cute.make_tiled_copy_B(copy_atom_n, tiled_mma)
+ tiled_copy_K = cute.make_tiled_copy_B(copy_atom_n, tiled_mma)
+ tiled_copy_Kr = cute.make_tiled_copy_B(copy_atom_t, tiled_mma)
+ tiled_copy_V = cute.make_tiled_copy_A(copy_atom_t, tiled_mma)
+ tiled_copy_Mqk = cute.make_tiled_copy_B(copy_atom_n, tiled_mma)
+ tiled_copy_Mkk = cute.make_tiled_copy_B(copy_atom_n, tiled_mma)
+ tiled_copy_O = cute.make_tiled_copy_C(copy_atom_o, tiled_mma)
+
+ smem_layout_atom_sw128 = cute.nvgpu.warpgroup.make_smem_layout_atom(
+ cute.nvgpu.warpgroup.SmemLayoutAtomKind.K_SW128,
+ mQ.element_type,
+ )
+ smem_layout_atom_sw32 = cute.nvgpu.warpgroup.make_smem_layout_atom(
+ cute.nvgpu.warpgroup.SmemLayoutAtomKind.K_SW32,
+ mMqk.element_type,
+ )
+
+ sQK_layout_staged = cute.tile_to_shape(
+ smem_layout_atom_sw128,
+ (self.TT, self.D, self.load_stages),
+ order=(0, 1, 2),
+ )
+ sQKt_layout_staged = cute.select(sQK_layout_staged, mode=[1, 0, 2])
+ sQK_layout = cute.slice_(sQK_layout_staged, (None, None, 0))
+
+ sV_layout_staged = cute.tile_to_shape(
+ smem_layout_atom_sw128,
+ (self.TT, self.BV, self.load_stages),
+ order=(0, 1, 2),
+ )
+ sVt_layout_staged = cute.select(sV_layout_staged, mode=[1, 0, 2])
+ sV_layout = cute.slice_(sV_layout_staged, (None, None, 0))
+
+ sM_layout_staged = cute.tile_to_shape(
+ smem_layout_atom_sw32,
+ (self.BT, self.BT, self.CT, self.load_stages),
+ order=(0, 1, 2, 3),
+ )
+ sM_layout = cute.slice_(sM_layout_staged, (None, None, None, 0))
+
+ sGk_layout_staged = cute.make_layout(
+ (self.D, self.CT, self.load_stages),
+ stride=(1, self.D, self.CT * self.D),
+ )
+ sGk_layout = cute.slice_(sGk_layout_staged, (None, None, 0))
+
+ sO_layout_staged = cute.tile_to_shape(
+ smem_layout_atom_sw128,
+ (self.TT, self.BV, self.store_stages),
+ order=(0, 1, 2),
+ )
+ sOt_layout_staged = cute.select(sO_layout_staged, mode=[1, 0, 2])
+ sO_layout = cute.slice_(sO_layout_staged, (None, None, 0))
+
+ buffer_align_bytes = 1024
+
+ @cute.struct
+ class SharedStorage:
+ sQ: cute.struct.Align[
+ cute.struct.MemRange[mQ.element_type, cute.cosize(sQK_layout_staged)],
+ buffer_align_bytes,
+ ]
+ sK: cute.struct.Align[
+ cute.struct.MemRange[mK.element_type, cute.cosize(sQK_layout_staged)],
+ buffer_align_bytes,
+ ]
+ sKr: cute.struct.Align[
+ cute.struct.MemRange[mKr.element_type, cute.cosize(sQK_layout_staged)],
+ buffer_align_bytes,
+ ]
+ sV: cute.struct.Align[
+ cute.struct.MemRange[mV.element_type, cute.cosize(sV_layout_staged)],
+ buffer_align_bytes,
+ ]
+ sMqk: cute.struct.Align[
+ cute.struct.MemRange[mMqk.element_type, cute.cosize(sM_layout_staged)],
+ buffer_align_bytes,
+ ]
+ sMkk: cute.struct.Align[
+ cute.struct.MemRange[mMkk.element_type, cute.cosize(sM_layout_staged)],
+ buffer_align_bytes,
+ ]
+ sGk: cute.struct.Align[
+ cute.struct.MemRange[mGk.element_type, cute.cosize(sGk_layout_staged)],
+ buffer_align_bytes,
+ ]
+ sO: cute.struct.Align[
+ cute.struct.MemRange[mO.element_type, cute.cosize(sO_layout_staged)],
+ buffer_align_bytes,
+ ]
+ tma_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_stages * 2]
+ tma_O_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.store_stages * 2]
+
+ self.num_tma_load_bytes = (
+ cute.size_in_bytes(mQ.element_type, sQK_layout)
+ + cute.size_in_bytes(mK.element_type, sQK_layout)
+ + cute.size_in_bytes(mKr.element_type, sQK_layout)
+ + cute.size_in_bytes(mV.element_type, sV_layout)
+ + cute.size_in_bytes(mMqk.element_type, sM_layout)
+ + cute.size_in_bytes(mMkk.element_type, sM_layout)
+ + cute.size_in_bytes(mGk.element_type, sGk_layout)
+ )
+
+ tma_atom_Q, tma_tensor_Q = cpasync.make_tiled_tma_atom(
+ cpasync.CopyBulkTensorTileG2SOp(),
+ mQ,
+ sQK_layout,
+ (self.TT, self.D),
+ )
+ tma_atom_K, tma_tensor_K = cpasync.make_tiled_tma_atom(
+ cpasync.CopyBulkTensorTileG2SOp(),
+ mK,
+ sQK_layout,
+ (self.TT, self.D),
+ )
+ tma_atom_Kr, tma_tensor_Kr = cpasync.make_tiled_tma_atom(
+ cpasync.CopyBulkTensorTileG2SOp(),
+ mKr,
+ sQK_layout,
+ (self.TT, self.D),
+ )
+ tma_atom_V, tma_tensor_V = cpasync.make_tiled_tma_atom(
+ cpasync.CopyBulkTensorTileG2SOp(),
+ mV,
+ sV_layout,
+ (self.TT, self.BV),
+ )
+ tma_atom_Mqk, tma_tensor_Mqk = cpasync.make_tiled_tma_atom(
+ cpasync.CopyBulkTensorTileG2SOp(),
+ mMqk,
+ sM_layout,
+ (self.BT, self.BT, self.CT),
+ )
+ tma_atom_Mkk, tma_tensor_Mkk = cpasync.make_tiled_tma_atom(
+ cpasync.CopyBulkTensorTileG2SOp(),
+ mMkk,
+ sM_layout,
+ (self.BT, self.BT, self.CT),
+ )
+ tma_atom_Gk, tma_tensor_Gk = cpasync.make_tiled_tma_atom(
+ cpasync.CopyBulkTensorTileG2SOp(),
+ mGk,
+ sGk_layout,
+ (self.D, self.CT),
+ )
+ tma_atom_O, tma_tensor_O = cpasync.make_tiled_tma_atom(
+ cpasync.CopyBulkTensorTileS2GOp(),
+ mO,
+ sO_layout,
+ (self.TT, self.BV),
+ )
+
+ universal_copy_bits = 128
+ copy_atom_O = cute.make_copy_atom(
+ cute.nvgpu.CopyUniversalOp(),
+ mO.element_type,
+ num_bits_per_copy=universal_copy_bits,
+ )
+ copy_elements = universal_copy_bits // mO.element_type.width
+ row_threads = self.BV // copy_elements
+ thr_layout_O = cute.make_layout(
+ (self.warp_threads // row_threads, row_threads), stride=(row_threads, 1)
+ )
+ val_layout_O = cute.make_layout((1, copy_elements))
+ tiled_copy_O_s2g = cute.make_tiled_copy_tv(
+ copy_atom_O, thr_layout_O, val_layout_O
+ )
+
+ grid = (self.D // self.BV, H, B)
+ self.kernel(
+ tma_atom_Q,
+ tma_tensor_Q,
+ tma_atom_K,
+ tma_tensor_K,
+ tma_atom_Kr,
+ tma_tensor_Kr,
+ tma_atom_V,
+ tma_tensor_V,
+ tma_atom_Mqk,
+ tma_tensor_Mqk,
+ tma_atom_Mkk,
+ tma_tensor_Mkk,
+ tma_atom_Gk,
+ tma_tensor_Gk,
+ tma_atom_O,
+ tma_tensor_O,
+ mO,
+ mS_in,
+ mS_out,
+ mCuSeqlens,
+ mStateIndices,
+ sQK_layout_staged,
+ sQKt_layout_staged,
+ sV_layout_staged,
+ sVt_layout_staged,
+ sM_layout_staged,
+ sGk_layout_staged,
+ sO_layout_staged,
+ sOt_layout_staged,
+ tiled_mma,
+ tiled_copy_Q,
+ tiled_copy_K,
+ tiled_copy_Kr,
+ tiled_copy_V,
+ tiled_copy_Mqk,
+ tiled_copy_Mkk,
+ tiled_copy_O,
+ tiled_copy_O_s2g,
+ SharedStorage,
+ ).launch(
+ grid=grid,
+ block=(self.threads_per_cta, 1, 1),
+ stream=stream,
+ )
+
+ @cute.kernel
+ def kernel(
+ self,
+ tma_atom_Q: cute.CopyAtom,
+ tma_tensor_Q: cute.Tensor,
+ tma_atom_K: cute.CopyAtom,
+ tma_tensor_K: cute.Tensor,
+ tma_atom_Kr: cute.CopyAtom,
+ tma_tensor_Kr: cute.Tensor,
+ tma_atom_V: cute.CopyAtom,
+ tma_tensor_V: cute.Tensor,
+ tma_atom_Mqk: cute.CopyAtom,
+ tma_tensor_Mqk: cute.Tensor,
+ tma_atom_Mkk: cute.CopyAtom,
+ tma_tensor_Mkk: cute.Tensor,
+ tma_atom_Gk: cute.CopyAtom,
+ tma_tensor_Gk: cute.Tensor,
+ tma_atom_O: cute.CopyAtom,
+ tma_tensor_O: cute.Tensor,
+ mO: cute.Tensor,
+ mS_in: Optional[cute.Tensor],
+ mS_out: Optional[cute.Tensor],
+ mCuSeqlens: Optional[cute.Tensor],
+ mStateIndices: Optional[cute.Tensor],
+ sQK_layout_staged: cute.ComposedLayout,
+ sQKt_layout_staged: cute.ComposedLayout,
+ sV_layout_staged: cute.ComposedLayout,
+ sVt_layout_staged: cute.ComposedLayout,
+ sM_layout_staged: cute.ComposedLayout,
+ sGk_layout_staged: cute.Layout,
+ sO_layout_staged: cute.ComposedLayout,
+ sOt_layout_staged: cute.ComposedLayout,
+ tiled_mma: cute.TiledMma,
+ tiled_copy_Q: cute.TiledCopy,
+ tiled_copy_K: cute.TiledCopy,
+ tiled_copy_Kr: cute.TiledCopy,
+ tiled_copy_V: cute.TiledCopy,
+ tiled_copy_Mqk: cute.TiledCopy,
+ tiled_copy_Mkk: cute.TiledCopy,
+ tiled_copy_O: cute.TiledCopy,
+ tiled_copy_O_s2g: cute.TiledCopy,
+ SharedStorage: cutlass.Constexpr,
+ ):
+ tidx, _, _ = cute.arch.thread_idx()
+ lane_idx = cute.arch.lane_idx()
+ warp_idx = cute.arch.make_warp_uniform(tidx // self.warp_threads)
+ warpgroup_idx = cute.arch.make_warp_uniform(tidx // self.warpgroup_threads)
+
+ v_idx, h_idx, b_idx = cute.arch.block_idx()
+ state_idx = b_idx
+ if cutlass.const_expr(mStateIndices is not None):
+ state_idx = mStateIndices[b_idx]
+
+ seq_len = cute.size(tma_tensor_Q, mode=[0])
+ seq_offset = 0
+ seq_offset_padded = 0
+ chunk_offset = 0
+ if cutlass.const_expr(mCuSeqlens is not None):
+ seq_len = get_seq_len(mCuSeqlens, b_idx)
+ seq_offset = get_seq_offset(mCuSeqlens, b_idx)
+ seq_offset_padded = get_seq_offset(mCuSeqlens, b_idx, padding=self.TT)
+ chunk_offset = get_seq_offset(
+ mCuSeqlens, b_idx, padding=self.TT, index_unit=self.BT
+ )
+ num_tiles = cute.ceil_div(seq_len, self.TT)
+
+ smem = cutlass.utils.SmemAllocator()
+ storage = smem.allocate(SharedStorage)
+ sQ = storage.sQ.get_tensor(
+ sQK_layout_staged.outer, swizzle=sQK_layout_staged.inner
+ )
+ sK = storage.sK.get_tensor(
+ sQK_layout_staged.outer, swizzle=sQK_layout_staged.inner
+ )
+ sKr = storage.sKr.get_tensor(
+ sQK_layout_staged.outer, swizzle=sQK_layout_staged.inner
+ )
+ sKrt = storage.sKr.get_tensor(
+ sQKt_layout_staged.outer, swizzle=sQKt_layout_staged.inner
+ )
+ sV = storage.sV.get_tensor(
+ sV_layout_staged.outer, swizzle=sV_layout_staged.inner
+ )
+ sVt = storage.sV.get_tensor(
+ sVt_layout_staged.outer, swizzle=sVt_layout_staged.inner
+ )
+ sMqk = storage.sMqk.get_tensor(
+ sM_layout_staged.outer, swizzle=sM_layout_staged.inner
+ )
+ sMkk = storage.sMkk.get_tensor(
+ sM_layout_staged.outer, swizzle=sM_layout_staged.inner
+ )
+ sGk = storage.sGk.get_tensor(sGk_layout_staged)
+ sO = storage.sO.get_tensor(
+ sO_layout_staged.outer, swizzle=sO_layout_staged.inner
+ )
+ sOt = storage.sO.get_tensor(
+ sOt_layout_staged.outer, swizzle=sOt_layout_staged.inner
+ )
+
+ tma_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+ tma_consumer_group = pipeline.CooperativeGroup(
+ pipeline.Agent.Thread, self.compute_warps
+ )
+ tma_producer, tma_consumer = pipeline.PipelineTmaAsync.create(
+ num_stages=self.load_stages,
+ producer_group=tma_producer_group,
+ consumer_group=tma_consumer_group,
+ barrier_storage=storage.tma_mbar_ptr.data_ptr(),
+ tx_count=self.num_tma_load_bytes,
+ ).make_participants()
+
+ compute_group = pipeline.CooperativeGroup(
+ pipeline.Agent.Thread, self.compute_threads
+ )
+ store_group = pipeline.CooperativeGroup(
+ pipeline.Agent.Thread, self.warp_threads
+ )
+ tma_producer_O, tma_consumer_O = pipeline.PipelineAsync.create(
+ num_stages=self.store_stages,
+ producer_group=compute_group,
+ consumer_group=store_group,
+ barrier_storage=storage.tma_O_mbar_ptr.data_ptr(),
+ ).make_participants()
+ store_pipeline = pipeline.PipelineTmaStore.create(
+ num_stages=self.store_stages,
+ producer_group=store_group,
+ )
+
+ if warp_idx == self.load_warp_id:
+ if cutlass.const_expr(mCuSeqlens is not None):
+ mQ_seq = cute.domain_offset((seq_offset_padded, 0, 0), tma_tensor_Q)
+ mK_seq = cute.domain_offset((seq_offset_padded, 0, 0), tma_tensor_K)
+ mKr_seq = cute.domain_offset((seq_offset_padded, 0, 0), tma_tensor_Kr)
+ mV_seq = cute.domain_offset((seq_offset, 0, 0), tma_tensor_V)
+ mMqk_seq = cute.domain_offset((0, 0, chunk_offset, 0), tma_tensor_Mqk)
+ mMkk_seq = cute.domain_offset((0, 0, chunk_offset, 0), tma_tensor_Mkk)
+ mGk_seq = cute.domain_offset((0, chunk_offset, 0), tma_tensor_Gk)
+ gQ = cute.local_tile(
+ mQ_seq[None, None, h_idx], (self.TT, self.D), (None, None)
+ )
+ gK = cute.local_tile(
+ mK_seq[None, None, h_idx], (self.TT, self.D), (None, None)
+ )
+ gKr = cute.local_tile(
+ mKr_seq[None, None, h_idx], (self.TT, self.D), (None, None)
+ )
+ gV = cute.local_tile(
+ mV_seq[None, None, h_idx], (self.TT, self.BV), (None, None)
+ )
+ gMqk = cute.local_tile(
+ mMqk_seq[None, None, None, h_idx],
+ (self.BT, self.BT, self.CT),
+ (None, None, None),
+ )
+ gMkk = cute.local_tile(
+ mMkk_seq[None, None, None, h_idx],
+ (self.BT, self.BT, self.CT),
+ (None, None, None),
+ )
+ gGk = cute.local_tile(
+ mGk_seq[None, None, h_idx], (self.D, self.CT), (None, None)
+ )
+ else:
+ gQ = cute.local_tile(
+ tma_tensor_Q[None, None, h_idx, b_idx],
+ (self.TT, self.D),
+ (None, None),
+ )
+ gK = cute.local_tile(
+ tma_tensor_K[None, None, h_idx, b_idx],
+ (self.TT, self.D),
+ (None, None),
+ )
+ gKr = cute.local_tile(
+ tma_tensor_Kr[None, None, h_idx, b_idx],
+ (self.TT, self.D),
+ (None, None),
+ )
+ gV = cute.local_tile(
+ tma_tensor_V[None, None, h_idx, b_idx],
+ (self.TT, self.BV),
+ (None, None),
+ )
+ gMqk = cute.local_tile(
+ tma_tensor_Mqk[None, None, None, h_idx, b_idx],
+ (self.BT, self.BT, self.CT),
+ (None, None, None),
+ )
+ gMkk = cute.local_tile(
+ tma_tensor_Mkk[None, None, None, h_idx, b_idx],
+ (self.BT, self.BT, self.CT),
+ (None, None, None),
+ )
+ gGk = cute.local_tile(
+ tma_tensor_Gk[None, None, h_idx, b_idx],
+ (self.D, self.CT),
+ (None, None),
+ )
+
+ tQsQ, tQgQ = cpasync.tma_partition(
+ tma_atom_Q,
+ 0,
+ cute.make_layout(1),
+ cute.group_modes(sQ, 0, 2),
+ cute.group_modes(gQ, 0, 2),
+ )
+ tKsK, tKgK = cpasync.tma_partition(
+ tma_atom_K,
+ 0,
+ cute.make_layout(1),
+ cute.group_modes(sK, 0, 2),
+ cute.group_modes(gK, 0, 2),
+ )
+ tKrsKr, tKrgKr = cpasync.tma_partition(
+ tma_atom_Kr,
+ 0,
+ cute.make_layout(1),
+ cute.group_modes(sKr, 0, 2),
+ cute.group_modes(gKr, 0, 2),
+ )
+ tVsV, tVgV = cpasync.tma_partition(
+ tma_atom_V,
+ 0,
+ cute.make_layout(1),
+ cute.group_modes(sV, 0, 2),
+ cute.group_modes(gV, 0, 2),
+ )
+ tMqksMqk, tMqkgMqk = cpasync.tma_partition(
+ tma_atom_Mqk,
+ 0,
+ cute.make_layout(1),
+ cute.group_modes(sMqk, 0, 3),
+ cute.group_modes(gMqk, 0, 3),
+ )
+ tMkksMkk, tMkkgMkk = cpasync.tma_partition(
+ tma_atom_Mkk,
+ 0,
+ cute.make_layout(1),
+ cute.group_modes(sMkk, 0, 3),
+ cute.group_modes(gMkk, 0, 3),
+ )
+ tGkGk, tGkgGk = cpasync.tma_partition(
+ tma_atom_Gk,
+ 0,
+ cute.make_layout(1),
+ cute.group_modes(sGk, 0, 2),
+ cute.group_modes(gGk, 0, 2),
+ )
+
+ for t in cutlass.range(num_tiles, unroll=0):
+ tma_empty = tma_producer.acquire_and_advance()
+ cute.copy(
+ tma_atom_Q,
+ tQgQ[None, t, 0],
+ tQsQ[None, tma_empty.index],
+ tma_bar_ptr=tma_empty.barrier,
+ )
+ cute.copy(
+ tma_atom_K,
+ tKgK[None, t, 0],
+ tKsK[None, tma_empty.index],
+ tma_bar_ptr=tma_empty.barrier,
+ )
+ cute.copy(
+ tma_atom_Kr,
+ tKrgKr[None, t, 0],
+ tKrsKr[None, tma_empty.index],
+ tma_bar_ptr=tma_empty.barrier,
+ )
+ cute.copy(
+ tma_atom_V,
+ tVgV[None, t, v_idx],
+ tVsV[None, tma_empty.index],
+ tma_bar_ptr=tma_empty.barrier,
+ )
+ cute.copy(
+ tma_atom_Mqk,
+ tMqkgMqk[None, 0, 0, t],
+ tMqksMqk[None, tma_empty.index],
+ tma_bar_ptr=tma_empty.barrier,
+ )
+ cute.copy(
+ tma_atom_Mkk,
+ tMkkgMkk[None, 0, 0, t],
+ tMkksMkk[None, tma_empty.index],
+ tma_bar_ptr=tma_empty.barrier,
+ )
+ cute.copy(
+ tma_atom_Gk,
+ tGkgGk[None, 0, t],
+ tGkGk[None, tma_empty.index],
+ tma_bar_ptr=tma_empty.barrier,
+ )
+
+ elif warp_idx == self.store_warp_id:
+ if cutlass.const_expr(mCuSeqlens is not None):
+ mO_seq = cute.domain_offset((seq_offset, 0, 0), tma_tensor_O)
+ gO = cute.local_tile(
+ mO_seq[None, None, h_idx], (self.TT, self.BV), (None, None)
+ )
+ mO_seq_s2g = cute.domain_offset((seq_offset, 0, 0), mO)
+ gO_s2g = cute.local_tile(
+ mO_seq_s2g[None, None, h_idx], (self.TT, self.BV), (None, None)
+ )
+ else:
+ gO = cute.local_tile(
+ tma_tensor_O[None, None, h_idx, b_idx],
+ (self.TT, self.BV),
+ (None, None),
+ )
+
+ tOsO, tOgO = cpasync.tma_partition(
+ tma_atom_O,
+ 0,
+ cute.make_layout(1),
+ cute.group_modes(sO, 0, 2),
+ cute.group_modes(gO, 0, 2),
+ )
+
+ for t in cutlass.range(num_tiles, unroll=0):
+ tma_O_full = tma_consumer_O.wait_and_advance()
+ cute.arch.fence_view_async_shared()
+ if cutlass.const_expr(mCuSeqlens is not None):
+ valid_tokens = min(self.TT, seq_len - t * self.TT)
+ if valid_tokens == self.TT:
+ cute.copy(
+ tma_atom_O,
+ tOsO[None, tma_O_full.index],
+ tOgO[None, t, v_idx],
+ )
+ store_pipeline.producer_commit()
+ store_pipeline.producer_acquire()
+ else:
+ thr_copy_O_s2g = tiled_copy_O_s2g.get_slice(lane_idx)
+ tOsO_s2g = thr_copy_O_s2g.partition_S(
+ sO[None, None, tma_O_full.index]
+ )
+ tOrO_s2g = cute.make_rmem_tensor_like(tOsO_s2g)
+ tOgO_s2g = thr_copy_O_s2g.partition_D(
+ gO_s2g[None, None, t, v_idx]
+ )
+
+ cO = cute.make_identity_tensor((self.TT, self.BV))
+ tOcO_s2g = thr_copy_O_s2g.partition_S(cO)
+ pred_shape = (1, *tOcO_s2g.shape[1:])
+ tOpO_s2g = cute.make_rmem_tensor(pred_shape, cutlass.Boolean)
+ for m_idx in cutlass.range_constexpr(tOcO_s2g.shape[1]):
+ for n_idx in cutlass.range_constexpr(tOcO_s2g.shape[2]):
+ m, n = tOcO_s2g[(0, 0), m_idx, n_idx]
+ tOpO_s2g[0, m_idx, n_idx] = m < valid_tokens
+
+ cute.copy(tiled_copy_O_s2g, tOsO_s2g, tOrO_s2g)
+ cute.copy(tiled_copy_O_s2g, tOrO_s2g, tOgO_s2g, pred=tOpO_s2g)
+ else:
+ cute.copy(
+ tma_atom_O,
+ tOsO[None, tma_O_full.index],
+ tOgO[None, t, v_idx],
+ )
+ store_pipeline.producer_commit()
+ store_pipeline.producer_acquire()
+ tma_O_full.release()
+ store_pipeline.producer_tail()
+
+ elif warpgroup_idx in self.compute_warpgroup_ids:
+ tidx_in_compute = tidx % self.compute_threads
+
+ thr_mma = tiled_mma.get_slice(tidx_in_compute)
+
+ S_shape = thr_mma.partition_shape_C((self.BV, self.D))
+ tSrS = thr_mma.make_fragment_C(S_shape) # 64 regs
+ S_shape = thr_mma.partition_shape_A((self.BV, self.D))
+ tArS = thr_mma.make_fragment_A(S_shape) # 32 regs
+
+ # tSrS always represents S^T; transpose_S describes the external buffers.
+ if cutlass.const_expr(mS_in is not None):
+ if cutlass.const_expr(self.transpose_S):
+ gS_in = cute.local_tile(
+ mS_in[None, None, h_idx, state_idx], (self.BV, self.D), (v_idx, 0)
+ )
+ else:
+ gS_in = cute.local_tile(
+ mS_in[None, None, h_idx, state_idx], (self.D, self.BV), (0, v_idx)
+ )
+ gS_in = cute.make_tensor(
+ gS_in.iterator, cute.select(gS_in.layout, mode=[1, 0])
+ )
+ tSgS_in = thr_mma.partition_C(gS_in)
+ if cutlass.const_expr(tSgS_in.element_type.width == 16):
+ tSrS_f16 = cute.make_rmem_tensor_like(tSrS, tSgS_in.element_type)
+ cute.autovec_copy(tSgS_in, tSrS_f16)
+ tSrS.store(tSrS_f16.load().to(tSrS.element_type))
+ else:
+ cute.autovec_copy(tSgS_in, tSrS)
+ else:
+ tSrS.fill(0.0)
+
+ for _ in cutlass.range(num_tiles, unroll=0):
+ tma_full = tma_consumer.wait_and_advance()
+ tma_O_empty = None
+
+ sQ_tile = cute.local_tile(
+ sQ[None, None, tma_full.index], (self.BT, self.D), (None, 0)
+ )
+ sK_tile = cute.local_tile(
+ sK[None, None, tma_full.index], (self.BT, self.D), (None, 0)
+ )
+ sKr_tile = cute.local_tile(
+ sKrt[None, None, tma_full.index], (self.D, self.BT), (0, None)
+ )
+ sVt_tile = cute.local_tile(
+ sVt[None, None, tma_full.index], (self.BV, self.BT), (0, None)
+ )
+ sMqk_tile = sMqk[None, None, None, tma_full.index]
+ sMkk_tile = sMkk[None, None, None, tma_full.index]
+ sGk_tile = cute.make_tensor(
+ sGk[None, None, tma_full.index].iterator,
+ cute.make_layout((self.BV, self.D, self.CT), stride=(0, 1, self.D)),
+ )
+ sOt_tile = cute.local_tile(
+ sOt[None, None, 0], (self.BV, self.BT), (0, None)
+ ) # dummy
+
+ tVrK = thr_mma.make_fragment_B(thr_mma.partition_B(sK_tile))
+ thr_copy_K = tiled_copy_K.get_slice(tidx_in_compute)
+ tVsK_copy_view = thr_copy_K.partition_S(sK_tile)
+ tVrK_copy_view = thr_copy_K.retile(tVrK)
+
+ tOrQ = thr_mma.make_fragment_B(thr_mma.partition_B(sQ_tile))
+ thr_copy_Q = tiled_copy_Q.get_slice(tidx_in_compute)
+ tOsQ_copy_view = thr_copy_Q.partition_S(sQ_tile)
+ tOrQ_copy_view = thr_copy_Q.retile(tOrQ)
+
+ tPrV = thr_mma.make_fragment_A(thr_mma.partition_A(sVt_tile))
+ tPrMkk = thr_mma.make_fragment_B(thr_mma.partition_B(sMkk_tile))
+ thr_copy_V = tiled_copy_V.get_slice(tidx_in_compute)
+ tPsV_copy_view = thr_copy_V.partition_S(sVt_tile)
+ tPrV_copy_view = thr_copy_V.retile(tPrV)
+ thr_copy_Mkk = tiled_copy_Mkk.get_slice(tidx_in_compute)
+ tPsMkk_copy_view = thr_copy_Mkk.partition_S(sMkk_tile)
+ tPrMkk_copy_view = thr_copy_Mkk.retile(tPrMkk)
+
+ tSrKr = thr_mma.make_fragment_B(thr_mma.partition_B(sKr_tile))
+ tSsGk = thr_mma.partition_C(sGk_tile)
+ thr_copy_Kr = tiled_copy_Kr.get_slice(tidx_in_compute)
+ tSsKr_copy_view = thr_copy_Kr.partition_S(sKr_tile)
+ tSrKr_copy_view = thr_copy_Kr.retile(tSrKr)
+
+ tOrO = thr_mma.make_fragment_C(
+ thr_mma.partition_C(sOt_tile)
+ ) # 8 regs per BT
+ tOrO.fill(0.0)
+ tOrO_f16 = cute.make_rmem_tensor_like(tOrO, sOt_tile.element_type)
+ thr_copy_O = tiled_copy_O.get_slice(tidx_in_compute)
+ tOrO_copy_view = thr_copy_O.retile(tOrO_f16)
+ tOsO_copy_view = thr_copy_O.partition_D(sOt_tile)
+
+ tOrMqk = thr_mma.make_fragment_B(thr_mma.partition_B(sMqk_tile))
+ thr_copy_Mqk = tiled_copy_Mqk.get_slice(tidx_in_compute)
+ tOsMqk_copy_view = thr_copy_Mqk.partition_S(sMqk_tile)
+ tOrMqk_copy_view = thr_copy_Mqk.retile(tOrMqk)
+
+ for ti in cutlass.range_constexpr(self.CT):
+ tArS_f32 = cute.make_tensor(tSrS.iterator, tArS.layout)
+ if cutlass.const_expr(ti == 0):
+ cute.copy(
+ tiled_copy_K,
+ tVsK_copy_view[None, None, None, ti],
+ tVrK_copy_view[None, None, None, ti],
+ ) # 32 regs, 8 instructions
+ tma_O_empty = tma_producer_O.acquire_and_advance()
+
+ ### S f32 -> f16, dV = K @ S, load Q, V, Mkk, and first half of Gk
+
+ V_shape = thr_mma.partition_shape_C((self.BV, self.BT))
+ tVrV = thr_mma.make_fragment_C(V_shape) # 8 regs
+ tVrV.fill(0.0)
+
+ tSsGk_mn = cute.make_tensor(
+ tSsGk[None, None, None, ti].iterator,
+ layout_acc_mn(tiled_mma, tSsGk[None, None, None, ti].layout),
+ )
+ tSsGk_mn = cute.flatten(tSsGk_mn[0, None])
+ tSrGk_mn = cute.make_rmem_tensor_like(tSsGk_mn)
+ tSrS_mn = cute.make_tensor(
+ tSrS.iterator, layout_acc_mn(tiled_mma, tSrS.layout)
+ ) # ((2,1),(2,16)):((2,0),(1,4))
+
+ total_k = cute.size(tArS, mode=[2]) # 8
+ tArS[None, None, 0].store(
+ tArS_f32[None, None, 0].load().to(tArS.element_type)
+ )
+ for k in cutlass.range_constexpr(total_k): # 8
+ if cutlass.const_expr(k < total_k - 1):
+ tArS[None, None, k + 1].store(
+ tArS_f32[None, None, k + 1].load().to(tArS.element_type)
+ )
+
+ cute.gemm(
+ tiled_mma,
+ tVrV,
+ tArS[None, None, k],
+ tVrK[None, None, k, ti],
+ tVrV,
+ )
+
+ if cutlass.const_expr(k < total_k // 2):
+ cute.copy(
+ tiled_copy_Q,
+ tOsQ_copy_view[None, None, 2 * k, ti],
+ tOrQ_copy_view[None, None, 2 * k, ti],
+ ) # 32 regs
+ cute.copy(
+ tiled_copy_Q,
+ tOsQ_copy_view[None, None, 2 * k + 1, ti],
+ tOrQ_copy_view[None, None, 2 * k + 1, ti],
+ ) # 32 regs
+ else:
+ # 16 regs, 8 instructions
+ cute.autovec_copy(
+ tSsGk_mn[None, 2 * k - total_k],
+ tSrGk_mn[None, 2 * k - total_k],
+ )
+ cute.autovec_copy(
+ tSsGk_mn[None, 2 * k + 1 - total_k],
+ tSrGk_mn[None, 2 * k + 1 - total_k],
+ )
+
+ if cutlass.const_expr(k == 6):
+ cute.copy(
+ tiled_copy_V,
+ tPsV_copy_view[None, None, None, ti],
+ tPrV_copy_view[None, None, None, ti],
+ ) # 4 regs, 1 instruction
+ cute.copy(
+ tiled_copy_Mkk,
+ tPsMkk_copy_view[None, None, None, ti],
+ tPrMkk_copy_view[None, None, None, ti],
+ ) # 4 regs, 1 instruction
+
+ ### O = Q @ S, P = Mkk @ (V - dV), load second half of Gk and K^T
+
+ P_shape = thr_mma.partition_shape_C((self.BV, self.BT))
+ tPrP = thr_mma.make_fragment_C(P_shape) # 8 regs
+ tPrP.fill(0.0)
+
+ total_k = cute.size(tArS, mode=[2]) # 8
+ for k in cutlass.range_constexpr(total_k): # 8
+ cute.gemm(
+ tiled_mma,
+ tOrO[None, None, None, ti],
+ tArS[None, None, k],
+ tOrQ[None, None, k, ti],
+ tOrO[None, None, None, ti],
+ )
+
+ if cutlass.const_expr(k == total_k // 2):
+ tPrV_tile = tPrV[None, None, None, ti]
+ for i in cutlass.range_constexpr(cute.size(tPrV_tile)):
+ tPrV_tile[i] -= tVrV[i].to(tPrV_tile.element_type)
+ cute.gemm(
+ tiled_mma,
+ tPrP,
+ tPrV_tile,
+ tPrMkk[None, None, None, ti],
+ tPrP,
+ )
+
+ # 16 regs, 8 instructions
+ cute.autovec_copy(
+ tSsGk_mn[None, total_k + k], tSrGk_mn[None, total_k + k]
+ )
+
+ cute.copy(
+ tiled_copy_Kr,
+ tSsKr_copy_view[None, k, None, ti],
+ tSrKr_copy_view[None, k, None, ti],
+ ) # 32 regs, 8 instructions
+
+ if cutlass.const_expr(k >= total_k // 2):
+ for j in cutlass.range_constexpr(2):
+ c = ((k - total_k // 2) * 2 + j) * 2
+ for r in cutlass.range_constexpr(2):
+ tSrS_mn[r, c], tSrS_mn[r, c + 1] = (
+ cute.arch.mul_packed_f32x2(
+ (tSrGk_mn[c], tSrGk_mn[c + 1]),
+ (tSrS_mn[r, c], tSrS_mn[r, c + 1]),
+ )
+ )
+
+ P_shape = thr_mma.partition_shape_A((self.BV, self.BT))
+ tSrP = thr_mma.make_fragment_A(P_shape)
+ tPrP_f16 = cute.make_tensor(tSrP.iterator, tPrP.layout)
+ tPrP_f16.store(tPrP.load().to(tPrP_f16.element_type))
+
+ ### S = Gk * S + K^T @ P, load Mqk and K, O += Mqk @ P, store O
+
+ cute.copy(
+ tiled_copy_Mqk,
+ tOsMqk_copy_view[None, None, None, ti],
+ tOrMqk_copy_view[None, None, None, ti],
+ ) # 4 regs, 1 instruction
+
+ if cutlass.const_expr(ti == self.CT - 1):
+ tma_full.release()
+
+ assert (
+ cute.size(tSrP, mode=[1]) == 1
+ and cute.size(tSrP, mode=[2]) == 1
+ ) # rest_M = 1, rest_K = 1
+ total_n = cute.size(tSrKr, mode=[1])
+ prefetch_gk_s = total_n // 2
+ for n in cutlass.range_constexpr(total_n): # 16
+ if cutlass.const_expr(n == total_n // 2):
+ cute.gemm(
+ tiled_mma,
+ tOrO[None, None, None, ti],
+ tSrP,
+ tOrMqk[None, None, None, ti],
+ tOrO[None, None, None, ti],
+ ) # 2 instructions
+
+ if cutlass.const_expr(n < total_n - prefetch_gk_s):
+ c = 2 * (n + prefetch_gk_s)
+ for r in cutlass.range_constexpr(2):
+ tSrS_mn[r, c], tSrS_mn[r, c + 1] = (
+ cute.arch.mul_packed_f32x2(
+ (tSrGk_mn[c], tSrGk_mn[c + 1]),
+ (tSrS_mn[r, c], tSrS_mn[r, c + 1]),
+ )
+ )
+ cute.gemm(
+ tiled_mma,
+ self.append_unit_mode(tSrS[None, 0, n], 2),
+ self.append_unit_mode(tSrP[None, 0, 0]),
+ self.append_unit_mode(tSrKr[None, n, 0, ti]),
+ self.append_unit_mode(tSrS[None, 0, n], 2),
+ )
+ if cutlass.const_expr(ti < self.CT - 1 and n % 2 == 0):
+ cute.copy(
+ tiled_copy_K,
+ tVsK_copy_view[None, None, n // 2, ti + 1],
+ tVrK_copy_view[None, None, n // 2, ti + 1],
+ ) # 32 regs, 8 instructions
+
+ if cutlass.const_expr(ti == 0):
+ sOt_tile = cute.local_tile(
+ sOt[None, None, tma_O_empty.index],
+ (self.BV, self.BT),
+ (0, None),
+ )
+ tOsO_copy_view = thr_copy_O.partition_D(sOt_tile)
+ tOrO_f16[None, None, None, ti].store(
+ tOrO[None, None, None, ti].load().to(tOrO_f16.element_type)
+ )
+ cute.copy(
+ tiled_copy_O,
+ tOrO_copy_view[None, None, None, ti],
+ tOsO_copy_view[None, None, None, ti],
+ )
+
+ tma_O_empty.commit()
+
+ if cutlass.const_expr(mS_out is not None):
+ if cutlass.const_expr(self.transpose_S):
+ gS_out = cute.local_tile(
+ mS_out[None, None, h_idx, state_idx], (self.BV, self.D), (v_idx, 0)
+ )
+ else:
+ gS_out = cute.local_tile(
+ mS_out[None, None, h_idx, state_idx], (self.D, self.BV), (0, v_idx)
+ )
+ gS_out = cute.make_tensor(
+ gS_out.iterator, cute.select(gS_out.layout, mode=[1, 0])
+ )
+ tSgS_out = thr_mma.partition_C(gS_out)
+ if cutlass.const_expr(tSgS_out.element_type.width == 16):
+ tSrS_f16 = cute.make_rmem_tensor_like(tSrS, tSgS_out.element_type)
+ tSrS_f16.store(tSrS.load().to(tSrS_f16.element_type))
+ cute.autovec_copy(tSrS_f16, tSgS_out)
+ else:
+ cute.autovec_copy(tSrS, tSgS_out)
+
+ @staticmethod
+ def append_unit_mode(t: cute.Tensor, mode: int = 1):
+ unit = cute.make_layout(1, stride=0)
+ layout = t.layout
+ for _ in range(mode):
+ layout = cute.append(layout, unit)
+ return cute.make_tensor(t.iterator, layout)
+
+
+_BT = 16
+_CT = 4
+_TT = _CT * _BT
+
+
+def _cute_dtype(dtype: torch.dtype):
+ if dtype == torch.bfloat16:
+ return cutlass.BFloat16
+ if dtype == torch.float16:
+ return cutlass.Float16
+ if dtype == torch.float32:
+ return cutlass.Float32
+ if dtype == torch.int32:
+ return cutlass.Int32
+ if dtype == torch.int64:
+ return cutlass.Int64
+ raise ValueError(f"unsupported dtype: {dtype}")
+
+
+def _fake(dtype, shape, assumed_align=16):
+ return cute.runtime.make_fake_compact_tensor(
+ dtype,
+ shape,
+ stride_order=tuple(reversed(range(len(shape)))),
+ assumed_align=assumed_align,
+ )
+
+
+def _fake_stride(dtype, shape):
+ stride = tuple(cute.sym_int() for _ in range(len(shape) - 1))
+ stride = stride + (1,)
+ return cute.runtime.make_fake_tensor(
+ dtype,
+ shape,
+ stride=stride,
+ assumed_align=16,
+ )
+
+
+@functools.cache
+def _compile_k1(
+ dtype: torch.dtype,
+ seqlen_dtype: torch.dtype,
+ d: int,
+ h: int,
+ gate_lower_bound: bool,
+ is_varlen: bool,
+ max_active_clusters: int,
+):
+ if torch.cuda.is_current_stream_capturing():
+ raise RuntimeError("Warm Small-BH KDA before CUDA graph capture")
+ cu_dtype = _cute_dtype(dtype)
+ b, t, tp, nt = (cute.sym_int() for _ in range(4))
+ feature: tuple[object, ...]
+ padded_feature: tuple[object, ...]
+ beta_shape: tuple[object, ...]
+ matrix: tuple[object, ...]
+ gate: tuple[object, ...]
+ if is_varlen:
+ feature = (t, h, d)
+ padded_feature = (tp, h, d)
+ beta_shape = (t, h)
+ matrix = (nt, h, _BT, _BT)
+ gate = (nt, h, d)
+ cu_seqlens = _fake(_cute_dtype(seqlen_dtype), (b,))
+ else:
+ feature = (b, t, h, d)
+ padded_feature = feature
+ beta_shape = (b, t, h)
+ matrix = (b, nt, h, _BT, _BT)
+ gate = (b, nt, h, d)
+ cu_seqlens = None
+
+ q, k, g = [_fake_stride(cu_dtype, feature) for _ in range(3)]
+ qd, kd, kr = [_fake(cu_dtype, padded_feature) for _ in range(3)]
+ kernel = ChunkKdaFwdK1(
+ D=d,
+ H=h,
+ BT=_BT,
+ gate_lower_bound=gate_lower_bound,
+ max_active_clusters=max_active_clusters,
+ )
+ stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+ return cute.compile(
+ kernel,
+ q,
+ k,
+ g,
+ _fake_stride(cu_dtype, beta_shape),
+ _fake(cutlass.Float32, (h, d)),
+ _fake(cutlass.Float32, (h,)),
+ qd,
+ kd,
+ kr,
+ _fake(cu_dtype, matrix),
+ _fake(cu_dtype, matrix),
+ _fake(cutlass.Float32, gate),
+ cutlass.Float32(1),
+ cutlass.Float32(1),
+ cu_seqlens,
+ cutlass.Int32(1),
+ stream,
+ options="--enable-tvm-ffi",
+ )
+
+
+@functools.cache
+def _compile_k2(
+ dtype: torch.dtype,
+ seqlen_dtype: torch.dtype,
+ state_dtype: torch.dtype,
+ d: int,
+ bv: int,
+ has_state_input: bool,
+ has_state_output: bool,
+ transpose_state: bool,
+ is_varlen: bool,
+ state_row_stride: int | None,
+ has_state_indices: bool,
+):
+ if torch.cuda.is_current_stream_capturing():
+ raise RuntimeError("Warm Small-BH KDA before CUDA graph capture")
+ cu_dtype = _cute_dtype(dtype)
+ b, bs, t, tp, h, nt = (cute.sym_int() for _ in range(6))
+ feature: tuple[object, ...]
+ value: tuple[object, ...]
+ output: tuple[object, ...]
+ matrix: tuple[object, ...]
+ gate: tuple[object, ...]
+ if is_varlen:
+ feature = (tp, h, d)
+ value = output = (t, h, d)
+ matrix = (nt, h, _BT, _BT)
+ gate = (nt, h, d)
+ cu_seqlens = _fake(_cute_dtype(seqlen_dtype), (bs,))
+ else:
+ feature = value = output = (b, t, h, d)
+ matrix = (b, nt, h, _BT, _BT)
+ gate = (b, nt, h, d)
+ cu_seqlens = None
+ state = (b, h, d, d)
+
+ def fake_state():
+ if state_row_stride is None:
+ return _fake(_cute_dtype(state_dtype), state)
+ return cute.runtime.make_fake_tensor(
+ _cute_dtype(state_dtype), state,
+ stride=(state_row_stride, d * d, d, 1), assumed_align=16,
+ )
+
+ q, k, kr = [_fake(cu_dtype, feature) for _ in range(3)]
+ kernel = ChunkKdaFwdK2(D=d, BT=_BT, BV=bv, transpose_S=transpose_state)
+ stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+ return cute.compile(
+ kernel,
+ q,
+ k,
+ kr,
+ _fake(cu_dtype, value),
+ _fake(cu_dtype, matrix),
+ _fake(cu_dtype, matrix),
+ _fake(cutlass.Float32, gate),
+ _fake(cu_dtype, output),
+ fake_state() if has_state_input else None,
+ fake_state() if has_state_output else None,
+ cu_seqlens,
+ _fake(cutlass.Int32, (cute.sym_int(),), assumed_align=4) if has_state_indices else None,
+ stream,
+ options="--enable-tvm-ffi",
+ )
+
+
+def chunk_kda_fwd_k1(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ g: torch.Tensor,
+ beta: torch.Tensor,
+ dt_bias: torch.Tensor,
+ a_log: torch.Tensor,
+ *,
+ scale: float | None = None,
+ gate_scale: float = 1.0,
+ gate_lower_bound: bool = False,
+ cu_seqlens: torch.Tensor | None = None,
+ max_seqlen: int | None = None,
+ outputs: tuple[torch.Tensor, ...] | None = None,
+ max_active_clusters: int = 152,
+) -> tuple[torch.Tensor, ...]:
+ is_varlen = cu_seqlens is not None
+ matrix_shape: tuple[int, ...]
+ gate_shape: tuple[int, ...]
+ if is_varlen:
+ t, h, d = q.shape
+ b = len(cu_seqlens) - 1
+ t_pad = t + b * _TT
+ nt = (t_pad + _BT - 1) // _BT
+ if max_seqlen is None:
+ max_seqlen = t
+ seq_shape = (t_pad, h, d)
+ matrix_shape = (nt, h, _BT, _BT)
+ gate_shape = (nt, h, d)
+ else:
+ b, t, h, d = q.shape
+ nt = ((t + _TT - 1) // _TT) * _CT
+ max_seqlen = t
+ seq_shape = q.shape
+ matrix_shape = (b, nt, h, _BT, _BT)
+ gate_shape = (b, nt, h, d)
+
+ if d != 128:
+ raise ValueError(f"unsupported d: {d}")
+
+ if outputs is None:
+ outputs = (
+ *(torch.empty(seq_shape, dtype=q.dtype, device=q.device) for _ in range(3)),
+ torch.empty(matrix_shape, dtype=q.dtype, device=q.device),
+ torch.empty(matrix_shape, dtype=q.dtype, device=q.device),
+ torch.empty(gate_shape, dtype=torch.float32, device=q.device),
+ )
+ qd, kd, kr, mqk, mkk, gk = outputs
+
+ compiled = _compile_k1(
+ q.dtype,
+ cu_seqlens.dtype if cu_seqlens is not None else torch.int32,
+ d,
+ h,
+ gate_lower_bound,
+ is_varlen,
+ max_active_clusters,
+ )
+ compiled(
+ q,
+ k,
+ g,
+ beta,
+ dt_bias,
+ a_log,
+ qd,
+ kd,
+ kr,
+ mqk,
+ mkk,
+ gk,
+ d**-0.5 if scale is None else scale,
+ gate_scale,
+ cu_seqlens,
+ max_seqlen,
+ )
+ return outputs
+
+
+def chunk_kda_fwd_k2(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ kr: torch.Tensor,
+ v: torch.Tensor,
+ mqk: torch.Tensor,
+ mkk: torch.Tensor,
+ gk: torch.Tensor,
+ *,
+ return_state: bool = False,
+ transpose_state: bool = False,
+ cu_seqlens: torch.Tensor | None = None,
+ output: torch.Tensor | None = None,
+ state_input: torch.Tensor | None = None,
+ state_output: torch.Tensor | None = None,
+ state_indices: torch.Tensor | None = None,
+ max_active_clusters: int = 152,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+ is_varlen = cu_seqlens is not None
+ if is_varlen:
+ _, h, d = q.shape
+ b = len(cu_seqlens) - 1
+ else:
+ b, _, h, d = q.shape
+ bv = d if b * h * 2 > max_active_clusters else d // 2
+
+ output = torch.empty_like(v) if output is None else output
+ if return_state and state_output is None:
+ state_output = torch.empty((b, h, d, d), dtype=torch.float32, device=q.device)
+
+ has_input, has_output = state_input is not None, state_output is not None
+ state_dtype = (
+ state_input.dtype
+ if state_input is not None
+ else state_output.dtype
+ if state_output is not None
+ else q.dtype
+ )
+ compiled = _compile_k2(
+ q.dtype,
+ cu_seqlens.dtype if cu_seqlens is not None else torch.int32,
+ state_dtype,
+ d,
+ bv,
+ has_input,
+ has_output,
+ transpose_state,
+ is_varlen,
+ state_input.stride(0) if state_indices is not None else None,
+ state_indices is not None,
+ )
+ compiled(
+ q,
+ k,
+ kr,
+ v,
+ mqk,
+ mkk,
+ gk,
+ output,
+ state_input,
+ state_output,
+ cu_seqlens,
+ state_indices,
+ )
+ return output, state_output
+
+
+def chunk_kda_fwd(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ g: torch.Tensor,
+ beta: torch.Tensor,
+ dt_bias: torch.Tensor,
+ a_log: torch.Tensor,
+ *,
+ scale: float | None = None,
+ gate_scale: float = 1.0,
+ gate_lower_bound: bool = False,
+ return_state: bool = False,
+ transpose_state: bool = False,
+ cu_seqlens: torch.Tensor | None = None,
+ max_seqlen: int | None = None,
+ output: torch.Tensor | None = None,
+ state_input: torch.Tensor | None = None,
+ state_output: torch.Tensor | None = None,
+ state_indices: torch.Tensor | None = None,
+ workspace_tensors: tuple[torch.Tensor, ...] | None = None,
+ max_active_clusters: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+ if max_active_clusters <= 0:
+ max_active_clusters = torch.cuda.get_device_properties(q.device).multi_processor_count
+ qd, kd, kr, mqk, mkk, gk = chunk_kda_fwd_k1(
+ q,
+ k,
+ g,
+ beta,
+ dt_bias,
+ a_log,
+ scale=scale,
+ gate_scale=gate_scale,
+ gate_lower_bound=gate_lower_bound,
+ cu_seqlens=cu_seqlens,
+ max_seqlen=max_seqlen,
+ outputs=workspace_tensors,
+ max_active_clusters=max_active_clusters,
+ )
+ return chunk_kda_fwd_k2(
+ qd,
+ kd,
+ kr,
+ v,
+ mqk,
+ mkk,
+ gk,
+ return_state=return_state,
+ transpose_state=transpose_state,
+ cu_seqlens=cu_seqlens,
+ output=output,
+ state_input=state_input,
+ state_output=state_output,
+ state_indices=state_indices,
+ max_active_clusters=max_active_clusters,
+ )
+
+
+def clear_compilation_cache() -> None:
+ _compile_k1.cache_clear()
+ _compile_k2.cache_clear()
+
+
+__all__ = [
+ "chunk_kda_fwd_k1",
+ "chunk_kda_fwd_k2",
+ "chunk_kda_fwd",
+ "clear_compilation_cache",
+]
diff --git a/tensorrt_llm/_torch/custom_ops/__init__.py b/tensorrt_llm/_torch/custom_ops/__init__.py
index 03c214b201b0..eed9430cf90f 100644
--- a/tensorrt_llm/_torch/custom_ops/__init__.py
+++ b/tensorrt_llm/_torch/custom_ops/__init__.py
@@ -94,9 +94,9 @@ def inplace_slice_copy(dest: torch.Tensor, src: torch.Tensor, dim1_start: int,
cute_dsl_megamoe_nvfp4_blackwell
__all__ += ['cute_dsl_megamoe_nvfp4_blackwell']
-if IS_CUTLASS_DSL_AVAILABLE and IS_FLASHINFER_AVAILABLE:
- from .cute_dsl_kimi_k3_custom_ops import kda_prefill
- __all__ += ['kda_prefill']
+if IS_CUTLASS_DSL_AVAILABLE:
+ from .cute_dsl_kimi_k3_custom_ops import kda_prefill, kda_prefill_conv
+ __all__ += ['kda_prefill', 'kda_prefill_conv']
if IS_CUDA_TILE_AVAILABLE:
from .cuda_tile_custom_ops import (cuda_tile_rms_norm,
diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py
index 958ceee0b7ff..5210664f89c3 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py
@@ -1,905 +1,42 @@
-# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""FlashInfer BT16/Small-BH KDA prefill and fused short convolution.
-"""CuTe DSL custom op for Kimi K3 KDA prefill.
-
-The implementation source-integrates the production orchestration from the
-KDA ``chunk_kda_fwd`` benchmark. The default pipeline launches fused K123,
-Akk inverse, and persistent K4 kernels.
-
-Supported:
- - Equal-length sequences (B ≥ 1, with the existing B=1 padding path)
- - Variable-length sequences (B=1 with cu_seqlens)
- - safe_gate mode (sigmoid + lower_bound) and softplus mode
- - dt_bias
- - use_gate_in_kernel=True with A_log
- - optional in-kernel beta sigmoid
- - chunk_size=64
-
-The ``trtllm::kda_prefill`` operator updates selected recurrent-state pool rows
-in place. Intermediate matrices remain private runner workspace.
+Inputs are raw BF16 Q/K/V, gate and beta logits. BT16 performs Q/K
+normalization and gate/beta activation. Recurrent state is updated directly
+in its resident BF16 or FP32 pool, with no FP32 staging/writeback.
"""
-import os
-import weakref
-from typing import Any, Hashable, Optional, TypeVar
+import math
+from functools import lru_cache
import torch
-from tensorrt_llm.logger import logger
-
-from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE
-from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE
-if IS_CUTLASS_DSL_AVAILABLE and IS_FLASHINFER_AVAILABLE:
- import cuda.bindings.driver as cuda_driver
+@lru_cache(maxsize=32)
+def _compile_prefill(device: torch.device, state_dtype: torch.dtype, lower_bound: float, mode: str):
+ if torch.cuda.is_current_stream_capturing():
+ raise RuntimeError("Warm KDA prefill before CUDA graph capture")
import cutlass
- import cutlass.cute as cute
- from cutlass.cute.runtime import from_dlpack
-
- from ..attention.kernels.blackwell.kimi_k3_kda.akk_inverse import akk_inv_host as _akk_inv_host
- from ..attention.kernels.blackwell.kimi_k3_kda.fused_k123 import (
- make_host_function as _fused_make_host,
- )
- from ..attention.kernels.blackwell.kimi_k3_kda.k4_persistent import (
- BYTES_PER_TENSORMAP as _K4P_BTM,
- )
- from ..attention.kernels.blackwell.kimi_k3_kda.k4_persistent import NUM_TENSORMAPS as _K4P_NTM
- from ..attention.kernels.blackwell.kimi_k3_kda.k4_persistent import (
- make_host_fn as _k4p_make_host,
- )
- from ..modules.fla.index import prepare_chunk_indices
-else:
- raise ImportError("Kimi K3 KDA prefill requires NVIDIA CUTLASS DSL and FlashInfer")
-
-
-def _ct(t, etype):
- """Create a CuTe tensor from PyTorch tensor."""
- r = from_dlpack(t, assumed_align=16)
- r.element_type = etype
- return r
-
-
-def _fits_32bit_stride(tensor: torch.Tensor) -> bool:
- int32_max = 2**31 - 1
- max_offset = int(tensor.storage_offset())
- for size, stride in zip(tensor.shape, tensor.stride()):
- stride = abs(int(stride))
- if stride > int32_max:
- return False
- if size:
- max_offset += (int(size) - 1) * stride
- if max_offset > int32_max:
- return False
- return True
-
-
-def _dynamic_dlpack_arg(tensor: torch.Tensor):
- wrapper = from_dlpack(
- tensor,
- assumed_align=16,
- use_32bit_stride=_fits_32bit_stride(tensor),
- )
- for dim, stride in enumerate(tensor.stride()):
- if stride == 1:
- return wrapper.mark_layout_dynamic(dim)
- return wrapper
-
-
-def _layout_key(tensor: torch.Tensor):
- return (
- tensor.dtype,
- tuple(tensor.shape),
- tuple(tensor.stride()),
- _fits_32bit_stride(tensor),
- )
-
-
-def _dynamic_vector_layout_key(tensor: torch.Tensor):
- """Key a 1D CuTe argument whose sole extent is runtime dynamic."""
- assert tensor.ndim == 1
- return (
- tensor.dtype,
- tuple(tensor.stride()),
- _fits_32bit_stride(tensor),
- )
-
-
-def _current_cu_stream(device):
- """CUstream handle of torch's CURRENT stream, queried per call.
-
- The executor runs the model on a dedicated non-blocking
- ``torch.cuda.Stream`` (``py_executor.execution_stream``), not the default
- stream. Every kernel launch must go to this stream: launching on the DSL
- default stream (what ``.launch()`` does without a ``stream`` argument)
- races with the projections that produce q/k/v/g/beta and with the
- consumers of O/final_state — silent, intermittent corruption in the
- runtime while single-stream unit tests pass. This is also why the value
- must never be cached alongside the scratch buffers: the buffers may be
- created under a different stream (e.g. warmup) than later forwards.
- """
- return cuda_driver.CUstream(torch.cuda.current_stream(device).cuda_stream)
-
-
-# Cached eqlen dummy cu/ci cute wrappers — shared by K123 and akk_inv.
-# Avoids per-call torch.empty + from_dlpack overhead (~10-12us each).
-_eqlen_dummy_cache = {}
-
-
-def _get_eqlen_dummies(device, idx_dtype=torch.int64):
- """Returns cached (cu_ct, ci_ct) cute wrappers for eqlen (B+1=2, NT+1=2)."""
- key = (device.index if device.index is not None else 0, idx_dtype)
- if key not in _eqlen_dummy_cache:
- cu_t = torch.empty(2, dtype=idx_dtype, device=device)
- ci_t = torch.empty(1, 2, dtype=idx_dtype, device=device)
- cu_etype = cutlass.Int64 if idx_dtype == torch.int64 else cutlass.Int32
- _eqlen_dummy_cache[key] = (_ct(cu_t, cu_etype), _ct(ci_t, cu_etype))
- return _eqlen_dummy_cache[key]
-
-
-def _cute_int_type(dtype):
- """Map PyTorch integer dtype to CUTLASS element type."""
- if dtype == torch.int32:
- return cutlass.Int32
- elif dtype == torch.int64:
- return cutlass.Int64
- else:
- raise ValueError(f"Unsupported integer dtype: {dtype}")
-
-
-# ========== Fused K1+K2+K3 compilation cache ==========
-_fused_k123_cache = {}
-# id(cu_seqlens) -> bool. Skips per-call GPU->CPU sync on subsequent calls
-# when the same cu_seqlens tensor is reused (typical training/inference loop).
-_varlen_pure_cache = {}
-# id(cu_seqlens) -> int seqlen, populated alongside _varlen_pure_cache for
-# single-seq cu_seqlens.
-_varlen_single_seqlen_cache = {}
-
-
-def _prune_on_gc(cache, key, *keyobjs):
- """Drop ``cache[key]`` when any of ``keyobjs`` is garbage-collected.
-
- The id()-keyed caches in this module are only sound while the keyed
- Python objects stay alive — CPython reuses an object's address (its id)
- the moment it is freed, so a dead entry could otherwise be served for an
- unrelated new tensor with recycled id but different contents/storage.
- A finalizer runs at dealloc, before the address can be recycled, so a
- pruned entry can never alias a new object. This also bounds cache growth
- when callers pass fresh tensors every call.
- """
- for o in keyobjs:
- weakref.finalize(o, cache.pop, key, None)
-
-
-def _scratch_ct(cute_wrappers, name, t, etype):
- """`_ct(t, etype)` for the ``_get_buffers`` scratch tensor ``t``, cached
- under ``name`` in ``cute_wrappers``, the wrapper dict of the same
- ``_buf_cache`` entry, so a launch skips its ``from_dlpack`` (~5-10us each).
-
- A CuTe wrapper keeps the tensor it wraps alive (its DLPack capsule holds a
- reference), so a wrapper is cached only beside its tensor: evicting the
- entry from ``_buf_cache`` releases the scratch and its wrappers together.
- Per-call tensors are never cached; they are wrapped with plain ``_ct``.
- Without an entry (``cute_wrappers`` is None) ``t`` is wrapped afresh.
- """
- if cute_wrappers is None:
- return _ct(t, etype)
- key = ("scratch_ct", name, etype)
- cached = cute_wrappers.get(key)
- if cached is None or cached[0] is not t:
- cached = cute_wrappers[key] = (t, _ct(t, etype))
- return cached[1]
-
-
-# Cache for dt_bias `.float().contiguous().view(H, K)` + cute wrapper.
-# dt_bias is typically a nn.Parameter — same object across iterations.
-_dt_bias_cache = {}
-
-
-def _get_dt_bias_ct(dt_bias, H, K):
- """Returns cached cute wrapper for dt_bias.float().view(H, K)."""
- key = (id(dt_bias), H, K)
- entry = _dt_bias_cache.get(key)
- if entry is None:
- # Copy (never view) the keyed object: a view would pin it, making
- # the entry immortal (weakref pruning could never fire). Callers
- # pass a fresh .detach() per call, so an aliasing entry would leak.
- bias_t = dt_bias.detach().float().reshape(H, K).clone()
- entry = (bias_t, _ct(bias_t, cutlass.Float32))
- _dt_bias_cache[key] = entry
- _prune_on_gc(_dt_bias_cache, key, dt_bias)
- return entry[1]
-
-
-# Cache for the empty 1x1 fp32 bias tensor used when dt_bias is None.
-_empty_bias_cache = {}
-
-
-def _get_empty_bias_ct(device):
- idx = device.index if device.index is not None else 0
- if idx not in _empty_bias_cache:
- t = torch.empty(1, 1, dtype=torch.float32, device=device)
- _empty_bias_cache[idx] = _ct(t, cutlass.Float32)
- return _empty_bias_cache[idx]
-
-
-# K4 varlen cu_seqlens / chunk_offsets cute wrappers. Caches the int32-cast
-# tensor and its mark_layout_dynamic wrapper so they survive multiple calls
-# with the same input objects.
-_k4_varlen_cu_co_cache = {}
-
-def _get_k4_varlen_cu_co(cu_seqlens, chunk_offsets):
- key = (id(cu_seqlens), id(chunk_offsets))
- entry = _k4_varlen_cu_co_cache.get(key)
- if entry is None:
- # Always materialize copies: if the cached value aliased the keyed
- # object, the entry would pin it and the weakref pruning could
- # never fire (immortal entry).
- cu_int32 = cu_seqlens.to(torch.int32).contiguous()
- if cu_int32 is cu_seqlens:
- cu_int32 = cu_seqlens.clone()
- co_int32 = chunk_offsets.to(torch.int32).contiguous()
- if co_int32 is chunk_offsets:
- co_int32 = chunk_offsets.clone()
- cu_ct = from_dlpack(cu_int32, assumed_align=4).mark_layout_dynamic()
- cu_ct.element_type = cutlass.Int32
- co_ct = from_dlpack(co_int32, assumed_align=4).mark_layout_dynamic()
- co_ct.element_type = cutlass.Int32
- # Hold refs to the int32 tensors so they don't get GC'd and the
- # underlying storage remain valid as long as cu_ct / co_ct live.
- entry = (cu_int32, co_int32, cu_ct, co_ct)
- _k4_varlen_cu_co_cache[key] = entry
- _prune_on_gc(_k4_varlen_cu_co_cache, key, cu_seqlens, chunk_offsets)
- return entry[2], entry[3]
-
-
-# Cache the (cu_for_k4, chunk_offsets_for_k4) pair keyed by id(cu_seqlens).
-# Both tensors are computed on-GPU with no host sync — replaces the previous
-# `cu_seqlens.cpu().tolist() + Python loop + torch.tensor` chain that forced
-# a 50-200us GPU->CPU stall on the K4 prep path every varlen call.
-_varlen_k4_input_cache = {}
-
-
-def _get_varlen_k4_inputs(cu_seqlens, BT):
- key = id(cu_seqlens)
- entry = _varlen_k4_input_cache.get(key)
- if entry is None:
- # Copy (never alias) the keyed object — see _get_k4_varlen_cu_co.
- cu_int32 = cu_seqlens.to(torch.int32).contiguous()
- if cu_int32 is cu_seqlens:
- cu_int32 = cu_seqlens.clone()
- seq_lens = cu_int32[1:] - cu_int32[:-1]
- chunk_counts = (seq_lens + (BT - 1)) // BT
- zero = torch.zeros(1, dtype=torch.int32, device=cu_int32.device)
- co_int32 = torch.cat([zero, torch.cumsum(chunk_counts, dim=0).to(torch.int32)])
- co_int32 = co_int32.contiguous()
- _varlen_k4_input_cache[key] = (cu_int32, co_int32)
- _prune_on_gc(_varlen_k4_input_cache, key, cu_seqlens)
- return _varlen_k4_input_cache[key]
-
-
-# ========== BF16 akk_inv compilation cache ==========
-_akk_inv_cache = {}
-# ========== K4 persistent (varlen via cu_seqlens) compilation cache ==========
-_k4p_cache = {}
-_k4p_tm_ws = {}
-
-# ========== Shape-keyed scratch caches ==========
-# The scratch caches below are keyed by the packed prefill token count. In the
-# executor that is attn_metadata.num_ctx_tokens, which the scheduler does not
-# quantize, so almost every prefill iteration brings a new key. Every such
-# cache must therefore be LRU-bounded through _lru_put / _lru_touch.
-
-
-_V = TypeVar("_V")
-
-
-def _lru_put(cache: dict, key: Hashable, value: _V, max_entries: int) -> _V:
- """Insert ``value`` under ``key``, first evicting least-recently-used entries."""
- while len(cache) >= max_entries:
- cache.pop(next(iter(cache)))
- cache[key] = value
- return value
-
-
-def _lru_touch(cache: dict, key: Hashable) -> Any:
- """Mark ``key`` most recently used and return its value."""
- cache[key] = cache.pop(key)
- return cache[key]
-
-
-_DEFAULT_SCRATCH_CACHE_ENTRIES = 2
-
-
-def _scratch_cache_max_entries() -> int:
- """Entry cap of the shape-keyed scratch caches: TLLM_KDA_BUF_CACHE_ENTRIES, default 2.
-
- A non-integer value falls back to the default with a warning instead of
- failing the import of this module.
- """
- raw = os.environ.get("TLLM_KDA_BUF_CACHE_ENTRIES")
- if raw is None:
- return _DEFAULT_SCRATCH_CACHE_ENTRIES
- try:
- return max(1, int(raw))
- except ValueError:
- logger.warning(
- f"Ignoring invalid TLLM_KDA_BUF_CACHE_ENTRIES={raw!r}; "
- f"using the default of {_DEFAULT_SCRATCH_CACHE_ENTRIES}."
+ from ..attention.kernels.blackwell.kimi_k3_kda.kda_chunked_bt16 import compile
+
+ with torch.cuda.device(device):
+ return compile(
+ dtype=cutlass.BFloat16,
+ state_dtype=cutlass.BFloat16 if state_dtype == torch.bfloat16 else cutlass.Float32,
+ gate_dtype=cutlass.BFloat16,
+ safe_gate=True,
+ gate_lower_bound=lower_bound,
+ has_state_in=True,
+ has_state_out=True,
+ has_state_indices=True,
+ mode=mode,
)
- return _DEFAULT_SCRATCH_CACHE_ENTRIES
-
-
-# Buffer cache: avoid re-allocating ~67us of intermediate tensors per call.
-# Also caches cute.Tensor wrappers (saves ~7us each call from from_dlpack).
-#
-# One entry owns ~1.25 KiB per token per KDA head, i.e. ~1 GiB at H=96 and
-# T=8192. This scratch is allocated after the KV-cache pool has been sized, so
-# every resident entry is memory the KV-cache accounting never saw: keep the
-# cap small. A small cap costs almost nothing, because the reuse that matters
-# is within one forward: all KDA layers of an iteration share one T, so even a
-# single entry hits on every layer but the first. Cross-iteration hits are
-# incidental, since T usually changes from one iteration to the next.
-# Override with TLLM_KDA_BUF_CACHE_ENTRIES (minimum 1).
-_buf_cache = {}
-_BUF_CACHE_MAX_ENTRIES = _scratch_cache_max_entries()
-
-# Padded-input scratch cache for the eqlen partial-chunk path. Keyed by
-# (B, T_padded, H, K, dtype_qkv, dtype_g, dtype_beta, device, real_T).
-# real_T is part of the key so the g sentinel tail [real_T:T_padded] = -1e3
-# is set once and reused across calls with the same shape.
-_padded_input_cache = {}
-_PADDED_INPUT_CACHE_MAX_ENTRIES = _BUF_CACHE_MAX_ENTRIES
-
-# Sentinel-padded g scratch for varlen single-seq Phase 2.1 path. Keyed by
-# (B, T_padded, H, K, dtype, device, real_T). The tail [real_T:T_padded] is
-# pre-set to -1e3 once at cache init; subsequent calls only overwrite the
-# valid prefix [0, real_T). real_T is the request's true prompt length, so
-# the key space is as wide as the prompt-length distribution.
-_g_sentinel_cache = {}
-_G_SENTINEL_CACHE_MAX_ENTRIES = _BUF_CACHE_MAX_ENTRIES
-
-
-def _get_g_sentinel_buffer(B, T_padded, H, K, dtype_g, device, real_T):
- key = (B, T_padded, H, K, dtype_g, device.index if device.index is not None else 0, real_T)
- if key in _g_sentinel_cache:
- return _lru_touch(_g_sentinel_cache, key)
- e = torch.zeros(B, T_padded, H, K, dtype=dtype_g, device=device)
- if real_T < T_padded:
- e[:, real_T:] = -1000.0
- return _lru_put(_g_sentinel_cache, key, e, _G_SENTINEL_CACHE_MAX_ENTRIES)
-
-
-def _get_padded_input_buffers(B, T_padded, H, K, dtype_qkv, dtype_g, dtype_beta, device, real_T):
- key = (
- B,
- T_padded,
- H,
- K,
- dtype_qkv,
- dtype_g,
- dtype_beta,
- device.index if device.index is not None else 0,
- real_T,
- )
- if key in _padded_input_cache:
- return _lru_touch(_padded_input_cache, key)
- q_pad = torch.zeros(B, T_padded, H, K, dtype=dtype_qkv, device=device)
- k_pad = torch.zeros(B, T_padded, H, K, dtype=dtype_qkv, device=device)
- v_pad = torch.zeros(B, T_padded, H, K, dtype=dtype_qkv, device=device)
- beta_pad = torch.zeros(B, T_padded, H, dtype=dtype_beta, device=device)
- # g: zero in [0, real_T), sentinel -1e3 in [real_T, T_padded). Caller's
- # data overwrites the prefix each call; the sentinel tail never moves.
- g_pad = torch.zeros(B, T_padded, H, K, dtype=dtype_g, device=device)
- if real_T < T_padded:
- g_pad[:, real_T:] = -1000.0
- e = (q_pad, k_pad, v_pad, g_pad, beta_pad)
- return _lru_put(_padded_input_cache, key, e, _PADDED_INPUT_CACHE_MAX_ENTRIES)
-
-
-def _get_buffers(dev, dtype_k, B, T, H, K_dim, V_dim, NT, N_seqs, BT, varlen=False):
- """Return cached intermediate tensors for the KDA prefill pipeline."""
- key = (dev.index or 0, B, T, H, K_dim, V_dim, NT, N_seqs, varlen)
- if key not in _buf_cache:
- bf16 = cutlass.BFloat16
- fp32 = cutlass.Float32
- # Varlen chunk-tile kernels (akk_inv's Stage-0 cp.async A-tile load,
- # K4's A-tile loads) transfer the full BT-row tile of every chunk
- # and neutralize invalid rows only after the access; when the
- # batch's FINAL chunk is partial they read up to BT-1 rows past the
- # logical T. These are driver-owned scratch buffers, so honor the
- # kernel's boundary-at-the-data contract by allocating one chunk of
- # zeroed slack past T (the eqlen path gets the same guarantee from
- # its 256-multiple input padding). The slack rows are never
- # consumed — OOB rows are zeroed in SMEM or masked at stores.
- # Input tensors (beta) get the opposite treatment: the K123 kernel
- # bounds-checks those loads, since the driver doesn't own them.
- if varlen:
- assert B == 1, f"varlen expects packed B=1 input, got B={B}"
- T_alloc = T + BT if varlen else T
-
- def _with_slack(ctor, *shape, dtype):
- full = ctor(*shape, device=dev, dtype=dtype)
- return full[:, :T]
-
- k_scaled = _with_slack(torch.empty, B, T_alloc, H, K_dim, dtype=dtype_k) # raw, no beta
- kg = _with_slack(torch.empty, B, T_alloc, H, K_dim, dtype=dtype_k)
- q_scaled = _with_slack(torch.empty, B, T_alloc, H, K_dim, dtype=dtype_k)
- beta_activated = _with_slack(torch.empty, B, T_alloc, H, dtype=torch.bfloat16)
- gk_last_exp = torch.empty(B, NT, H, K_dim, device=dev, dtype=torch.float32)
- A_qk = _with_slack(torch.zeros, B, T_alloc, H, BT, dtype=dtype_k)
- A_kk = _with_slack(torch.zeros, B, T_alloc, H, BT, dtype=dtype_k)
- cu_eqlen = torch.arange(0, (B + 1) * T, T, dtype=torch.int32, device=dev)
- # co_eqlen is consumed only on the eqlen path (varlen K4 derives its
- # chunk offsets from cu_seqlens). Eqlen inputs are pre-padded to a
- # 256-multiple upstream, so T // BT >= 4 there; varlen calls can
- # carry T < BT (short-prompt batches), where arange(step=T // BT)
- # would raise "step must be nonzero" building this dead buffer.
- # Clamp the step so the buffer stays constructible.
- nt_eq = max(T // BT, 1)
- co_eqlen = torch.arange(0, (B + 1) * nt_eq, nt_eq, dtype=torch.int32, device=dev)
-
- T_total = B * T
- A_kk_flat = A_kk.reshape(T_total, H, BT)
- A_qk_flat = A_qk.reshape(T_total, H, BT)
- KS_flat = k_scaled.reshape(T_total, H, K_dim)
- QS_flat = q_scaled.reshape(T_total, H, K_dim)
- KG_flat = kg.reshape(T_total, H, K_dim)
- gk_flat = gk_last_exp.reshape(-1, H, K_dim)
-
- def _wrap(t, etype):
- r = from_dlpack(t, assumed_align=16).mark_layout_dynamic()
- r.element_type = etype
- return r
-
- a_ct = _wrap(A_kk_flat, bf16)
- aqc_ct = _wrap(A_qk_flat, bf16)
- ks_ct = _wrap(KS_flat, bf16) # raw k_scaled (beta absorbed in akk_inv)
- qs_ct = _wrap(QS_flat, bf16)
- kg_ct = _wrap(KG_flat, bf16)
- gk_ct = _wrap(gk_flat, fp32)
- cu_eqlen_ct = from_dlpack(cu_eqlen, assumed_align=4).mark_layout_dynamic()
- cu_eqlen_ct.element_type = cutlass.Int32
- co_eqlen_ct = from_dlpack(co_eqlen, assumed_align=4).mark_layout_dynamic()
- co_eqlen_ct.element_type = cutlass.Int32
-
- # akk_inv views: bf16 storage reinterpreted as fp32 (packed 2x bf16 -> 1x fp32).
- # Layout-dynamic so the single compiled akk_inv (shape-independent
- # key) accepts them regardless of T; the host rebuilds its own
- # runtime-shaped views from the iterator anyway.
- akk_in_view = from_dlpack(A_kk, assumed_align=16).mark_layout_dynamic()
- akk_in_view.element_type = fp32
- akk_out_view = from_dlpack(A_kk, assumed_align=16).mark_layout_dynamic()
- akk_out_view.element_type = fp32
-
- cute_wrappers = dict(
- a_ct=a_ct,
- aqc_ct=aqc_ct,
- ks_ct=ks_ct,
- qs_ct=qs_ct,
- kg_ct=kg_ct,
- gk_ct=gk_ct,
- cu_eqlen_ct=cu_eqlen_ct,
- co_eqlen_ct=co_eqlen_ct,
- akk_in_view=akk_in_view,
- akk_out_view=akk_out_view,
- # Filled lazily on first launch — saves cache_key tuple build +
- # outer dict lookup on subsequent calls.
- _k123_fns={},
- _akk_inv_fn=None,
- )
-
- return _lru_put(
- _buf_cache,
- key,
- (
- k_scaled,
- kg,
- q_scaled,
- beta_activated,
- gk_last_exp,
- A_qk,
- A_kk,
- cu_eqlen,
- co_eqlen,
- cute_wrappers,
- ),
- _BUF_CACHE_MAX_ENTRIES,
- )
- return _lru_touch(_buf_cache, key)
-
-
-def _new_output(dev, dtype, B, T, H, V_dim, BT, varlen=False):
- """The prefill's output ``o`` [B, T, H, V] and K4's wrapper of it, new for every call: ``o`` belongs to the
- caller, so no later call may write it. On the varlen path the buffer has one chunk of slack past T, like the
- scratch."""
- T_alloc = T + BT if varlen else T
- o = torch.empty(B, T_alloc, H, V_dim, device=dev, dtype=dtype)[:, :T]
- o_ct = from_dlpack(o.reshape(B * T, H, V_dim), assumed_align=16).mark_layout_dynamic()
- o_ct.element_type = cutlass.BFloat16
- return o, o_ct
-
-
-def _launch_k4_persistent(
- cute_wrappers,
- o_ct,
- v_beta,
- state_pool,
- state_indices,
- cu_seqlens,
- chunk_offsets,
- cu_eqlen_passed=False,
- num_sm=148,
- H=None,
- V_dim=None,
-):
- """Launch persistent K4 with cached CuTe wrappers.
-
- Precondition: all sequence lengths in cu_seqlens must be > 0 (a
- zero-length sequence deadlocks the kernel's chunk-loop barriers; see the
- k4_persistent module docstring). Not validated here: cu_seqlens is on the
- GPU and a host-side check would sync the hot path, and the prefill
- runtime never emits zero-length sequences.
-
- No fast-launch (args-tuple) cache here: such a cache pins the per-call
- activation and state-pool tensors via their CuTe wrappers (the wrapper holds the
- storage, so the keyed object never dies and weakref pruning never
- fires) — a per-call GPU memory leak in the executor runtime. The
- per-call cost is re-wrapping a handful of tensors (~10us each).
- """
- bf16 = cutlass.BFloat16
- N_seqs = cu_seqlens.shape[0] - 1
- dev = v_beta.device
- dev_idx = dev.index or 0
- s_ct = _dynamic_dlpack_arg(state_pool)
- a_ct = cute_wrappers["a_ct"]
- b_ct = cute_wrappers["ks_ct"] # raw k_scaled (beta absorbed in akk_inv)
- q_ct = cute_wrappers["qs_ct"]
- aqc_ct = cute_wrappers["aqc_ct"]
- kg_ct = cute_wrappers["kg_ct"]
- gk_ct = cute_wrappers["gk_ct"]
- # v is a per-call activation — wrap fresh every call (never cache; see
- # _scratch_ct: only the _get_buffers scratch keeps cached wrappers, in its
- # own _buf_cache entry).
- v_view = v_beta.reshape(-1, H, V_dim) if v_beta.dim() == 4 else v_beta
- v_ct = from_dlpack(v_view, assumed_align=16).mark_layout_dynamic()
- v_ct.element_type = bf16
-
- if cu_eqlen_passed:
- cu_ct = cute_wrappers["cu_eqlen_ct"]
- co_ct = cute_wrappers["co_eqlen_ct"]
- else:
- cu_ct, co_ct = _get_k4_varlen_cu_co(cu_seqlens, chunk_offsets)
-
- state_indices_ct = _dynamic_dlpack_arg(state_indices)
-
- # Allocate for the maximum persistent grid once. The scheduler applies
- # min(N_seqs * H, num_sm) at runtime, so pre-minimizing here would turn
- # every distinct request count into another compiled host function.
- tm_key = (dev_idx, num_sm)
- if tm_key not in _k4p_tm_ws:
- tm_ws_t = torch.zeros(
- num_sm * _K4P_NTM * _K4P_BTM,
- dtype=torch.uint8,
- device=dev,
- )
- tm_ct = from_dlpack(tm_ws_t, assumed_align=16)
- tm_ct.element_type = cutlass.Uint8
- _k4p_tm_ws[tm_key] = (tm_ws_t, tm_ct)
- else:
- tm_ws_t, tm_ct = _k4p_tm_ws[tm_key]
-
- # Launch on torch's CURRENT stream — see _current_cu_stream.
- stream = _current_cu_stream(dev)
-
- # H and the state V/K dims MUST be in the key: s_ct bakes the inner
- # [H, V, K] shape and all strides at compile time, so the compiled
- # function is head-count- and layout-specific. N_seqs and the
- # token/chunk counts remain runtime values.
- cache_key = (
- dev_idx,
- num_sm,
- H,
- state_pool.shape[-1],
- V_dim,
- _layout_key(state_pool),
- # The sole dimension of this vector is runtime dynamic. Excluding
- # its extent lets warmup serve full and underfilled context batches.
- _dynamic_vector_layout_key(state_indices),
- )
- k4_fn = _k4p_cache.get(cache_key)
- if k4_fn is None:
- host_fn = _k4p_make_host(num_sm=num_sm)
- k4_fn = cute.compile(
- host_fn,
- a_ct,
- b_ct,
- v_ct,
- q_ct,
- aqc_ct,
- kg_ct,
- o_ct,
- gk_ct,
- s_ct,
- state_indices_ct,
- cu_ct,
- co_ct,
- tm_ct,
- N_seqs,
- stream,
- )
- _k4p_cache[cache_key] = k4_fn
-
- args = (
- a_ct,
- b_ct,
- v_ct,
- q_ct,
- aqc_ct,
- kg_ct,
- o_ct,
- gk_ct,
- s_ct,
- state_indices_ct,
- cu_ct,
- co_ct,
- tm_ct,
- N_seqs,
- stream,
- )
- k4_fn(*args)
-
-
-def _launch_fused_k123_inv(
- q,
- k,
- g,
- A_log,
- beta,
- beta_activated,
- scale,
- k_scaled,
- kg,
- q_scaled,
- gk_last_exp,
- A_qk,
- A_kk_inv,
- cu_seqlens,
- chunk_indices,
- is_varlen,
- NT,
- valid_tokens,
- dt_bias=None,
- safe_gate=False,
- lower_bound=None,
- use_beta_sigmoid_in_kernel=True,
- akk_in_view=None,
- akk_out_view=None,
- cute_wrappers=None,
- varlen_pure_override=False,
- varlen_is_aligned=None,
-):
- """Persistent K1+K2 (writes A_kk in I+L format with diag=1) chained with
- BF16 akk_inv (in-place inversion). Final A_kk_inv = (I+L)^-1."""
-
- # No fast-launch (args-tuple) cache: it would pin the per-call q/k/g/beta
- # activations via their cute wrappers (the wrapper holds the storage, so
- # the keyed objects never die and weakref pruning never fires) — a
- # per-call GPU memory leak in the executor runtime. Wrapper gathering
- # below costs ~10us per tensor.
- B, T, H, K = q.shape
- BT = 64
- dev = q.device.index or 0
- T_padded = T if is_varlen else None
- has_bias = dt_bias is not None
-
- # Prefer the host-derived alignment property supplied by the model runtime.
- # Direct low-level callers that omit it retain the cached device inference
- # fallback for backward compatibility.
- varlen_pure = False
- if is_varlen and cu_seqlens is not None:
- if varlen_pure_override:
- # Caller (Phase 2.1 single-seq path) sentinel-padded the data for
- # this call; the cached per-object verdict stays untouched.
- varlen_pure = True
- elif varlen_is_aligned is not None:
- varlen_pure = varlen_is_aligned
- else:
- _vp_key = id(cu_seqlens)
- if _vp_key not in _varlen_pure_cache:
- cu_cpu = cu_seqlens.cpu().tolist()
- seq_lens = [cu_cpu[i + 1] - cu_cpu[i] for i in range(len(cu_cpu) - 1)]
- _varlen_pure_cache[_vp_key] = all((sl % BT) == 0 for sl in seq_lens)
- _prune_on_gc(_varlen_pure_cache, _vp_key, cu_seqlens)
- varlen_pure = _varlen_pure_cache[_vp_key]
- # B/NT/T are runtime args of the compiled host_fn (see make_host_function),
- # avoiding shape-dependent grid and view specialization. Varlen remains
- # fully shape-independent; eqlen retains T only as a cache discriminator
- # for the raw tensor batch strides described below.
- #
- # cu/ci DTYPE must be part of the key: the kernel reads
- # mCuSeqlens/mChunkIndices with the element type baked at compile time
- # (int64 elements are addressed with stride 8, int32 with stride 4).
- # Reusing an int64-compiled kernel on int32 tensors (or vice versa)
- # misaddresses every cu/ci element — garbage seq ids/chunk starts ->
- # OOB reads (cudaErrorIllegalAddress) or silent corruption. Observed as
- # a crash when a dump-replay batch (int64 cu/ci) preceded a synthetic
- # batch (int32 cu/ci) in one process.
- cu_ci_dtypes = (cu_seqlens.dtype, chunk_indices.dtype) if is_varlen else None
- # Eqlen kernels directly index mBeta/mAqk/mAkk using the input wrappers'
- # compile-time batch stride, which depends on T. Varlen fixes B=1, so that
- # batch stride is never used and all T shapes can keep sharing one compile.
- eqlen_t_key = T if not is_varlen else -1
- cache_key = (
- H,
- is_varlen,
- dev,
- has_bias,
- safe_gate,
- varlen_pure,
- cu_ci_dtypes,
- eqlen_t_key,
- use_beta_sigmoid_in_kernel,
- beta.dtype,
- )
-
- # Inputs are guaranteed contiguous by upstream linear projections.
- # A_log is fp32 model param; .float() is no-op when dtype already matches.
- # q/k/g/beta/cu/ci are per-call activations: plain _ct, never cached.
- # The _get_buffers scratch below is wrapped once per _buf_cache entry, and
- # its wrappers live in that entry's cute_wrappers (see _scratch_ct).
- q_ct = _ct(q, cutlass.BFloat16)
- k_ct = _ct(k, cutlass.BFloat16)
- g_ct = _ct(g, cutlass.BFloat16)
- alog_ct = _ct(A_log if A_log.dtype == torch.float32 else A_log.float(), cutlass.Float32)
- if beta.dtype == torch.float32:
- beta_etype = cutlass.Float32
- elif beta.dtype == torch.bfloat16:
- beta_etype = cutlass.BFloat16
- else:
- raise ValueError(f"Kimi K3 KDA prefill beta must be float32 or bfloat16, got {beta.dtype}")
- beta_ct = _ct(beta, beta_etype)
- beta_activated_ct = _scratch_ct(
- cute_wrappers, "beta_activated", beta_activated, cutlass.BFloat16
- )
-
- ks_ct = _scratch_ct(cute_wrappers, "k_scaled", k_scaled, cutlass.BFloat16)
- kg_ct = _scratch_ct(cute_wrappers, "kg", kg, cutlass.BFloat16)
- qs_ct = _scratch_ct(cute_wrappers, "q_scaled", q_scaled, cutlass.BFloat16)
- gk_ct = _scratch_ct(cute_wrappers, "gk_last_exp", gk_last_exp, cutlass.Float32)
- aqk_ct = _scratch_ct(cute_wrappers, "A_qk", A_qk, cutlass.BFloat16)
- akk_ct = _scratch_ct(cute_wrappers, "A_kk", A_kk_inv, cutlass.BFloat16)
-
- if is_varlen:
- cu_ct = _ct(cu_seqlens, _cute_int_type(cu_seqlens.dtype))
- ci_ct = _ct(chunk_indices, _cute_int_type(chunk_indices.dtype))
- else:
- cu_ct, ci_ct = _get_eqlen_dummies(q.device, torch.int64)
-
- if dt_bias is not None:
- bias_ct = _get_dt_bias_ct(dt_bias, H, K)
- else:
- bias_ct = _get_empty_bias_ct(q.device)
- lb_val = float(lower_bound) if lower_bound is not None else 0.0
-
- # Launch on torch's CURRENT stream — see _current_cu_stream.
- stream = _current_cu_stream(q.device)
-
- ct_args = (
- q_ct,
- k_ct,
- g_ct,
- alog_ct,
- beta_ct,
- beta_activated_ct,
- scale,
- ks_ct,
- kg_ct,
- qs_ct,
- gk_ct,
- aqk_ct,
- akk_ct,
- cu_ct,
- ci_ct,
- bias_ct,
- lb_val,
- # Runtime shape scalars (rt_nt / rt_b / rt_t_total in host_fn).
- NT,
- B,
- B * (T_padded if is_varlen else NT * BT),
- valid_tokens,
- stream,
- )
-
- if cache_key not in _fused_k123_cache:
- host_fn = _fused_make_host(
- B,
- NT,
- H,
- is_varlen=is_varlen,
- T_padded=T_padded,
- has_bias=has_bias,
- use_safe_gate=safe_gate,
- use_beta_sigmoid=use_beta_sigmoid_in_kernel,
- varlen_pure=varlen_pure,
- )
- _fused_k123_cache[cache_key] = cute.compile(host_fn, *ct_args)
- k123_fn = _fused_k123_cache[cache_key]
- k123_fn(*ct_args)
-
- # ===== Chained BF16 akk_inv (in-place: A_kk_inv = (I+L)^-1) =====
- # Views are layout-dynamic so one compiled akk_inv serves every batch
- # shape (the host builds its own runtime-shaped views from the iterator;
- # the wrapper layout only matters for the call signature).
- if akk_in_view is None:
- akk_in_view = from_dlpack(A_kk_inv, assumed_align=16).mark_layout_dynamic()
- akk_in_view.element_type = cutlass.Float32
- akk_out_view = from_dlpack(A_kk_inv, assumed_align=16).mark_layout_dynamic()
- akk_out_view.element_type = cutlass.Float32
-
- if is_varlen:
- # Reuse the cached cute wrappers from K123 (same tensor objects).
- akk_cu_ct = cu_ct
- akk_ci_ct = ci_ct
- is_varlen_int = 1
- T_val = T
- else:
- akk_cu_ct, akk_ci_ct = cu_ct, ci_ct
- is_varlen_int = 0
- T_val = NT * BT
-
- beta_for_akk = beta_activated if use_beta_sigmoid_in_kernel else beta
- if beta_for_akk.dtype == torch.float32:
- akk_beta_etype = cutlass.Float32
- else:
- akk_beta_etype = cutlass.BFloat16
- akk_beta_ct = (
- _scratch_ct(cute_wrappers, "beta_activated", beta_for_akk, akk_beta_etype)
- if use_beta_sigmoid_in_kernel
- else _ct(beta_for_akk, akk_beta_etype)
- )
-
- # B/NT/T_val remain runtime args of akk_inv_host. Eqlen T is still part of
- # the cache key because the beta wrapper bakes its batch stride.
- # cu/ci dtype is a specializer for the same reason as in the K123 key above
- # (element type baked into the compiled reader).
- akk_cache_key = (H, is_varlen, dev, cu_ci_dtypes, eqlen_t_key, beta_for_akk.dtype)
- if akk_cache_key not in _akk_inv_cache:
- _akk_inv_cache[akk_cache_key] = cute.compile(
- _akk_inv_host,
- akk_in_view,
- akk_out_view,
- akk_beta_ct,
- B,
- NT,
- H,
- akk_cu_ct,
- akk_ci_ct,
- is_varlen_int,
- T_val,
- stream,
- )
- akk_fn = _akk_inv_cache[akk_cache_key]
- akk_args = (akk_in_view, akk_out_view, akk_beta_ct, B, NT, akk_cu_ct, akk_ci_ct, T_val, stream)
- akk_fn(*akk_args)
-
-
-def _chunk_kda_fwd(
+@torch.library.custom_op("trtllm::kda_prefill", mutates_args=("state_pool",), device_types="cuda")
+def kda_prefill(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
@@ -907,417 +44,241 @@ def _chunk_kda_fwd(
beta: torch.Tensor,
state_pool: torch.Tensor,
state_indices: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ A_log: torch.Tensor,
+ dt_bias: torch.Tensor,
scale: float,
- cu_seqlens: torch.LongTensor | None = None,
- chunk_indices: torch.LongTensor | None = None,
- chunk_size: int = 64,
- safe_gate: bool = False,
- lower_bound: float | None = None,
- use_gate_in_kernel: bool = False,
- use_beta_sigmoid_in_kernel: bool = True,
- A_log: torch.Tensor | None = None,
- dt_bias: torch.Tensor | None = None,
- varlen_is_aligned: bool | None = None,
- single_sequence_length: int | None = None,
+ lower_bound: float,
+ mode: str = "auto",
+ max_seqlen: int = 0,
+ seq_order: torch.Tensor | None = None,
+ cu_chunks: torch.Tensor | None = None,
) -> torch.Tensor:
- """Run optimized KDA prefill against indexed recurrent-state rows."""
- if safe_gate and lower_bound is None:
- lower_bound = -5.0
-
- is_varlen = cu_seqlens is not None
-
- if q.shape[1] == 0:
- raise ValueError("Indexed KDA prefill does not support an empty token batch.")
-
- B, T, H, K = q.shape
- V_dim = v.shape[-1]
+ """Update unique selected V-first state rows using packed raw Q/K/V.
+
+ The scheduler must prepare monotone int64 offsets and unique, nonnegative
+ int32 slots outside execution. Empty sequences are allowed. Fresh rows
+ must be reset by the caller. Optional int32 ``seq_order`` (stable descending
+ length order) and ``cu_chunks`` (prefix sums of ceil(length / 16) in request
+ order) are prepared once per batch and shared across layers. Both must be
+ refreshed before replay when packed lengths change, retaining their storage.
+ Without them, BT16 regenerates metadata on the device each invocation.
+
+ Auto selects Small-BH for BF16 state when twice the sequence/head count
+ fits the device SM count. ``bt16`` retains BT16's engine/decomposition
+ selection; ``small-bh`` forces Small-BH. ``max_seqlen`` is an optional
+ upper bound for all replays, or zero to use the total token capacity.
+ Small-BH and BT16 decomposition consume aligned, head-contiguous beta
+ views directly; the engine requires compact beta storage.
+ """
device = q.device
- BT = 64
-
- # Phase 1: handle eqlen with T % 64 != 0 entirely on the eqlen path —
- # never borrow varlen's mask code. Pad inputs to a CHUNKS_PER_BLOCK*BT
- # = 256 multiple so the persistent scheduler's `cgs_per_head = NT // 4`
- # divides cleanly and every chunk has 64 valid rows of (zero-padded /
- # sentinel-padded) data. K123 eqlen kernel runs unchanged — it doesn't
- # know partial chunks exist, the boundary is handled at the data
- # boundary (K4's bounded-TMA principle, just expressed via host pad
- # since K123 stores via autovec_copy not TMA).
- #
- # Varlen with non-aligned seq lengths is a separate problem (Phase 2):
- # multi-seq varlen can't be host-padded without repacking memory.
- real_T = T
- CPB_BT = 4 * BT # CHUNKS_PER_BLOCK * BT, the cgs_per_head divisibility unit
- # The trigger is CPB_BT, not BT: the eqlen scheduler floors `cgs_per_head`
- # and has no `chunk_idx < num_chunks` guard, so a BT-aligned T whose chunk
- # count is not a multiple of CHUNKS_PER_BLOCK (T = 64, 128, 192, 320, ...)
- # would silently drop its trailing chunks.
- needs_eqlen_pad = (not is_varlen) and (T % CPB_BT != 0)
- if needs_eqlen_pad:
- if B != 1:
- raise NotImplementedError(
- f"eqlen with B>1 and T % {CPB_BT} != 0 not supported (got B={B}, T={T})."
- )
- T_padded = ((T + CPB_BT - 1) // CPB_BT) * CPB_BT
- # Pre-allocated padded scratch buffers (per (B,T_padded,H,K,dtype) cache
- # key). torch.cat would reallocate + copy the full 200MB q tensor every
- # call — caching the destination buffer drops that to a single slice
- # copy of the valid prefix (caller already lives in our buffer for
- # subsequent calls reusing the same id, but we re-copy unconditionally
- # since the caller may have updated the data in-place).
- q_pad, k_pad, v_pad, g_pad, beta_pad = _get_padded_input_buffers(
- B, T_padded, H, K, q.dtype, g.dtype, beta.dtype, q.device, real_T
- )
- # q/k/v/beta zero-padded → K1/K2 MMAs naturally produce 0 for OOB rows.
- # Tail [real_T:] of q_pad/k_pad/v_pad/beta_pad is pre-zeroed at cache
- # init and never written, so we only copy the valid prefix.
- q_pad[:, :real_T].copy_(q)
- k_pad[:, :real_T].copy_(k)
- v_pad[:, :real_T].copy_(v)
- beta_pad[:, :real_T].copy_(beta)
- # g uses a -1e3 sentinel so the gate activation saturates to 0 for OOB
- # rows (both safe_gate sigmoid and softplus paths). Plain g=0 gives
- # nonzero activation that would corrupt the cumsum past seq end. The
- # tail is set to -1e3 once at cache init; we only copy the valid prefix.
- g_pad[:, :real_T].copy_(g)
- q, k, v, g, beta = q_pad, k_pad, v_pad, g_pad, beta_pad
- T = T_padded # downstream buffer alloc + kernel layout use T_padded
-
- # Phase 2.1: varlen with a SINGLE non-aligned sequence — caller already
- # zero-padded q/k/v/beta to a 64-multiple (FLA convention), but g's tail
- # is also zero, which causes the gate activation to be non-zero past seq
- # end and corrupts the cumsum / GkLast. We sentinel-pad g (cheap: ~5MB
- # copy) and force VARLEN_PURE=1 so all 4 mask sites compile-elide. Same
- # K4 "boundary at the data" principle as the eqlen path.
- #
- # Multi-seq varlen is NOT handled here — its OOB regions overlap with
- # adjacent seqs' data, so sentinel-pad on g would corrupt the next seq.
- # Multi-seq optimization needs per-seq dynamic tensormap (Phase 2.2).
- needs_varlen_single_pad = False
- if is_varlen and cu_seqlens is not None and cu_seqlens.shape[0] == 2 and B == 1:
- if single_sequence_length is not None:
- single_sequence_is_aligned = single_sequence_length % BT == 0
- if varlen_is_aligned is not None and varlen_is_aligned != single_sequence_is_aligned:
- raise ValueError("varlen_is_aligned disagrees with single_sequence_length")
- varlen_is_aligned = single_sequence_is_aligned
- if not single_sequence_is_aligned:
- real_T = single_sequence_length
- needs_varlen_single_pad = True
- else:
- _vl_key = id(cu_seqlens)
- if _vl_key not in _varlen_pure_cache:
- cu_cpu = cu_seqlens.cpu().tolist()
- sl = cu_cpu[1] - cu_cpu[0]
- _varlen_pure_cache[_vl_key] = sl % BT == 0
- _varlen_single_seqlen_cache[_vl_key] = sl
- _prune_on_gc(_varlen_pure_cache, _vl_key, cu_seqlens)
- _prune_on_gc(_varlen_single_seqlen_cache, _vl_key, cu_seqlens)
- varlen_is_aligned = _varlen_pure_cache[_vl_key]
- if not varlen_is_aligned:
- real_T = _varlen_single_seqlen_cache[_vl_key]
- needs_varlen_single_pad = True
- varlen_pure_override = False
- if needs_varlen_single_pad:
- # q/k/v/beta already zero-padded by caller (FLA convention). Re-build g
- # with -1000 sentinel in the tail so VARLEN_PURE=1 path is correct.
- # Cache the resulting g buffer so repeated calls with same input ids
- # don't re-allocate.
- cur_T = q.shape[1] # caller's padded T
- assert cur_T % BT == 0 and cur_T >= real_T, (
- f"varlen single-seq path expects caller-padded input "
- f"(T={cur_T}, seqlen={real_T}); see "
- "KDAKernelDispatch.prefill_chunk_kda"
+ if torch.cuda.get_device_capability(device) not in ((10, 0), (10, 3), (10, 7)):
+ raise ValueError("FlashInfer KDA prefill requires SM100, SM103 or SM107")
+ if (
+ mode not in ("auto", "bt16", "small-bh", "engine", "decomp")
+ or not math.isfinite(lower_bound)
+ or lower_bound >= 0
+ ):
+ raise ValueError("KDA prefill requires a finite negative gate bound and a supported mode")
+ if q.ndim != 4 or q.shape[0] != 1 or q.shape[-1] != 128 or q.shape[1] < 1:
+ raise ValueError("KDA prefill requires packed [1, T, H, 128] inputs")
+ if any(
+ t.shape != q.shape or t.dtype != torch.bfloat16 or not t.is_contiguous()
+ for t in (q, k, v, g)
+ ):
+ raise ValueError("KDA Q/K/V and raw gates must be contiguous BF16")
+ if beta.shape != q.shape[:-1] or beta.dtype != torch.bfloat16:
+ raise ValueError("KDA beta logits must be BF16 [1, T, H]")
+ n, heads = state_indices.numel(), q.shape[2]
+ if (
+ state_pool.dtype not in (torch.bfloat16, torch.float32)
+ or state_pool.ndim != 4
+ or state_pool.shape[1:] != (heads, 128, 128)
+ or state_pool.stride()[1:] != (128 * 128, 128, 1)
+ or state_pool.stride(0) < heads * 128 * 128
+ or state_pool.stride(0) * state_pool.element_size() % 16
+ or state_pool.data_ptr() % 16
+ ):
+ raise ValueError("KDA requires aligned BF16/FP32 V-first state rows")
+ if (
+ state_indices.dtype != torch.int32
+ or state_indices.shape != (n,)
+ or not state_indices.is_contiguous()
+ or cu_seqlens.dtype != torch.int64
+ or cu_seqlens.shape != (n + 1,)
+ or not cu_seqlens.is_contiguous()
+ or n == 0
+ ):
+ raise ValueError("Prepare contiguous int32 indices and int64 packed offsets")
+ if (
+ A_log.dtype != torch.float32
+ or A_log.numel() != heads
+ or not A_log.is_contiguous()
+ or dt_bias.dtype != torch.float32
+ or dt_bias.numel() != heads * 128
+ or not dt_bias.is_contiguous()
+ ):
+ raise ValueError("KDA gate parameters must be contiguous FP32")
+ if any(
+ t.device != device
+ for t in (k, v, g, beta, state_pool, state_indices, cu_seqlens, A_log, dt_bias)
+ ):
+ raise ValueError("KDA inputs and state must share one CUDA device")
+ if any(t.data_ptr() % 16 for t in (q, k, v, g, A_log, dt_bias)):
+ raise ValueError("KDA inputs and gate parameters must be 16-byte aligned")
+ if (seq_order is None) != (cu_chunks is None):
+ raise ValueError("Provide both KDA sequence order and chunk offsets")
+ if seq_order is not None:
+ for tensor, shape in ((seq_order, (n,)), (cu_chunks, (n + 1,))):
+ if (
+ tensor.device != device
+ or tensor.dtype != torch.int32
+ or tensor.shape != shape
+ or not tensor.is_contiguous()
+ ):
+ raise ValueError("KDA scheduling metadata must be contiguous CUDA int32 [N]/[N+1]")
+ if tensor.untyped_storage().data_ptr() == state_pool.untyped_storage().data_ptr():
+ raise ValueError("KDA state must not alias its scheduling metadata")
+ if any(
+ state_pool.untyped_storage().data_ptr() == t.untyped_storage().data_ptr()
+ for t in (q, k, v, g, beta, state_indices, cu_seqlens, A_log, dt_bias)
+ ):
+ raise ValueError("KDA state must not alias its immutable inputs")
+ if max_seqlen < 0 or max_seqlen > q.shape[1]:
+ raise ValueError("max_seqlen must be zero or a sequence-length bound within total tokens")
+ sm_count = torch.cuda.get_device_properties(device).multi_processor_count
+ small_bh_eligible = state_pool.dtype == torch.bfloat16 and q.shape[1] > 1
+ if mode == "small-bh" and not small_bh_eligible:
+ raise ValueError("Small-BH requires BF16 state and more than one packed token")
+ use_small_bh = mode == "small-bh" or (
+ mode == "auto" and small_bh_eligible and 2 * n * heads <= sm_count
+ )
+ # Match BT16's occupancy router, including auto with FP32 state where
+ # Small-BH is unavailable. Forced engine calls must still pack beta.
+ use_decomp = mode == "decomp" or (mode in ("auto", "bt16") and 2 * n * heads <= sm_count)
+ # The merged [f_a | beta] projection has unit head stride and a padded,
+ # 16-byte-aligned token pitch. Preserve packing for other input layouts.
+ strided_beta_supported = (
+ (use_small_bh or use_decomp)
+ and beta.stride(-1) == 1
+ and beta.stride(1) >= heads
+ and beta.stride(1) * beta.element_size() % 16 == 0
+ )
+ aligned_beta = beta if strided_beta_supported else beta.contiguous()
+ if aligned_beta.data_ptr() % 16:
+ aligned_beta = aligned_beta.clone(memory_format=torch.contiguous_format)
+ output = torch.empty_like(q)
+ if use_small_bh:
+ from ..attention.kernels.blackwell.kimi_k3_kda.kda_chunked_bt16_small_bh import (
+ chunk_kda_fwd,
)
- g_pad = _get_g_sentinel_buffer(B, cur_T, H, K, g.dtype, g.device, real_T)
- g_pad[:, :real_T].copy_(g[:, :real_T])
- g = g_pad
- # Force VARLEN_PURE=1 for THIS call only (the data is sentinel-padded
- # now). Never write the override into _varlen_pure_cache: the cached
- # verdict must keep saying "not pure" so a later call with the SAME
- # cu_seqlens object re-runs this sentinel pad — poisoning the cache
- # made the second call skip the pad while still compiling the
- # mask-free variant (silent tail corruption).
- varlen_pure_override = True
- # Phase 2.2 (multi-seq via host repack) was attempted but isn't net positive:
- # the scatter/gather memcpy cost (~800us GPU bandwidth per call) exceeds
- # the kernel mask-elision savings (~250us). Keeping multi-seq non-pure on
- # the original masked path. The right fix is kernel-level dynamic
- # tensormap (K4-style per-tile bounded TMA) but that requires a major
- # K123 kernel refactor — left as future work.
- # multiseq_info = None
-
- if is_varlen:
- if chunk_indices is None:
- chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
- NT = len(chunk_indices)
- # The varlen scheduler supports any chunk count.
- N_seqs = len(cu_seqlens) - 1
- else:
- NT = T // BT
- N_seqs = B
-
- # ===== Cached buffers + cute wrappers (avoid alloc + from_dlpack overhead per call) =====
- (
- k_scaled,
- kg,
- q_scaled,
- beta_activated,
- gk_last_exp,
- A_qk,
- A_kk,
- cu_eqlen,
- co_eqlen,
- cute_wrappers,
- ) = _get_buffers(device, k.dtype, B, T, H, K, V_dim, NT, N_seqs, BT, varlen=is_varlen)
- o, o_ct = _new_output(device, k.dtype, B, T, H, V_dim, BT, varlen=is_varlen)
-
- # Beta is fused entirely in akk_inv kernel epilogue (post-inv column-scale).
- # No host v*beta and no K1 k_scaled*beta any more.
- _launch_fused_k123_inv(
+ with torch.cuda.device(device):
+ chunk_kda_fwd(
+ q[0],
+ k[0],
+ v[0],
+ g[0],
+ aligned_beta[0],
+ dt_bias.reshape(heads, 128),
+ A_log.reshape(heads),
+ scale=scale,
+ gate_scale=lower_bound,
+ gate_lower_bound=True,
+ return_state=True,
+ transpose_state=True,
+ cu_seqlens=cu_seqlens,
+ max_seqlen=max_seqlen or q.shape[1],
+ output=output[0],
+ state_input=state_pool,
+ state_output=state_pool,
+ state_indices=state_indices,
+ max_active_clusters=sm_count,
+ )
+ return output
+ # Torch allocations belong to the CUDA graph's private memory pool during
+ # capture. Keeping scratch per layer would otherwise retain O(layers*T*H)
+ # memory; eager execution reuses allocator storage between layer calls.
+ compiled = _compile_prefill(
+ device, state_pool.dtype, lower_bound, "auto" if mode == "bt16" else mode
+ )
+ nonempty = min(q.shape[1], n)
+ capacity = nonempty + (q.shape[1] - nonempty) // 16
+ size = compiled.workspace_size_from_total_chunks(n, heads, capacity, device)
+ workspace = torch.empty(size, dtype=torch.uint8, device=device) if size else None
+ generate_metadata = seq_order is None
+ if generate_metadata:
+ seq_order = torch.empty(n, dtype=torch.int32, device=device)
+ cu_chunks = torch.empty(n + 1, dtype=torch.int32, device=device)
+ compiled(
q,
k,
+ v,
g,
A_log,
- beta,
- beta_activated,
- scale,
- k_scaled,
- kg,
- q_scaled,
- gk_last_exp,
- A_qk,
- A_kk,
+ dt_bias.reshape(heads, 128),
+ aligned_beta,
cu_seqlens,
- chunk_indices,
- is_varlen,
- NT,
- real_T,
- dt_bias=dt_bias,
- safe_gate=safe_gate,
- lower_bound=lower_bound,
- use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel,
- akk_in_view=cute_wrappers["akk_in_view"],
- akk_out_view=cute_wrappers["akk_out_view"],
- cute_wrappers=cute_wrappers,
- varlen_pure_override=varlen_pure_override,
- varlen_is_aligned=varlen_is_aligned,
- )
-
- # ===== K4: persistent kernel (eqlen + varlen via cu_seqlens) =====
- if is_varlen:
- # GPU-side cumsum + cache by id(cu_seqlens). No host sync.
- cu_for_k4, chunk_offsets_for_k4 = _get_varlen_k4_inputs(cu_seqlens, BT)
- else:
- cu_for_k4 = cu_eqlen
- chunk_offsets_for_k4 = co_eqlen
-
- _launch_k4_persistent(
- cute_wrappers,
- o_ct,
- v,
state_pool,
- state_indices,
- cu_for_k4,
- chunk_offsets_for_k4,
- cu_eqlen_passed=(not is_varlen),
- H=H,
- V_dim=V_dim,
+ output,
+ state_pool,
+ workspace,
+ torch.cuda.current_stream(device).cuda_stream,
+ scale=scale,
+ seq_order=seq_order,
+ planned_cu_chunks=cu_chunks,
+ planned_total_chunks=capacity,
+ generate_planned_metadata=generate_metadata,
+ state_indices=state_indices,
)
+ return output
- if needs_eqlen_pad:
- # Caller called with original T = real_T; their downstream code expects
- # outputs at that shape. Slice the padded tail back off.
- o = o[:, :real_T]
- return o
-
-
-class KdaPrefillRunner:
- """Runs the source-integrated Kimi K3 KDA prefill pipeline."""
- @staticmethod
- def forward(
- q: torch.Tensor,
- k: torch.Tensor,
- v: torch.Tensor,
- g: torch.Tensor,
- beta: torch.Tensor,
- state_pool: torch.Tensor,
- state_indices: torch.Tensor,
- scale: float,
- cu_seqlens: Optional[torch.Tensor] = None,
- chunk_indices: Optional[torch.Tensor] = None,
- chunk_size: int = 64,
- safe_gate: bool = False,
- lower_bound: Optional[float] = None,
- use_gate_in_kernel: bool = False,
- use_beta_sigmoid_in_kernel: bool = True,
- A_log: Optional[torch.Tensor] = None,
- dt_bias: Optional[torch.Tensor] = None,
- varlen_is_aligned: Optional[bool] = None,
- single_sequence_length: Optional[int] = None,
- ) -> torch.Tensor:
- """Run KDA prefill and update selected recurrent-state rows."""
- if not IS_CUTLASS_DSL_AVAILABLE:
- raise RuntimeError("Kimi K3 KDA prefill requires NVIDIA CUTLASS DSL")
- if chunk_size != 64:
- raise ValueError(f"Kimi K3 KDA prefill requires chunk_size=64, got {chunk_size}")
- if A_log is None:
- raise ValueError("Kimi K3 KDA prefill requires A_log")
-
- return _chunk_kda_fwd(
- q=q,
- k=k,
- v=v,
- g=g,
- beta=beta,
- state_pool=state_pool,
- state_indices=state_indices,
- scale=scale,
- cu_seqlens=cu_seqlens,
- chunk_indices=chunk_indices,
- chunk_size=chunk_size,
- safe_gate=safe_gate,
- lower_bound=lower_bound,
- use_gate_in_kernel=use_gate_in_kernel,
- use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel,
- A_log=A_log,
- dt_bias=dt_bias,
- varlen_is_aligned=varlen_is_aligned,
- single_sequence_length=single_sequence_length,
- )
-
-
-def _validate_indexed_state_pool(
- q: torch.Tensor,
- v: torch.Tensor,
- state_pool: torch.Tensor,
- state_indices: torch.Tensor,
- cu_seqlens: Optional[torch.Tensor],
-) -> None:
- num_sequences = q.shape[0] if cu_seqlens is None else cu_seqlens.shape[0] - 1
- H, K, V = q.shape[2], q.shape[3], v.shape[3]
- if state_pool.dtype != torch.float32 or state_pool.ndim != 4:
- raise ValueError(
- "Expected state_pool to be a rank-4 fp32 tensor with shape [slots, H, V, K]."
- )
- if state_pool.shape[1:] != (H, V, K):
- raise ValueError(
- f"Expected state_pool shape [slots, {H}, {V}, {K}], got {tuple(state_pool.shape)}."
- )
- expected_inner_strides = (V * K, K, 1)
- if state_pool.stride()[1:] != expected_inner_strides:
- raise ValueError(
- "Indexed KDA prefill requires dense H/V/K inner dimensions; "
- f"expected strides {expected_inner_strides}, got "
- f"{state_pool.stride()[1:]}."
- )
- if state_pool.stride(0) < H * V * K or state_pool.stride(0) % 4 != 0:
- raise ValueError(
- "Indexed KDA prefill requires a non-overlapping, 16-byte-aligned slot stride."
- )
- if state_pool.data_ptr() % 16 != 0:
- raise ValueError("Indexed KDA prefill requires a 16-byte-aligned state pool.")
- if state_indices.ndim != 1 or state_indices.shape[0] != num_sequences:
- raise ValueError(
- f"Expected state_indices shape [{num_sequences}], got {tuple(state_indices.shape)}."
- )
- if state_indices.dtype not in (torch.int32, torch.int64):
- raise ValueError("state_indices must have dtype int32 or int64.")
- if not state_indices.is_contiguous():
- raise ValueError("state_indices must be contiguous.")
- if state_indices.data_ptr() % 16 != 0:
- raise ValueError("Indexed KDA prefill requires 16-byte-aligned state_indices.")
- for name, tensor in (
- ("q", q),
- ("v", v),
- ("state_indices", state_indices),
- ):
- if tensor.device != state_pool.device:
- raise ValueError(f"{name} and state_pool must be on the same device.")
-
-
-if IS_CUTLASS_DSL_AVAILABLE:
-
- @torch.library.custom_op(
- "trtllm::kda_prefill",
- mutates_args=("state_pool",),
- device_types="cuda",
- )
- def kda_prefill(
- q: torch.Tensor,
- k: torch.Tensor,
- v: torch.Tensor,
- g: torch.Tensor,
- beta: torch.Tensor,
- state_pool: torch.Tensor,
- state_indices: torch.Tensor,
- scale: float,
- cu_seqlens: Optional[torch.Tensor] = None,
- chunk_indices: Optional[torch.Tensor] = None,
- chunk_size: int = 64,
- safe_gate: bool = False,
- lower_bound: Optional[float] = None,
- use_gate_in_kernel: bool = False,
- use_beta_sigmoid_in_kernel: bool = True,
- A_log: Optional[torch.Tensor] = None,
- dt_bias: Optional[torch.Tensor] = None,
- varlen_is_aligned: Optional[bool] = None,
- single_sequence_length: Optional[int] = None,
- ) -> torch.Tensor:
- """Run KDA prefill directly against a V-first recurrent-state pool."""
- _validate_indexed_state_pool(
- q,
- v,
- state_pool,
- state_indices,
- cu_seqlens,
- )
- return KdaPrefillRunner.forward(
- q=q,
- k=k,
- v=v,
- g=g,
- beta=beta,
- state_pool=state_pool,
- state_indices=state_indices,
- A_log=A_log,
- scale=scale,
- dt_bias=dt_bias,
- cu_seqlens=cu_seqlens,
- chunk_indices=chunk_indices,
- chunk_size=chunk_size,
- safe_gate=safe_gate,
- lower_bound=lower_bound,
- use_gate_in_kernel=use_gate_in_kernel,
- use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel,
- varlen_is_aligned=varlen_is_aligned,
- single_sequence_length=single_sequence_length,
- )
-
- @kda_prefill.register_fake
- def _(
- q: torch.Tensor,
- k: torch.Tensor,
- v: torch.Tensor,
- g: torch.Tensor,
- beta: torch.Tensor,
- state_pool: torch.Tensor,
- state_indices: torch.Tensor,
- scale: float,
- cu_seqlens: Optional[torch.Tensor] = None,
- chunk_indices: Optional[torch.Tensor] = None,
- chunk_size: int = 64,
- safe_gate: bool = False,
- lower_bound: Optional[float] = None,
- use_gate_in_kernel: bool = False,
- use_beta_sigmoid_in_kernel: bool = True,
- A_log: Optional[torch.Tensor] = None,
- dt_bias: Optional[torch.Tensor] = None,
- varlen_is_aligned: Optional[bool] = None,
- single_sequence_length: Optional[int] = None,
- ) -> torch.Tensor:
- del q, k, g, beta, state_pool, state_indices
- del scale, cu_seqlens, chunk_indices, chunk_size, safe_gate
- del lower_bound, use_gate_in_kernel, use_beta_sigmoid_in_kernel, A_log, dt_bias
- del varlen_is_aligned, single_sequence_length
- return v.new_empty(v.shape)
+@kda_prefill.register_fake
+def _kda_prefill_fake(
+ q,
+ k,
+ v,
+ g,
+ beta,
+ state_pool,
+ state_indices,
+ cu_seqlens,
+ A_log,
+ dt_bias,
+ scale,
+ lower_bound,
+ mode="auto",
+ max_seqlen=0,
+ seq_order=None,
+ cu_chunks=None,
+):
+ return torch.empty_like(q)
+
+
+@torch.library.custom_op(
+ "trtllm::kda_prefill_conv", mutates_args=("conv_state",), device_types="cuda"
+)
+def kda_prefill_conv(
+ src: torch.Tensor,
+ weight: torch.Tensor,
+ conv_state: torch.Tensor,
+ starts: torch.Tensor,
+ slots: torch.Tensor,
+ has_init: torch.Tensor,
+ max_seqlen: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Fuse short convolution, SiLU and raw Q/K/V layout conversion."""
+ from ..attention.kernels.blackwell.kimi_k3_kda.fused_prefill_conv import fused_prefill_conv
+
+ shape = (1, src.shape[0], weight.shape[0] // 384, 128)
+ q, k, v = (src.new_empty(shape) for _ in range(3))
+ with torch.cuda.device(src.device):
+ fused_prefill_conv(src, weight, conv_state, starts, slots, has_init, max_seqlen, q, k, v)
+ return q, k, v
+
+
+@kda_prefill_conv.register_fake
+def _kda_prefill_conv_fake(src, weight, conv_state, starts, slots, has_init, max_seqlen):
+ shape = (1, src.shape[0], weight.shape[0] // 384, 128)
+ return tuple(src.new_empty(shape) for _ in range(3))
diff --git a/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py
index 8d81fb9fe6b0..555bf4b28cae 100644
--- a/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py
+++ b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py
@@ -1,24 +1,6 @@
# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
-"""Kernel dispatch for the in-tree KDA module.
-
-Both optimized KDA kernels are source-integrated into TensorRT-LLM: the
-chunked prefill (CuTe DSL ``trtllm::kda_prefill``, see
-``tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py``) and the
-fused CUDA C++ single-token decode (``trtllm::kda_decode`` thop op wrapped
-by ``_kda_decode``). Neither requires the external
-``exisiting_optimization_work`` collection at runtime.
-
-On Blackwell (sm_100/sm_103) with the CuTe DSL toolchain available the
-dispatch runs the optimized kernels. On any other arch (or when the in-tree
-prefill op is unavailable) the dispatch falls back to FLA's ``chunk_kda`` /
-``fused_recurrent_kda`` on-device references.
-
-Callers are the ``KimiKDALinearAttention`` module. The module owns the
-HF-parity semantics (Q/K/V projections, conv, gating, ``o_norm``,
-``o_proj``); the dispatch here only wraps the kernel-level computation of
-the delta-rule inner loop plus its state update.
-"""
+"""KDA FlashInfer prefill, native decode/verification and portable FLA dispatch."""
from __future__ import annotations
@@ -68,133 +50,49 @@ def is_kda_optimized_supported() -> bool:
return get_kda_sm_version() in (100, 103, 107)
-@triton.jit(do_not_specialize=["num_tokens", "feature_stride"])
-def _fused_kda_post_conv_kernel(
- packed_ptr,
- q_out_ptr,
- k_out_ptr,
- v_out_ptr,
- num_tokens,
- token_stride,
- feature_stride,
- l2_norm_eps,
- NUM_HEADS: tl.constexpr,
- HEAD_DIM: tl.constexpr,
- BLOCK_TOKENS: tl.constexpr,
- BLOCK_DIM: tl.constexpr,
- TOKEN_MAJOR: tl.constexpr,
-):
- """Normalize and transpose packed Q/K/V in one launch.
-
- ``packed`` is ``[3 * NUM_HEADS * HEAD_DIM, num_tokens]`` with one unit
- stride: along tokens (channel-major) or, with ``TOKEN_MAJOR``, along
- features. ``token_stride`` stays specialized so Triton sees when
- token-major rows start 16-byte aligned and vectorizes their loads.
- ``feature_stride`` is the token count in channel-major mode, where
- specializing it selects a slower layout.
- """
- token_offsets = tl.program_id(0) * BLOCK_TOKENS + tl.arange(0, BLOCK_TOKENS)
- head_idx = tl.program_id(1)
- dim_offsets = tl.arange(0, BLOCK_DIM)
- token_mask = token_offsets < num_tokens
- dim_mask = dim_offsets < HEAD_DIM
- mask = token_mask[:, None] & dim_mask[None, :]
-
- feature_offsets = head_idx * HEAD_DIM + dim_offsets
- projection_size = NUM_HEADS * HEAD_DIM
- if TOKEN_MAJOR:
- outer_stride_i64 = token_stride.to(tl.int64)
- source_offsets = token_offsets[:, None].to(tl.int64) * outer_stride_i64 + feature_offsets[
- None, :
- ].to(tl.int64)
- section_stride = projection_size
- else:
- outer_stride_i64 = feature_stride.to(tl.int64)
- source_offsets = feature_offsets[None, :].to(tl.int64) * outer_stride_i64 + token_offsets[
- :, None
- ].to(tl.int64)
- section_stride = projection_size * outer_stride_i64
- output_offsets = (
- token_offsets[:, None].to(tl.int64) * projection_size + feature_offsets[None, :]
- )
-
- q = tl.load(packed_ptr + source_offsets, mask=mask, other=0.0).to(tl.float32)
- q /= tl.sqrt(tl.sum(q * q, axis=1) + l2_norm_eps)[:, None]
- tl.store(q_out_ptr + output_offsets, q, mask=mask)
-
- k = tl.load(
- packed_ptr + section_stride + source_offsets,
- mask=mask,
- other=0.0,
- ).to(tl.float32)
- k /= tl.sqrt(tl.sum(k * k, axis=1) + l2_norm_eps)[:, None]
- tl.store(k_out_ptr + output_offsets, k, mask=mask)
-
- v = tl.load(
- packed_ptr + 2 * section_stride + source_offsets,
- mask=mask,
- other=0.0,
- )
- tl.store(v_out_ptr + output_offsets, v, mask=mask)
-
-
def fused_kda_post_conv(
packed: torch.Tensor,
num_heads: int,
head_dim: int,
l2_norm_eps: float = 1e-6,
+ *,
+ normalize_qk: bool = True,
+ out: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
- """Convert packed convolution output to KDA Q/K/V.
+ """Portable layout conversion for the unfused convolution fallback.
- ``packed`` has shape ``[3 * num_heads * head_dim, tokens]`` and is either
- channel-major (unit token stride) or a transposed view of token-major
- rows (unit feature stride). The returned tensors are contiguous
- ``[1, tokens, num_heads, head_dim]``; Q and K are L2-normalized along the
- head dimension.
+ The production FlashInfer path produces this layout directly in the
+ convolution kernel and normalizes Q/K within the recurrent kernel.
"""
- projection_size = num_heads * head_dim
- if packed.ndim != 2 or packed.shape[0] != 3 * projection_size:
- raise ValueError(
- "Packed KDA post-conv expected shape "
- f"[{3 * projection_size}, tokens], got {tuple(packed.shape)}"
- )
- num_tokens = packed.shape[1]
- feature_stride, token_stride = packed.stride()
- token_major = feature_stride == 1 and num_tokens > 1
- if not token_major and token_stride != 1 and num_tokens > 1:
- raise ValueError(
- "Packed KDA post-conv requires a unit stride along tokens or features, "
- f"got strides {packed.stride()}"
- )
-
- output_shape = (1, num_tokens, num_heads, head_dim)
- q_out = torch.empty(output_shape, dtype=packed.dtype, device=packed.device)
- k_out = torch.empty_like(q_out)
- v_out = torch.empty_like(q_out)
- if num_tokens == 0:
- return q_out, k_out, v_out
-
- block_tokens = 16
- block_dim = triton.next_power_of_2(head_dim)
- grid = (triton.cdiv(num_tokens, block_tokens), num_heads)
- _fused_kda_post_conv_kernel[grid](
- packed,
- q_out,
- k_out,
- v_out,
- num_tokens,
- token_stride,
- feature_stride,
- l2_norm_eps,
- num_heads,
- head_dim,
- block_tokens,
- block_dim,
- token_major,
- num_warps=8,
- num_stages=3,
+ if packed.ndim != 2 or packed.shape[0] != 3 * num_heads * head_dim:
+ raise ValueError("Expected packed [3 * H * D, T] convolution output")
+ if packed.shape[1] > 1 and 1 not in packed.stride():
+ raise ValueError("Packed KDA post-conv requires a unit stride along tokens or features")
+ qkv = tuple(
+ t.t().reshape(1, packed.shape[1], num_heads, head_dim).contiguous()
+ for t in packed.chunk(3, dim=0)
)
- return q_out, k_out, v_out
+ if normalize_qk:
+ q, k, v = qkv
+
+ def normalize(t):
+ f = t.float()
+ return (f * torch.rsqrt((f * f).sum(-1, keepdim=True) + l2_norm_eps)).to(t.dtype)
+
+ qkv = (normalize(q), normalize(k), v)
+ if out is not None:
+ if len(out) != 3 or any(
+ dst.shape != src.shape
+ or dst.dtype != src.dtype
+ or dst.device != src.device
+ or not dst.is_contiguous()
+ for dst, src in zip(out, qkv)
+ ):
+ raise ValueError("Invalid KDA layout output buffers")
+ for dst, src in zip(out, qkv):
+ dst.copy_(src)
+ return out
+ return qkv
@triton.jit
@@ -551,7 +449,7 @@ def _load_prefill_module() -> ModuleType:
module = importlib.import_module(
"tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_custom_ops"
)
- except Exception as exc: # typically ImportError when CuTe DSL is unavailable
+ except ImportError as exc:
_PREFILL_IMPORT_ERROR = exc
raise
_PREFILL_MODULE = module
@@ -559,11 +457,15 @@ def _load_prefill_module() -> ModuleType:
def is_intree_prefill_available() -> bool:
- """True when the in-tree CuTe DSL prefill op can be imported."""
+ """True when the vendored prefill toolchain and operator are available."""
+ from ...cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE
+
+ if not IS_CUTLASS_DSL_AVAILABLE:
+ return False
try:
_load_prefill_module()
return True
- except Exception:
+ except ImportError:
return False
@@ -741,27 +643,33 @@ def can_use_indexed_prefill(
"""Return whether prefill can update this V-first pool directly."""
if self.prefill_kernel_path != "optimized" or num_tokens == 0:
return False
- if cu_seqlens is not None:
- if cu_seqlens.shape[0] - 1 != num_sequences:
- return False
- if state_pool.dtype != torch.float32 or state_pool.ndim != 4:
+ if (
+ cu_seqlens is None
+ or cu_seqlens.dtype != torch.int64
+ or cu_seqlens.shape != (num_sequences + 1,)
+ or not cu_seqlens.is_contiguous()
+ ):
+ return False
+ if state_pool.dtype not in (torch.bfloat16, torch.float32) or state_pool.ndim != 4:
return False
_, H, V, K = state_pool.shape
if state_pool.stride()[1:] != (V * K, K, 1):
return False
- if state_pool.stride(0) < H * V * K or state_pool.stride(0) % 4:
+ if (
+ (V, K) != (128, 128)
+ or state_pool.stride(0) < H * V * K
+ or state_pool.stride(0) * state_pool.element_size() % 16
+ ):
return False
if state_pool.data_ptr() % 16:
return False
if (
state_indices.ndim != 1
or state_indices.shape[0] != num_sequences
- or state_indices.dtype not in (torch.int32, torch.int64)
+ or state_indices.dtype != torch.int32
or not state_indices.is_contiguous()
):
return False
- if state_indices.data_ptr() % 16:
- return False
if (
has_initial_states.ndim != 1
or has_initial_states.shape[0] != num_sequences
@@ -797,84 +705,39 @@ def prefill_chunk_kda(
state_indices: Optional[torch.Tensor] = None,
varlen_is_aligned: Optional[bool] = None,
single_sequence_length: Optional[int] = None,
+ max_seqlen: int = 0,
+ seq_order: Optional[torch.Tensor] = None,
+ cu_chunks: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
- """Run KDA chunked prefill.
-
- Q and K must be L2-normalized by the caller. On the optimized path,
- beta sigmoid runs inside the in-tree CuTe DSL op.
-
- On the FLA path it calls ``fla.ops.kda.chunk_kda`` directly with the
- matching flags so the semantics are byte-equivalent.
+ """Run prefill on raw Q/K/V and gate/beta logits.
- The optimized kernel always reads and writes indexed V-first
- ``[slots, H, V, K]`` state-pool rows. Callers provide the executor
- pool and reset fresh rows before entering the kernel. The FLA fallback
- retains the batch-dense ``initial_state``/``final_state`` contract.
-
- ``chunk_indices``, ``varlen_is_aligned``, and
- ``single_sequence_length`` are prepared once from Kimi K3 runtime
- metadata. They keep the optimized path from reading ``cu_seqlens``
- back from the GPU in every KDA layer.
+ FlashInfer updates indexed BF16/FP32 V-first state directly; FLA uses
+ dense FP32 initial/final states. Both normalize Q/K internally.
"""
- use_indexed_state = state_pool is not None
- if use_indexed_state and (initial_state is not None or state_indices is None):
- raise ValueError(
- "Indexed KDA prefill requires state_indices and does not accept initial_state."
- )
- use_optimized = self.prefill_kernel_path == "optimized" and use_indexed_state
-
- if use_indexed_state and not use_optimized:
- raise RuntimeError(
- "Indexed KDA prefill requires the optimized prefill path; "
- "use the FLA state path for this batch."
- )
-
- if use_optimized:
- import torch.nn.functional as F
-
- real_T = q.shape[1]
- if cu_seqlens is not None:
- # The op's varlen single-seq path (Phase 2.1) expects the
- # caller to zero-pad the packed tensors to a chunk multiple
- # (FLA convention) while cu_seqlens keeps the real length;
- # the op re-sentinels g's tail itself. Without this the op
- # would run the mask-free kernel on a partial final chunk.
- # Multi-seq varlen runs the masked path and needs no pad.
- if cu_seqlens.shape[0] == 2 and real_T % chunk_size != 0:
- pad = chunk_size - real_T % chunk_size
- q = F.pad(q, (0, 0, 0, 0, 0, pad))
- k = F.pad(k, (0, 0, 0, 0, 0, pad))
- v = F.pad(v, (0, 0, 0, 0, 0, pad))
- g = F.pad(g, (0, 0, 0, 0, 0, pad))
- beta = F.pad(beta, (0, 0, 0, pad))
-
- A_log_kernel = A_log.detach() if A_log is not None else None
- dt_bias_kernel = dt_bias.detach() if dt_bias is not None else None
-
- _load_prefill_module() # registers trtllm::kda_prefill
+ if state_pool is not None:
+ if self.prefill_kernel_path != "optimized" or initial_state is not None:
+ raise ValueError("Indexed prefill requires FlashInfer and no dense initial state")
+ if not safe_gate or lower_bound is None:
+ raise ValueError("FlashInfer requires a finite negative gate bound")
+ _load_prefill_module()
out = torch.ops.trtllm.kda_prefill(
- q=q,
- k=k,
- v=v,
- g=g,
- beta=beta,
- state_pool=state_pool,
- state_indices=state_indices,
- scale=scale,
- cu_seqlens=cu_seqlens,
- chunk_indices=chunk_indices,
- chunk_size=chunk_size,
- safe_gate=safe_gate,
- lower_bound=lower_bound,
- use_gate_in_kernel=True,
- use_beta_sigmoid_in_kernel=True,
- A_log=A_log_kernel,
- dt_bias=dt_bias_kernel,
- varlen_is_aligned=varlen_is_aligned,
- single_sequence_length=single_sequence_length,
+ q,
+ k,
+ v,
+ g,
+ beta.to(torch.bfloat16),
+ state_pool,
+ state_indices,
+ cu_seqlens,
+ A_log.float().contiguous(),
+ dt_bias.float().contiguous(),
+ scale,
+ lower_bound,
+ getattr(self, "prefill_mode", "auto"),
+ max_seqlen,
+ seq_order,
+ cu_chunks,
)
- if out.shape[1] != real_T:
- out = out[:, :real_T]
return out, None
from fla.ops.kda import chunk_kda
@@ -890,7 +753,7 @@ def prefill_chunk_kda(
scale=scale,
initial_state=initial_state,
output_final_state=True,
- use_qk_l2norm_in_kernel=False,
+ use_qk_l2norm_in_kernel=True,
use_gate_in_kernel=True,
use_beta_sigmoid_in_kernel=True,
safe_gate=safe_gate,
diff --git a/tensorrt_llm/_torch/modules/kimi_kda/kimi_k3_mamba_metadata.py b/tensorrt_llm/_torch/modules/kimi_kda/kimi_k3_mamba_metadata.py
index 3f15c77647f0..8e3ca5aeaf12 100644
--- a/tensorrt_llm/_torch/modules/kimi_kda/kimi_k3_mamba_metadata.py
+++ b/tensorrt_llm/_torch/modules/kimi_kda/kimi_k3_mamba_metadata.py
@@ -1,55 +1,31 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
-"""Kimi K3-specific recurrent-state metadata."""
+"""Host-prepared scheduling metadata for packed FlashInfer KDA prefill."""
+
+from itertools import accumulate
import torch
from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata
from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata
-
-KDA_PREFILL_CHUNK_SIZE = 64
-
-
-@torch.compile(dynamic=True, fullgraph=True)
-def _prepare_kda_chunk_indices(
- sequence_lengths: torch.Tensor,
- chunk_indices: torch.Tensor,
-) -> torch.Tensor:
- """Fill FLA-compatible chunk indices without reading CUDA counts on CPU."""
- chunk_counts = (sequence_lengths + KDA_PREFILL_CHUNK_SIZE - 1).div(
- KDA_PREFILL_CHUNK_SIZE, rounding_mode="floor"
- )
- num_chunks = chunk_indices.shape[0]
- # Supplying output_size keeps repeat_interleave from synchronizing to read
- # sum(chunk_counts) from the device.
- sequence_indices = torch.repeat_interleave(
- torch.arange(sequence_lengths.shape[0], dtype=torch.long, device=sequence_lengths.device),
- chunk_counts,
- output_size=num_chunks,
- )
- chunk_starts = torch.cumsum(chunk_counts, dim=0) - chunk_counts
- local_chunk_indices = (
- torch.arange(num_chunks, dtype=torch.long, device=sequence_lengths.device)
- - chunk_starts[sequence_indices]
- )
- chunk_indices[:, 0].copy_(sequence_indices)
- chunk_indices[:, 1].copy_(local_chunk_indices)
- return chunk_indices
+from tensorrt_llm._utils import prefer_pinned
class KimiK3MambaMetadata(Mamba2Metadata):
- """Mamba metadata extended with Kimi K3 KDA preparation."""
+ """Reuse stable int32 slots/starts and int64 offsets from Mamba metadata."""
- # The KDA prefill kernel keys its compile cache on chunk alignment.
- prefill_chunk_alignment = KDA_PREFILL_CHUNK_SIZE
+ prefill_chunk_alignment = 16
def __init__(self, max_batch_size: int, chunk_size: int, max_num_tokens: int) -> None:
super().__init__(max_batch_size, chunk_size, max_num_tokens)
- # Every chunk contains at least one token, so this is a hard upper bound.
- self._kda_chunk_indices = torch.empty((max_num_tokens, 2), dtype=torch.long, device="cuda")
- self.kda_chunk_indices: torch.Tensor | None = None
- self.kda_varlen_is_aligned: bool | None = None
- self.kda_single_sequence_length: int | None = None
+ self.kda_max_seqlen = 0
+ # Align both rows to 16 bytes. One transfer refreshes the metadata shared
+ # by all layers; device storage stays fixed for CUDA graph replay.
+ width = ((max_batch_size + 1 + 3) // 4) * 4
+ self._kda_schedule = torch.zeros((2, width), dtype=torch.int32, device="cuda")
+ self.kda_seq_order = self._kda_schedule[0, :0]
+ self.kda_cu_chunks = self._kda_schedule[1, :1]
+
# KDA's CuTe verify kernel requires a 16-byte-aligned DLPack pointer.
# Keep the backing allocation stable so metadata preparation remains
# safe for CUDA graph capture and replay.
@@ -80,28 +56,21 @@ def prepare(self, attn_metadata: AttentionMetadata) -> None:
aligned_state_indices.copy_(generation_state_indices)
self.generation_state_indices = aligned_state_indices
- self.kda_varlen_is_aligned = (
- all(length % KDA_PREFILL_CHUNK_SIZE == 0 for length in context_lengths)
- if context_lengths
- else None
- )
- self.kda_single_sequence_length = context_lengths[0] if len(context_lengths) == 1 else None
-
- if not context_lengths:
- self.kda_chunk_indices = None
+ self.kda_max_seqlen = max(context_lengths, default=0)
+ n = len(context_lengths)
+ self.kda_seq_order = self._kda_schedule[0, :n]
+ self.kda_cu_chunks = self._kda_schedule[1, : n + 1]
+ if not n:
return
-
- chunk_counts = [
- (length + KDA_PREFILL_CHUNK_SIZE - 1) // KDA_PREFILL_CHUNK_SIZE
- for length in context_lengths
- ]
- num_chunks = sum(chunk_counts)
- assert num_chunks <= self._kda_chunk_indices.shape[0], (
- f"Expected at most {self._kda_chunk_indices.shape[0]} KDA chunks, got {num_chunks}"
- )
- self.kda_chunk_indices = self._kda_chunk_indices[:num_chunks]
-
- _prepare_kda_chunk_indices(
- attn_metadata.seq_lens_cuda[: len(context_lengths)],
- self.kda_chunk_indices,
+ order = sorted(range(n), key=lambda i: (-context_lengths[i], i))
+ chunks = list(accumulate(((length + 15) // 16 for length in context_lengths), initial=0))
+ width = self._kda_schedule.shape[1]
+ # Use a fresh pinned source so another prepare cannot overwrite an
+ # in-flight transfer. The pinned allocator retains it until DMA ends.
+ host_schedule = torch.tensor(
+ [order + [0] * (width - n), chunks + [0] * (width - n - 1)],
+ dtype=torch.int32,
+ device="cpu",
+ pin_memory=prefer_pinned(),
)
+ self._kda_schedule.copy_(host_schedule, non_blocking=True)
diff --git a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py
index 30231196346d..aa31750d45b1 100644
--- a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py
+++ b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py
@@ -20,7 +20,8 @@
import math
import os
import weakref
-from typing import List, NamedTuple, Optional, Tuple, Union
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, List, NamedTuple, Optional, Tuple, Union
import torch
from fla.modules import ShortConvolution
@@ -41,6 +42,9 @@
from ..multi_stream_utils import maybe_execute_in_parallel
from ._kda_kernels import KDAKernelDispatch, fused_kda_post_conv
+if TYPE_CHECKING:
+ from ..mamba.mamba2_metadata import Mamba2Metadata
+
_KDA_INDEXED_STATE_POOL_ENABLED = os.environ.get("TLLM_KDA_ENABLE_INDEXED_STATE_POOL", "1") == "1"
# Heuristic ported from SGLang's Blackwell cutoff:
# https://github.com/sgl-project/sglang/blob/e84bbf68efb683c9e2eef4168c5198042544599d/python/sglang/srt/models/kimi_k3.py#L946-L954
@@ -62,6 +66,15 @@ class _KdaPrequantizedCore(NamedTuple):
scale: torch.Tensor
+@dataclass
+class _KDAPrefillState:
+ kernel_pool: torch.Tensor
+ kernel_indices: torch.Tensor
+ recurrent_in: torch.Tensor | None
+ staged_state: torch.Tensor | None
+ use_indexed_state: bool
+
+
def _meta_safe_cast_dtype(module: nn.Module, dtype: torch.dtype) -> None:
"""Cast floating parameters while preserving meta-device construction.
@@ -93,8 +106,8 @@ def _cast(tensor: torch.Tensor) -> torch.Tensor:
def _stage_state_rows(ssm_pool: torch.Tensor, slot_indices: torch.Tensor) -> torch.Tensor:
"""Dense fp32 copy of the addressed recurrent-state rows.
- The fp32-only consumers (indexed prefill kernel, fused decode kernel) run on
- this copy when the pool is bf16; ``_writeback_state_rows`` rounds it back.
+ The fp32-only consumers (native decode/verification and FLA cores) run on
+ this copy when the pool is bf16; ``_writeback_state_rows`` rounds it back into the pool.
"""
return ssm_pool.index_select(0, slot_indices.long()).float()
@@ -283,10 +296,20 @@ def projection(name: str, in_features: int, out_features: int) -> nn.Module:
# K == V == 128 shape. Other shapes must use the portable fallback.
kernel_shape_ok = self.head_k_dim == 128 and self.head_dim == 128
self._dispatch = KDAKernelDispatch(
+ use_optimized_prefill=(
+ os.getenv("TLLM_KDA_ENABLE_OPT_PREFILL", "1") == "1"
+ and kernel_shape_ok
+ and self.conv_size == 4
+ and self.gate_lower_bound is not None
+ and math.isfinite(self.gate_lower_bound)
+ and self.gate_lower_bound < 0
+ ),
use_optimized_decode=kernel_shape_ok,
use_optimized_verify=kernel_shape_ok,
)
+ self._fuse_prefill_conv = True
+
# Fused prefill/decode/verify projection weights, built after checkpoint
# load. BF16 uses separate fused [q | k | v | g] and [f_a | b]
# GEMMs; FP8 fuses QKVG from checkpoint codes/scales and keeps BFA BF16.
@@ -592,16 +615,25 @@ def _project_gate_inputs(
def _project_packed_conv_input(
self, x: torch.Tensor, x2d: torch.Tensor
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
- """Project the block input into the ``[3D, T]`` the convolution needs.
+ """Project packed inputs for fused or portable convolution.
- The result is channel-major when the transposed GEMM can write it
- directly, and otherwise a transposed view of token-major rows, which
- the convolution and ``fused_kda_post_conv`` read through strides.
+ Fused convolution consumes token-major ``[T, 3D]`` or ``[T, 4D]``
+ projections. The portable path consumes ``[3D, T]``, channel-major
+ when the transposed GEMM is aligned and otherwise a transposed view.
Returns ``(packed_conv, onorm_g)``; ``onorm_g`` is None for
configurations without a full-rank gate.
"""
d = self.proj_size
+ if self._dispatch.prefill_kernel_path == "optimized" and self._fuse_prefill_conv:
+ if self._qkvg_proj_weight is not None:
+ qkvg = torch.nn.functional.linear(x, self._qkvg_proj_weight)
+ elif self.qkvg_proj is not None:
+ qkvg = self.qkvg_proj(x)
+ else:
+ qkvg = torch.cat((self.q_proj(x), self.k_proj(x), self.v_proj(x)), dim=-1)
+ onorm_g = qkvg[..., 3 * d : 4 * d] if qkvg.shape[-1] == 4 * d else None
+ return qkvg.squeeze(0), onorm_g
if self._qkvg_proj_weight is not None:
weight = self._qkvg_proj_weight
# The transposed GEMM writes rows of T elements; cuBLAS's fast
@@ -654,20 +686,18 @@ def forward_prefill(
layer_cache=None,
output: Optional[torch.Tensor] = None,
) -> torch.Tensor:
- """Run KDA prefill for ``num_prefills`` packed context sequences.
-
- ``x2d`` is ``[T, hidden]`` with sequence boundaries in ``cu_seqlens``.
- Convolution windows in ``conv_pool`` and recurrent states in
- ``ssm_pool`` are read and updated in place at ``slot_indices``.
- Returns the gated core output before ``o_proj``, written into
- ``output`` when one is given.
- """
- chunk_indices = getattr(mamba_metadata, "kda_chunk_indices", None)
- varlen_is_aligned = getattr(mamba_metadata, "kda_varlen_is_aligned", None)
- single_sequence_length = getattr(mamba_metadata, "kda_single_sequence_length", None)
from einops import rearrange
+ if x2d.shape[0] == 0:
+ reset_recurrent_state_rows(
+ ssm_pool, slot_indices, mamba_metadata.has_initial_states[:num_prefills], conv_pool
+ )
+ self._sync_kda_replay_conv_window(layer_cache, slot_indices, conv_pool)
+ return self._store_core(x2d.new_empty((0, self.num_heads, self.head_dim)), output)
x = x2d.unsqueeze(0) # [1, T, hidden]
+ if self._dispatch.prefill_kernel_path == "optimized":
+ # Metadata owns stable int32 indices; never convert per layer.
+ slot_indices = mamba_metadata.state_indices[:num_prefills]
packed_conv, onorm_g = self._project_packed_conv_input(x, x2d)
@@ -675,29 +705,102 @@ def forward_prefill(
# and for prefix-cache hits (block reuse), where the previous
# conv/recurrent state was onboarded into this request's slot.
has_init = mamba_metadata.has_initial_states[:num_prefills]
- slot_indices_long = slot_indices.long()
- # The indexed prefill op updates fp32 V-first pool rows in place, so
- # ``can_use_indexed_prefill`` rejects a bf16 pool outright. Staging the
- # addressed rows through a dense fp32 copy keeps prefill on that
- # fixed-configuration kernel instead of dropping it onto the FLA
- # fallback, whose Triton kernels are autotuned per process and
- # therefore do not reproduce run to run. The copy is rounded back into
- # the pool once the kernel has finished with it.
+ prepared_state = self._prepare_prefill_state(
+ cu_seqlens, mamba_metadata, num_prefills, ssm_pool, slot_indices, x2d.shape[0]
+ )
+ packed_conv = self._convolve_prefill(
+ packed_conv,
+ cu_seqlens,
+ slot_indices,
+ has_init,
+ conv_pool,
+ mamba_metadata,
+ num_prefills,
+ )
+
+ if self._bfa_proj_weight is not None:
+ beta, g, onorm_lowrank = self._project_gate_inputs(x)
+ if onorm_g is None:
+ onorm_g = onorm_lowrank
+ else:
+ g = self.f_b_proj(self.f_a_proj(x))
+ beta = self.b_proj(x)
+ g = rearrange(g, "... (h d) -> ... h d", d=self.head_dim)
+
+ o = self.forward_prefill_core(
+ packed_conv,
+ g,
+ beta,
+ cu_seqlens,
+ mamba_metadata,
+ num_prefills,
+ ssm_pool,
+ slot_indices,
+ _prepared_state=prepared_state,
+ )
+ # Fused-verify replay caches: seed the committed conv window so the
+ # first verify round convolves the correct history (pending drafts
+ # are zero for a fresh request, so the tail columns are unused).
+ self._sync_kda_replay_conv_window(layer_cache, slot_indices, conv_pool)
+
+ core = self._output_gate(x, o, onorm_g)
+ return self._store_core(core, output)
+
+ def _convolve_prefill(
+ self,
+ packed_conv: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ slot_indices: torch.Tensor,
+ has_init: torch.Tensor,
+ conv_pool: torch.Tensor,
+ mamba_metadata: Mamba2Metadata,
+ num_prefills: int,
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Convolve packed inputs and update the live convolution cache."""
+ assert self._packed_conv_weight is not None
+ if self._dispatch.prefill_kernel_path == "optimized" and self._fuse_prefill_conv:
+ return torch.ops.trtllm.kda_prefill_conv(
+ packed_conv,
+ self._packed_conv_weight,
+ conv_pool,
+ mamba_metadata.query_start_loc[: num_prefills + 1],
+ slot_indices,
+ has_init,
+ getattr(mamba_metadata, "kda_max_seqlen", packed_conv.shape[0]),
+ )
+ return causal_conv1d_fn(
+ packed_conv,
+ self._packed_conv_weight,
+ query_start_loc=mamba_metadata.query_start_loc[: num_prefills + 1],
+ cache_indices=mamba_metadata.state_indices[:num_prefills],
+ has_initial_state=has_init,
+ conv_states=conv_pool,
+ activation="silu",
+ )
+
+ def _prepare_prefill_state(
+ self,
+ cu_seqlens: torch.Tensor,
+ mamba_metadata,
+ num_prefills: int,
+ ssm_pool: torch.Tensor,
+ slot_indices: torch.Tensor,
+ num_tokens: int,
+ ) -> _KDAPrefillState:
+ """Prepare state at its original position before convolution/gate projection."""
+ has_init = mamba_metadata.has_initial_states[:num_prefills]
staged_state = None
kernel_pool, kernel_indices = ssm_pool, slot_indices
- if ssm_pool.dtype != torch.float32:
- staged_state = _stage_state_rows(ssm_pool, slot_indices_long)
- kernel_pool = staged_state
- # Staged rows are dense, so the kernel addresses them by position.
- kernel_indices = mamba_metadata._arange_buffer[:num_prefills]
use_indexed_state = self._dispatch.can_use_indexed_prefill(
- state_pool=kernel_pool,
- state_indices=kernel_indices,
+ state_pool=ssm_pool,
+ state_indices=slot_indices,
has_initial_states=has_init,
cu_seqlens=cu_seqlens,
num_sequences=num_prefills,
- num_tokens=x2d.shape[0],
+ num_tokens=num_tokens,
)
+ if self._dispatch.prefill_kernel_path == "optimized" and not use_indexed_state:
+ raise ValueError("Unsupported FlashInfer KDA state layout or prepared metadata")
recurrent_in = None
if use_indexed_state:
# The packed convolution clears fresh convolution rows itself.
@@ -707,41 +810,56 @@ def forward_prefill(
recurrent_in = (
staged_state
if staged_state is not None
- else _stage_state_rows(ssm_pool, slot_indices_long)
+ else _stage_state_rows(ssm_pool, slot_indices)
)
recurrent_in[~has_init] = 0
- # Reuse GDN's packed variable-length causal convolution. It reads
- # and writes the live [slots, 3D, W - 1] pool directly and honors
- # has_initial_state for fresh versus continuation requests. It works in
- # place on channel-major input and returns a new token-major tensor
- # for a transposed view.
- assert self._packed_conv_weight is not None
- packed_conv = causal_conv1d_fn(
- packed_conv,
- self._packed_conv_weight,
- query_start_loc=mamba_metadata.query_start_loc[: num_prefills + 1],
- cache_indices=mamba_metadata.state_indices[:num_prefills],
- has_initial_state=has_init,
- conv_states=conv_pool,
- activation="silu",
+ return _KDAPrefillState(
+ kernel_pool, kernel_indices, recurrent_in, staged_state, use_indexed_state
)
- if self._bfa_proj_weight is not None:
- beta, g, onorm_lowrank = self._project_gate_inputs(x)
- beta = beta.float()
- if onorm_g is None:
- onorm_g = onorm_lowrank
- else:
- g = self.f_b_proj(self.f_a_proj(x))
- beta = self.b_proj(x).float()
- g = rearrange(g, "... (h d) -> ... h d", d=self.head_dim)
+ def forward_prefill_core(
+ self,
+ packed_conv: torch.Tensor,
+ g: torch.Tensor,
+ beta: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ mamba_metadata,
+ num_prefills: int,
+ ssm_pool: torch.Tensor,
+ slot_indices: torch.Tensor,
+ *,
+ _prepared_state: _KDAPrefillState | None = None,
+ ) -> torch.Tensor:
+ """Run post-convolution KDA and update the selected recurrent-state rows.
- q, k, v = fused_kda_post_conv(
- packed_conv,
- num_heads=self.num_heads,
- head_dim=self.head_dim,
- )
+ ``packed_conv`` is BF16 ``[3 * H * D, tokens]``; ``g`` and ``beta``
+ are raw logits with shapes ``[1, tokens, H, D]`` and ``[1, tokens, H]``.
+ The state pool uses value-first ``[slots, H, D, D]`` storage.
+ """
+ state = _prepared_state
+ if state is None:
+ state = self._prepare_prefill_state(
+ cu_seqlens,
+ mamba_metadata,
+ num_prefills,
+ ssm_pool,
+ slot_indices,
+ packed_conv[0].shape[1] if isinstance(packed_conv, tuple) else packed_conv.shape[1],
+ )
+ kernel_pool, kernel_indices = state.kernel_pool, state.kernel_indices
+ recurrent_in, staged_state = state.recurrent_in, state.staged_state
+ use_indexed_state = state.use_indexed_state
+
+ if isinstance(packed_conv, tuple):
+ q, k, v = packed_conv
+ else:
+ q, k, v = fused_kda_post_conv(
+ packed_conv,
+ num_heads=self.num_heads,
+ head_dim=self.head_dim,
+ normalize_qk=False,
+ )
# The optimized kernel reads and writes the V-first pool directly;
# the FLA core exchanges batch-dense recurrent states.
@@ -752,35 +870,28 @@ def forward_prefill(
v=v,
g=g,
beta=beta,
- A_log=self.A_log,
- dt_bias=self.dt_bias,
+ A_log=self._A_log_f32 if self._A_log_f32 is not None else self.A_log,
+ dt_bias=self._dt_bias_f32 if self._dt_bias_f32 is not None else self.dt_bias,
scale=self.head_k_dim**-0.5,
initial_state=None if use_indexed_state else recurrent_in,
safe_gate=lower_bound is not None,
lower_bound=lower_bound,
cu_seqlens=cu_seqlens,
- chunk_indices=chunk_indices,
state_pool=kernel_pool if use_indexed_state else None,
state_indices=kernel_indices if use_indexed_state else None,
- varlen_is_aligned=varlen_is_aligned,
- single_sequence_length=single_sequence_length,
+ max_seqlen=getattr(mamba_metadata, "kda_max_seqlen", 0),
+ seq_order=getattr(mamba_metadata, "kda_seq_order", None),
+ cu_chunks=getattr(mamba_metadata, "kda_cu_chunks", None),
)
# The packed convolution persisted the live convolution pool in place.
if final_state is not None:
- _writeback_state_rows(ssm_pool, slot_indices_long, final_state)
+ _writeback_state_rows(ssm_pool, slot_indices, final_state)
else:
assert use_indexed_state
if staged_state is not None:
- _writeback_state_rows(ssm_pool, slot_indices_long, staged_state)
- # Fused-verify replay caches: seed the committed conv window so the
- # first verify round convolves the correct history (pending drafts
- # are zero for a fresh request, so the tail columns are unused).
- self._sync_kda_replay_conv_window(layer_cache, slot_indices, conv_pool)
-
- core = self._output_gate(x, o, onorm_g)
- # Can be removed once attention writes the core in place.
- return self._store_core(core, output)
+ _writeback_state_rows(ssm_pool, slot_indices, staged_state)
+ return o
def forward_decode(
self,
diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml
index d95c426ee9e3..5515860466e6 100644
--- a/tests/integration/test_lists/test-db/l0_b200.yml
+++ b/tests/integration/test_lists/test-db/l0_b200.yml
@@ -147,18 +147,18 @@ l0_b200:
- unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py
- unittest/_torch/modeling/test_kimi_kda_fp8_packed_prefill.py
- unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py
- - unittest/_torch/modules/kimi_kda/test_kda_buf_cache_bounded.py
- unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py
- unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py
+ - unittest/_torch/modules/kimi_kda/test_kda_flashinfer_prefill.py
+ - unittest/_torch/modules/kimi_kda/test_kda_fused_prefill_conv.py
- unittest/_torch/modules/kimi_kda/test_kda_decode_op.py
+ - unittest/_torch/modules/kimi_kda/test_kimi_kda_decode_state_parity.py
- unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py
# CuTe MTP verify kernel parity (trtllm::kda_mtp_decode).
- unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py
- unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.py
- unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_rmsnorm_op.py
- # KDA runtime: host-derived prefill metadata and the bf16 state pool
- # round-trip. Both are single-device cases that use GPU 0 only.
- - unittest/_torch/modules/kimi_kda/test_kda_host_metadata.py
+ # KDA runtime state-pool and mixed-batch core coverage.
- unittest/_torch/modules/kimi_kda/test_kimi_kda_bf16_state_pool.py
- unittest/_torch/modules/kimi_kda/test_kda_mixed_batch_core_layout.py
# GPU KDA disagg transfer + peer-validation (cpu_only cases skipped by the
diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml
index 0519a1d91874..6d9feb4cafa8 100644
--- a/tests/integration/test_lists/test-db/l0_cpu.yml
+++ b/tests/integration/test_lists/test-db/l0_cpu.yml
@@ -53,7 +53,6 @@ l0_cpu:
# file without a pytest.mark.cpu_only when the stage runs -m cpu_only, so this entry
# contributes only the marked files under attention/, in these sub-trees or elsewhere.
- unittest/_torch/attention
- - unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py
- unittest/_torch/compilation/test_auto_multi_stream.py
- unittest/_torch/compilation/test_phase_selective_forward.py
- unittest/_torch/compilation/test_remove_copy_pass.py::test_remove_copy_preserves_minimax_producer_outputs
diff --git a/tests/integration/test_lists/test-db/l0_gb300.yml b/tests/integration/test_lists/test-db/l0_gb300.yml
index 4173fa0b8141..1908fd1db4ce 100644
--- a/tests/integration/test_lists/test-db/l0_gb300.yml
+++ b/tests/integration/test_lists/test-db/l0_gb300.yml
@@ -43,6 +43,18 @@ l0_gb300:
- accuracy/test_disaggregated_serving.py::TestQwen3_8_Flash_Next::test_fp8_nixl_python[prefix_cache]
- unittest/_torch/moe/kernels/parallel
- unittest/_torch/kernels/parallel TIMEOUT (90)
+ # Kimi K3 KDA unit tests (single GPU, random weights, no checkpoint).
+ - unittest/_torch/modules/kimi_kda/test_kda_flashinfer_prefill.py
+ - unittest/_torch/modules/kimi_kda/test_kda_fused_prefill_conv.py
+ - unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py
+ - unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py
+ - unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py
+ - unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py
+ - unittest/_torch/modules/kimi_kda/test_kda_decode_op.py
+ - unittest/_torch/modules/kimi_kda/test_kimi_kda_decode_state_parity.py
+ - unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py
+ - unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py
+ - unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py
- unittest/_torch/visual_gen/kernels/parallel
- condition:
ranges:
diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
index 10d787106189..5ac740d94e1c 100644
--- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
+++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml
@@ -60,22 +60,9 @@ l0_gb300_multi_gpus:
- unittest/_torch/moe/test_moe_routing.py
- unittest/_torch/modeling -k "modeling_llama"
- unittest/_torch/modeling -k "modeling_gpt_oss"
- # Kimi K3 speculative decoding: exact kernel-level verify parity (random
- # weights, no checkpoint).
- - unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py
- # Fused KDA multi-token verify (trtllm::kda_mtp_decode): kernel-level
- # parity vs CPU golden + FLA sequential, and runtime-level fused-vs-
- # sequential two-round replay parity.
- - unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py
- - unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py
- # Kimi K3 op parity tests (random weights, no checkpoint; SM100/SM103).
- - unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py
- - unittest/_torch/modules/kimi_kda/test_kda_decode_op.py
- unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.py
- unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_rmsnorm_op.py
- # KDA runtime: host-derived prefill metadata and the bf16 state pool
- # round-trip. Both are single-device cases that use GPU 0 only.
- - unittest/_torch/modules/kimi_kda/test_kda_host_metadata.py
+ # KDA runtime state-pool and mixed-batch core coverage.
- unittest/_torch/modules/kimi_kda/test_kimi_kda_bf16_state_pool.py
- unittest/_torch/modules/kimi_kda/test_kda_mixed_batch_core_layout.py
# Kimi K3 MoE routing shape on the TRTLLM backend (random weights, no
diff --git a/tests/unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py b/tests/unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py
deleted file mode 100644
index 026db0130c21..000000000000
--- a/tests/unittest/_torch/cute_dsl/test_kimi_k3_kda_ptx_patch.py
+++ /dev/null
@@ -1,174 +0,0 @@
-# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Contract tests for the k4_persistent ptx-options monkey-patch.
-
-The patch's exception handling encodes two constraints that are invisible in
-its behavior and have already round-tripped through review once:
-
- * The handler must stay a single ``except Exception`` clause: the
- nvidia-cutlass-dsl 4.5.0 AST preprocessor rejects tuple except handlers
- anywhere in a kernel module ("'Tuple' object has no attribute 'id'"),
- breaking every ``cute.compile`` of the file. A narrowing to
- ``except (AttributeError, ImportError)`` had to be reverted for exactly
- this (commit f79c6ef3af6).
- * Only the expected "DSL absent / API surface moved" kinds
- (``AttributeError``, ``ImportError``) may be swallowed. Anything else
- must fail the import loudly: silently skipping the patch defers the
- blow-up to the first K4 compile (ptxas over WG0's 40-register budget
- without ``--uumn``) with no trail back to the patch.
-
-CPU-only: the import scenarios run the module top level in subprocesses (no
-kernel is compiled); the parser constraint is checked structurally on the AST.
-"""
-
-import ast
-import importlib.util
-import subprocess
-import sys
-from pathlib import Path
-
-import pytest
-
-# The L0 CPU-Generic stages run `pytest -m cpu_only`, and their conftest only
-# collects files containing the literal string "pytest.mark.cpu_only"; without
-# this marker every test here is deselected and pytest exits 5 (no tests
-# collected), which the test_unittests_v2 wrapper reports as a failure.
-pytestmark = pytest.mark.cpu_only
-
-try:
- import cutlass # noqa: F401
- import flashinfer.gdn_kernels # noqa: F401
-
- _KERNEL_DEPS_AVAILABLE = True
-except ImportError:
- _KERNEL_DEPS_AVAILABLE = False
-
-needs_kernel_deps = pytest.mark.skipif(
- not _KERNEL_DEPS_AVAILABLE, reason="requires nvidia-cutlass-dsl and flashinfer (gdn_kernels)"
-)
-
-
-def _k4_path() -> Path:
- """The k4_persistent.py that is actually in use (installed package if
- importable, else the source tree this test file lives in)."""
- try:
- spec = importlib.util.find_spec(
- "tensorrt_llm._torch.attention.kernels.blackwell.kimi_k3_kda.k4_persistent"
- )
- if spec is not None and spec.origin:
- return Path(spec.origin)
- except ModuleNotFoundError:
- pass
- return (
- Path(__file__).resolve().parents[4]
- / "tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/"
- "k4_persistent.py"
- )
-
-
-def test_no_tuple_except_handlers():
- """Guard against reintroducing ``except (A, B):`` anywhere in the module.
-
- Nothing else in CI compiles these kernels, so without this check the
- breakage only shows up as a runtime cute.compile failure on Blackwell.
- """
- tree = ast.parse(_k4_path().read_text())
- bad = [
- handler.lineno
- for node in ast.walk(tree)
- for handler in getattr(node, "handlers", [])
- if isinstance(handler.type, ast.Tuple)
- ]
- assert not bad, (
- f"tuple except handler(s) at line(s) {bad} of k4_persistent.py: the "
- f"nvidia-cutlass-dsl 4.5.0 AST preprocessor cannot parse these and "
- f"every cute.compile of the module fails. Use a single `except "
- f"Exception` with an isinstance re-raise in the body instead."
- )
-
-
-_LOAD_K4 = """
-import importlib.util, sys
-spec = importlib.util.spec_from_file_location("k4_under_test", {k4!r})
-mod = importlib.util.module_from_spec(spec)
-spec.loader.exec_module(mod)
-"""
-
-_FAKE_DSL = """
-import sys, types
-import cutlass.cutlass_dsl # materialize the real parent package first
-fake = types.ModuleType("cutlass.cutlass_dsl.cutlass")
-{fake_body}
-sys.modules["cutlass.cutlass_dsl.cutlass"] = fake
-"""
-
-
-def _run_import(inject: str = "") -> subprocess.CompletedProcess:
- code = _FAKE_DSL.format(fake_body=inject) if inject else ""
- code += _LOAD_K4.format(k4=str(_k4_path()))
- return subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=600)
-
-
-@needs_kernel_deps
-def test_import_applies_patch():
- proc = subprocess.run(
- [
- sys.executable,
- "-c",
- _LOAD_K4.format(k4=str(_k4_path()))
- + """
-from cutlass.cutlass_dsl.cutlass import CuTeDSL
-assert CuTeDSL._get_pipeline.__name__ == "_patched_get_pipeline", \\
- f"ptx-options patch not applied: {CuTeDSL._get_pipeline}"
-""",
- ],
- capture_output=True,
- text=True,
- timeout=600,
- )
- assert proc.returncode == 0, proc.stderr[-2000:]
-
-
-@needs_kernel_deps
-def test_expected_patch_failure_is_soft():
- """AttributeError (API surface moved) logs a warning and continues."""
- proc = _run_import(
- inject="""
-class CuTeDSL: # no _get_pipeline -> AttributeError in the patch block
- pass
-fake.CuTeDSL = CuTeDSL
-"""
- )
- assert proc.returncode == 0, proc.stderr[-2000:]
-
-
-@needs_kernel_deps
-def test_unexpected_patch_failure_raises():
- """Anything but AttributeError/ImportError must fail the import."""
- proc = _run_import(
- inject="""
-class _Meta(type):
- def __getattr__(cls, name):
- raise RuntimeError("simulated unexpected patch failure")
-class CuTeDSL(metaclass=_Meta):
- pass
-fake.CuTeDSL = CuTeDSL
-"""
- )
- assert proc.returncode != 0, (
- "an unexpected exception in the ptx-options patch was swallowed at "
- "import; it must propagate (a silently skipped patch surfaces later "
- "as an unexplained ptxas register-budget failure)"
- )
- assert "simulated unexpected patch failure" in proc.stderr
diff --git a/tests/unittest/_torch/modeling/test_kimi_kda_fp8_packed_prefill.py b/tests/unittest/_torch/modeling/test_kimi_kda_fp8_packed_prefill.py
index 66273d618b3d..93219a9d7837 100644
--- a/tests/unittest/_torch/modeling/test_kimi_kda_fp8_packed_prefill.py
+++ b/tests/unittest/_torch/modeling/test_kimi_kda_fp8_packed_prefill.py
@@ -31,12 +31,16 @@ class _Cfg:
def _has_supported_gpu() -> bool:
- return torch.cuda.is_available() and torch.cuda.get_device_capability(0) in {(10, 0), (10, 3)}
+ return torch.cuda.is_available() and torch.cuda.get_device_capability(0) in {
+ (10, 0),
+ (10, 3),
+ (10, 7),
+ }
pytestmark = pytest.mark.skipif(
not _has_supported_gpu(),
- reason="Kimi K3 FP8 projection is supported only on SM100/SM103",
+ reason="Kimi K3 FP8 projection is supported only on SM100/SM103/SM107",
)
diff --git a/tests/unittest/_torch/modules/kimi_kda/kda_module_test_utils.py b/tests/unittest/_torch/modules/kimi_kda/kda_module_test_utils.py
new file mode 100644
index 000000000000..ea2816626653
--- /dev/null
+++ b/tests/unittest/_torch/modules/kimi_kda/kda_module_test_utils.py
@@ -0,0 +1,165 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Module, packed-batch, and numerical parity fixtures for KDA GPU tests."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+ import torch
+
+ from tensorrt_llm._torch.modules.kimi_kda import KimiKDALinearAttention
+
+
+def make_attention(
+ heads: int,
+ hidden_size: int = 7168,
+ *,
+ weight_seed: int = 17,
+ bf16_weights: bool = False,
+) -> KimiKDALinearAttention:
+ """Build an attention module with reproducible synthetic checkpoint values."""
+ import torch
+
+ from tensorrt_llm._torch.model_config import ModelConfig
+ from tensorrt_llm._torch.modules.kimi_kda import KimiKDALinearAttention
+ from tensorrt_llm._torch.modules.linear import Linear
+ from tensorrt_llm.models.modeling_utils import QuantConfig
+ from tensorrt_llm.quantization import QuantAlgo
+
+ cfg = SimpleNamespace(
+ hidden_size=hidden_size,
+ rms_norm_eps=1e-5,
+ linear_attn_config=dict(
+ num_heads=heads,
+ head_dim=128,
+ short_conv_kernel_size=4,
+ use_full_rank_gate=True,
+ gate_lower_bound=-5.0,
+ ),
+ )
+ config = ModelConfig(
+ quant_config=QuantConfig()
+ if bf16_weights
+ else QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES)
+ )
+ with torch.random.fork_rng(devices=[torch.cuda.current_device()]):
+ torch.manual_seed(weight_seed)
+ with torch.device("cuda"):
+ attention = KimiKDALinearAttention(cfg, layer_idx=0, model_config=config)
+ generator = torch.Generator(device="cuda").manual_seed(weight_seed)
+ for module in attention.children():
+ if isinstance(module, Linear):
+ codes = (
+ (torch.randn(module.weight.shape, device="cuda", generator=generator) * 16)
+ .clamp(-64, 64)
+ .to(torch.float8_e4m3fn)
+ )
+ scales = torch.full(
+ module.weight_scale.shape, 1 / (16 * module.in_features**0.5), device="cuda"
+ )
+ module.load_weights([{"weight": codes, "weight_scale": scales}])
+ attention.dt_bias.data.zero_()
+ attention.finalize_decode_weights()
+ for module in attention.modules():
+ if isinstance(module, Linear):
+ module.post_load_weights()
+ attention.eval()
+ if attention._dispatch.prefill_kernel_path != "optimized":
+ raise RuntimeError("The test requires optimized production prefill")
+ return attention
+
+
+def make_batch(
+ heads: int,
+ lengths: list[int],
+ hidden_size: int = 7168,
+ *,
+ initial: str = "fresh",
+ input_seed: int = 23,
+ padded: bool = False,
+) -> SimpleNamespace:
+ """Prepare all index/offset conversions and synthetic inputs outside execution."""
+ import torch
+
+ n, total = len(lengths), sum(lengths)
+ generator = torch.Generator(device="cuda").manual_seed(input_seed)
+ offsets = [0]
+ for length in lengths:
+ offsets.append(offsets[-1] + length)
+ cu = torch.tensor(offsets, dtype=torch.int64, device="cuda")
+ indices = torch.arange(n, 0, -1, dtype=torch.int32, device="cuda")
+ has_init = [initial == "continuation" or (initial == "mixed" and i % 2 == 0) for i in range(n)]
+ metadata = SimpleNamespace(
+ has_initial_states=torch.tensor(has_init, dtype=torch.bool, device="cuda"),
+ use_initial_states=any(has_init),
+ state_indices=indices,
+ query_start_loc=cu.to(torch.int32),
+ query_start_loc_long=cu,
+ kda_max_seqlen=max(lengths, default=0),
+ _arange_buffer=torch.arange(n + 1, dtype=torch.int32, device="cuda"),
+ )
+ hidden = (
+ torch.randn(total, hidden_size, device="cuda", dtype=torch.bfloat16, generator=generator)
+ * 0.1
+ )
+ conv = (
+ torch.randn(
+ n + 2, 3 * heads * 128, 3, device="cuda", dtype=torch.bfloat16, generator=generator
+ )
+ * 0.05
+ )
+ state = torch.empty_strided(
+ (n + 2, heads, 128, 128),
+ (heads * 16384 + (8 if padded else 0), 16384, 128, 1),
+ device="cuda",
+ dtype=torch.bfloat16,
+ )
+ state.copy_(torch.randn(state.shape, device="cuda", generator=generator) * 0.05)
+ return SimpleNamespace(
+ hidden=hidden,
+ cu=cu,
+ indices=indices,
+ metadata=metadata,
+ conv=conv,
+ state=state,
+ n=n,
+ lengths=lengths,
+ )
+
+
+def clone_pool(pool: torch.Tensor) -> torch.Tensor:
+ """Copy pool values while preserving padding between rows."""
+ import torch
+
+ return torch.empty_strided(
+ pool.shape, pool.stride(), device=pool.device, dtype=pool.dtype
+ ).copy_(pool)
+
+
+def compare(actual: torch.Tensor, expected: torch.Tensor) -> dict:
+ """Report the existing KDA parity criteria and elementwise diagnostics."""
+ import torch
+
+ contract = actual.shape == expected.shape and actual.dtype == expected.dtype
+ a, b = actual.float().flatten(), expected.float().flatten()
+ finite = bool(torch.isfinite(a).all() & torch.isfinite(b).all())
+ error = a - b
+ norm = b.norm().item()
+ relative = error.norm().item() / max(norm, 1e-12)
+ cosine = torch.nn.functional.cosine_similarity(a, b, dim=0).item() if norm else 1.0
+ return dict(
+ passed=contract and finite and relative < 0.03 and cosine > 0.999,
+ shape_dtype_match=contract,
+ finite=finite,
+ max_abs=error.abs().max().item() if finite else None,
+ relative_l2=relative if finite else None,
+ cosine=cosine if finite else None,
+ mismatches=int((~torch.isclose(a, b, rtol=0.01, atol=0.01)).sum()),
+ rtol=0.01,
+ atol=0.01,
+ relative_l2_limit=0.03,
+ cosine_min=0.999,
+ )
diff --git a/tests/unittest/_torch/modules/kimi_kda/kda_prefill_test_utils.py b/tests/unittest/_torch/modules/kimi_kda/kda_prefill_test_utils.py
index c52d9fde6425..a3d2e179e4cb 100644
--- a/tests/unittest/_torch/modules/kimi_kda/kda_prefill_test_utils.py
+++ b/tests/unittest/_torch/modules/kimi_kda/kda_prefill_test_utils.py
@@ -3,7 +3,6 @@
"""Helpers for indexed KDA prefill parity tests."""
import torch
-from fla.modules.l2norm import l2norm_fwd
from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import KDAKernelDispatch
from tensorrt_llm._torch.modules.mamba.recurrent_state_cache import reset_recurrent_state_rows
@@ -27,10 +26,16 @@ def run_indexed_prefill(
"""Run optimized prefill against a pool and return its selected rows."""
assert dispatch.prefill_kernel_path == "optimized"
q, k, v, g, beta = inputs
- q, _ = l2norm_fwd(q.clone())
- k, _ = l2norm_fwd(k.clone())
a_log, dt_bias = gate_params
num_sequences = q.shape[0] if cu_seqlens is None else cu_seqlens.shape[0] - 1
+ if cu_seqlens is None:
+ cu_seqlens = (
+ torch.arange(num_sequences + 1, device=q.device, dtype=torch.int64) * q.shape[1]
+ )
+ q, k, v, g = (t.reshape(1, -1, *t.shape[2:]) for t in (q, k, v, g))
+ beta = beta.reshape(1, -1, beta.shape[-1])
+ else:
+ cu_seqlens = cu_seqlens.to(torch.int64)
pool_was_provided = state_pool is not None
if state_indices is None:
@@ -80,7 +85,7 @@ def run_indexed_prefill(
)
assert final_state is None
selected_state = state_pool.index_select(0, state_indices_long)
- return output, selected_state, state_pool
+ return output.reshape_as(inputs[0]), selected_state, state_pool
def run_fla_prefill(
@@ -95,8 +100,6 @@ def run_fla_prefill(
"""Run the independent FLA state path."""
assert dispatch.prefill_kernel_path == "fla"
q, k, v, g, beta = inputs
- q, _ = l2norm_fwd(q.clone())
- k, _ = l2norm_fwd(k.clone())
a_log, dt_bias = gate_params
output, final_state = dispatch.prefill_chunk_kda(
q=q,
diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_buf_cache_bounded.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_buf_cache_bounded.py
deleted file mode 100644
index 9f2292345beb..000000000000
--- a/tests/unittest/_torch/modules/kimi_kda/test_kda_buf_cache_bounded.py
+++ /dev/null
@@ -1,163 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Bounds of the shape-keyed KDA prefill scratch caches.
-
-``_buf_cache``, ``_padded_input_cache`` and ``_g_sentinel_cache`` are keyed by
-the packed prefill token count, which changes on almost every executor
-iteration, so each of them is LRU-bounded by ``_lru_put`` / ``_lru_touch``.
-One ``_buf_cache`` entry holds about 1 GiB at Kimi K3 sizes and is allocated
-after the KV-cache pool is sized, so the cap (``TLLM_KDA_BUF_CACHE_ENTRIES``,
-default 2) is a real memory knob.
-
-These tests call the cache getters directly and launch no KDA kernel, so they
-run on any CUDA GPU. Releasing the scratch of entries evicted after a real
-prefill launch is covered by ``test_kda_cache_soundness.py``.
-"""
-
-import gc
-
-import pytest
-import torch
-
-# Small on purpose: the bound does not depend on the shape, and K3's H=96 at
-# T=8192 would take about 1 GiB per entry.
-H = 4
-K_DIM = 128
-V_DIM = 128
-BT = 64
-T0 = 1024
-N_SHAPES = 16
-
-
-def _op_module():
- # The module raises ImportError at import time without CuTe DSL / FlashInfer.
- return pytest.importorskip(
- "tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_custom_ops", exc_type=ImportError
- )
-
-
-def _settle() -> None:
- gc.collect()
- torch.cuda.synchronize()
-
-
-def test_lru_put_evicts_least_recently_used() -> None:
- module = _op_module()
- cache = {}
- module._lru_put(cache, "a", 1, max_entries=2)
- module._lru_put(cache, "b", 2, max_entries=2)
- assert module._lru_touch(cache, "a") == 1
- assert module._lru_put(cache, "c", 3, max_entries=2) == 3
- assert list(cache) == ["a", "c"]
-
-
-@pytest.mark.parametrize(
- ("value", "expected"), [(None, 2), ("5", 5), ("1", 1), ("0", 1), ("abc", 2), ("", 2)]
-)
-def test_scratch_cache_cap_from_env(
- monkeypatch: pytest.MonkeyPatch, value: str | None, expected: int
-) -> None:
- module = _op_module()
- if value is None:
- monkeypatch.delenv("TLLM_KDA_BUF_CACHE_ENTRIES", raising=False)
- else:
- monkeypatch.setenv("TLLM_KDA_BUF_CACHE_ENTRIES", value)
- assert module._scratch_cache_max_entries() == expected
-
-
-@pytest.fixture
-def clean_caches():
- if not torch.cuda.is_available():
- pytest.skip("needs a CUDA device")
- module = _op_module()
- caches = (module._buf_cache, module._padded_input_cache, module._g_sentinel_cache)
- saved = [dict(cache) for cache in caches]
- for cache in caches:
- cache.clear()
- _settle()
- yield module
- for cache, entries in zip(caches, saved):
- cache.clear()
- cache.update(entries)
- _settle()
-
-
-def _get_buffers(module, device: torch.device, t: int):
- return module._get_buffers(
- device, torch.bfloat16, 1, t, H, K_DIM, V_DIM, t // BT, 1, BT, varlen=True
- )
-
-
-def _buf_entry_bytes(t: int) -> int:
- """Bytes of the large tensors of one varlen ``_get_buffers`` entry.
-
- k_scaled, kg and q_scaled are [T + BT, H, K]; A_qk and A_kk are
- [T + BT, H, BT]; O_flat is [T + BT, H, V]; all bf16. The smaller tensors
- are covered by the slack of the budgets below.
- """
- return (t + BT) * H * 2 * (3 * K_DIM + 2 * BT + V_DIM)
-
-
-@torch.no_grad()
-def test_buf_cache_holds_at_most_cap_entries(clean_caches) -> None:
- module = clean_caches
- device = torch.device("cuda", torch.cuda.current_device())
- base = torch.cuda.memory_allocated(device)
-
- for i in range(N_SHAPES):
- _get_buffers(module, device, T0 + BT * i)
- _settle()
-
- cap = module._BUF_CACHE_MAX_ENTRIES
- assert len(module._buf_cache) <= cap
- live = torch.cuda.memory_allocated(device) - base
- budget = int(1.25 * cap * _buf_entry_bytes(T0 + BT * (N_SHAPES - 1))) + (8 << 20)
- assert live <= budget, (
- f"{live / 2**20:.1f} MiB live after {N_SHAPES} shapes, "
- f"budget {budget / 2**20:.1f} MiB for {cap} entries"
- )
-
- module._buf_cache.clear()
- _settle()
- leaked = torch.cuda.memory_allocated(device) - base
- assert leaked < (1 << 20), f"{leaked / 2**20:.1f} MiB survived clearing _buf_cache"
-
-
-@torch.no_grad()
-def test_buf_cache_keeps_recently_used_entry(clean_caches) -> None:
- module = clean_caches
- device = torch.device("cuda", torch.cuda.current_device())
- cap = module._BUF_CACHE_MAX_ENTRIES
- if cap < 2:
- pytest.skip("a recently used entry survives an insertion only with a cap of 2 or more")
-
- hot = _get_buffers(module, device, T0)
- for i in range(1, 2 * cap + 1):
- assert _get_buffers(module, device, T0) is hot
- _get_buffers(module, device, T0 + BT * i)
- assert _get_buffers(module, device, T0) is hot
-
-
-@torch.no_grad()
-def test_padded_input_and_g_sentinel_caches_are_bounded(clean_caches) -> None:
- module = clean_caches
- device = torch.device("cuda", torch.cuda.current_device())
-
- for i in range(N_SHAPES):
- t_padded = T0 + BT * i
- real_t = t_padded - 7
- module._get_g_sentinel_buffer(1, t_padded, H, K_DIM, torch.bfloat16, device, real_t)
- module._get_padded_input_buffers(
- 1,
- t_padded,
- H,
- K_DIM,
- torch.bfloat16,
- torch.bfloat16,
- torch.float32,
- device,
- real_t,
- )
-
- assert len(module._g_sentinel_cache) <= module._G_SENTINEL_CACHE_MAX_ENTRIES
- assert len(module._padded_input_cache) <= module._PADDED_INPUT_CACHE_MAX_ENTRIES
diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py
index cda5e71ed7e3..0a0a5d1e73f6 100644
--- a/tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py
+++ b/tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py
@@ -1,42 +1,12 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
-"""Runtime-integration soundness tests for indexed optimized KDA prefill.
-
-Covers the two classes of bug that ordinary op-math parity tests cannot see:
-
-1. STREAM soundness. The executor runs the model on a dedicated non-blocking
- ``torch.cuda.Stream`` (``py_executor.execution_stream``); the CuTe DSL
- kernels must launch on torch's current stream, not the DSL default stream.
- A default-stream launch races with the projections producing q/k/v/g/beta
- and with consumers of the output and mutated state-pool rows — silent,
- intermittent corruption e2e (GSM8K 68.99 vs 97.01 on the FLA control)
- while single-stream unit tests all pass.
- ``test_indexed_prefill_uses_current_stream`` reproduces this
- deterministically by delaying the execution stream with ``_sleep`` so any
- kernel launched on the default stream reads not-yet-written inputs.
-
-2. id()-keyed cache soundness under the runtime's fresh-tensor-per-call
- pattern. CPython recycles an object's address (= its id) as soon as it is
- freed; the ``_prune_on_gc`` weakref guard must pop the cache entry at
- dealloc, before a new tensor can alias the key. The recycled-ID test
- attempts to engineer exactly that aliasing: cache a verdict under id(t1),
- free t1, re-allocate until a new tensor lands on the same id, refill it
- with data for which the stale verdict would be WRONG, and parity-check
- against FLA. ``test_repeated_single_sequence_metadata_matches_fla``
- additionally covers the Phase 2.1 poisoning case (same cu_seqlens object,
- two calls).
-
-3. Bounded scratch. The prefill scratch cache ``_buf_cache`` keeps at most
- ``_BUF_CACHE_MAX_ENTRIES`` token-count shapes, and the runtime calls with a
- new token count for most prefill batches, so entries are evicted all the
- time. ``test_evicted_prefill_scratch_is_released`` checks that an evicted
- entry's scratch is freed: nothing else, such as a cached CuTe wrapper, may
- keep it alive. ``test_prefill_scratch_plateaus_over_distinct_token_counts``
- checks that memory stays bounded when every call brings a new token count.
+"""Prefill must use the caller's stream and current packed sequence metadata.
+
+Delayed producer operations exercise stream ordering. Repeated calls with
+changed sequence boundaries exercise metadata regeneration.
"""
import gc
-import weakref
import pytest
import torch
@@ -56,12 +26,14 @@
def _has_supported_gpu() -> bool:
- return torch.cuda.is_available() and torch.cuda.get_device_capability(0) in {(10, 0), (10, 3)}
+ if not torch.cuda.is_available():
+ return False
+ return torch.cuda.get_device_capability(0) in {(10, 0), (10, 3), (10, 7)}
pytestmark = pytest.mark.skipif(
not _has_supported_gpu(),
- reason="Kimi K3 is supported only on Blackwell (SM100/SM103)",
+ reason="FlashInfer KDA prefill requires SM100/SM103/SM107",
)
@@ -117,47 +89,6 @@ def _make_cu_seqlens(sequence_lengths: list[int]) -> torch.Tensor:
)
-def _op_module():
- from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_custom_ops
-
- return cute_dsl_kimi_k3_custom_ops
-
-
-def _flush_tensor_cache_pins() -> None:
- from fla.ops.utils.index import prepare_chunk_indices as fla_prepare_chunk_indices
- from fla.ops.utils.index import prepare_chunk_offsets
-
- from tensorrt_llm._torch.modules.fla.index import prepare_chunk_indices
-
- for _ in range(5):
- dummy = torch.tensor([0, 64], dtype=torch.long, device="cuda")
- fla_prepare_chunk_indices(dummy, 64)
- prepare_chunk_offsets(dummy, 64)
- prepare_chunk_indices(dummy, 64)
-
-
-def _release_cu_seqlens(module, holder: list[torch.Tensor]) -> tuple[int, bool]:
- old_id = id(holder[0])
- holder.clear()
- _flush_tensor_cache_pins()
- gc.collect()
- return old_id, old_id not in module._varlen_pure_cache
-
-
-def _allocate_with_recycled_id(target_id: int, make, attempts: int = 512):
- retained = []
- for _ in range(attempts):
- candidate = make()
- if id(candidate) == target_id:
- return candidate
- retained.append(candidate)
- return None
-
-
-def _tensor_weakrefs(values) -> list[weakref.ref]:
- return [weakref.ref(value) for value in values if isinstance(value, torch.Tensor)]
-
-
@pytest.mark.parametrize("regime", ["eqlen", "varlen"])
@torch.no_grad()
def test_indexed_prefill_uses_current_stream(
@@ -209,106 +140,31 @@ def test_indexed_prefill_uses_current_stream(
ids=["varlen", "single_sequence"],
)
@torch.no_grad()
-def test_recycled_cu_seqlens_id_matches_fla(
+def test_replaced_packed_metadata_matches_fla(
dispatch_pair: tuple[KDAKernelDispatch, KDAKernelDispatch],
gate_params: tuple[torch.Tensor, torch.Tensor],
aligned_lengths: list[int],
unaligned_lengths: list[int],
) -> None:
- """A recycled tensor ID cannot reuse stale alignment or length metadata."""
+ """Each invocation must use its current packed sequence boundaries."""
optimized, reference = dispatch_pair
- module = _op_module()
- cu_holder = [_make_cu_seqlens(aligned_lengths)]
- inputs = _make_inputs(sum(aligned_lengths), seed=11)
- run_indexed_prefill(optimized, gate_params, inputs, cu_holder[0])
- torch.cuda.synchronize()
- assert module._varlen_pure_cache.get(id(cu_holder[0])) is True
-
- old_id, pruned = _release_cu_seqlens(module, cu_holder)
- if not pruned:
- collision = _allocate_with_recycled_id(
- old_id,
- lambda: _make_cu_seqlens(unaligned_lengths),
- attempts=64,
+ for iteration, lengths in enumerate((aligned_lengths, unaligned_lengths, aligned_lengths)):
+ inputs = _make_inputs(sum(lengths), seed=11 + iteration)
+ cu = _make_cu_seqlens(lengths)
+ actual_output, actual_state, _ = run_indexed_prefill(
+ optimized,
+ gate_params,
+ inputs,
+ cu,
)
- assert collision is None, "a live cache entry allowed its tensor ID to be recycled"
- pytest.skip("cu_seqlens remains pinned, so its ID cannot be recycled")
-
- cu_seqlens = _allocate_with_recycled_id(
- old_id,
- lambda: _make_cu_seqlens(unaligned_lengths),
- )
- if cu_seqlens is None:
- pytest.skip("could not obtain a recycled cu_seqlens tensor ID")
-
- inputs = _make_inputs(sum(unaligned_lengths), seed=12)
- actual_output, actual_state, _ = run_indexed_prefill(
- optimized,
- gate_params,
- inputs,
- cu_seqlens,
- )
- expected_output, expected_state = run_fla_prefill(
- reference,
- gate_params,
- inputs,
- _make_cu_seqlens(unaligned_lengths),
- )
- assert_kda_close("recycled_id/output", actual_output, expected_output)
- assert_kda_close("recycled_id/state", actual_state, expected_state)
- assert module._varlen_pure_cache.get(id(cu_seqlens)) is False
-
-
-@torch.no_grad()
-def test_evicted_prefill_scratch_is_released(
- dispatch_pair: tuple[KDAKernelDispatch, KDAKernelDispatch],
- gate_params: tuple[torch.Tensor, torch.Tensor],
-) -> None:
- """Scratch evicted from the bounded prefill buffer cache is freed."""
- optimized, _ = dispatch_pair
- module = _op_module()
- # Every call has its own token count, and one pass has more shapes than the
- # cache holds, so the cyclic passes below miss and evict on every call.
- # Two sequences per call keep the batch at four or more chunks.
- first_lengths = [300 + 64 * i for i in range(module._BUF_CACHE_MAX_ENTRIES + 4)]
-
- def run_pass() -> None:
- for i, first_length in enumerate(first_lengths):
- sequence_lengths = [first_length, 200]
- inputs = _make_inputs(sum(sequence_lengths), seed=500 + i)
- run_indexed_prefill(optimized, gate_params, inputs, _make_cu_seqlens(sequence_lengths))
- torch.cuda.synchronize()
- _flush_tensor_cache_pins()
- gc.collect()
-
- # The first pass compiles the kernels and leaves the cache holding the last
- # shapes of the cycle; every later pass evicts all of those entries and
- # ends with the same shapes cached again.
- run_pass()
- baseline = torch.cuda.memory_allocated()
- evicted_scratch = _tensor_weakrefs(next(iter(module._buf_cache.values())))
- assert evicted_scratch
- for iteration in range(2):
- run_pass()
- assert len(module._buf_cache) <= module._BUF_CACHE_MAX_ENTRIES
- grown = torch.cuda.memory_allocated() - baseline
- assert grown <= 16 << 20, (
- f"pass {iteration}: {grown / (1 << 20):.1f} MiB more allocated after the evictions"
+ expected_output, expected_state = run_fla_prefill(
+ reference,
+ gate_params,
+ inputs,
+ cu.clone(),
)
- alive = sum(ref() is not None for ref in evicted_scratch)
- assert alive == 0, f"pass {iteration}: {alive} tensors of an evicted entry are still alive"
-
-
-def _prefill_scratch_set_bytes(total_tokens: int) -> int:
- """Upper bound on the large tensors of one varlen ``_get_buffers`` entry.
-
- k_scaled, kg and q_scaled are [T + BT, H, K] and A_qk and A_kk are
- [T + BT, H, BT], all bf16. One extra [T + BT, H, V] term covers the
- smaller tensors and the prefill output with room to spare.
- """
- chunk = 64
- rows = total_tokens + chunk
- return rows * NUM_HEADS * 2 * (3 * HEAD_DIM + 2 * chunk + HEAD_DIM)
+ assert_kda_close("metadata/output", actual_output, expected_output)
+ assert_kda_close("metadata/state", actual_state, expected_state)
@torch.no_grad()
@@ -316,79 +172,34 @@ def test_prefill_scratch_plateaus_over_distinct_token_counts(
dispatch_pair: tuple[KDAKernelDispatch, KDAKernelDispatch],
gate_params: tuple[torch.Tensor, torch.Tensor],
) -> None:
- """Prefill memory plateaus when every call brings a token count never seen before.
-
- The executor's packed prefill token count is the sum of the iteration's
- context chunks and is not quantized, so a long run keeps presenting new
- shapes. The eviction test above cycles through a fixed set of shapes, so a
- shape-keyed cache that is not bounded would be filled on its first pass and
- never grow after the baseline. Here every call uses a new token count, so
- any such cache, and any state pinning evicted scratch, grows by one buffer
- set per call and fails the budget.
- """
+ """Prefill must not retain GPU scratch for every distinct packed token count."""
optimized, _ = dispatch_pair
- module = _op_module()
- num_calls = 24
def run(first_length: int, seed: int) -> None:
- sequence_lengths = [first_length, 200]
- inputs = _make_inputs(sum(sequence_lengths), seed=seed)
- run_indexed_prefill(optimized, gate_params, inputs, _make_cu_seqlens(sequence_lengths))
+ lengths = [first_length, 200]
+ inputs = _make_inputs(sum(lengths), seed=seed)
+ run_indexed_prefill(optimized, gate_params, inputs, _make_cu_seqlens(lengths))
def settle() -> None:
torch.cuda.synchronize()
- _flush_tensor_cache_pins()
gc.collect()
- # Compile the kernels and fill the shape-independent compile caches before
- # the baseline, at token counts the measured loop does not use.
+ # Warm compilation before measuring retained tensor storage. The measured
+ # calls each use a new token count, so a cache retaining per-shape scratch
+ # cannot hide behind reuse of the warmup allocations.
for i in range(2):
run(100 + 64 * i, seed=700 + i)
settle()
- module._buf_cache.clear()
- settle()
baseline = torch.cuda.memory_allocated()
- first_lengths = [300 + 64 * i for i in range(num_calls)]
- for i, first_length in enumerate(first_lengths):
- run(first_length, seed=800 + i)
+ num_calls = 24
+ for i in range(num_calls):
+ run(300 + 64 * i, seed=800 + i)
settle()
grown = torch.cuda.memory_allocated() - baseline
-
- cap = module._BUF_CACHE_MAX_ENTRIES
- assert len(module._buf_cache) <= cap
- # Clamp the cap used for the budget, so raising TLLM_KDA_BUF_CACHE_ENTRIES
- # cannot hide a leak; a deliberately larger cap should revisit this bound.
- max_tokens = first_lengths[-1] + 200
- budget = (min(cap, 4) + 2) * _prefill_scratch_set_bytes(max_tokens) + (128 << 20)
+ # Allow fixed runtime overhead, but not one large scratch set per shape.
+ budget = 128 << 20
assert grown <= budget, (
- f"prefill scratch grows with the token count instead of plateauing: "
- f"{grown / (1 << 20):.0f} MiB allocated after {num_calls} new token counts, "
- f"budget {budget / (1 << 20):.0f} MiB for a cache of {cap} entries"
+ f"prefill retains {grown / (1 << 20):.0f} MiB after {num_calls} new token counts; "
+ f"budget is {budget / (1 << 20):.0f} MiB"
)
-
-
-@torch.no_grad()
-def test_repeated_single_sequence_metadata_matches_fla(
- dispatch_pair: tuple[KDAKernelDispatch, KDAKernelDispatch],
- gate_params: tuple[torch.Tensor, torch.Tensor],
-) -> None:
- """Repeated unaligned calls cannot poison cached single-sequence metadata."""
- optimized, reference = dispatch_pair
- cu_seqlens = _make_cu_seqlens([300])
- for iteration in range(2):
- inputs = _make_inputs(300, seed=40 + iteration)
- actual_output, actual_state, _ = run_indexed_prefill(
- optimized,
- gate_params,
- inputs,
- cu_seqlens,
- )
- expected_output, expected_state = run_fla_prefill(
- reference,
- gate_params,
- inputs,
- _make_cu_seqlens([300]),
- )
- assert_kda_close(f"iteration_{iteration}/output", actual_output, expected_output)
- assert_kda_close(f"iteration_{iteration}/state", actual_state, expected_state)
diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_flashinfer_prefill.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_flashinfer_prefill.py
new file mode 100644
index 000000000000..b683582dc1b2
--- /dev/null
+++ b/tests/unittest/_torch/modules/kimi_kda/test_kda_flashinfer_prefill.py
@@ -0,0 +1,335 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Production module/core behavior for FlashInfer prefill and BF16 state ownership."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytest.importorskip("fla")
+
+from kda_module_test_utils import compare, make_attention, make_batch
+
+pytestmark = pytest.mark.skipif(
+ not torch.cuda.is_available()
+ or torch.cuda.get_device_capability() not in ((10, 0), (10, 3), (10, 7)),
+ reason="FlashInfer KDA requires SM100/SM103/SM107",
+)
+
+
+def assert_parity(actual: torch.Tensor, expected: torch.Tensor) -> None:
+ metrics = compare(actual, expected)
+ assert metrics["passed"], metrics
+
+
+@pytest.mark.parametrize("heads", [6, 12, 24, 48, 96])
+@pytest.mark.parametrize("lengths", [[0, 15, 16, 17], [63, 64, 65], [150]])
+@pytest.mark.parametrize("distribution", ["normal", "uniform", "sparse"])
+@torch.inference_mode()
+def test_indexed_prefill_matches_reference(
+ heads: int, lengths: list[int], distribution: str
+) -> None:
+ from tensorrt_llm._torch.attention.kernels.blackwell.kimi_k3_kda.kda_chunked_bt16 import (
+ kda_chunked_mma_reference,
+ )
+ from tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_custom_ops import kda_prefill
+
+ torch.manual_seed(1901)
+ total, n = sum(lengths), len(lengths)
+
+ def sample(shape: tuple[int, ...]) -> torch.Tensor:
+ x = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
+ if distribution == "uniform":
+ x.uniform_(-1, 1)
+ elif distribution == "sparse":
+ x *= torch.rand(shape, device="cuda") < 0.1
+ return x * 0.2
+
+ q, k, v, g = [sample((1, total, heads, 128)) for _ in range(4)]
+ projection_width = ((128 + heads + 7) // 8) * 8
+ beta = sample((1, total, projection_width))[..., 128 : 128 + heads]
+ a_log = torch.linspace(-2, 1, heads, device="cuda")
+ bias = torch.linspace(-1, 1, heads * 128, device="cuda").reshape(heads, 128)
+ slots_backing = torch.arange(n + 1, 0, -1, device="cuda", dtype=torch.int32)
+ slots = slots_backing[1:] # Contiguous, with only natural int32 alignment.
+ offsets = torch.tensor([0] + lengths, device="cuda", dtype=torch.int64).cumsum(0)
+ row_size = heads * 128 * 128
+ backing = sample((n + 2, row_size + 128))
+ state = backing[:, :row_size].view(n + 2, heads, 128, 128)
+ state[slots[::2].long()] = 0 # Mixed fresh and continuation rows.
+ before = backing.clone()
+ immutable = [q, k, v, g, beta, a_log, bias, slots, offsets]
+ snapshots = [x.clone() for x in immutable]
+ expected_out, expected_state = kda_chunked_mma_reference(
+ *(x.transpose(1, 2) for x in (q, k, v, g, beta)),
+ initial_state=state[slots.long()].transpose(-1, -2),
+ a_log=a_log,
+ dt_bias=bias,
+ cu_seqlens=offsets,
+ )
+ for mode in ("engine", "decomp", "bt16", "small-bh", "auto"):
+ backing.copy_(before)
+ output = kda_prefill(
+ q,
+ k,
+ v,
+ g,
+ beta,
+ state,
+ slots,
+ offsets,
+ a_log,
+ bias,
+ 128**-0.5,
+ -5.0,
+ mode,
+ max(lengths),
+ )
+ torch.testing.assert_close(output, expected_out.transpose(1, 2), rtol=0.01, atol=0.01)
+ torch.testing.assert_close(
+ state[slots.long()],
+ expected_state.transpose(-1, -2).to(torch.bfloat16),
+ rtol=0.01,
+ atol=0.01,
+ )
+ torch.testing.assert_close(backing[[0, n + 1]], before[[0, n + 1]], rtol=0, atol=0)
+ torch.testing.assert_close(backing[:, row_size:], before[:, row_size:], rtol=0, atol=0)
+ for actual, expected in zip(immutable, snapshots):
+ torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize(
+ "mode,state_dtype",
+ [(mode, torch.bfloat16) for mode in ("engine", "decomp", "auto", "bt16", "small-bh")]
+ + [("auto", torch.float32)],
+)
+@pytest.mark.parametrize("heads", [6, 96])
+@pytest.mark.parametrize("beta_layout", ["projection", "head_stride", "unaligned", "token_pitch"])
+@torch.inference_mode()
+def test_prefill_beta_views_match_compact_eager_and_graph(
+ mode: str, state_dtype: torch.dtype, heads: int, beta_layout: str
+) -> None:
+ from tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_custom_ops import kda_prefill
+
+ torch.manual_seed(1901)
+ total = 33
+ q, k, v, g = [
+ torch.randn(1, total, heads, 128, device="cuda", dtype=torch.bfloat16) * 0.2
+ for _ in range(4)
+ ]
+ if beta_layout == "projection":
+ width = ((128 + heads + 7) // 8) * 8
+ beta = torch.randn(1, total, width, device="cuda", dtype=torch.bfloat16)[
+ ..., 128 : 128 + heads
+ ]
+ elif beta_layout == "head_stride":
+ beta = torch.randn(1, total, 2 * heads, device="cuda", dtype=torch.bfloat16)[..., ::2]
+ elif beta_layout == "unaligned":
+ beta = torch.randn(total * heads + 1, device="cuda", dtype=torch.bfloat16)[1:].view(
+ 1, total, heads
+ )
+ else:
+ beta = torch.randn(1, total, heads + 1, device="cuda", dtype=torch.bfloat16)[..., :heads]
+ compact_beta = beta.clone(memory_format=torch.contiguous_format)
+ slots = torch.tensor([2, 1], device="cuda", dtype=torch.int32)
+ offsets = torch.tensor([0, 16, total], device="cuda", dtype=torch.int64)
+ a_log = torch.linspace(-2, 1, heads, device="cuda")
+ bias = torch.linspace(-1, 1, heads * 128, device="cuda").reshape(heads, 128)
+ row_size = heads * 128 * 128
+ initial = torch.randn(4, row_size + 128, device="cuda", dtype=state_dtype) * 0.1
+ initial[2, :row_size] = 0
+ backing = initial.clone()
+ state = backing[:, :row_size].view(4, heads, 128, 128)
+ immutable = [q, k, v, g, beta, slots, a_log, bias]
+ snapshots = [tensor.clone() for tensor in immutable]
+
+ def run(beta_input: torch.Tensor) -> torch.Tensor:
+ return kda_prefill(
+ q, k, v, g, beta_input, state, slots, offsets, a_log, bias, 128**-0.5, -5.0, mode
+ )
+
+ run(beta)
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ graph_output = run(beta)
+ for split in (16, 15):
+ offsets[1] = split
+ backing.copy_(initial)
+ expected = run(compact_beta).clone()
+ expected_state = backing.clone()
+ backing.copy_(initial)
+ actual = run(beta)
+ torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+ torch.testing.assert_close(backing, expected_state, rtol=0, atol=0)
+ backing.copy_(initial)
+ graph.replay()
+ torch.testing.assert_close(graph_output, expected, rtol=0, atol=0)
+ torch.testing.assert_close(backing, expected_state, rtol=0, atol=0)
+ torch.testing.assert_close(backing[[0, 3]], initial[[0, 3]], rtol=0, atol=0)
+ torch.testing.assert_close(backing[:, row_size:], initial[:, row_size:], rtol=0, atol=0)
+ for tensor, snapshot in zip(immutable, snapshots):
+ torch.testing.assert_close(tensor, snapshot, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("conv_mode", ["separate", "fused_all"])
+@torch.inference_mode()
+def test_module_continuation_decode_and_verify(conv_mode: str) -> None:
+ heads, hidden = 6, 768
+ baseline = make_attention(heads, hidden_size=hidden)
+ baseline._dispatch.prefill_kernel_path = "fla"
+ candidate = make_attention(heads, hidden_size=hidden)
+ candidate._fuse_prefill_conv = conv_mode == "fused_all"
+ batch = make_batch(heads, [17, 65], hidden_size=hidden, initial="mixed")
+ states = [batch.state.clone(), batch.state.clone()]
+ convs = [batch.conv.clone(), batch.conv.clone()]
+ for _ in range(3):
+ outputs = [
+ module.forward_prefill(
+ batch.hidden, batch.cu, batch.metadata, batch.n, conv, state, batch.indices
+ )
+ for module, conv, state in zip((baseline, candidate), convs, states)
+ ]
+ assert_parity(outputs[1], outputs[0])
+ assert_parity(states[1], states[0])
+ torch.testing.assert_close(convs[1], convs[0], rtol=0, atol=0)
+ batch.metadata.has_initial_states.fill_(True)
+ batch.metadata.use_initial_states = True
+ for _ in range(3):
+ x = torch.randn(batch.n, hidden, device="cuda", dtype=torch.bfloat16) * 0.1
+ outputs = [
+ module.forward_decode(
+ x,
+ conv,
+ state,
+ batch.indices,
+ batch.metadata,
+ SimpleNamespace(kda_k_cache=None, has_kda_replay_caches=False),
+ ssm_state_indices=batch.indices,
+ )
+ for module, conv, state in zip((baseline, candidate), convs, states)
+ ]
+ assert_parity(outputs[1], outputs[0])
+ assert_parity(states[1], states[0])
+ steps = 3
+ x = torch.randn(batch.n * steps, hidden, device="cuda", dtype=torch.bfloat16) * 0.1
+ outputs, caches = [], []
+ for module, conv, state in zip((baseline, candidate), convs, states):
+ cache = SimpleNamespace(
+ kda_k_cache=None,
+ has_kda_replay_caches=False,
+ intermediate_conv_window=torch.empty(
+ batch.n, steps, 3 * heads * 128, 3, device="cuda", dtype=torch.bfloat16
+ ),
+ intermediate_ssm=torch.empty(
+ batch.n, steps, heads, 128, 128, device="cuda", dtype=torch.bfloat16
+ ),
+ )
+ outputs.append(module.forward_verify(x, steps, cache, conv, state, batch.indices))
+ caches.append(cache)
+ assert_parity(outputs[1], outputs[0])
+ assert_parity(caches[1].intermediate_ssm, caches[0].intermediate_ssm)
+
+
+@pytest.mark.parametrize("mode", ["engine", "decomp", "auto", "small-bh"])
+@pytest.mark.parametrize("conv_mode", ["separate", "fused_all"])
+@torch.inference_mode()
+def test_graph_replay_refreshes_packed_lengths(mode: str, conv_mode: str) -> None:
+ candidate = make_attention(6, hidden_size=768)
+ candidate._dispatch.prefill_mode = mode
+ candidate._fuse_prefill_conv = conv_mode == "fused_all"
+ batch = make_batch(6, [17, 65], hidden_size=768, initial="mixed")
+ batch.metadata.kda_max_seqlen = sum([17, 65])
+ state, conv = batch.state.clone(), batch.conv.clone()
+
+ def run() -> torch.Tensor:
+ return candidate.forward_prefill(
+ batch.hidden, batch.cu, batch.metadata, 2, conv, state, batch.indices
+ )
+
+ for _ in range(3):
+ run()
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ output = run()
+ for lengths in ([17, 65], [33, 49], [0, 82], [64, 18]):
+ batch.cu.copy_(torch.tensor([0, lengths[0], 82], device="cuda", dtype=torch.int64))
+ batch.metadata.query_start_loc.copy_(batch.cu)
+ state.copy_(batch.state)
+ conv.copy_(batch.conv)
+ eager = run().clone()
+ expected_state, expected_conv = state.clone(), conv.clone()
+ state.copy_(batch.state)
+ conv.copy_(batch.conv)
+ graph.replay()
+ torch.testing.assert_close(output, eager, rtol=0, atol=0)
+ torch.testing.assert_close(state, expected_state, rtol=0, atol=0)
+ torch.testing.assert_close(conv, expected_conv, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32])
+@torch.inference_mode()
+def test_production_forward_updates_live_cache(state_dtype: torch.dtype) -> None:
+ """The model entry point updates recurrent, convolution and replay caches."""
+ heads, hidden = 6, 768
+ attention = make_attention(heads, hidden_size=hidden)
+ reference = make_attention(heads, hidden_size=hidden)
+ reference._dispatch.prefill_kernel_path = "fla"
+ batch = make_batch(heads, [17, 65], hidden_size=hidden, initial="mixed", padded=True)
+ # Packed int32 metadata may be a contiguous slice with only natural
+ # 4-byte alignment; FlashInfer does not require a 16-byte slot pointer.
+ slots = torch.empty(batch.n + 1, dtype=torch.int32, device="cuda")
+ slots[1:].copy_(batch.indices)
+ batch.indices = slots[1:]
+ batch.metadata.state_indices = batch.indices
+ state = batch.state.to(state_dtype)
+ initial_state = state.clone()
+ ref_state, ref_conv = state.clone(), batch.conv.clone()
+ from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import PythonMambaCacheManager
+
+ replay_caches = {}
+ for label in ("q", "k", "v"):
+ replay_caches[f"kda_conv_{label}"] = torch.randn(
+ state.shape[0], 7, heads * 128, device="cuda", dtype=torch.bfloat16
+ ).transpose(-1, -2)
+ cache = PythonMambaCacheManager.SpeculativeState(
+ conv=batch.conv.clone(),
+ temporal=state,
+ kda_k_cache=torch.empty(1, device="cuda"),
+ **replay_caches,
+ )
+ old_replay = [getattr(cache, f"kda_conv_{label}").clone() for label in ("q", "k", "v")]
+ metadata = SimpleNamespace(
+ mamba_metadata=batch.metadata,
+ num_contexts=batch.n,
+ num_ctx_tokens=batch.hidden.shape[0],
+ num_tokens=batch.hidden.shape[0],
+ seq_lens=torch.tensor([17, 65]),
+ kv_cache_manager=SimpleNamespace(mamba_layer_cache=lambda _: cache),
+ )
+ original = batch.hidden.clone()
+ actual = attention(batch.hidden, metadata)
+ expected = reference.forward_prefill(
+ batch.hidden,
+ batch.cu,
+ batch.metadata,
+ batch.n,
+ ref_conv,
+ ref_state,
+ batch.indices,
+ )
+ expected = reference._project_output(expected)
+ assert_parity(actual, expected)
+ assert_parity(state[batch.indices.long()], ref_state[batch.indices.long()])
+ torch.testing.assert_close(cache.conv, ref_conv, rtol=0, atol=0)
+ torch.testing.assert_close(batch.hidden, original, rtol=0, atol=0)
+ untouched = [0, batch.n + 1]
+ torch.testing.assert_close(state[untouched], initial_state[untouched], rtol=0, atol=0)
+ for i, label in enumerate(("q", "k", "v")):
+ replay = getattr(cache, f"kda_conv_{label}")
+ selected = cache.conv[batch.indices.long(), i * heads * 128 : (i + 1) * heads * 128]
+ torch.testing.assert_close(replay[batch.indices.long(), :, :3], selected, rtol=0, atol=0)
+ torch.testing.assert_close(replay[untouched], old_replay[i][untouched], rtol=0, atol=0)
+ torch.testing.assert_close(replay[..., 3:], old_replay[i][..., 3:], rtol=0, atol=0)
diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_fused_prefill_conv.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_fused_prefill_conv.py
new file mode 100644
index 000000000000..b04c4ba3bd67
--- /dev/null
+++ b/tests/unittest/_torch/modules/kimi_kda/test_kda_fused_prefill_conv.py
@@ -0,0 +1,180 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Fused convolution raw Q/K/V, state ownership, and CUDA graph parity."""
+
+from itertools import accumulate
+
+import pytest
+import torch
+from kda_module_test_utils import clone_pool, make_batch
+
+from tensorrt_llm._torch.attention.kernels.blackwell.kimi_k3_kda.fused_prefill_conv import (
+ fused_prefill_conv,
+)
+from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import fused_kda_post_conv
+from tensorrt_llm._torch.modules.mamba.causal_conv1d import causal_conv1d_fn
+from tensorrt_llm._torch.modules.mamba.fuse_elementwise_ops import extract_transpose_prefill_slice
+
+pytestmark = pytest.mark.skipif(
+ not torch.cuda.is_available()
+ or torch.cuda.get_device_capability() not in ((10, 0), (10, 3), (10, 7)),
+ reason="BT16 requires SM100/SM103/SM107",
+)
+
+
+@pytest.mark.parametrize("heads", [1, 3, 6, 12, 96])
+@pytest.mark.parametrize("initial", ["fresh", "continuation", "mixed"])
+@torch.inference_mode()
+def test_fused_conv_raw_outputs_and_dynamic_graph(heads: int, initial: str) -> None:
+ # Include empty/short sequences and tile-boundary crossings. Deliberately
+ # offset the input and pad pool rows to exercise the actual stride contract.
+ lengths = [0, 1, 2, 3, 15, 16, 17, 31, 32, 33, 63, 64, 65]
+ batch = make_batch(heads, lengths, hidden_size=8, initial=initial)
+ total, channels = sum(lengths), 3 * heads * 128
+ backing = (
+ torch.randn(total * (4 * heads * 128 + 8) + 1, device="cuda", dtype=torch.bfloat16) * 0.1
+ )
+ src = torch.as_strided(backing, (total, 4 * heads * 128), (4 * heads * 128 + 8, 1), 1)
+ weight = torch.randn(channels, 4, device="cuda", dtype=torch.bfloat16) * 0.1
+ state = torch.empty_strided(
+ batch.conv.shape, (channels * 3 + 8, 3, 1), device="cuda", dtype=torch.bfloat16
+ ).copy_(batch.conv)
+ seed = clone_pool(state)
+ qkv = tuple(
+ torch.empty(1, total, heads, 128, device="cuda", dtype=torch.bfloat16) for _ in range(3)
+ )
+ immutable = [t.clone() for t in (src, weight)]
+
+ def run() -> None:
+ fused_prefill_conv(
+ src,
+ weight,
+ state,
+ batch.metadata.query_start_loc,
+ batch.indices,
+ batch.metadata.has_initial_states,
+ total,
+ *qkv,
+ )
+
+ run()
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ run()
+ for current in (lengths, list(reversed(lengths)), [total] + [0] * (len(lengths) - 1)):
+ offsets = torch.tensor([0] + list(accumulate(current)), device="cuda", dtype=torch.int64)
+ batch.cu.copy_(offsets)
+ batch.metadata.query_start_loc.copy_(offsets)
+ expected_state = clone_pool(seed)
+ expected_packed = extract_transpose_prefill_slice(src, total, 0, channels)
+ causal_conv1d_fn(
+ expected_packed,
+ weight,
+ query_start_loc=batch.metadata.query_start_loc,
+ cache_indices=batch.indices,
+ has_initial_state=batch.metadata.has_initial_states,
+ conv_states=expected_state,
+ activation="silu",
+ )
+ expected = fused_kda_post_conv(expected_packed, heads, 128, normalize_qk=False)
+ for launch in (run, graph.replay):
+ state.copy_(seed)
+ launch()
+ for actual, ref in zip(qkv, expected):
+ assert torch.isfinite(actual).all()
+ torch.testing.assert_close(actual, ref, rtol=0.008, atol=0.0001)
+ torch.testing.assert_close(state, expected_state, rtol=0, atol=0)
+ for tensor, original in zip((src, weight), immutable):
+ torch.testing.assert_close(tensor, original, rtol=0, atol=0)
+
+
+@torch.inference_mode()
+def test_fused_conv_rejects_output_aliasing() -> None:
+ heads, total = 6, 16
+ batch = make_batch(heads, [total], hidden_size=8)
+ src = torch.randn(total, 3 * heads * 128, dtype=torch.bfloat16, device="cuda")
+ weight = torch.randn(3 * heads * 128, 4, dtype=torch.bfloat16, device="cuda")
+ # A contiguous output in the source allocation must be rejected before launch.
+ q = src.view(-1)[: total * heads * 128].view(1, total, heads, 128)
+ with pytest.raises(ValueError, match="disjoint storage"):
+ fused_prefill_conv(
+ src,
+ weight,
+ batch.conv,
+ batch.metadata.query_start_loc,
+ batch.indices,
+ batch.metadata.has_initial_states,
+ total,
+ q,
+ torch.empty_like(q),
+ torch.empty_like(q),
+ )
+
+
+@pytest.mark.parametrize("source_offset,padded_state", [(0, False), (1, True)])
+@torch.inference_mode()
+def test_fused_conv_warmed_layout_accepts_new_buffers_in_graph(
+ source_offset: int, padded_state: bool
+) -> None:
+ heads, lengths = 1, [17, 31]
+ total, channels = sum(lengths), 384
+ previous_outputs = None
+ previous_values = None
+ for seed in (701, 702):
+ torch.manual_seed(seed)
+ batch = make_batch(heads, lengths, hidden_size=8, initial="mixed", input_seed=seed)
+ storage = torch.randn(total * channels + source_offset, device="cuda", dtype=torch.bfloat16)
+ src = storage[source_offset:].view(total, channels)
+ weight = torch.randn(channels, 4, device="cuda", dtype=torch.bfloat16) * 0.1
+ state = torch.empty_strided(
+ batch.conv.shape,
+ (channels * 3 + (8 if padded_state else 0), 3, 1),
+ device="cuda",
+ dtype=torch.bfloat16,
+ ).copy_(batch.conv)
+ qkv = tuple(
+ torch.empty(1, total, heads, 128, device="cuda", dtype=torch.bfloat16) for _ in range(3)
+ )
+ expected_state = clone_pool(state)
+ packed = extract_transpose_prefill_slice(src, total, 0, channels)
+ causal_conv1d_fn(
+ packed,
+ weight,
+ query_start_loc=batch.metadata.query_start_loc,
+ cache_indices=batch.indices,
+ has_initial_state=batch.metadata.has_initial_states,
+ conv_states=expected_state,
+ activation="silu",
+ )
+ expected = fused_kda_post_conv(packed, heads, 128, normalize_qk=False)
+
+ def run() -> None:
+ fused_prefill_conv(
+ src,
+ weight,
+ state,
+ batch.metadata.query_start_loc,
+ batch.indices,
+ batch.metadata.has_initial_states,
+ max(lengths),
+ *qkv,
+ )
+
+ if previous_outputs is None:
+ run()
+ else:
+ # Only the first allocation was warmed. New pointers with the same
+ # layout must reuse that specialization without compiling in capture.
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ run()
+ state.copy_(batch.conv)
+ graph.replay()
+ for old, snapshot in zip(previous_outputs, previous_values):
+ torch.testing.assert_close(old, snapshot, rtol=0, atol=0)
+ for actual, reference in zip(qkv, expected):
+ torch.testing.assert_close(actual, reference, rtol=0.008, atol=0.0001)
+ torch.testing.assert_close(state, expected_state, rtol=0, atol=0)
+ previous_outputs = qkv
+ previous_values = tuple(t.clone() for t in qkv)
diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_host_metadata.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_host_metadata.py
deleted file mode 100644
index 29da696514ef..000000000000
--- a/tests/unittest/_torch/modules/kimi_kda/test_kda_host_metadata.py
+++ /dev/null
@@ -1,181 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-"""Parity coverage for host-derived KDA prefill metadata."""
-
-import pytest
-import torch
-
-pytest.importorskip("fla")
-
-from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import KDAKernelDispatch # noqa: E402
-
-NUM_HEADS = 6
-HEAD_DIM = 128
-LOWER_BOUND = -5.0
-
-
-def _has_supported_gpu() -> bool:
- # SM100/SM103 (Blackwell) and SM107 (Rubin) -- must track
- # _kda_kernels.is_kda_optimized_supported(), which accepts all three.
- if not torch.cuda.is_available():
- return False
- return torch.cuda.get_device_capability(0) in {(10, 0), (10, 3), (10, 7)}
-
-
-pytestmark = pytest.mark.skipif(
- not _has_supported_gpu(),
- reason="Kimi K3 is supported only on Blackwell SM100/SM103 or Rubin SM107",
-)
-
-
-@pytest.fixture(scope="module")
-def dispatch_pair() -> tuple[KDAKernelDispatch, KDAKernelDispatch]:
- optimized = KDAKernelDispatch(use_optimized_prefill=True, use_optimized_decode=False)
- reference = KDAKernelDispatch(use_optimized_prefill=False, use_optimized_decode=False)
- assert optimized.prefill_kernel_path == "optimized"
- assert reference.prefill_kernel_path == "fla"
- return optimized, reference
-
-
-@pytest.fixture(scope="module")
-def gate_params() -> tuple[torch.Tensor, torch.Tensor]:
- generator = torch.Generator(device="cuda").manual_seed(0)
- a_log = torch.randn(NUM_HEADS, generator=generator, dtype=torch.float32, device="cuda") * 0.5
- dt_bias = (
- torch.randn(
- NUM_HEADS * HEAD_DIM,
- generator=generator,
- dtype=torch.float32,
- device="cuda",
- )
- * 0.1
- )
- return a_log, dt_bias
-
-
-def _make_inputs(total_tokens: int, seed: int) -> tuple[torch.Tensor, ...]:
- generator = torch.Generator(device="cuda").manual_seed(seed)
-
- def random_tensor(*shape: int, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor:
- return torch.randn(*shape, generator=generator, dtype=torch.float32, device="cuda").to(
- dtype
- )
-
- q = random_tensor(1, total_tokens, NUM_HEADS, HEAD_DIM)
- k = random_tensor(1, total_tokens, NUM_HEADS, HEAD_DIM)
- v = random_tensor(1, total_tokens, NUM_HEADS, HEAD_DIM)
- g = random_tensor(1, total_tokens, NUM_HEADS, HEAD_DIM)
- beta = random_tensor(1, total_tokens, NUM_HEADS, dtype=torch.float32)
- return q, k, v, g, beta
-
-
-def _run(
- dispatch: KDAKernelDispatch,
- gate_params: tuple[torch.Tensor, torch.Tensor],
- inputs: tuple[torch.Tensor, ...],
- cu_seqlens: torch.Tensor,
- *,
- chunk_indices: torch.Tensor | None = None,
- varlen_is_aligned: bool | None = None,
- single_sequence_length: int | None = None,
-) -> tuple[torch.Tensor, torch.Tensor | None]:
- q, k, v, g, beta = inputs
- a_log, dt_bias = gate_params
- return dispatch.prefill_chunk_kda(
- q=q.clone(),
- k=k.clone(),
- v=v.clone(),
- g=g.clone(),
- beta=beta.clone(),
- A_log=a_log,
- dt_bias=dt_bias,
- scale=HEAD_DIM**-0.5,
- initial_state=None,
- safe_gate=True,
- lower_bound=LOWER_BOUND,
- cu_seqlens=cu_seqlens,
- chunk_indices=chunk_indices,
- varlen_is_aligned=varlen_is_aligned,
- single_sequence_length=single_sequence_length,
- )
-
-
-def _assert_close(actual: torch.Tensor, expected: torch.Tensor) -> None:
- actual_float = actual.float()
- expected_float = expected.float()
- cosine = torch.nn.functional.cosine_similarity(
- actual_float.flatten(), expected_float.flatten(), dim=0
- ).item()
- relative_l2 = ((actual_float - expected_float).norm() / (expected_float.norm() + 1e-12)).item()
- assert cosine > 0.999
- assert relative_l2 < 3e-2
-
-
-@pytest.mark.parametrize(
- "sequence_lengths",
- ([128, 256, 192], [100, 257, 219], [300]),
- ids=("aligned", "mixed", "single_unaligned"),
-)
-@torch.no_grad()
-def test_host_metadata_bypasses_device_inference(
- dispatch_pair: tuple[KDAKernelDispatch, KDAKernelDispatch],
- gate_params: tuple[torch.Tensor, torch.Tensor],
- sequence_lengths: list[int],
- monkeypatch: pytest.MonkeyPatch,
-) -> None:
- import fla.ops.utils.index as fla_index
-
- from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_custom_ops
-
- optimized, reference = dispatch_pair
- cu_seqlens = torch.tensor(
- [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()],
- dtype=torch.long,
- device="cuda",
- )
- cache_key = id(cu_seqlens)
- assert cache_key not in cute_dsl_kimi_k3_custom_ops._varlen_pure_cache
- assert cache_key not in cute_dsl_kimi_k3_custom_ops._varlen_single_seqlen_cache
-
- inputs = _make_inputs(sum(sequence_lengths), seed=700 + len(sequence_lengths))
- chunk_indices = torch.tensor(
- [
- [sequence_index, chunk_index]
- for sequence_index, length in enumerate(sequence_lengths)
- for chunk_index in range((length + 63) // 64)
- ],
- dtype=torch.long,
- device="cuda",
- )
- varlen_is_aligned = all(length % 64 == 0 for length in sequence_lengths)
- single_sequence_length = sequence_lengths[0] if len(sequence_lengths) == 1 else None
- with monkeypatch.context() as context:
- context.setattr(
- fla_index,
- "prepare_chunk_indices",
- lambda *_args, **_kwargs: pytest.fail(
- "host metadata should bypass FLA chunk-index preparation"
- ),
- )
- actual_output, actual_state = _run(
- optimized,
- gate_params,
- inputs,
- cu_seqlens,
- chunk_indices=chunk_indices,
- varlen_is_aligned=varlen_is_aligned,
- single_sequence_length=single_sequence_length,
- )
- reference_output, reference_state = _run(
- reference,
- gate_params,
- inputs,
- cu_seqlens.clone(),
- )
-
- assert actual_state is not None
- assert reference_state is not None
- _assert_close(actual_output, reference_output)
- _assert_close(actual_state, reference_state)
- assert cache_key not in cute_dsl_kimi_k3_custom_ops._varlen_pure_cache
- assert cache_key not in cute_dsl_kimi_k3_custom_ops._varlen_single_seqlen_cache
diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py
index ad848fca5512..489dca8ef7d5 100644
--- a/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py
+++ b/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py
@@ -2,10 +2,7 @@
# SPDX-License-Identifier: Apache-2.0
"""Parity tests for the optimized Kimi K3 KDA prefill op."""
-import importlib.util
-import sys
from collections.abc import Mapping
-from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock
@@ -35,10 +32,6 @@
def _has_supported_gpu() -> bool:
- # Defer to the predicate the runtime actually dispatches on rather than
- # restating its capability set. It accepts SM100/SM103 only; SM107 takes
- # the FLA fallback in KDAKernelDispatch, so admitting it here would run
- # these parity assertions against the unoptimized path.
if not torch.cuda.is_available():
return False
return is_kda_optimized_supported()
@@ -46,7 +39,7 @@ def _has_supported_gpu() -> bool:
pytestmark = pytest.mark.skipif(
not _has_supported_gpu(),
- reason="KDA optimized prefill kernels require Blackwell SM100/SM103",
+ reason="FlashInfer KDA prefill requires SM100/SM103/SM107",
)
@@ -147,7 +140,7 @@ def _run_production_prefill(
)
core = attention.forward_prefill(
hidden_states.reshape(-1, HIDDEN_SIZE),
- cu_seqlens,
+ cu_seqlens.to(torch.int64),
metadata,
batch_size,
conv_pool,
@@ -173,10 +166,7 @@ def _assert_close(actual: torch.Tensor, expected: torch.Tensor) -> None:
@pytest.mark.parametrize("sequence_length", [1, 127, 256])
@torch.no_grad()
def test_fused_kda_post_conv_matches_reference(sequence_length: int, token_major: bool) -> None:
- """Packed convolution output is normalized and transposed correctly.
-
- Token-major input is the transposed view of ``[T, 3P]`` rows.
- """
+ """The unfused convolution path produces contiguous raw Q/K/V without normalization."""
torch.manual_seed(1)
num_heads, head_dim = 8, 128
projection_size = num_heads * head_dim
@@ -187,21 +177,18 @@ def test_fused_kda_post_conv_matches_reference(sequence_length: int, token_major
if token_major:
packed = packed.t()
- actual_q, actual_k, actual_v = fused_kda_post_conv(packed, num_heads, head_dim)
+ actual_q, actual_k, actual_v = fused_kda_post_conv(
+ packed, num_heads, head_dim, normalize_qk=False
+ )
unpacked = packed.reshape(3, projection_size, sequence_length)
expected_q = unpacked[0].transpose(0, 1).reshape(1, sequence_length, num_heads, head_dim)
expected_k = unpacked[1].transpose(0, 1).reshape(1, sequence_length, num_heads, head_dim)
expected_v = unpacked[2].transpose(0, 1).reshape(1, sequence_length, num_heads, head_dim)
- expected_q = expected_q.float()
- expected_k = expected_k.float()
- expected_q *= torch.rsqrt((expected_q * expected_q).sum(dim=-1, keepdim=True) + 1e-6)
- expected_k *= torch.rsqrt((expected_k * expected_k).sum(dim=-1, keepdim=True) + 1e-6)
-
assert actual_q.is_contiguous()
assert actual_k.is_contiguous()
assert actual_v.is_contiguous()
- torch.testing.assert_close(actual_q, expected_q.to(packed.dtype), rtol=1e-2, atol=1e-2)
- torch.testing.assert_close(actual_k, expected_k.to(packed.dtype), rtol=1e-2, atol=1e-2)
+ torch.testing.assert_close(actual_q, expected_q, rtol=0, atol=0)
+ torch.testing.assert_close(actual_k, expected_k, rtol=0, atol=0)
assert torch.equal(actual_v, expected_v)
@@ -236,6 +223,7 @@ def test_project_packed_conv_input_layout(
torch.manual_seed(3)
# Only the projection layout is checked, so two heads keep the weights small.
kda = _make_kda(num_heads=2).to(dtype)
+ kda._fuse_prefill_conv = False
kda.finalize_decode_weights()
assert kda._qkvg_proj_weight is not None
x2d = torch.randn(num_tokens, HIDDEN_SIZE, dtype=dtype, device="cuda") * 0.05
@@ -411,8 +399,7 @@ def test_optimized_prefill_matches_fla_reference(num_heads: int, use_full_rank_g
optimized.state_dict(), num_heads=num_heads, use_full_rank_gate=use_full_rank_gate
)
- # Keep B=2 across a T transition: eqlen mBeta/mAqk/mAkk batch strides
- # depend on T and therefore require distinct compiled kernel variants.
+ # Reuse the same module across changing token counts and batch sizes.
for batch_size, sequence_length in [(2, 256), (2, 512), (1, 1024)]:
hidden_states = (
torch.randn(
@@ -454,233 +441,6 @@ def test_optimized_prefill_matches_fla_reference(num_heads: int, use_full_rank_g
_assert_close(actual, expected)
-@torch.no_grad()
-@pytest.mark.parametrize(
- "sequence_lengths",
- [None, [100, 257, 63]],
- ids=["padded_eqlen", "varlen"],
-)
-def test_fused_prefill_beta_sigmoid_matches_unfused_kernel(
- sequence_lengths: list[int] | None,
-) -> None:
- """In-kernel beta sigmoid matches the legacy host launch."""
- from fla.modules.l2norm import l2norm_fwd
- from fla.ops.common.gate import fused_beta_sigmoid
-
- torch.manual_seed(10)
- batch_size = 1
- if sequence_lengths is None:
- sequence_length = 300
- num_sequences = batch_size
- cu_seqlens = None
- else:
- sequence_length = sum(sequence_lengths)
- num_sequences = len(sequence_lengths)
- cu_seqlens = torch.tensor(
- [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()],
- dtype=torch.long,
- device="cuda",
- )
- A_log = torch.randn(NUM_HEADS, dtype=torch.float32, device="cuda") * 0.1
- dt_bias = torch.randn(NUM_HEADS * HEAD_DIM, dtype=torch.float32, device="cuda") * 0.1
-
- for seed in (11, 12):
- generator = torch.Generator(device="cuda").manual_seed(seed)
-
- def randn(
- *shape: int, dtype: torch.dtype = torch.bfloat16, scale: float = 0.05
- ) -> torch.Tensor:
- return (
- torch.randn(*shape, generator=generator, dtype=torch.float32, device="cuda") * scale
- ).to(dtype)
-
- q = randn(batch_size, sequence_length, NUM_HEADS, HEAD_DIM)
- k = randn(batch_size, sequence_length, NUM_HEADS, HEAD_DIM)
- v = randn(batch_size, sequence_length, NUM_HEADS, HEAD_DIM)
- g = randn(batch_size, sequence_length, NUM_HEADS, HEAD_DIM)
- if seed == 11:
- beta = randn(batch_size, sequence_length, NUM_HEADS, dtype=torch.float32)
- else:
- # Exercise the L2 epsilon path and saturated sigmoid tails.
- q[:, ::17].zero_()
- k[:, ::19].zero_()
- beta = torch.linspace(
- -20.0,
- 20.0,
- batch_size * sequence_length * NUM_HEADS,
- dtype=torch.float32,
- device="cuda",
- ).reshape(batch_size, sequence_length, NUM_HEADS)
- initial_state = randn(
- num_sequences, NUM_HEADS, HEAD_DIM, HEAD_DIM, dtype=torch.float32, scale=0.01
- )
- state_indices = torch.arange(num_sequences, dtype=torch.int32, device="cuda")
-
- common = dict(
- v=v,
- g=g,
- scale=HEAD_DIM**-0.5,
- state_indices=state_indices,
- safe_gate=True,
- lower_bound=-5.0,
- use_gate_in_kernel=True,
- A_log=A_log,
- dt_bias=dt_bias,
- cu_seqlens=cu_seqlens,
- )
-
- q_unfused, _ = l2norm_fwd(q)
- k_unfused, _ = l2norm_fwd(k)
- beta_unfused = fused_beta_sigmoid(beta, scale=1.0).to(torch.bfloat16)
- state_unfused = initial_state.clone()
- output_unfused = torch.ops.trtllm.kda_prefill(
- q=q_unfused,
- k=k_unfused,
- beta=beta_unfused,
- state_pool=state_unfused,
- use_beta_sigmoid_in_kernel=False,
- **common,
- )
- # The op reuses shape-keyed output buffers; retain the first result
- # before the fused call overwrites it.
- output_unfused = output_unfused.clone()
-
- state_fused = initial_state.clone()
- output_fused = torch.ops.trtllm.kda_prefill(
- q=q_unfused,
- k=k_unfused,
- beta=beta,
- state_pool=state_fused,
- use_beta_sigmoid_in_kernel=True,
- **common,
- )
-
- _assert_close(output_fused, output_unfused)
- _assert_close(state_fused, state_unfused)
-
-
-# fused_k123's K1 warps meet at this barrier after the cross-warp scan, then
-# read the scanned prefix rows and the chunk total from sPartialLast.
-_K1_SCAN_BARRIER = " k1_internal_barrier()\n"
-_K1_SCAN_READS = (
- _K1_SCAN_BARRIER
- + "\n"
- + " for vi in cutlass.range_constexpr(VEC):\n"
- + " rGkLast[vi] = sPartialLast[K1_ROW_GROUPS - 1, col_base + vi]\n"
-)
-# K1 warp 1 reads the prefix row 0 and the chunk total.
-_LATE_K1_WARP = (
- " if k1_warp == 1:\n"
- " late_t = _test_globaltimer()\n"
- " late_end = late_t + cutlass.Int64(20000)\n"
- " while late_t < late_end:\n"
- " late_t = _test_globaltimer()\n"
-)
-_GLOBALTIMER = """
-
-@dsl_user_op
-def _test_globaltimer(*, loc=None, ip=None):
- return cutlass.Int64(
- llvm.inline_asm(
- T.i64(),
- [],
- "mov.u64 $0, %globaltimer;",
- "=l",
- has_side_effects=True,
- is_align_stack=False,
- asm_dialect=llvm.AsmDialect.AD_ATT,
- loc=loc,
- ip=ip,
- )
- )
-"""
-
-
-@torch.no_grad()
-@pytest.mark.parametrize(
- "sequence_lengths",
- [None, [100, 257, 63]],
- ids=["padded_eqlen", "varlen"],
-)
-def test_fused_prefill_does_not_depend_on_k1_warp_timing(
- sequence_lengths: list[int] | None, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
-) -> None:
- """A K1 warp that reads the scanned partial sums late gets the same bits.
-
- A copy of fused_k123 stalls K1 warp 1 for 20 us right after the scan
- barrier, so the other K1 warps run ahead into the next chunk before it
- reads sPartialLast. Outputs and states must equal the unpatched kernel's.
- """
- from tensorrt_llm._torch.attention.kernels.blackwell.kimi_k3_kda import fused_k123
- from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_custom_ops as kimi_k3_ops
-
- source = Path(fused_k123.__file__).read_text()
- assert source.count(_K1_SCAN_READS) == 1, "fused_k123's K1 scan reads moved"
- late_source = source.replace(
- _K1_SCAN_READS, _K1_SCAN_BARRIER + _LATE_K1_WARP + _K1_SCAN_READS[len(_K1_SCAN_BARRIER) :]
- )
- late_path = tmp_path / "fused_k123_late_k1_warp.py"
- late_path.write_text(late_source + _GLOBALTIMER)
- spec = importlib.util.spec_from_file_location(late_path.stem, late_path)
- late_module = importlib.util.module_from_spec(spec)
- monkeypatch.setitem(sys.modules, late_path.stem, late_module)
- spec.loader.exec_module(late_module)
-
- torch.manual_seed(10)
- if sequence_lengths is None:
- sequence_length = 300
- num_sequences = 1
- cu_seqlens = None
- else:
- sequence_length = sum(sequence_lengths)
- num_sequences = len(sequence_lengths)
- cu_seqlens = torch.tensor(
- [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()],
- dtype=torch.long,
- device="cuda",
- )
-
- def randn(*shape: int, dtype: torch.dtype = torch.bfloat16, scale: float = 0.05):
- return (torch.randn(*shape, dtype=torch.float32, device="cuda") * scale).to(dtype)
-
- def l2norm(x: torch.Tensor) -> torch.Tensor:
- return torch.nn.functional.normalize(x.float(), dim=-1).to(torch.bfloat16)
-
- shape = (1, sequence_length, NUM_HEADS, HEAD_DIM)
- inputs = dict(
- q=l2norm(randn(*shape)),
- k=l2norm(randn(*shape)),
- v=randn(*shape),
- g=randn(*shape),
- beta=randn(1, sequence_length, NUM_HEADS, dtype=torch.float32),
- scale=HEAD_DIM**-0.5,
- state_indices=torch.arange(num_sequences, dtype=torch.int32, device="cuda"),
- safe_gate=True,
- lower_bound=-5.0,
- use_gate_in_kernel=True,
- A_log=torch.randn(NUM_HEADS, dtype=torch.float32, device="cuda") * 0.1,
- dt_bias=torch.randn(NUM_HEADS * HEAD_DIM, dtype=torch.float32, device="cuda") * 0.1,
- cu_seqlens=cu_seqlens,
- use_beta_sigmoid_in_kernel=True,
- )
- initial_state = randn(
- num_sequences, NUM_HEADS, HEAD_DIM, HEAD_DIM, dtype=torch.float32, scale=0.01
- )
-
- def run() -> tuple[torch.Tensor, torch.Tensor]:
- state = initial_state.clone()
- output = torch.ops.trtllm.kda_prefill(state_pool=state, **inputs).clone()
- return output, state
-
- expected_output, expected_state = run()
- monkeypatch.setattr(kimi_k3_ops, "_fused_make_host", late_module.make_host_function)
- monkeypatch.setattr(kimi_k3_ops, "_fused_k123_cache", {})
- output, state = run()
-
- torch.testing.assert_close(output, expected_output, rtol=0, atol=0)
- torch.testing.assert_close(state, expected_state, rtol=0, atol=0)
-
-
@torch.no_grad()
def test_kda_mixer_empty_prefill():
"""An empty token payload still publishes the conv state for its slot."""
@@ -727,155 +487,11 @@ def test_kda_mixer_empty_prefill_preserves_initial_conv_state():
assert torch.equal(conv_pool, conv_seed)
-@torch.no_grad()
-def test_indexed_prefill_routing_validates_equal_length_state_indices() -> None:
- dispatch = _kda_kernels.KDAKernelDispatch(
- use_optimized_prefill=True,
- use_optimized_decode=False,
- )
- assert dispatch.prefill_kernel_path == "optimized"
- state_pool = torch.zeros(3, 1, 2, 2, dtype=torch.float32, device="cuda")
- state_indices = torch.tensor([0, 1], dtype=torch.int32, device="cuda")
- has_initial_states = torch.zeros(2, dtype=torch.bool, device="cuda")
- common = {
- "state_pool": state_pool,
- "has_initial_states": has_initial_states,
- "cu_seqlens": None,
- "num_tokens": 256,
- }
-
- assert dispatch.can_use_indexed_prefill(
- state_indices=state_indices,
- num_sequences=2,
- **common,
- )
- assert not dispatch.can_use_indexed_prefill(
- state_indices=state_indices,
- num_sequences=1,
- **common,
- )
-
- misaligned_storage = torch.tensor([-1, 0, 1], dtype=torch.int32, device="cuda")
- misaligned_indices = misaligned_storage[1:]
- assert misaligned_indices.is_contiguous()
- assert misaligned_indices.data_ptr() % 16 != 0
- assert not dispatch.can_use_indexed_prefill(
- state_indices=misaligned_indices,
- num_sequences=2,
- **common,
- )
-
-
-@torch.no_grad()
-def test_kda_prefill_op_rejects_misaligned_state_indices() -> None:
- from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_custom_ops # noqa: F401
-
- q = torch.zeros(1, 256, 1, HEAD_DIM, dtype=torch.bfloat16, device="cuda")
- beta = torch.zeros(1, 256, 1, dtype=torch.float32, device="cuda")
- state_pool = torch.zeros(1, 1, HEAD_DIM, HEAD_DIM, dtype=torch.float32, device="cuda")
- index_storage = torch.tensor([-1, 0], dtype=torch.int32, device="cuda")
- state_indices = index_storage[1:]
- assert state_indices.is_contiguous()
- assert state_indices.data_ptr() % 16 != 0
-
- with pytest.raises(ValueError, match="16-byte-aligned state_indices"):
- torch.ops.trtllm.kda_prefill(
- q=q,
- k=q,
- v=q,
- g=q,
- beta=beta,
- state_pool=state_pool,
- state_indices=state_indices,
- scale=HEAD_DIM**-0.5,
- )
-
-
-@torch.no_grad()
-@pytest.mark.parametrize(
- "sequence_lengths",
- [[100, 200], None],
- ids=["varlen", "equal-length-padded"],
-)
-def test_kda_prefill_op_output_survives_the_next_call(sequence_lengths) -> None:
- """The op's output belongs to the caller, as its fake (a new tensor) declares.
-
- A second call with the same shapes and other inputs must leave the first
- call's output as it was. ``None`` runs one 300-token row without
- cu_seqlens, which the op pads to a chunk multiple and slices back.
- """
- from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_custom_ops # noqa: F401
- from tensorrt_llm._torch.modules.fla.index import prepare_chunk_indices
-
- heads = 2
- tokens = sum(sequence_lengths) if sequence_lengths else 300
- rows = len(sequence_lengths) if sequence_lengths else 1
- state_pool = torch.zeros(rows, heads, HEAD_DIM, HEAD_DIM, dtype=torch.float32, device="cuda")
- state_indices = torch.arange(rows, dtype=torch.int32, device="cuda")
- cu_seqlens = None
- chunk_indices = None
- if sequence_lengths:
- cu_seqlens = torch.tensor(
- [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()],
- dtype=torch.long,
- device="cuda",
- )
- chunk_indices = prepare_chunk_indices(cu_seqlens, 64)
-
- def call(seed: int) -> torch.Tensor:
- gen = torch.Generator(device="cuda").manual_seed(seed)
- shape = (1, tokens, heads, HEAD_DIM)
-
- def unit() -> torch.Tensor:
- t = torch.randn(*shape, generator=gen, device="cuda")
- return (t / t.norm(dim=-1, keepdim=True)).to(torch.bfloat16)
-
- state_pool.zero_()
- return torch.ops.trtllm.kda_prefill(
- q=unit(),
- k=unit(),
- v=(torch.randn(*shape, generator=gen, device="cuda") * 0.5).to(torch.bfloat16),
- g=torch.randn(*shape, generator=gen, device="cuda").to(torch.bfloat16),
- beta=torch.randn(1, tokens, heads, generator=gen, device="cuda"),
- state_pool=state_pool,
- state_indices=state_indices,
- scale=HEAD_DIM**-0.5,
- cu_seqlens=cu_seqlens,
- chunk_indices=chunk_indices,
- safe_gate=True,
- lower_bound=-5.0,
- use_gate_in_kernel=True,
- use_beta_sigmoid_in_kernel=True,
- A_log=torch.zeros(heads, device="cuda"),
- dt_bias=torch.zeros(heads * HEAD_DIM, device="cuda"),
- )
-
- first = call(seed=0)
- kept = first.clone()
- second = call(seed=1)
- torch.cuda.synchronize()
- assert first.shape == second.shape == (1, tokens, heads, HEAD_DIM)
- assert not torch.equal(second, kept), "the two calls must compute different outputs"
- assert torch.equal(first, kept), "a later call wrote an earlier call's output"
-
-
@torch.no_grad()
def test_kda_prefill_op_partial_final_chunk_large_batch():
- """Regression: varlen batches whose FINAL chunk is partial.
-
- The chunk-tile kernels access the full 64-row tile of every chunk and
- neutralize invalid rows only after the access, so the batch's final
- partial chunk touches up to 63 rows past the logical packed length —
- OOB reads on the beta input (now bounds-checked in fused_k123) and on
- the A_kk/A_qk scratch (now allocated with one chunk of slack). The
- runtime's autotuner-warmup shape [max_seq_len - 1, 1] = [8191, 1] hit
- this as CUDA_ERROR_ILLEGAL_ADDRESS whenever the following page was
- unmapped.
-
- - [8191, 1]: the exact autotuner-warmup composition (one max_seq_len-1
- context plus a 1-token remainder).
- - [8000, 150, 42]: interior partial chunks (cross-sequence rows) plus
- a partial final chunk, at eval-like scale.
+ """Large packed batches with short final sequences match the FLA reference.
+
+ Cover the autotuner warmup split [8191, 1] and a mixed-length 8192-token batch.
"""
optimized = _make_kda()
reference = _make_reference(optimized.state_dict())
@@ -905,23 +521,21 @@ def test_kda_prefill_op_partial_final_chunk_large_batch():
"sequence_lengths",
([1], [30], [65], [6, 12], [64, 64], [1, 2, 3], [64, 64, 63], [6, 12, 20, 25]),
ids=(
- # The only shape taking the sentinel-padded, forced-VARLEN_PURE branch.
"single-token",
- "one-chunk",
- "two-chunk-single-seq",
- "two-chunk-varlen",
- # All-aligned, so the tile guard alone bounds the kernel's stores.
- "two-chunk-varlen-aligned",
+ "short-single-seq",
+ "partial-single-seq",
+ "two-short-seqs",
+ "two-aligned-seqs",
"three-short-seqs",
- "three-chunk-boundary",
- "four-chunk",
+ "three-seqs-partial-tail",
+ "four-short-seqs",
),
)
def test_kda_prefill_small_varlen_stays_on_optimized_op(
monkeypatch: pytest.MonkeyPatch,
sequence_lengths: list[int],
) -> None:
- """The optimized op owns every small varlen chunk count, not just NT >= 4.
+ """Short packed prompts stay on the optimized prefill path.
Routing these short-prompt shapes to FLA instead would make a real serving
run pay FLA's cold Triton autotune on its first short context.
@@ -953,20 +567,22 @@ def test_kda_prefill_small_varlen_stays_on_optimized_op(
actual = _run_production_prefill(optimized, hidden_states, cumulative_lengths)
assert fla_spy.call_count == 0, (
f"prefill fell back to FLA for sequence lengths {sequence_lengths}; "
- "the optimized op must handle every varlen chunk count"
+ "the optimized op must handle short packed prompts"
)
expected = reference.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths)
_assert_close(actual, expected)
-@pytest.mark.parametrize("misaligned_indices", [False, True], ids=["indexed", "fallback"])
+@pytest.mark.parametrize(
+ "misaligned_indices", [False, True], ids=["aligned-indices", "offset-indices"]
+)
@torch.no_grad()
def test_kda_prefill_state_pool_matches_fallback_with_mixed_initial_states(
monkeypatch: pytest.MonkeyPatch,
misaligned_indices: bool,
) -> None:
- """Indexed and unsupported-index paths match FLA state handling."""
+ """Aligned and offset int32 slot views both match the explicit FLA fallback."""
torch.manual_seed(3)
optimized = _make_kda()
monkeypatch.setattr(_kda_kernels, "is_intree_prefill_available", lambda: False)
@@ -1048,6 +664,66 @@ def test_kda_prefill_state_pool_matches_fallback_with_mixed_initial_states(
)
+@torch.no_grad()
+@pytest.mark.parametrize(
+ "sequence_lengths",
+ [[100, 200], None],
+ ids=["varlen", "single-sequence"],
+)
+def test_kda_prefill_op_output_survives_the_next_call(sequence_lengths) -> None:
+ """The op's output belongs to the caller, as its fake (a new tensor) declares.
+
+ A second call with the same shapes and other inputs must leave the first
+ call's output as it was. ``None`` runs one packed 300-token sequence.
+ """
+ from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_custom_ops # noqa: F401
+
+ heads = 2
+ tokens = sum(sequence_lengths) if sequence_lengths else 300
+ rows = len(sequence_lengths) if sequence_lengths else 1
+ state_pool = torch.zeros(rows, heads, HEAD_DIM, HEAD_DIM, dtype=torch.float32, device="cuda")
+ state_indices = torch.arange(rows, dtype=torch.int32, device="cuda")
+ cu_seqlens = torch.tensor([0, tokens], dtype=torch.int64, device="cuda")
+ if sequence_lengths:
+ cu_seqlens = torch.tensor(
+ [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()],
+ dtype=torch.long,
+ device="cuda",
+ )
+
+ def call(seed: int) -> torch.Tensor:
+ gen = torch.Generator(device="cuda").manual_seed(seed)
+ shape = (1, tokens, heads, HEAD_DIM)
+
+ def unit() -> torch.Tensor:
+ t = torch.randn(*shape, generator=gen, device="cuda")
+ return (t / t.norm(dim=-1, keepdim=True)).to(torch.bfloat16)
+
+ state_pool.zero_()
+ return torch.ops.trtllm.kda_prefill(
+ q=unit(),
+ k=unit(),
+ v=(torch.randn(*shape, generator=gen, device="cuda") * 0.5).to(torch.bfloat16),
+ g=torch.randn(*shape, generator=gen, device="cuda").to(torch.bfloat16),
+ beta=torch.randn(1, tokens, heads, generator=gen, device="cuda", dtype=torch.bfloat16),
+ state_pool=state_pool,
+ state_indices=state_indices,
+ scale=HEAD_DIM**-0.5,
+ cu_seqlens=cu_seqlens,
+ lower_bound=-5.0,
+ A_log=torch.zeros(heads, device="cuda"),
+ dt_bias=torch.zeros(heads * HEAD_DIM, device="cuda"),
+ )
+
+ first = call(seed=0)
+ kept = first.clone()
+ second = call(seed=1)
+ torch.cuda.synchronize()
+ assert first.shape == second.shape == (1, tokens, heads, HEAD_DIM)
+ assert not torch.equal(second, kept), "the two calls must compute different outputs"
+ assert torch.equal(first, kept), "a later call wrote an earlier call's output"
+
+
@torch.no_grad()
def test_token_major_prefill_matches_channel_major_with_initial_states(
monkeypatch: pytest.MonkeyPatch,
@@ -1061,6 +737,7 @@ def test_token_major_prefill_matches_channel_major_with_initial_states(
"""
torch.manual_seed(5)
kda = _make_kda()
+ kda._fuse_prefill_conv = False
sequence_lengths = [64, 64, 64, 63]
assert sum(sequence_lengths) * torch.bfloat16.itemsize % 16 != 0
cumulative_lengths = torch.tensor(
@@ -1138,30 +815,18 @@ def channel_major_project(self, x, x2d):
@torch.no_grad()
def test_kda_prefill_op_shape_growth_and_cu_dtype_transitions():
- """Cross-call transitions through one process's compile caches.
-
- Regression for the cu/ci-dtype cache-key bug: the K123/akk_inv compile
- caches were keyed shape-independently but NOT on the cu_seqlens /
- chunk_indices dtype, while the compiled kernels bake the element type
- (int64 reads use stride 8, int32 stride 4). Reusing an int64-compiled
- kernel on int32 cu/ci misaddressed every cu/ci element — garbage seq
- ids / chunk starts -> cudaErrorIllegalAddress on the first call after
- the flip (memcheck: 4-byte read one element past the 2-entry int32 cu);
- the reverse direction (int32-compiled, int64 passed) corrupted
- silently. Shape growth alone (same dtype) was already sound.
-
- The sequence below covers, in one process: buffer-cache growth
- (T 1171 -> 8191), int64 -> int32 flip on the grown shape, shrink with
- a flip back, and a multi-seq int32 batch. Every call is parity-checked
- against FLA (catches the silent-corruption direction too).
+ """Changing token counts and caller offset dtypes preserve reference parity.
+
+ The production adapter prepares int64 packed offsets for FlashInfer from
+ either caller dtype. Repeated growth and shrinkage exercise cached kernels.
"""
torch.manual_seed(0)
optimized = _make_kda()
reference = _make_reference(optimized.state_dict())
cases = [
([517, 654], torch.long), # small batch, int64 cu (dump-replay-like)
- ([8191], torch.int32), # buffer growth + dtype flip (crashed pre-fix)
- ([1171], torch.long), # shrink + flip back (silent corruption pre-fix)
+ ([8191], torch.int32), # growth + dtype conversion
+ ([1171], torch.long), # shrink + flip back
([150, 900, 333, 640], torch.int32), # multi-seq int32
]
for sequence_lengths, cu_dtype in cases:
diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py
index b864ef63bc96..b15fe4774957 100644
--- a/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py
+++ b/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py
@@ -1,23 +1,10 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
-"""State parity tests: indexed optimized KDA prefill vs FLA, kernel level.
-
-Complements test_kda_prefill_op.py (module outputs, no recurrent-state pool)
-with the cases the executor runtime actually exercises through
-KDAKernelDispatch:
-
- * selected state-pool row parity in the pool's V-first [N, H, V, K] layout
- — K == V == 128 for Kimi K3, so a layout mix-up is invisible to shape
- checks and only caught numerically;
- * a LARGE carried initial state — the K4 sign error fixed by dev-tech commit
- e45ae259 (NV = U - W@S, was U + W@S) is masked by tiny/zero initial states
- and random inputs, and only blows up when a real-magnitude state is
- carried in;
- * non-64-aligned varlen (single-sequence pad path and multi-sequence masked
- path);
- * fresh input tensors on every call — the runtime never reuses tensor
- objects, so wrapper and compile-cache behavior sees runtime-style tensor
- lifetimes instead of benchmark-style stable-object reuse.
+"""Indexed FlashInfer state parity against the independent FLA path.
+
+Cover V-first state layout, large carried states, partial packed sequences,
+changing head counts, chunked continuation, and repeated eager/graph calls.
+Inputs are raw Q/K/V; both recurrent kernels normalize Q/K internally.
"""
import pytest
@@ -38,12 +25,14 @@
def _has_supported_gpu() -> bool:
- return torch.cuda.is_available() and torch.cuda.get_device_capability(0) in {(10, 0), (10, 3)}
+ if not torch.cuda.is_available():
+ return False
+ return torch.cuda.get_device_capability(0) in {(10, 0), (10, 3), (10, 7)}
pytestmark = pytest.mark.skipif(
not _has_supported_gpu(),
- reason="Kimi K3 is supported only on Blackwell (SM100/SM103)",
+ reason="FlashInfer KDA prefill requires SM100/SM103/SM107",
)
@@ -147,7 +136,7 @@ def _expand_initial_state(initial_state: torch.Tensor | None, num_sequences: int
],
ids=[
"eqlen_fresh",
- "eqlen_padded_carried",
+ "eqlen_partial_carried",
"varlen_aligned_carried",
"varlen_unaligned_carried",
"single_unaligned_carried",
@@ -240,65 +229,6 @@ def test_runtime_scale_state_pool_matches_fla(
assert_kda_close("runtime_scale/state", actual_state, expected_state)
-def _num_chunks(sequence_lengths: list[int]) -> int:
- return sum((length + 63) // 64 for length in sequence_lengths)
-
-
-@pytest.mark.parametrize(
- "victim_lengths,poison_lengths",
- [([8191, 1], [8127, 65]), ([8000, 150, 42], [7999, 129, 64])],
- ids=["autotuner_boundary", "mixed_partial_chunks"],
-)
-@torch.no_grad()
-def test_final_state_ignores_poisoned_scratch(
- dispatch_pair: tuple[KDAKernelDispatch, KDAKernelDispatch],
- gate_params: tuple[torch.Tensor, torch.Tensor],
- victim_lengths: list[int],
- poison_lengths: list[int],
-) -> None:
- """Partial final chunks cannot consume stale rows from a cached scratch buffer."""
- optimized, reference = dispatch_pair
- assert sum(victim_lengths) == sum(poison_lengths)
- assert _num_chunks(victim_lengths) == _num_chunks(poison_lengths)
- total_tokens = sum(victim_lengths)
-
- poison_inputs, _ = _make_inputs(1, total_tokens, seed=31)
- q, k, v, g, beta = poison_inputs
- poison_inputs = q * 100, k * 100, v * 1000, g, beta * 10
- run_indexed_prefill(
- optimized,
- gate_params,
- poison_inputs,
- _make_cu_seqlens(poison_lengths),
- )
- torch.cuda.synchronize()
-
- victim_inputs, initial_state = _make_inputs(
- 1,
- total_tokens,
- seed=32,
- initial_state_scale=1.0,
- )
- initial_state = _expand_initial_state(initial_state, len(victim_lengths))
- cu_seqlens = _make_cu_seqlens(victim_lengths)
- actual_output, actual_state, _ = run_indexed_prefill(
- optimized,
- gate_params,
- victim_inputs,
- cu_seqlens,
- initial_state=initial_state,
- )
- expected_output, expected_state = run_fla_prefill(
- reference,
- gate_params,
- victim_inputs,
- cu_seqlens,
- initial_state=initial_state,
- )
- assert_kda_close("poisoned_scratch/output", actual_output, expected_output)
- assert_kda_close("poisoned_scratch/state", actual_state, expected_state)
-
-
def _run_head_count_case(
dispatch_pair: tuple[KDAKernelDispatch, KDAKernelDispatch],
num_heads: int,
@@ -401,8 +331,8 @@ def test_chunked_continuation_matches_fla(
@pytest.mark.parametrize("use_cuda_graph", [False, True], ids=["eager", "cuda_graph"])
@torch.no_grad()
-def test_prefill_stage_reuse_is_repeatable(use_cuda_graph: bool) -> None:
- """TMA must wait for K1's Q/K reads as well as MMA before reusing a stage."""
+def test_prefill_is_repeatable(use_cuda_graph: bool) -> None:
+ """Repeated prefill calls must agree in eager execution and graph replay."""
generator = torch.Generator(device="cuda").manual_seed(20260921)
length = 1186
shape = (1, length, NUM_HEADS, HEAD_DIM)
@@ -410,24 +340,16 @@ def test_prefill_stage_reuse_is_repeatable(use_cuda_graph: bool) -> None:
def random_tensor(shape: tuple[int, ...]) -> torch.Tensor:
return torch.randn(shape, device="cuda", generator=generator)
- q = torch.nn.functional.normalize(random_tensor(shape), dim=-1).bfloat16()
- k = torch.nn.functional.normalize(random_tensor(shape), dim=-1).bfloat16()
+ q = random_tensor(shape).bfloat16()
+ k = random_tensor(shape).bfloat16()
v = random_tensor(shape).bfloat16()
g = random_tensor(shape).bfloat16()
beta = random_tensor(shape[:-1])
a_log = torch.zeros(NUM_HEADS, device="cuda")
dt_bias = torch.zeros(NUM_HEADS * HEAD_DIM, device="cuda")
state = torch.empty(1, NUM_HEADS, HEAD_DIM, HEAD_DIM, device="cuda")
- indices = torch.zeros(1, dtype=torch.int64, device="cuda")
+ indices = torch.zeros(1, dtype=torch.int32, device="cuda")
cu_seqlens = _make_cu_seqlens([length])
- chunk_count = (length + 63) // 64
- chunk_indices = torch.stack(
- (
- torch.zeros(chunk_count, device="cuda", dtype=torch.int64),
- torch.arange(chunk_count, device="cuda"),
- ),
- dim=-1,
- )
dispatch = KDAKernelDispatch(use_optimized_prefill=True, use_optimized_decode=False)
def run() -> torch.Tensor:
@@ -445,11 +367,8 @@ def run() -> torch.Tensor:
safe_gate=True,
lower_bound=-5.0,
cu_seqlens=cu_seqlens,
- chunk_indices=chunk_indices,
state_pool=state,
state_indices=indices,
- varlen_is_aligned=False,
- single_sequence_length=length,
)
return output
diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py b/tests/unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py
index 2553fe8f23ce..a7a4c201e2d5 100644
--- a/tests/unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py
+++ b/tests/unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py
@@ -2,24 +2,20 @@
# SPDX-License-Identifier: Apache-2.0
"""Tests for Kimi K3-specific Mamba metadata."""
+from itertools import accumulate
from types import SimpleNamespace
import pytest
import torch
from tensorrt_llm._device_copy import DeviceCopyBuffer
-from tensorrt_llm._torch.modules.fla.index import prepare_chunk_indices
from tensorrt_llm._torch.modules.kimi_kda.kimi_k3_mamba_metadata import KimiK3MambaMetadata
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-def _make_attention_metadata(
- context_lengths: list[int],
- generation_lengths: list[int] | None = None,
-) -> SimpleNamespace:
- generation_lengths = generation_lengths or []
- seq_lens = torch.tensor(context_lengths + generation_lengths or [1], dtype=torch.int)
+def _make_attention_metadata(context_lengths: list[int]) -> SimpleNamespace:
+ seq_lens = torch.tensor(context_lengths or [1], dtype=torch.int)
num_contexts = len(context_lengths)
return SimpleNamespace(
seq_lens=seq_lens,
@@ -29,32 +25,222 @@ def _make_attention_metadata(
kv_cache_manager=None,
request_ids=None,
kv_cache_params=SimpleNamespace(
- num_cached_tokens_per_seq=torch.zeros(len(seq_lens), dtype=torch.int),
+ num_cached_tokens_per_seq=torch.zeros(num_contexts, dtype=torch.int),
),
)
-def test_k3_prepare_metadata_match_chunk_indices() -> None:
+@pytest.mark.parametrize("lengths", [[128, 256], [1, 64, 65, 129], [17, 1, 17, 16], [300], []])
+def test_prepare_packed_kda_metadata(lengths: list[int]) -> None:
metadata = KimiK3MambaMetadata(max_batch_size=4, chunk_size=128, max_num_tokens=512)
- cases = (
- ([128, 256], [], True, None),
- ([1, 64, 65, 129], [], False, None),
- ([300], [], False, 300),
- ([65, 129], [1, 1], False, None),
+ slots = torch.tensor([3, 1, 2, 0], dtype=torch.int32, device="cuda")
+ metadata.state_indices.copy_(slots)
+ metadata.prepare(_make_attention_metadata(lengths))
+ assert metadata.kda_max_seqlen == max(lengths, default=0)
+ torch.testing.assert_close(
+ metadata.kda_seq_order,
+ torch.tensor(
+ sorted(range(len(lengths)), key=lambda i: (-lengths[i], i)),
+ device="cuda",
+ dtype=torch.int32,
+ ),
+ )
+ torch.testing.assert_close(
+ metadata.kda_cu_chunks,
+ torch.tensor(
+ list(accumulate(((length + 15) // 16 for length in lengths), initial=0)),
+ device="cuda",
+ dtype=torch.int32,
+ ),
+ )
+ torch.testing.assert_close(metadata.state_indices, slots)
+ if lengths:
+ expected = torch.tensor(
+ [0, *torch.tensor(lengths).cumsum(0).tolist()], dtype=torch.int64, device="cuda"
+ )
+ torch.testing.assert_close(metadata.query_start_loc_long, expected)
+ torch.testing.assert_close(metadata.query_start_loc, expected.int())
+ assert metadata.query_start_loc_long.is_contiguous()
+ assert metadata.query_start_loc.is_contiguous()
+
+
+def test_prepare_refreshes_offsets_without_changing_storage() -> None:
+ metadata = KimiK3MambaMetadata(max_batch_size=4, chunk_size=128, max_num_tokens=512)
+ metadata.prepare(_make_attention_metadata([17, 65]))
+ ptrs = (
+ metadata.query_start_loc.data_ptr(),
+ metadata.query_start_loc_long.data_ptr(),
+ metadata.state_indices.data_ptr(),
+ metadata.kda_seq_order.data_ptr(),
+ metadata.kda_cu_chunks.data_ptr(),
+ )
+ metadata.prepare(_make_attention_metadata([64, 18]))
+ assert metadata.kda_max_seqlen == 64
+ expected = torch.tensor([0, 64, 82], dtype=torch.int64, device="cuda")
+ torch.testing.assert_close(metadata.query_start_loc_long, expected)
+ torch.testing.assert_close(metadata.query_start_loc, expected.int())
+ assert ptrs == (
+ metadata.query_start_loc.data_ptr(),
+ metadata.query_start_loc_long.data_ptr(),
+ metadata.state_indices.data_ptr(),
+ metadata.kda_seq_order.data_ptr(),
+ metadata.kda_cu_chunks.data_ptr(),
+ )
+ metadata.prepare(_make_attention_metadata([]))
+ assert metadata.kda_max_seqlen == 0
+
+
+def test_prepare_schedule_excludes_decode_requests() -> None:
+ metadata = KimiK3MambaMetadata(max_batch_size=4, chunk_size=128, max_num_tokens=512)
+ batch = _make_attention_metadata([17, 65])
+ batch.seq_lens = torch.tensor([17, 65, 1, 1], dtype=torch.int)
+ batch.seq_lens_cuda = batch.seq_lens.cuda()
+ metadata.prepare(batch)
+ torch.testing.assert_close(
+ metadata.kda_seq_order, torch.tensor([1, 0], dtype=torch.int32, device="cuda")
+ )
+ torch.testing.assert_close(
+ metadata.kda_cu_chunks, torch.tensor([0, 2, 7], dtype=torch.int32, device="cuda")
+ )
+
+
+@pytest.mark.parametrize("heads", [6, 96])
+@torch.inference_mode()
+def test_module_with_prepared_schedule_matches_device_generation(heads: int) -> None:
+ if torch.cuda.get_device_capability() not in ((10, 0), (10, 3), (10, 7)):
+ pytest.skip("FlashInfer KDA requires SM100/SM103/SM107")
+ pytest.importorskip("fla")
+ from kda_module_test_utils import clone_pool, make_attention, make_batch
+
+ module = make_attention(heads, hidden_size=768, bf16_weights=True)
+ lengths = [17, 65]
+ batch = make_batch(heads, lengths, hidden_size=768, initial="mixed", padded=True)
+ prepared = KimiK3MambaMetadata(max_batch_size=4, chunk_size=128, max_num_tokens=512)
+ prepared.prepare(_make_attention_metadata(lengths))
+ prepared.state_indices[: batch.n].copy_(batch.indices)
+ prepared.has_initial_states[: batch.n].copy_(batch.metadata.has_initial_states)
+ prepared.use_initial_states = batch.metadata.use_initial_states
+ states = [clone_pool(batch.state), clone_pool(batch.state)]
+ convs = [batch.conv.clone(), batch.conv.clone()]
+ outputs = [
+ module.forward_prefill(
+ batch.hidden, batch.cu, metadata, batch.n, conv, state, batch.indices
+ )
+ for metadata, conv, state in zip((batch.metadata, prepared), convs, states)
+ ]
+ torch.testing.assert_close(outputs[0], outputs[1], rtol=0, atol=0)
+ torch.testing.assert_close(states[0], states[1], rtol=0, atol=0)
+ torch.testing.assert_close(convs[0], convs[1], rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("mode", ["engine", "decomp", "bt16", "auto", "small-bh"])
+@pytest.mark.parametrize("heads", [6, 96])
+@torch.inference_mode()
+def test_prepared_metadata_matches_device_generation_across_layers_and_replays(
+ mode: str, heads: int
+) -> None:
+ if torch.cuda.get_device_capability() not in ((10, 0), (10, 3), (10, 7)):
+ pytest.skip("FlashInfer KDA requires SM100/SM103/SM107")
+ from tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_custom_ops import kda_prefill
+
+ torch.manual_seed(1901)
+ metadata = KimiK3MambaMetadata(max_batch_size=4, chunk_size=128, max_num_tokens=512)
+ metadata.prepare(_make_attention_metadata([1, 15, 47, 65]))
+ offsets = metadata.query_start_loc_long
+ slots = torch.tensor([4, 1, 3, 2], dtype=torch.int32, device="cuda")
+ q, k, v, g = [
+ torch.randn(1, 128, heads, 128, dtype=torch.bfloat16, device="cuda") * 0.2 for _ in range(4)
+ ]
+ beta = torch.randn(1, 128, heads, dtype=torch.bfloat16, device="cuda")
+ a_log = torch.linspace(-2, 1, heads, device="cuda")
+ bias = torch.linspace(-1, 1, heads * 128, device="cuda")
+ row_size = heads * 128 * 128
+ initial = torch.randn(6, row_size + 128, dtype=torch.bfloat16, device="cuda") * 0.1
+ initial[slots[::2].long(), :row_size] = 0
+ # Two independent layer pools consume the same prepared scheduling buffers.
+ backing = [initial.clone(), initial.clone()]
+ states = [tensor[:, :row_size].view(6, heads, 128, 128) for tensor in backing]
+ pointers = (
+ offsets.data_ptr(),
+ metadata.kda_seq_order.data_ptr(),
+ metadata.kda_cu_chunks.data_ptr(),
)
- for (
- context_lengths,
- generation_lengths,
- expected_alignment,
- expected_single_length,
- ) in cases:
- metadata.prepare(_make_attention_metadata(context_lengths, generation_lengths))
-
- assert metadata.kda_varlen_is_aligned is expected_alignment
- assert metadata.kda_single_sequence_length == expected_single_length
- cu_seqlens = metadata.query_start_loc_long[: len(context_lengths) + 1]
- legacy_chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size=64)
- torch.testing.assert_close(metadata.kda_chunk_indices, legacy_chunk_indices)
+
+ def run(prepared: bool) -> list[torch.Tensor]:
+ return [
+ kda_prefill(
+ q,
+ k,
+ v,
+ g,
+ beta,
+ state,
+ slots,
+ offsets,
+ a_log,
+ bias,
+ 128**-0.5,
+ -5.0,
+ mode,
+ 128,
+ metadata.kda_seq_order if prepared else None,
+ metadata.kda_cu_chunks if prepared else None,
+ )
+ for state in states
+ ]
+
+ run(True)
+ run(False)
+ torch.cuda.synchronize()
+ graph = torch.cuda.CUDAGraph()
+ with torch.cuda.graph(graph):
+ captured = run(True)
+ # Fixed token/sequence capacity, but different ordering and chunk counts.
+ for lengths in ([1, 15, 47, 65], [17, 15, 64, 32], [62, 1, 1, 64]):
+ metadata.prepare(_make_attention_metadata(lengths))
+ assert pointers == (
+ metadata.query_start_loc_long.data_ptr(),
+ metadata.kda_seq_order.data_ptr(),
+ metadata.kda_cu_chunks.data_ptr(),
+ )
+ immutable = [
+ q,
+ k,
+ v,
+ g,
+ beta,
+ offsets,
+ slots,
+ a_log,
+ bias,
+ metadata.kda_seq_order,
+ metadata.kda_cu_chunks,
+ ]
+ snapshots = [tensor.clone() for tensor in immutable]
+ for tensor in backing:
+ tensor.copy_(initial)
+ expected = run(False)
+ expected_states = [tensor.clone() for tensor in backing]
+ for replay in (False, True):
+ for tensor in backing:
+ tensor.copy_(initial)
+ if replay:
+ graph.replay()
+ actual = captured
+ else:
+ actual = run(True)
+ for output, reference, tensor, reference_state in zip(
+ actual, expected, backing, expected_states
+ ):
+ assert torch.isfinite(output).all()
+ torch.testing.assert_close(output, reference, rtol=0, atol=0)
+ torch.testing.assert_close(tensor, reference_state, rtol=0, atol=0)
+ torch.testing.assert_close(tensor[[0, 5]], initial[[0, 5]], rtol=0, atol=0)
+ torch.testing.assert_close(
+ tensor[:, row_size:], initial[:, row_size:], rtol=0, atol=0
+ )
+ for tensor, snapshot in zip(immutable, snapshots):
+ torch.testing.assert_close(tensor, snapshot, rtol=0, atol=0)
def test_k3_prepare_materializes_stable_aligned_generation_indices() -> None:
diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_bf16_state_pool.py b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_bf16_state_pool.py
index d0a27a505632..053966bf48cb 100644
--- a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_bf16_state_pool.py
+++ b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_bf16_state_pool.py
@@ -1,18 +1,6 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
-"""A bf16 KDA recurrent-state pool must round-trip through the fp32-only paths.
-
-The fused T=1 decode kernel and the indexed prefill kernel read and write the
-state in fp32 only, so a bf16 pool goes through an fp32 copy of the addressed
-rows: plain decode takes ``forward_decode_fallback`` (gather, widen, run the
-kernel, round back) and prefill stages the rows for the same indexed kernel the
-fp32 pool uses. The fp32 pool is the reference in both cases: decode must agree
-up to bf16 rounding of the committed state, prefill must be bit-identical
-(same kernel, same fp32 inputs) with the state rounded once, and rows the batch
-does not address must stay untouched.
-
-Needs 1 GPU + fla-core; runs with random weights (no checkpoint).
-"""
+"""BF16 recurrent-state parity for direct prefill and staged native decode."""
from types import SimpleNamespace
@@ -123,46 +111,27 @@ def test_bf16_pool_decode_matches_fp32_pool():
torch.testing.assert_close(ssm_bf16[untouched], initial[untouched], rtol=0, atol=0)
-def _prefill_metadata(
- sequence_lengths, use_initial_states, has_initial_states, cu_seqlens, slot_indices
-):
- """The Kimi K3 runtime metadata forward_prefill reads (chunk indices precomputed)."""
- chunk_indices = torch.tensor(
- [
- [sequence_index, chunk_index]
- for sequence_index, length in enumerate(sequence_lengths)
- for chunk_index in range((length + 63) // 64)
- ],
- device="cuda",
- dtype=torch.long,
- )
+def _prefill_metadata(use_initial_states, has_initial_states, cu_seqlens, slot_indices):
+ """Prepare packed offsets and state ownership for production prefill."""
return SimpleNamespace(
use_initial_states=use_initial_states,
has_initial_states=torch.tensor(has_initial_states, device="cuda", dtype=torch.bool),
- kda_chunk_indices=chunk_indices,
- kda_varlen_is_aligned=all(length % 64 == 0 for length in sequence_lengths),
- kda_single_sequence_length=(sequence_lengths[0] if len(sequence_lengths) == 1 else None),
# the packed causal convolution takes int32 sequence starts and the conv-pool slots
query_start_loc=cu_seqlens.to(torch.int32),
state_indices=slot_indices,
- # the runtime metadata's preallocated 0..max_batch range (staged rows are dense)
- _arange_buffer=torch.arange(len(sequence_lengths) + 1, dtype=torch.int32, device="cuda"),
)
@pytest.mark.parametrize(
"sequence_lengths,use_initial_states,has_initial_states",
[
- # >= 4 chunks of 64 tokens in total: below that the optimized prefill
- # op would launch a zero-size grid, so the dispatch keeps such batches
- # on the FLA reference path and neither pool reaches the indexed kernel.
([17, 31, 64, 40], False, [False, False, False, False]),
([1, 129], True, [True, False]),
],
)
@torch.no_grad()
def test_bf16_pool_prefill_matches_fp32_pool(
- sequence_lengths, use_initial_states, has_initial_states, monkeypatch
+ sequence_lengths, use_initial_states, has_initial_states
):
torch.manual_seed(0)
runtime = _make_runtime()
@@ -178,9 +147,7 @@ def test_bf16_pool_prefill_matches_fp32_pool(
cu_seqlens = torch.tensor(
[0, *torch.tensor(sequence_lengths).cumsum(0).tolist()], device="cuda", dtype=torch.long
)
- metadata = _prefill_metadata(
- sequence_lengths, use_initial_states, has_initial_states, cu_seqlens, slot_indices
- )
+ metadata = _prefill_metadata(use_initial_states, has_initial_states, cu_seqlens, slot_indices)
hidden = torch.randn(num_tokens, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.05
conv_seed = torch.randn(slots, 3 * dim, CONV_WIDTH - 1, dtype=torch.bfloat16, device="cuda")
conv_seed *= 0.02
@@ -191,15 +158,6 @@ def test_bf16_pool_prefill_matches_fp32_pool(
ssm_fp32 = ssm_bf16.float()
initial = ssm_bf16.clone()
- indexed_calls = []
- prefill_chunk_kda = runtime._dispatch.prefill_chunk_kda
-
- def _record(**kwargs):
- indexed_calls.append(kwargs["state_pool"] is not None)
- return prefill_chunk_kda(**kwargs)
-
- monkeypatch.setattr(runtime._dispatch, "prefill_chunk_kda", _record)
-
ref_conv = conv_seed.clone()
expected = runtime.forward_prefill(
hidden, cu_seqlens, metadata, num_prefills, ref_conv, ssm_fp32, slot_indices
@@ -209,11 +167,8 @@ def _record(**kwargs):
hidden, cu_seqlens, metadata, num_prefills, actual_conv, ssm_bf16, slot_indices
)
- assert indexed_calls == [True, True], "both pools must take the indexed prefill kernel"
- torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+ assert _relative_error(actual, expected) < 0.03
torch.testing.assert_close(actual_conv, ref_conv, rtol=0, atol=0)
assert not torch.equal(ssm_bf16[addressed], initial[addressed]), "state was not written back"
- torch.testing.assert_close(
- ssm_bf16[addressed], ssm_fp32[addressed].to(torch.bfloat16), rtol=0, atol=0
- )
+ assert _relative_error(ssm_bf16[addressed], ssm_fp32[addressed]) < 0.03
torch.testing.assert_close(ssm_bf16[untouched], initial[untouched], rtol=0, atol=0)
diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_decode_state_parity.py b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_decode_state_parity.py
new file mode 100644
index 000000000000..f09c4929d988
--- /dev/null
+++ b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_decode_state_parity.py
@@ -0,0 +1,124 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Sequential KDA decode parity for live W-1 convolution and recurrent pools.
+
+Compare indexed decode with the production fallback using identical weights,
+changing tokens, and independently cloned BF16/FP32 recurrent state. Check the
+selected rows against the fallback and untouched rows exactly after every step.
+Requires one GPU and fla-core; no checkpoint.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytest.importorskip("fla")
+
+from kda_module_test_utils import compare # noqa: E402
+
+from tensorrt_llm._torch.modules.kimi_kda import KimiKDALinearAttention # noqa: E402
+
+HIDDEN_SIZE = 512
+NUM_HEADS = 4
+HEAD_DIM = 128
+CONV_WIDTH = 4
+
+
+class _Cfg:
+ hidden_size = HIDDEN_SIZE
+ rms_norm_eps = 1e-6
+ linear_attn_config = {
+ "num_heads": NUM_HEADS,
+ "head_dim": HEAD_DIM,
+ "short_conv_kernel_size": CONV_WIDTH,
+ "use_full_rank_gate": True,
+ "gate_lower_bound": -5.0,
+ }
+
+
+def _make_runtime() -> KimiKDALinearAttention:
+ if not torch.cuda.is_available():
+ pytest.skip("the fused KDA decode kernel requires a CUDA device")
+ if torch.cuda.get_device_capability(0) not in {(10, 0), (10, 3), (10, 7)}:
+ pytest.skip("Kimi K3 requires SM100/SM103/SM107")
+ runtime = KimiKDALinearAttention(_Cfg(), layer_idx=0).to("cuda")
+ # Checkpoints initialize dt_bias; seed it before finalizing kernel constants.
+ torch.nn.init.normal_(runtime.dt_bias, std=0.1)
+ runtime.finalize_decode_weights()
+ return runtime
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("batch", [1, 3])
+def test_decode_state_matches_fallback_over_many_steps(
+ batch: int, state_dtype: torch.dtype
+) -> None:
+ torch.manual_seed(0)
+ slots, steps = batch + 3, 6
+ dim = NUM_HEADS * HEAD_DIM
+ runtime = _make_runtime()
+ # Nonzero, reversed slots detect accidental batch-row addressing.
+ slot_indices = torch.arange(slots - 1, slots - batch - 1, -1, dtype=torch.int32, device="cuda")
+ slots_long = slot_indices.long()
+ untouched = torch.arange(3, device="cuda")
+ conv_pool = (
+ torch.randn(slots, 3 * dim, CONV_WIDTH - 1, dtype=torch.bfloat16, device="cuda") * 0.05
+ )
+ ssm_pool = (
+ torch.randn(slots, NUM_HEADS, HEAD_DIM, HEAD_DIM, dtype=state_dtype, device="cuda") * 0.05
+ )
+ direct_conv, direct_ssm = conv_pool.clone(), ssm_pool.clone()
+ fallback_conv, fallback_ssm = conv_pool.clone(), ssm_pool.clone()
+ hidden = torch.randn(batch, steps, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.1
+ hidden_before = hidden.clone()
+ metadata = SimpleNamespace(
+ _arange_buffer=torch.arange(batch + 1, dtype=torch.int32, device="cuda")
+ )
+
+ for step in range(steps):
+ actual = runtime.forward_decode(
+ hidden[:, step],
+ direct_conv,
+ direct_ssm,
+ slots_long,
+ mamba_metadata=metadata,
+ ssm_state_indices=slot_indices,
+ )
+ expected = runtime.forward_decode_fallback(
+ hidden[:, step],
+ fallback_conv,
+ fallback_ssm,
+ slots_long,
+ )
+ # Fused and separate projections can round differently; use the established
+ # module parity criteria, while op-only pool updates stay bit-exact.
+ for name, got, want in (
+ ("output", actual, expected),
+ (
+ "conv",
+ direct_conv.index_select(0, slots_long),
+ fallback_conv.index_select(0, slots_long),
+ ),
+ (
+ "state",
+ direct_ssm.index_select(0, slots_long),
+ fallback_ssm.index_select(0, slots_long),
+ ),
+ ):
+ metrics = compare(got, want)
+ assert metrics["passed"], (step, name, metrics)
+ for changed, original in (
+ (direct_conv, conv_pool),
+ (fallback_conv, conv_pool),
+ (direct_ssm, ssm_pool),
+ (fallback_ssm, ssm_pool),
+ ):
+ torch.testing.assert_close(
+ changed.index_select(0, untouched),
+ original.index_select(0, untouched),
+ rtol=0,
+ atol=0,
+ )
+ torch.testing.assert_close(hidden, hidden_before, rtol=0, atol=0)