From c4a0fb00f45370bc4a5caebadb5aab24d9493edb Mon Sep 17 00:00:00 2001 From: Hongyi Wu <62729549+Aharrypotter@users.noreply.github.com> Date: Fri, 21 Aug 2026 11:15:45 +0800 Subject: [PATCH 1/5] feat(gdn2): add Hopper SM90a CuTe DSL prefill kernel Fully fused packed variable-length Gated DeltaNet-2 forward prefill for SM90a: one 384-thread CTA per sequence, TMA/WGMMA, and a load-balanced longest-processing-time sequence order. Supports MHA, GVA2, and GVA4 with independently optional initial and final recurrent state. The intra-chunk factorization is blockwise rebased per 16-token sub-block, so the only operand carrying a positive exponent spans at most 15 in-block token gaps. The separable chunk-start alternative is cheaper -- one scaling of each operand serves the whole 64x64 tile -- but it overflows FP32 for inputs the public contract accepts, which is why it is not used. The supported decay contract is g in [-5, 0], matching the pinned FLA GDN2 safe_gate range. Private MLIR dialects are taken through cula.ops._mlir_compat, so this kernel is covered by the repository-wide CuTeDSL contract at import. --- cula/ops/gdn2/__init__.py | 4 + cula/ops/gdn2/sm90/__init__.py | 14 + cula/ops/gdn2/sm90/collective_inverse_hmma.py | 370 ++ cula/ops/gdn2/sm90/config.py | 37 + cula/ops/gdn2/sm90/inverse_helpers.py | 61 + cula/ops/gdn2/sm90/prefill.py | 198 ++ cula/ops/gdn2/sm90/prefill_kernel.py | 3043 +++++++++++++++++ 7 files changed, 3727 insertions(+) create mode 100644 cula/ops/gdn2/__init__.py create mode 100644 cula/ops/gdn2/sm90/__init__.py create mode 100644 cula/ops/gdn2/sm90/collective_inverse_hmma.py create mode 100644 cula/ops/gdn2/sm90/config.py create mode 100644 cula/ops/gdn2/sm90/inverse_helpers.py create mode 100644 cula/ops/gdn2/sm90/prefill.py create mode 100644 cula/ops/gdn2/sm90/prefill_kernel.py diff --git a/cula/ops/gdn2/__init__.py b/cula/ops/gdn2/__init__.py new file mode 100644 index 00000000..0ad6c45d --- /dev/null +++ b/cula/ops/gdn2/__init__.py @@ -0,0 +1,4 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Private GDN2 backend kernels.""" diff --git a/cula/ops/gdn2/sm90/__init__.py b/cula/ops/gdn2/sm90/__init__.py new file mode 100644 index 00000000..4a94f912 --- /dev/null +++ b/cula/ops/gdn2/sm90/__init__.py @@ -0,0 +1,14 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Private Hopper SM90a GDN2 backend.""" + +from .config import SM90_BACKEND_ID + +__all__ = ["get_sm90_gdn2_backend_identity"] + + +def get_sm90_gdn2_backend_identity() -> str: + """Return the stable product backend identity.""" + + return SM90_BACKEND_ID diff --git a/cula/ops/gdn2/sm90/collective_inverse_hmma.py b/cula/ops/gdn2/sm90/collective_inverse_hmma.py new file mode 100644 index 00000000..0d99f6f3 --- /dev/null +++ b/cula/ops/gdn2/sm90/collective_inverse_hmma.py @@ -0,0 +1,370 @@ +# Copyright 2025-2026 FlashInfer team. +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import warp + +from .inverse_helpers import SM80, select_tensor_10 + +# ─── CollectiveInverse ─────────────────────────────────────────────────────── +# Inverts a 64×64 lower-triangular smem matrix in-place (I + lower_tril(K)). +# Warp group: 128 threads (4 warps). Layout of sT: (64,64) col-major. + + +class CollectiveInverse: + def __init__( + self, + garbage_filled_diagonal: bool = True, + garbage_filled_upper_triangular: bool = False, + ): + self.garbage_filled_diagonal = garbage_filled_diagonal + self.garbage_filled_upper_triangular = garbage_filled_upper_triangular + + # ── Level 1: NxN Gauss elimination on diagonal blocks ─────────────────────── + + @cute.jit + def compute_diagonal_block_inverse( + self, + mat: cute.Tensor, + tid_in_group: cutlass.Int32, + block_size: cutlass.Constexpr, + ): + """Invert one square lower-triangular block in place. + + One row is assigned to each participating thread. + """ + shuffle_mask = (block_size - 1) | ((32 - block_size) << 8) + + # Vectorize the row load/store. Scalar U16 indexing creates excessive + # shared-memory wavefronts on the KK buffer. + row = cute.make_rmem_tensor(block_size, cutlass.Float16) + cute.autovec_copy(mat[tid_in_group, None], row) + + # Apply I + lower_tril masking in fp32. + frag = cute.make_rmem_tensor(block_size, cutlass.Float32) + for j in cutlass.range_constexpr(block_size): + raw = cutlass.Float32(row[j]) + if cutlass.const_expr(self.garbage_filled_diagonal or self.garbage_filled_upper_triangular): + if tid_in_group == cutlass.Int32(j): + frag[j] = cutlass.Float32(1.0) + elif tid_in_group < cutlass.Int32(j): + frag[j] = cutlass.Float32(0.0) + else: + frag[j] = raw + else: + frag[j] = raw + + # Gaussian elimination: row-reduce to produce inv(I + lower_tril(K)) + for src_row in cutlass.range_constexpr(block_size - 1): + row_scale = -frag[src_row] + for i in cutlass.range_constexpr(src_row): + src_val = cute.arch.shuffle_sync(frag[i], cutlass.Int32(src_row), mask_and_clamp=shuffle_mask) + if tid_in_group > cutlass.Int32(src_row): + frag[i] = frag[i] + row_scale * src_val + if tid_in_group > cutlass.Int32(src_row): + frag[src_row] = row_scale + + row_out = cute.make_rmem_tensor(block_size, cutlass.Float16) + for j in cutlass.range_constexpr(block_size): + row_out[j] = cutlass.Float16(frag[j]) + cute.autovec_copy(row_out, mat[tid_in_group, None]) + + # ── Level 2: 8×8 blocks → 16×16 blockwise inverse ─────────────────────────── + # 8×8-to-16×16 blockwise merge (column-major path). + # Called by all 32 threads of one warp. + # mat: (16, 16) smem slice (one 16×16 diagonal block). + + @cute.jit + def blockwise_8x8_to_16x16(self, mat: cute.Tensor): + lane_id = cute.arch.lane_idx() + + tiled_mma = cute.make_tiled_mma( + warp.MmaF16BF16Op(cutlass.Float16, cutlass.Float32, (16, 8, 8)), + (1, 1, 1), + permutation_mnk=(16, 8, 8), + ) + + mat_2x2 = cute.flat_divide(mat, (8, 8)) + sDinv = mat_2x2[None, None, 1, 1] + sC = select_tensor_10(mat_2x2[None, None, 1, 0]) + sAinv = select_tensor_10(mat_2x2[None, None, 0, 0]) + sO = mat_2x2[None, None, 1, 0] + + # Broadcast sDinv (8,8) → (16,8) by stride-0 on the extra M dimension + sDinv_bcast = cute.make_tensor( + sDinv.iterator, + cute.make_layout( + ((cute.size(sDinv, mode=[0]), 2), cute.size(sDinv, mode=[1])), + stride=((sDinv.layout.stride[0], 0), sDinv.layout.stride[1]), + ), + ) + sO_bcast = cute.make_tensor( + sO.iterator, + cute.make_layout( + ((cute.size(sO, mode=[0]), 2), cute.size(sO, mode=[1])), + stride=((sO.layout.stride[0], 0), sO.layout.stride[1]), + ), + ) + + thr_mma = tiled_mma.get_slice(lane_id) + tOrDinv = thr_mma.make_fragment_A(thr_mma.partition_A(sDinv_bcast)) + tOrC = thr_mma.make_fragment_B(thr_mma.partition_B(sC)) + tOrAinv = thr_mma.make_fragment_B(thr_mma.partition_B(sAinv)) + tDCrDC = cute.make_rmem_tensor(thr_mma.partition_shape_C((16, 8)), cutlass.Float32) + tOrO = cute.make_rmem_tensor(thr_mma.partition_shape_C((16, 8)), cutlass.Float32) + + # Row-major copy atoms: A→LdMatrix_N, B→LdMatrix_T, C→StMatrix_N + dinv_atom = cute.make_copy_atom(warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=1), cutlass.Float16) + b_atom = cute.make_copy_atom(warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=1), cutlass.Float16) + o_atom = cute.make_copy_atom(warp.StMatrix8x8x16bOp(transpose=False, num_matrices=1), cutlass.Float16) + + D_tiled_copy = cute.make_tiled_copy_A(dinv_atom, tiled_mma) + C_tiled_copy = cute.make_tiled_copy_B(b_atom, tiled_mma) + A_tiled_copy = cute.make_tiled_copy_B(b_atom, tiled_mma) + O_tiled_copy = cute.make_tiled_copy_C(o_atom, tiled_mma) + + D_thr_copy = D_tiled_copy.get_slice(lane_id) + C_thr_copy = C_tiled_copy.get_slice(lane_id) + A_thr_copy = A_tiled_copy.get_slice(lane_id) + O_thr_copy = O_tiled_copy.get_slice(lane_id) + + tOsDinv = D_thr_copy.partition_S(sDinv_bcast) + tOrDinv_cv = D_thr_copy.retile(tOrDinv) + tOsC = C_thr_copy.partition_S(sC) + tOrC_cv = C_thr_copy.retile(tOrC) + tOsAinv = A_thr_copy.partition_S(sAinv) + tOrAinv_cv = A_thr_copy.retile(tOrAinv) + tOsO = O_thr_copy.partition_D(sO_bcast) + tOrO_f16 = cute.make_fragment_like(tOrO, cutlass.Float16) + tOrO_cv = O_thr_copy.retile(tOrO_f16) + + # ── Step 1: tDCrDC = -inv(D) @ C ───────────────────────────────────────── + # Load only the first MMA_M slice of D_inv (rows 0-7 of the 16-row broadcast) + cute.copy(D_tiled_copy, tOsDinv[None, None, 0], tOrDinv_cv[None, None, 0]) + cute.copy(C_tiled_copy, tOsC, tOrC_cv) + tDCrDC.fill(0.0) + cute.gemm(tiled_mma, tDCrDC, tOrDinv, tOrC, tDCrDC) + for i in cutlass.range_constexpr(cute.size(tDCrDC)): + tDCrDC[i] = -tDCrDC[i] + + # ── Step 2: tOrO = tDCrDC @ inv(A) ─────────────────────────────────────── + tOrDC = SM80.make_acc_into_op(tDCrDC, tiled_mma, cutlass.Float16) + + cute.copy(A_tiled_copy, tOsAinv, tOrAinv_cv) + tOrO.fill(0.0) + cute.gemm(tiled_mma, tOrO, tOrDC, tOrAinv, tOrO) + + # ── Write output (first MMA_M slice → sO rows 0-7) ─────────────────────── + tOrO_f16.store(tOrO.load().to(cutlass.Float16)) + cute.copy(O_tiled_copy, tOrO_cv[None, None, 0], tOsO[None, None, 0]) + + # ── Level 3: 16×16 blocks → 32×32 blockwise inverse ───────────────────────── + # Called by one warp (thread_idx 0-31 or 32-63, i.e. thread_idx<64 in the WG). + # mat: (32, 32) smem slice. + + @cute.jit + def blockwise_16x16_to_32x32(self, mat: cute.Tensor): + lane_id = cute.arch.lane_idx() + + tiled_mma = cute.make_tiled_mma( + warp.MmaF16BF16Op(cutlass.Float16, cutlass.Float32, (16, 8, 16)), + (1, 1, 1), + permutation_mnk=(16, 16, 16), + ) + + mat_2x2 = cute.flat_divide(mat, (16, 16)) + sDinv = mat_2x2[None, None, 1, 1] + sC = select_tensor_10(mat_2x2[None, None, 1, 0]) + sAinv = select_tensor_10(mat_2x2[None, None, 0, 0]) + sO = mat_2x2[None, None, 1, 0] + + thr_mma = tiled_mma.get_slice(lane_id) + tOrDinv = thr_mma.make_fragment_A(thr_mma.partition_A(sDinv)) + tOrC = thr_mma.make_fragment_B(thr_mma.partition_B(sC)) + tOrAinv = thr_mma.make_fragment_B(thr_mma.partition_B(sAinv)) + tDCrDC = cute.make_rmem_tensor(thr_mma.partition_shape_C((16, 16)), cutlass.Float32) + tOrO = cute.make_rmem_tensor(thr_mma.partition_shape_C((16, 16)), cutlass.Float32) + + dinv_atom = cute.make_copy_atom(warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=2), cutlass.Float16) + b_atom = cute.make_copy_atom(warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=2), cutlass.Float16) + o_atom = cute.make_copy_atom(warp.StMatrix8x8x16bOp(transpose=False, num_matrices=2), cutlass.Float16) + + D_tiled_copy = cute.make_tiled_copy_A(dinv_atom, tiled_mma) + C_tiled_copy = cute.make_tiled_copy_B(b_atom, tiled_mma) + A_tiled_copy = cute.make_tiled_copy_B(b_atom, tiled_mma) + O_tiled_copy = cute.make_tiled_copy_C(o_atom, tiled_mma) + + D_thr_copy = D_tiled_copy.get_slice(lane_id) + C_thr_copy = C_tiled_copy.get_slice(lane_id) + A_thr_copy = A_tiled_copy.get_slice(lane_id) + O_thr_copy = O_tiled_copy.get_slice(lane_id) + + tOsDinv = D_thr_copy.partition_S(sDinv) + tOrDinv_cv = D_thr_copy.retile(tOrDinv) + tOsC = C_thr_copy.partition_S(sC) + tOrC_cv = C_thr_copy.retile(tOrC) + tOsAinv = A_thr_copy.partition_S(sAinv) + tOrAinv_cv = A_thr_copy.retile(tOrAinv) + tOsO = O_thr_copy.partition_D(sO) + tOrO_f16 = cute.make_fragment_like(tOrO, cutlass.Float16) + tOrO_cv = O_thr_copy.retile(tOrO_f16) + + # ── Step 1: tDCrDC = -inv(D) @ C ───────────────────────────────────────── + cute.copy(D_tiled_copy, tOsDinv, tOrDinv_cv) + cute.copy(C_tiled_copy, tOsC, tOrC_cv) + tDCrDC.fill(0.0) + cute.gemm(tiled_mma, tDCrDC, tOrDinv, tOrC, tDCrDC) + for i in cutlass.range_constexpr(cute.size(tDCrDC)): + tDCrDC[i] = -tDCrDC[i] + + # ── Step 2: tOrO = tDCrDC @ inv(A) ─────────────────────────────────────── + tOrDC = SM80.make_acc_into_op(tDCrDC, tiled_mma, cutlass.Float16) + + cute.copy(A_tiled_copy, tOsAinv, tOrAinv_cv) + tOrO.fill(0.0) + cute.gemm(tiled_mma, tOrO, tOrDC, tOrAinv, tOrO) + + tOrO_f16.store(tOrO.load().to(cutlass.Float16)) + cute.copy(O_tiled_copy, tOrO_cv, tOsO) + + # ── Level 4: 32×32 blocks → 64×64 blockwise inverse ───────────────────────── + # Called by all 4 warps (128 threads). + # sT: (64, 64) smem tensor (the full matrix). + + @cute.jit + def blockwise_32x32_to_64x64(self, sT: cute.Tensor, barrier_id: cutlass.Int32): + lane_id = cute.arch.lane_idx() + warp_id = cute.arch.warp_idx() % 4 # WG-local warp ID 0..3 + x = warp_id // 2 # 0 or 1 + y = warp_id % 2 # 0 or 1 + + # TiledMMA1: 16×16×32 (for -inv(D)@C, 1 K-pass of K=32) + tiled_mma1 = cute.make_tiled_mma( + warp.MmaF16BF16Op(cutlass.Float16, cutlass.Float32, (16, 8, 16)), + (1, 1, 1), + permutation_mnk=(16, 16, 32), + ) + # TiledMMA2: 16×32×16 (for (-inv(D)@C)@inv(A), N=32 output) + tiled_mma2 = cute.make_tiled_mma( + warp.MmaF16BF16Op(cutlass.Float16, cutlass.Float32, (16, 8, 16)), + (1, 1, 1), + permutation_mnk=(16, 32, 16), + ) + + mat_2x2 = cute.flat_divide(sT, (32, 32)) + mat_16x2_2x2 = cute.logical_divide(mat_2x2, (16, 16)) + + # Per-warp tile slices (each warp handles one 16×32 or 32×16 sub-tile) + sDinv = mat_16x2_2x2[(None, y), None, 1, 1] + sC = select_tensor_10(mat_16x2_2x2[None, (None, x), 1, 0]) + sAinv = select_tensor_10(mat_16x2_2x2[(None, x), None, 0, 0]) + sO = mat_16x2_2x2[(None, y), None, 1, 0] + + thr_mma1 = tiled_mma1.get_slice(lane_id) + thr_mma2 = tiled_mma2.get_slice(lane_id) + + tOrDinv = thr_mma1.make_fragment_A(thr_mma1.partition_A(sDinv)) + tOrC = thr_mma1.make_fragment_B(thr_mma1.partition_B(sC)) + tOrAinv = thr_mma2.make_fragment_B(thr_mma2.partition_B(sAinv)) + + tDCrDC = cute.make_rmem_tensor(thr_mma1.partition_shape_C((16, 16)), cutlass.Float32) + tOrO = cute.make_rmem_tensor(thr_mma2.partition_shape_C((16, 32)), cutlass.Float32) + + dinv_atom = cute.make_copy_atom(warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), cutlass.Float16) + c_atom = cute.make_copy_atom(warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=4), cutlass.Float16) + ainv_atom = cute.make_copy_atom(warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=2), cutlass.Float16) + O_atom_s2r = cute.make_copy_atom(warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), cutlass.Float16) + O_atom_r2s = cute.make_copy_atom(warp.StMatrix8x8x16bOp(transpose=False, num_matrices=4), cutlass.Float16) + + D_tiled_copy = cute.make_tiled_copy_A(dinv_atom, tiled_mma1) + C_tiled_copy = cute.make_tiled_copy_B(c_atom, tiled_mma1) + A_tiled_copy = cute.make_tiled_copy_B(ainv_atom, tiled_mma2) + O_tiled_s2r = cute.make_tiled_copy_C(O_atom_s2r, tiled_mma2) + O_tiled_r2s = cute.make_tiled_copy_C(O_atom_r2s, tiled_mma2) + + D_thr_copy = D_tiled_copy.get_slice(lane_id) + C_thr_copy = C_tiled_copy.get_slice(lane_id) + A_thr_copy = A_tiled_copy.get_slice(lane_id) + O_thr_s2r = O_tiled_s2r.get_slice(lane_id) + O_thr_r2s = O_tiled_r2s.get_slice(lane_id) + + tOsDinv = D_thr_copy.partition_S(sDinv) + tOrDinv_cv = D_thr_copy.retile(tOrDinv) + tOsC = C_thr_copy.partition_S(sC) + tOrC_cv = C_thr_copy.retile(tOrC) + tOsAinv = A_thr_copy.partition_S(sAinv) + tOrAinv_cv = A_thr_copy.retile(tOrAinv) + + # ── Step 1: tDCrDC = -inv(D) @ C ───────────────────────────────────────── + cute.copy(D_tiled_copy, tOsDinv, tOrDinv_cv) + cute.copy(C_tiled_copy, tOsC, tOrC_cv) + tDCrDC.fill(0.0) + cute.gemm(tiled_mma1, tDCrDC, tOrDinv, tOrC, tDCrDC) + for i in cutlass.range_constexpr(cute.size(tDCrDC)): + tDCrDC[i] = -tDCrDC[i] + + # ── Step 2: tOrO = tDCrDC @ inv(A) ─────────────────────────────────────── + tOrDC = SM80.make_acc_into_op(tDCrDC, tiled_mma2, cutlass.Float16) + + cute.copy(A_tiled_copy, tOsAinv, tOrAinv_cv) + tOrO.fill(0.0) + cute.gemm(tiled_mma2, tOrO, tOrDC, tOrAinv, tOrO) + + tOrO_f16 = cute.make_fragment_like(tOrO, cutlass.Float16) + tOrO_f16.store(tOrO.load().to(cutlass.Float16)) + + # ── Cross-warp reduction: warps with x=0 write first, x=1 reads+adds+writes + cute.arch.barrier(barrier_id=barrier_id, number_of_threads=128) + + tOsO = O_thr_r2s.partition_D(sO) + tOrO_cv = O_thr_r2s.retile(tOrO_f16) + if x == 0: + cute.copy(O_tiled_r2s, tOrO_cv, tOsO) + + cute.arch.barrier(barrier_id=barrier_id, number_of_threads=128) + + if x == 1: + tOrO_red = cute.make_fragment_like(tOrO_f16) + tOsO_s = O_thr_s2r.partition_S(sO) + tOrO_red_cv = O_thr_s2r.retile(tOrO_red) + cute.copy(O_tiled_s2r, tOsO_s, tOrO_red_cv) + for i in cutlass.range_constexpr(cute.size(tOrO_f16)): + tOrO_f16[i] = tOrO_f16[i] + tOrO_red[i] + cute.copy(O_tiled_r2s, tOrO_cv, tOsO) + + @cute.jit + def run(self, sT: cute.Tensor, barrier_id: cutlass.Int32): + """Invert a 64×64 column-major shared-memory tensor in place. + + The call requires one 128-thread warp group and an available barrier ID. + """ + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % 128 + + # ── Level 1: 8×8 Gauss on diagonal 8×8 blocks (threads 0-63) ──────────── + t8x8 = cute.flat_divide(sT, (8, 8)) + if thread_idx < 64: + blk = thread_idx // 8 + self.compute_diagonal_block_inverse(t8x8[None, None, blk, blk], thread_idx % 8, 8) + + cute.arch.barrier(barrier_id=barrier_id, number_of_threads=128) + + # ── Level 2: 16×16 blockwise inverse (all 4 warps → 4 diagonal blocks) ── + t16x16 = cute.flat_divide(sT, (16, 16)) + blk2 = thread_idx // 32 + self.blockwise_8x8_to_16x16(t16x16[None, None, blk2, blk2]) + + cute.arch.barrier(barrier_id=barrier_id, number_of_threads=128) + + # ── Level 3: 32×32 blockwise inverse (threads 0-63 → 2 diagonal blocks) ─ + t32x32 = cute.flat_divide(sT, (32, 32)) + if thread_idx < 64: + blk3 = thread_idx // 32 + self.blockwise_16x16_to_32x32(t32x32[None, None, blk3, blk3]) + + cute.arch.barrier(barrier_id=barrier_id, number_of_threads=128) + + # ── Level 4: 64×64 blockwise inverse (all 4 warps) ─────────────────────── + self.blockwise_32x32_to_64x64(sT, barrier_id) diff --git a/cula/ops/gdn2/sm90/config.py b/cula/ops/gdn2/sm90/config.py new file mode 100644 index 00000000..f7b4cffd --- /dev/null +++ b/cula/ops/gdn2/sm90/config.py @@ -0,0 +1,37 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Stable host-side contract for Hopper SM90a GDN2 prefill.""" + +CHUNK_SIZE = 64 +HEAD_SIZE = 128 +VALUE_SIZE = 128 +THREADS_PER_CTA = 384 +MAX_SEQUENCES = 32 +SUPPORTED_Q_HEADS = 16 +SUPPORTED_V_HEADS = (16, 32, 64) +SM90_BACKEND_ID = "sm90a_cutedsl_gdn2_prefill_v1" + +# Elementwise log-decay lower bound. The blockwise-rebased intra-chunk +# factorization keeps every stored exponent within a 15-token in-block span; +# g >= -5 leaves the largest such exponent (75 nats) a factor ~9e5 below the +# BF16/FP32 overflow boundary (~88.72 nats). See +# docs/gdn2_sm90_stable_factor.md. +SUPPORTED_G_MIN = -5.0 + +# The one supported nvidia-cutlass-dsl range for this backend. +# is_sm90_gdn2_available(), the runtime dispatch error, and +# docs/gdn2_sm90_api.md all derive from this string, and the installed +# version is read through cula.ops._mlir_compat so this backend and the +# shared gateway cannot disagree about which toolchain is in use. +# +# Relationship to the repository-wide contract in cula/ops/_mlir_compat.py +# (_SUPPORTED_MIN/_SUPPORTED_MAX/_EXCLUDED_VERSIONS, kept in sync with +# pyproject.toml): the upper bound is the same 4.7 and is enforced there on +# every private-dialect access, which this kernel triggers at import. GDN2 +# only raises the floor, from 4.4.2 to 4.5.1, because the kernel needs +# `cutlass.cute.nvgpu.OperandMajorMode`, which 4.4.x does not provide -- the +# backend cannot be imported there at all. Raising the shared floor also +# subsumes the shared 4.5.0 exclusion. Keep the upper bound in step with the +# gateway when it moves. +CUTLASS_DSL_REQUIREMENT = "nvidia-cutlass-dsl>=4.5.1,<4.7" diff --git a/cula/ops/gdn2/sm90/inverse_helpers.py b/cula/ops/gdn2/sm90/inverse_helpers.py new file mode 100644 index 00000000..6a078304 --- /dev/null +++ b/cula/ops/gdn2/sm90/inverse_helpers.py @@ -0,0 +1,61 @@ +# Copyright 2025-2026 FlashInfer team. +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""CuTe DSL helpers used by the GDN2 triangular-inverse collective.""" + +import cutlass +import cutlass.cute as cute + + +@cute.jit +def select_tensor_10(t: cute.Tensor) -> cute.Tensor: + """Swap the first two modes of a tensor without changing its storage.""" + + return cute.make_tensor( + t.iterator.align(t.iterator.max_alignment), + cute.make_layout( + (t.layout.shape[1], t.layout.shape[0]) + t.layout.shape[2:], + stride=(t.layout.stride[1], t.layout.stride[0]) + t.layout.stride[2:], + ), + ) + + +class SM80: + """Fragment-layout helpers for the HMMA inverse collective.""" + + @staticmethod + @cute.jit + def convert_c_layout_to_a_layout(c_layout, tiled_mma): + c_frag_atom_size = cute.size(c_layout, mode=[0]) + a_frag_atom_size = cute.size(tiled_mma.tv_layout_A, mode=[1]) + ratio = a_frag_atom_size // c_frag_atom_size + if cutlass.const_expr(ratio == 1): + return c_layout + + divided = cute.logical_divide(c_layout, (None, None, ratio)) + frag_layout = cute.flatten( + cute.make_layout( + (divided.shape[0], divided.shape[2][0]), + stride=(divided.stride[0], divided.stride[2][0]), + ), + ) + return cute.make_layout( + (frag_layout.shape, divided.shape[1], divided.shape[2][1]), + stride=( + frag_layout.stride, + divided.stride[1], + divided.stride[2][1], + ), + ) + + @staticmethod + @cute.jit + def make_acc_into_op(acc: cute.Tensor, tiled_mma, dtype) -> cute.Tensor: + operand = cute.make_fragment_like( + SM80.convert_c_layout_to_a_layout(acc.layout, tiled_mma), + dtype, + ) + operand_as_acc = cute.make_tensor(operand.iterator, acc.layout) + operand_as_acc.store(acc.load().to(dtype)) + return operand diff --git a/cula/ops/gdn2/sm90/prefill.py b/cula/ops/gdn2/sm90/prefill.py new file mode 100644 index 00000000..b807f140 --- /dev/null +++ b/cula/ops/gdn2/sm90/prefill.py @@ -0,0 +1,198 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Host launch adapter for the Hopper SM90a GDN2 prefill kernel.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import from_dlpack + +from .config import ( + MAX_SEQUENCES, + SM90_BACKEND_ID, + SUPPORTED_Q_HEADS, + SUPPORTED_V_HEADS, +) +from .prefill_kernel import GDN2PrefillKernel + +if TYPE_CHECKING: + from cula.gdn2.prefill import _GDN2Inputs + +_compiled: dict[tuple[int, int, bool, bool, bool, bool], object] = {} + + +@dataclass(frozen=True) +class GDN2ExecutionInfo: + """Metadata receipt for one product GDN2 launch.""" + + backend_id: str + total_tokens: int + num_sequences: int + num_q_heads: int + num_v_heads: int + has_initial_state: bool + store_final_state: bool + sequence_policy: str + compile_cache_entries: int + fallback: bool + + +def _device_key(device: torch.device) -> int: + return torch.cuda.current_device() if device.index is None else device.index + + +def _dynamic_mode0(tensor: torch.Tensor): + return from_dlpack( + tensor, + assumed_align=16, + ).mark_compact_shape_dynamic( + mode=0, + stride_order=tensor.dim_order(), + ) + + +def _resolve_support(inputs: _GDN2Inputs) -> None: + if inputs.num_q_heads != SUPPORTED_Q_HEADS: + raise NotImplementedError( + f"GDN2 SM90a prefill requires Hq={SUPPORTED_Q_HEADS}", + ) + if inputs.num_v_heads not in SUPPORTED_V_HEADS: + raise NotImplementedError( + f"GDN2 SM90a prefill requires Hv in {SUPPORTED_V_HEADS}", + ) + if not 1 <= inputs.num_sequences <= MAX_SEQUENCES: + raise NotImplementedError( + f"GDN2 SM90a prefill requires 1 <= N <= {MAX_SEQUENCES}", + ) + + +def _compile( + inputs: _GDN2Inputs, + initial_state: torch.Tensor, + final_state: torch.Tensor, + stream: cuda.CUstream, +): + has_initial_state = inputs.initial_state is not None + store_final_state = inputs.output_final_state + use_n1_hv16_v64 = ( + inputs.num_sequences == 1 + and inputs.num_v_heads == 16 + and has_initial_state + and store_final_state + and inputs.total_tokens > 64 + ) + retain_final_tail = store_final_state and not (inputs.num_sequences == 1 and inputs.total_tokens <= 64) + key = ( + _device_key(inputs.q.device), + inputs.num_v_heads, + has_initial_state, + store_final_state, + use_n1_hv16_v64, + retain_final_tail, + ) + compiled = _compiled.get(key) + if compiled is not None: + return compiled + + kernel = GDN2PrefillKernel( + has_initial_state=has_initial_state, + store_final_state=store_final_state, + value_tile=64 if use_n1_hv16_v64 else 128, + single_state_owner=use_n1_hv16_v64, + retain_final_tail=retain_final_tail, + ) + compiled = cute.compile( + kernel, + *( + _dynamic_mode0(tensor) + for tensor in ( + inputs.q, + inputs.k, + inputs.v, + inputs.b, + inputs.w, + inputs.cu_seqlens, + inputs.g, + inputs.q, + inputs.q, + initial_state, + inputs.output, + final_state, + ) + ), + cutlass.Int32(inputs.num_sequences), + cutlass.Int32(inputs.num_q_heads), + cutlass.Int32(inputs.num_v_heads), + cutlass.Int32(inputs.total_tokens), + cutlass.Float32(inputs.scale), + stream=stream, + options="--enable-tvm-ffi", + ) + _compiled[key] = compiled + return compiled + + +def launch_sm90_gdn2( + inputs: _GDN2Inputs, + *, + return_debug: bool = False, +) -> GDN2ExecutionInfo | None: + """Launch the product GDN2 backend without a fallback.""" + + _resolve_support(inputs) + initial_state = inputs.initial_state if inputs.initial_state is not None else inputs.q + final_state = inputs.output_state if inputs.output_state is not None else inputs.output + + device = inputs.q.device + with torch.cuda.device(device): + stream = cuda.CUstream( + torch.cuda.current_stream(device).cuda_stream, + ) + compiled = _compile( + inputs, + initial_state, + final_state, + stream, + ) + compiled( + inputs.q, + inputs.k, + inputs.v, + inputs.b, + inputs.w, + inputs.cu_seqlens, + inputs.g, + inputs.q, + inputs.q, + initial_state, + inputs.output, + final_state, + inputs.num_sequences, + inputs.num_q_heads, + inputs.num_v_heads, + inputs.total_tokens, + inputs.scale, + stream, + ) + + if not return_debug: + return None + return GDN2ExecutionInfo( + backend_id=SM90_BACKEND_ID, + total_tokens=inputs.total_tokens, + num_sequences=inputs.num_sequences, + num_q_heads=inputs.num_q_heads, + num_v_heads=inputs.num_v_heads, + has_initial_state=inputs.initial_state is not None, + store_final_state=inputs.output_final_state, + sequence_policy="stable_lpt32", + compile_cache_entries=len(_compiled), + fallback=False, + ) diff --git a/cula/ops/gdn2/sm90/prefill_kernel.py b/cula/ops/gdn2/sm90/prefill_kernel.py new file mode 100644 index 00000000..6da5c6c1 --- /dev/null +++ b/cula/ops/gdn2/sm90/prefill_kernel.py @@ -0,0 +1,3043 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Hopper SM90a Gated DeltaNet-2 forward-prefill kernel. + +WG1/WG2 retain distributed common transforms and resident FP32 state. +They construct the exact chunk-local FP32 prefix from public raw G in the +existing two-stage shared arena, then prepare Q/K/B/G-derived operands for +generation ``n + 1`` before executing recurrence generation ``n``. WG0 +computes causal QK, erase, and collective inverse for ``n + 1`` concurrently +with that recurrence. No global G-prefix workspace or second launch is used. +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils.hopper_helpers as sm90_utils +from cutlass.cute.nvgpu import OperandMajorMode, cpasync, warp, warpgroup +from cutlass.cutlass_dsl import T + +from cula.ops._mlir_compat import llvm + +from .collective_inverse_hmma import CollectiveInverse +from .config import CHUNK_SIZE, HEAD_SIZE, VALUE_SIZE + +_INV_LN2 = 1.4426950408889634 +_WARP_GROUP_SIZE = 128 +_THREADS_PER_CTA = 384 +_WGMMA_K = 16 +_RAW_KEY_TILES = HEAD_SIZE // _WGMMA_K +_VALUE_TILES = VALUE_SIZE // _WGMMA_K +_RAW_STAGES = 2 +_VW_PRIVATE_STAGES = 1 +_INPUT_STAGES = 2 +_FACTOR_WORKSPACE_STAGES = 1 +_WRITE_STAGES = 1 +_OUTPUT_STAGES = 2 +_PRODUCER_SIGNAL_WARPS = _WARP_GROUP_SIZE // 32 +_PRODUCER_REGISTER_TARGET = 72 +_STATE_REGISTER_TARGET = 216 +_STATE_VALUE_TILE = 64 +_QKBG_TRANSACTION_BYTES = 3 * CHUNK_SIZE * _WGMMA_K * 2 + CHUNK_SIZE * _WGMMA_K * 4 +_VW_TRANSACTION_BYTES = 2 * CHUNK_SIZE * _WGMMA_K * 2 +_STATE0_WRITE_BARRIER = 2 +_STATE1_WRITE_BARRIER = 3 +_STATE_COMMON_BARRIER = 4 +_STATE_ITERATION_DONE_BARRIER = 5 +_STATE0_PREFIX_BARRIER = 6 +_STATE1_PREFIX_BARRIER = 7 +_STORE_WG_BARRIER = 1 +_INVERSE_BARRIER = 13 +_QKB_STREAM_TILES = _RAW_KEY_TILES // 2 +_MAX_SEQUENCE_SCHEDULE = 32 + + +@cute.jit +def _device_fail_closed() -> None: + llvm.inline_asm( + None, + [], + "trap;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@cute.jit +def _convert_c_layout_to_a_layout( + c_layout: cute.Layout, + a_value_layout, +): + """Convert a Hopper C-fragment layout to its matching RS-A layout.""" + + return cute.make_layout( + ( + a_value_layout, + c_layout.shape[1], + ( + c_layout.shape[2], + cute.size(c_layout, mode=[0]) // cute.size(a_value_layout), + ), + ), + stride=( + c_layout.stride[0], + c_layout.stride[1], + ( + c_layout.stride[2], + cute.size(a_value_layout, mode=[2]) * c_layout.stride[0][2], + ), + ), + ) + + +@cute.jit +def _make_acc_into_op( + accumulator: cute.Tensor, + tiled_mma: cute.TiledMma, +) -> cute.Tensor: + operand = cute.make_rmem_tensor_like( + _convert_c_layout_to_a_layout( + accumulator.layout, + tiled_mma.tv_layout_A.shape[1], + ), + cutlass.BFloat16, + ) + operand_as_accumulator = cute.make_tensor( + operand.iterator, + accumulator.layout, + ) + operand_as_accumulator.store( + accumulator.load().to(cutlass.BFloat16), + ) + return operand + + +@cute.jit +def _fence_f32_register(reg: cutlass.Float32) -> cutlass.Float32: + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [reg.ir_value()], + "", + "=f,0", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ), + ) + + +@cute.jit +def _fence_u32_register(reg: cutlass.Uint32) -> cutlass.Uint32: + return cutlass.Uint32( + llvm.inline_asm( + T.i32(), + [reg.ir_value()], + "", + "=r,0", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ), + ) + + +@cute.jit +def _fence_register_fragment(fragment: cute.Tensor) -> None: + if cutlass.const_expr(fragment.element_type is cutlass.Float32): + values = cute.recast_tensor(fragment, cutlass.Float32) + for item in cutlass.range_constexpr(cute.size(values)): + values[item] = _fence_f32_register(values[item]) + else: + values = cute.recast_tensor(fragment, cutlass.Uint32) + for item in cutlass.range_constexpr(cute.size(values)): + values[item] = _fence_u32_register(values[item]) + + +@cute.jit +def _wgmma_gemm( + tiled_mma: cute.TiledMma, + accumulator: cute.Tensor, + operand_a: cute.Tensor, + operand_b: cute.Tensor, + accumulate: bool, +) -> None: + for k_block in cutlass.range( + cute.size(operand_a, mode=[2]), + unroll_full=True, + ): + tiled_mma.set( + warpgroup.Field.ACCUMULATE, + accumulate or k_block != 0, + ) + cute.gemm( + tiled_mma, + accumulator, + operand_a[(None, None, k_block)], + operand_b[(None, None, k_block)], + accumulator, + ) + + +@cute.jit +def _stable_lpt32_sequence( + cu_seqlens: cute.Tensor, + sequence_rank: cutlass.Int32, + num_sequences: cutlass.Int32, + lane: cutlass.Int32, +) -> cutlass.Int32: + """Return the stable descending chunk-count sequence for one rank.""" + + if num_sequences > cutlass.Int32(_MAX_SEQUENCE_SCHEDULE): + _device_fail_closed() + + chunk_count = cutlass.Int32(-1) + if lane < num_sequences: + start = cutlass.Int64(cu_seqlens[lane]) + end = cutlass.Int64( + cu_seqlens[lane + cutlass.Int32(1)], + ) + if start < cutlass.Int64(0) or end <= start: + _device_fail_closed() + chunk_count = cutlass.Int32( + (end - start + cutlass.Int64(CHUNK_SIZE - 1)) // cutlass.Int64(CHUNK_SIZE), + ) + + rank = cutlass.Int32(0) + for source_lane in cutlass.range(num_sequences, unroll=0): + other_chunk_count = cute.arch.shuffle_sync( + chunk_count, + source_lane, + ) + if other_chunk_count > chunk_count or (other_chunk_count == chunk_count and source_lane < lane): + rank = rank + cutlass.Int32(1) + + selected = cutlass.Int32(-1) + if lane < num_sequences and rank == sequence_rank: + selected = lane + for delta in (16, 8, 4, 2, 1): + other_selected = cute.arch.shuffle_sync_down( + selected, + delta, + ) + if other_selected > selected: + selected = other_selected + sequence = cute.arch.shuffle_sync(selected, 0) + if sequence < cutlass.Int32(0): + _device_fail_closed() + return sequence + + +_FACTOR_BLOCK = 16 +_FACTOR_SUB_BLOCKS = CHUNK_SIZE // _FACTOR_BLOCK +# Three precomputed pair-product rows: the distance-two products +# (delta1*delta2, delta2*delta3) and the distance-three product +# (delta1*delta2*delta3), stored FP16 so the factor fold reads one row per +# pair without growing the near-capacity V128 shared-memory budget. The +# FP16 rounding (2^-11 relative, subnormal floor 6e-8) is subdominant to +# the BF16 fold-operand rounding (2^-8) and to every comparison tolerance. +_GS_PAIR_ROWS = 3 +_FACTOR_LOWER_PAIRS = ( + # Ordered so the modulo-4 warp assignment balances pair count against + # fold count: the three-pair warps carry the diagonal (fold-free) pairs. + (0, 0), + (1, 1), + (2, 0), + (3, 1), + (2, 2), + (3, 3), + (3, 0), + (1, 0), + (2, 1), + (3, 2), +) +_FACTOR_UPPER_PAIRS = ( + (0, 1), + (0, 2), + (0, 3), + (1, 2), + (1, 3), + (2, 3), +) + + +def _factor_pair_scale( + gs_delta: cute.Tensor, + gs_pair: cute.Tensor, + channel: cutlass.Int32, + query_block: int, + key_block: int, +) -> cutlass.Float32: + """Per-channel exp(Gs(I) - Gs(J)) read as one precomputed row. + + Trace-time helper (deliberately not ``@cute.jit``): distance-one pairs + read the plain FP32 delta row; longer distances read the FP16 product + rows the gs_delta writer precomputes once per generation. + """ + + distance = query_block - key_block + if distance == 1: + return cutlass.Float32( + gs_delta[ + channel, + key_block + 1, + ], + ) + row = key_block if distance == 2 else 2 + return cutlass.Float32( + gs_pair[ + channel, + row, + ], + ) + + +@cute.jit +def _factor_block_pair( + thread_mma, + a_tiled_copy, + b_tiled_copy, + a_thread_copy, + b_thread_copy, + key_coordinates: cute.Tensor, + q_blocks: cute.Tensor, + erase_blocks: cute.Tensor, + k_blocks: cute.Tensor, + gs_delta: cute.Tensor, + gs_pair: cute.Tensor, + qk_accumulator: cute.Tensor, + erase_accumulator: cute.Tensor, + query_block: cutlass.Constexpr, + key_block: cutlass.Constexpr, +) -> None: + """Accumulate one 16x16 sub-block pair for both factor matrices. + + Operand fragments stream one 16-channel k-tile at a time through + LdMatrix tiled copies so the producer warp group stays inside its + register budget without scalar shared-memory traffic; the shared + ``k~'`` fragment is loaded once per tile and used by both matrices, + and off-diagonal pairs fold the bounded per-channel correction + ``exp(Gs(I) - Gs(J)) <= 1`` into the left fragments before each MMA. + """ + + q_tiles = cute.flat_divide( + q_blocks[None, None, query_block, 0], + (_FACTOR_BLOCK, _WGMMA_K), + ) + erase_tiles = cute.flat_divide( + erase_blocks[None, None, query_block, 0], + (_FACTOR_BLOCK, _WGMMA_K), + ) + key_tiles = cute.flat_divide( + k_blocks[None, None, key_block, 0], + (_FACTOR_BLOCK, _WGMMA_K), + ) + for key_tile in cutlass.range_constexpr(HEAD_SIZE // _WGMMA_K): + key_slice = key_tiles[None, None, 0, key_tile] + key_fragment = thread_mma.make_fragment_B( + thread_mma.partition_B(key_slice), + ) + cute.copy( + b_tiled_copy, + b_thread_copy.partition_S(key_slice), + b_thread_copy.retile(key_fragment), + ) + if cutlass.const_expr(query_block != key_block): + # Fold exp(Gs(I) - Gs(J)) into the shared right fragment once: + # the folded operand is k * exp(Gs(I) - G_j) <= |k|, and both + # matrices consume the same corrected fragment. + for element in cutlass.range_constexpr( + cute.size(key_fragment), + ): + _, tile_channel = key_coordinates[element] + channel = cutlass.Int32(key_tile * _WGMMA_K) + tile_channel + pair_scale = _factor_pair_scale( + gs_delta, + gs_pair, + channel, + query_block, + key_block, + ) + key_fragment[element] = cutlass.BFloat16( + cutlass.Float32(key_fragment[element]) * pair_scale, + ) + + query_slice = q_tiles[None, None, 0, key_tile] + query_fragment = thread_mma.make_fragment_A( + thread_mma.partition_A(query_slice), + ) + cute.copy( + a_tiled_copy, + a_thread_copy.partition_S(query_slice), + a_thread_copy.retile(query_fragment), + ) + erase_slice = erase_tiles[None, None, 0, key_tile] + erase_fragment = thread_mma.make_fragment_A( + thread_mma.partition_A(erase_slice), + ) + cute.copy( + a_tiled_copy, + a_thread_copy.partition_S(erase_slice), + a_thread_copy.retile(erase_fragment), + ) + cute.gemm( + thread_mma, + qk_accumulator, + query_fragment, + key_fragment, + qk_accumulator, + ) + cute.gemm( + thread_mma, + erase_accumulator, + erase_fragment, + key_fragment, + erase_accumulator, + ) + + +@cute.jit +def _store_factor_qk_block( + qk_store_tiled, + qk_store_thread, + block_coordinates: cute.Tensor, + aqk_blocks: cute.Tensor, + qk_accumulator: cute.Tensor, + valid_tokens: cutlass.Int32, + scale: cutlass.Float32, + query_block: cutlass.Constexpr, + key_block: cutlass.Constexpr, +) -> None: + masked = cute.make_fragment_like(qk_accumulator, cutlass.BFloat16) + for element in cutlass.range_constexpr(cute.size(qk_accumulator)): + local_row, local_column = block_coordinates[element] + global_row = cutlass.Int32(query_block * _FACTOR_BLOCK) + local_row + global_column = cutlass.Int32(key_block * _FACTOR_BLOCK) + local_column + qk_value = cutlass.Float32(0.0) + if cutlass.const_expr(query_block == key_block): + # Diagonal blocks: the causal test is live; the column bound is + # implied by column <= row < valid_tokens. + if global_row < valid_tokens and global_row >= global_column: + qk_value = qk_accumulator[element] * scale + else: + # Strictly lower blocks: row >= 16*I > 16*J + 15 >= column, so + # causality and the column bound both follow from the row bound. + if global_row < valid_tokens: + qk_value = qk_accumulator[element] * scale + masked[element] = cutlass.BFloat16(qk_value) + cute.copy( + qk_store_tiled, + qk_store_thread.retile(masked), + qk_store_thread.partition_D( + aqk_blocks[None, None, query_block, key_block], + ), + ) + + +@cute.jit +def _store_factor_erase_block( + erase_store_tiled, + erase_store_thread, + block_coordinates: cute.Tensor, + inverse_blocks: cute.Tensor, + erase_accumulator: cute.Tensor, + valid_tokens: cutlass.Int32, + query_block: cutlass.Constexpr, + key_block: cutlass.Constexpr, +) -> None: + masked = cute.make_fragment_like(erase_accumulator, cutlass.Float16) + for element in cutlass.range_constexpr(cute.size(erase_accumulator)): + local_row, local_column = block_coordinates[element] + global_row = cutlass.Int32(query_block * _FACTOR_BLOCK) + local_row + global_column = cutlass.Int32(key_block * _FACTOR_BLOCK) + local_column + erase_value = cutlass.Float32(0.0) + if cutlass.const_expr(query_block == key_block): + # Diagonal blocks: keep the strict-lower test; the column bound + # is implied by column < row < valid_tokens. + if global_row < valid_tokens and global_row > global_column: + erase_value = erase_accumulator[element] + else: + # Strictly lower blocks: row > column and the column bound both + # follow from the row bound. + if global_row < valid_tokens: + erase_value = erase_accumulator[element] + masked[element] = cutlass.Float16(erase_value) + cute.copy( + erase_store_tiled, + erase_store_thread.retile(masked), + erase_store_thread.partition_D( + inverse_blocks[None, None, query_block, key_block], + ), + ) + + +@cute.jit +def _publish_factor_blocks( + thread_in_group: cutlass.Int32, + q_tilde: cute.Tensor, + erase_tilde: cute.Tensor, + k_prime: cute.Tensor, + gs_delta: cute.Tensor, + gs_pair: cute.Tensor, + aqk: cute.Tensor, + inverse: cute.Tensor, + valid_tokens: cutlass.Int32, + scale: cutlass.Float32, + fill_static_upper: cutlass.Boolean, +) -> None: + """Publish causal QK and the strict-lower erase Gram per sub-block pair. + + The operands are blockwise rebased: ``q~``/``e~`` carry + ``exp(G_i - Gs(B(i)))`` and ``k~'`` carries ``exp(Gs(B(j)) - G_j)``, so + every stored exponent spans at most 15 in-block token gaps. Ten lower + 16x16 blocks per matrix run on four warps with warp-level m16n8k16 MMAs + sharing one ``k~'`` fragment per pair per k-tile; the six upper blocks + are zero-filled. The FP16 Gram destination aliases the dead raw-G + staging arena, so no factor operand overlaps it. + """ + + warp_index = thread_in_group // cutlass.Int32(32) + lane_index = thread_in_group % cutlass.Int32(32) + + tiled_mma = cute.make_tiled_mma( + warp.MmaF16BF16Op( + cutlass.BFloat16, + cutlass.Float32, + (_FACTOR_BLOCK, 8, _WGMMA_K), + ), + (1, 1, 1), + permutation_mnk=(_FACTOR_BLOCK, _FACTOR_BLOCK, _WGMMA_K), + ) + thread_mma = tiled_mma.get_slice(lane_index) + + q_blocks = cute.flat_divide(q_tilde, (_FACTOR_BLOCK, HEAD_SIZE)) + erase_blocks = cute.flat_divide(erase_tilde, (_FACTOR_BLOCK, HEAD_SIZE)) + k_blocks = cute.flat_divide(k_prime, (_FACTOR_BLOCK, HEAD_SIZE)) + aqk_blocks = cute.flat_divide(aqk, (_FACTOR_BLOCK, _FACTOR_BLOCK)) + inverse_blocks = cute.flat_divide( + inverse, + (_FACTOR_BLOCK, _FACTOR_BLOCK), + ) + + key_coordinates = thread_mma.partition_B( + cute.make_identity_tensor((_FACTOR_BLOCK, _WGMMA_K)), + ) + block_coordinates = thread_mma.partition_C( + cute.make_identity_tensor((_FACTOR_BLOCK, _FACTOR_BLOCK)), + ) + + a_atom = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=2), + cutlass.BFloat16, + ) + b_atom = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=2), + cutlass.BFloat16, + ) + qk_store_atom = cute.make_copy_atom( + warp.StMatrix8x8x16bOp(transpose=False, num_matrices=2), + cutlass.BFloat16, + ) + erase_store_atom = cute.make_copy_atom( + warp.StMatrix8x8x16bOp(transpose=False, num_matrices=2), + cutlass.Float16, + ) + a_tiled_copy = cute.make_tiled_copy_A(a_atom, tiled_mma) + b_tiled_copy = cute.make_tiled_copy_B(b_atom, tiled_mma) + qk_store_tiled = cute.make_tiled_copy_C(qk_store_atom, tiled_mma) + erase_store_tiled = cute.make_tiled_copy_C(erase_store_atom, tiled_mma) + a_thread_copy = a_tiled_copy.get_slice(lane_index) + b_thread_copy = b_tiled_copy.get_slice(lane_index) + qk_store_thread = qk_store_tiled.get_slice(lane_index) + erase_store_thread = erase_store_tiled.get_slice(lane_index) + + for pair in cutlass.range_constexpr(len(_FACTOR_LOWER_PAIRS)): + query_block = _FACTOR_LOWER_PAIRS[pair][0] + key_block = _FACTOR_LOWER_PAIRS[pair][1] + if warp_index == cutlass.Int32(pair % 4): + qk_accumulator = cute.make_rmem_tensor( + thread_mma.partition_shape_C((_FACTOR_BLOCK, _FACTOR_BLOCK)), + cutlass.Float32, + ) + erase_accumulator = cute.make_rmem_tensor( + thread_mma.partition_shape_C((_FACTOR_BLOCK, _FACTOR_BLOCK)), + cutlass.Float32, + ) + qk_accumulator.fill(0.0) + erase_accumulator.fill(0.0) + _factor_block_pair( + thread_mma, + a_tiled_copy, + b_tiled_copy, + a_thread_copy, + b_thread_copy, + key_coordinates, + q_blocks, + erase_blocks, + k_blocks, + gs_delta, + gs_pair, + qk_accumulator, + erase_accumulator, + query_block, + key_block, + ) + _store_factor_qk_block( + qk_store_tiled, + qk_store_thread, + block_coordinates, + aqk_blocks, + qk_accumulator, + valid_tokens, + scale, + query_block, + key_block, + ) + _store_factor_erase_block( + erase_store_tiled, + erase_store_thread, + block_coordinates, + inverse_blocks, + erase_accumulator, + valid_tokens, + query_block, + key_block, + ) + + for pair in cutlass.range_constexpr(len(_FACTOR_UPPER_PAIRS)): + query_block = _FACTOR_UPPER_PAIRS[pair][0] + key_block = _FACTOR_UPPER_PAIRS[pair][1] + for linear in cutlass.range( + thread_in_group, + _FACTOR_BLOCK * _FACTOR_BLOCK, + _WARP_GROUP_SIZE, + unroll=1, + ): + local_row = linear // cutlass.Int32(_FACTOR_BLOCK) + local_column = linear % cutlass.Int32(_FACTOR_BLOCK) + # The inverse arena aliases the raw-G staging that every chunk's + # G transaction rewrites, so its upper blocks are re-zeroed each + # generation. The aqk arena is private and nothing ever writes + # its upper blocks, so those zeros are published once per stage. + if fill_static_upper: + aqk_blocks[ + local_row, + local_column, + query_block, + key_block, + ] = cutlass.BFloat16(0.0) + inverse_blocks[ + local_row, + local_column, + query_block, + key_block, + ] = cutlass.Float16(0.0) + cute.arch.fence_proxy("async.shared", space="cta") + + +class GDN2PrefillKernel: + """Production raw-G prefill kernel with stable LPT32 sequence scheduling.""" + + value_tile = 128 + state_value_tile = _STATE_VALUE_TILE + threads_per_cta = _THREADS_PER_CTA + min_blocks_per_mp = 1 + + def __init__( + self, + *, + has_initial_state: bool, + store_final_state: bool, + value_tile: int = VALUE_SIZE, + single_state_owner: bool = False, + retain_final_tail: bool = False, + ) -> None: + if value_tile not in (64, VALUE_SIZE): + raise ValueError(f"unsupported GDN2 value tile: {value_tile}") + if single_state_owner != (value_tile == 64): + raise ValueError("V64 requires exactly one recurrent State WG") + self.has_initial_state = has_initial_state + self.store_final_state = store_final_state + self.value_tile = value_tile + self.single_state_owner = single_state_owner + self.retain_final_tail = retain_final_tail + self.subgroup_prefix = True + self.subgroup_exclusive_carry = True + self.elide_state_common_barrier = True + self.elide_state_iteration_done_barrier = True + self.sequence_wave_rotation = 0 + self.length_ranked_sequence_order = True + + @cute.jit + def __call__( + self, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + b: cute.Tensor, + w: cute.Tensor, + cu_seqlens: cute.Tensor, + g: cute.Tensor, + capsule_aqk: cute.Tensor, + capsule_akk: cute.Tensor, + initial_state: cute.Tensor, + output: cute.Tensor, + final_state: cute.Tensor, + num_sequences: cutlass.Int32, + num_q_heads: cutlass.Int32, + num_v_heads: cutlass.Int32, + total_tokens: cutlass.Int32, + scale: cutlass.Float32, + stream: cuda.CUstream, + ) -> None: + state_op = warpgroup.MmaF16BF16Op( + cutlass.BFloat16, + cutlass.Float32, + (self.state_value_tile, HEAD_SIZE, _WGMMA_K), + warpgroup.OperandSource.RMEM, + OperandMajorMode.K, + OperandMajorMode.K, + ) + # Token-dimension MMAs use an n16 atom so the state-side projections + # can issue per 16-token sub-block with the state fragments rescaled + # by the bounded block deltas between slices. + token_op = warpgroup.MmaF16BF16Op( + cutlass.BFloat16, + cutlass.Float32, + (self.state_value_tile, _FACTOR_BLOCK, _WGMMA_K), + warpgroup.OperandSource.RMEM, + OperandMajorMode.K, + OperandMajorMode.K, + ) + state_mma = cute.make_tiled_mma( + cute.make_mma_atom(state_op), + cute.make_layout((1, 1, 1)), + ) + token_mma = cute.make_tiled_mma( + cute.make_mma_atom(token_op), + cute.make_layout((1, 1, 1)), + ) + raw_layout = sm90_utils.make_smem_layout_a( + cutlass.utils.LayoutEnum.ROW_MAJOR, + (CHUNK_SIZE, CHUNK_SIZE, _WGMMA_K), + cutlass.BFloat16, + _RAW_STAGES, + ) + g_staging_layout = cute.make_layout( + (CHUNK_SIZE, _WGMMA_K, _RAW_STAGES), + stride=( + _WGMMA_K, + 1, + CHUNK_SIZE * _WGMMA_K, + ), + ) + operand_layout_atom = warpgroup.make_smem_layout_atom( + warpgroup.SmemLayoutAtomKind.K_SW32, + cutlass.BFloat16, + ) + key_operand_layout = cute.tile_to_shape( + operand_layout_atom, + (CHUNK_SIZE, HEAD_SIZE, _INPUT_STAGES), + (0, 1, 2), + ) + factor_workspace_layout = cute.tile_to_shape( + operand_layout_atom, + (CHUNK_SIZE, HEAD_SIZE, _FACTOR_WORKSPACE_STAGES), + (0, 1, 2), + ) + token_operand_layout = cute.tile_to_shape( + operand_layout_atom, + (CHUNK_SIZE, CHUNK_SIZE, _INPUT_STAGES), + (0, 1, 2), + ) + factor_inverse_layout = cute.make_layout( + (CHUNK_SIZE, CHUNK_SIZE, _FACTOR_WORKSPACE_STAGES), + stride=( + CHUNK_SIZE, + 1, + CHUNK_SIZE * CHUNK_SIZE, + ), + ) + state_update_layout = cute.tile_to_shape( + operand_layout_atom, + (HEAD_SIZE, CHUNK_SIZE, _INPUT_STAGES), + (0, 1, 2), + ) + write_layout = cute.make_layout( + (CHUNK_SIZE, self.value_tile, _WRITE_STAGES), + stride=( + self.value_tile, + 1, + CHUNK_SIZE * self.value_tile, + ), + ) + gamma_layout = cute.make_layout( + (HEAD_SIZE, _INPUT_STAGES), + stride=(1, HEAD_SIZE), + ) + gs_delta_layout = cute.make_layout( + (HEAD_SIZE, _FACTOR_SUB_BLOCKS, _INPUT_STAGES), + stride=( + _FACTOR_SUB_BLOCKS, + 1, + HEAD_SIZE * _FACTOR_SUB_BLOCKS, + ), + ) + gs_pair_layout = cute.make_layout( + (HEAD_SIZE, _GS_PAIR_ROWS, _INPUT_STAGES), + stride=( + _GS_PAIR_ROWS, + 1, + HEAD_SIZE * _GS_PAIR_ROWS, + ), + ) + output_layout_atom = warpgroup.make_smem_layout_atom( + warpgroup.SmemLayoutAtomKind.K_SW64, + cutlass.BFloat16, + ) + output_layout = cute.tile_to_shape( + output_layout_atom, + (CHUNK_SIZE, self.value_tile, _OUTPUT_STAGES), + (0, 1, 2), + ) + + q_heads = cute.size(q, mode=[1]) + v_heads = cute.size(v, mode=[1]) + token_extent = cute.size(q, mode=[0]) + raw_q_layout = cute.make_layout( + (token_extent, HEAD_SIZE, q_heads), + stride=(q_heads * HEAD_SIZE, 1, HEAD_SIZE), + ) + raw_v_layout = cute.make_layout( + (token_extent, VALUE_SIZE, v_heads), + stride=(v_heads * VALUE_SIZE, 1, VALUE_SIZE), + ) + q_global = cute.make_tensor(q.iterator, raw_q_layout) + k_global = cute.make_tensor(k.iterator, raw_q_layout) + b_global = cute.make_tensor(b.iterator, raw_q_layout) + g_global = cute.make_tensor(g.iterator, raw_q_layout) + v_global = cute.make_tensor(v.iterator, raw_v_layout) + w_global = cute.make_tensor(w.iterator, raw_v_layout) + + output_global_layout = cute.make_layout( + (token_extent, VALUE_SIZE, v_heads), + stride=(v_heads * VALUE_SIZE, 1, VALUE_SIZE), + ) + tma_output_global = cute.make_tensor( + output.iterator, + output_global_layout, + ) + + q_atom, q_tma = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + q_global, + cute.slice_(raw_layout, (None, None, 0)), + (CHUNK_SIZE, _WGMMA_K), + ) + k_atom, k_tma = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + k_global, + cute.slice_(raw_layout, (None, None, 0)), + (CHUNK_SIZE, _WGMMA_K), + ) + b_atom, b_tma = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + b_global, + cute.slice_(raw_layout, (None, None, 0)), + (CHUNK_SIZE, _WGMMA_K), + ) + v_atom, v_tma = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + v_global, + cute.slice_(raw_layout, (None, None, 0)), + (CHUNK_SIZE, _WGMMA_K), + ) + w_atom, w_tma = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + w_global, + cute.slice_(raw_layout, (None, None, 0)), + (CHUNK_SIZE, _WGMMA_K), + ) + g_atom, g_tma = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + g_global, + cute.slice_(g_staging_layout, (None, None, 0)), + (CHUNK_SIZE, _WGMMA_K), + ) + output_atom, output_tma = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + tma_output_global, + cute.slice_(output_layout, (None, None, 0)), + (CHUNK_SIZE, self.value_tile), + ) + + @cute.struct + class SharedStorage: + qkb0_barriers: cute.struct.MemRange[ + cutlass.Int64, + 2, + ] + qkb1_barriers: cute.struct.MemRange[ + cutlass.Int64, + 2, + ] + vw0_barriers: cute.struct.MemRange[ + cutlass.Int64, + 2 * _VW_PRIVATE_STAGES, + ] + vw1_barriers: cute.struct.MemRange[ + cutlass.Int64, + 2 * _VW_PRIVATE_STAGES, + ] + raw_handoff_barriers: cute.struct.MemRange[ + cutlass.Int64, + 2, + ] + factor_ready_barriers: cute.struct.MemRange[ + cutlass.Int64, + 2 * _INPUT_STAGES, + ] + factor_done_barriers: cute.struct.MemRange[ + cutlass.Int64, + 2 * _INPUT_STAGES, + ] + output_handoff_barriers: cute.struct.MemRange[ + cutlass.Int64, + 2 * _OUTPUT_STAGES, + ] + producer_value_work_by_warp: cute.struct.MemRange[ + cutlass.Int32, + _PRODUCER_SIGNAL_WARPS, + ] + raw_q: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.BFloat16, + cute.cosize(raw_layout), + ], + 128, + ] + raw_k: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.BFloat16, + cute.cosize(raw_layout), + ], + 128, + ] + raw_b: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.BFloat16, + cute.cosize(raw_layout), + ], + 128, + ] + raw_g: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.Float32, + cute.cosize(g_staging_layout), + ], + 128, + ] + raw_v: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.BFloat16, + cute.cosize(raw_layout), + ], + 128, + ] + raw_w: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.BFloat16, + cute.cosize(raw_layout), + ], + 128, + ] + q_bar: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.BFloat16, + cute.cosize(key_operand_layout), + ], + 128, + ] + erase_bar: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.BFloat16, + cute.cosize(key_operand_layout), + ], + 128, + ] + key_tail: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.BFloat16, + cute.cosize(state_update_layout), + ], + 128, + ] + aqk_scaled: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.BFloat16, + cute.cosize(token_operand_layout), + ], + 128, + ] + akk_inverse: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.BFloat16, + cute.cosize(token_operand_layout), + ], + 128, + ] + write_value: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.BFloat16, + cute.cosize(write_layout), + ], + 128, + ] + factor_workspace: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.BFloat16, + cute.cosize(factor_workspace_layout), + ], + 128, + ] + gs_delta: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.Float32, + cute.cosize(gs_delta_layout), + ], + 128, + ] + gs_pair: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.Float16, + cute.cosize(gs_pair_layout), + ], + 128, + ] + output: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.BFloat16, + cute.cosize(output_layout), + ], + 128, + ] + gamma_end: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.Float32, + cute.cosize(gamma_layout), + ], + 128, + ] + + self.shared_storage = SharedStorage + self.dynamic_smem_bytes = SharedStorage.size_in_bytes() + self.kernel( + q_atom, + q_tma, + k_atom, + k_tma, + b_atom, + b_tma, + g_atom, + g_tma, + v_atom, + v_tma, + w_atom, + w_tma, + output_atom, + output_tma, + output, + cu_seqlens, + g, + initial_state, + final_state, + num_sequences, + num_q_heads, + num_v_heads, + total_tokens, + scale, + state_mma, + token_mma, + raw_layout, + g_staging_layout, + key_operand_layout, + factor_workspace_layout, + token_operand_layout, + factor_inverse_layout, + state_update_layout, + write_layout, + gamma_layout, + gs_delta_layout, + gs_pair_layout, + output_layout, + ).launch( + grid=( + num_sequences * num_v_heads * cutlass.Int32(VALUE_SIZE // self.value_tile), + 1, + 1, + ), + block=(self.threads_per_cta, 1, 1), + cluster=(1, 1, 1), + smem=self.dynamic_smem_bytes, + stream=stream, + min_blocks_per_mp=self.min_blocks_per_mp, + ) + + @cute.kernel + def kernel( + self, + q_atom: cute.CopyAtom, + q_tma: cute.Tensor, + k_atom: cute.CopyAtom, + k_tma: cute.Tensor, + b_atom: cute.CopyAtom, + b_tma: cute.Tensor, + g_atom: cute.CopyAtom, + g_tma: cute.Tensor, + v_atom: cute.CopyAtom, + v_tma: cute.Tensor, + w_atom: cute.CopyAtom, + w_tma: cute.Tensor, + output_atom: cute.CopyAtom, + output_tma: cute.Tensor, + output: cute.Tensor, + cu_seqlens: cute.Tensor, + g: cute.Tensor, + initial_state: cute.Tensor, + final_state: cute.Tensor, + num_sequences: cutlass.Int32, + num_q_heads: cutlass.Int32, + num_v_heads: cutlass.Int32, + total_tokens: cutlass.Int32, + scale: cutlass.Float32, + state_mma: cute.TiledMma, + token_mma: cute.TiledMma, + raw_layout: cute.ComposedLayout, + g_staging_layout: cute.Layout, + key_operand_layout: cute.ComposedLayout, + factor_workspace_layout: cute.ComposedLayout, + token_operand_layout: cute.ComposedLayout, + factor_inverse_layout: cute.Layout, + state_update_layout: cute.ComposedLayout, + write_layout: cute.Layout, + gamma_layout: cute.Layout, + gs_delta_layout: cute.Layout, + gs_pair_layout: cute.Layout, + output_layout: cute.ComposedLayout, + ) -> None: + thread, _, _ = cute.arch.thread_idx() + work_index, _, _ = cute.arch.block_idx() + warp_group = cute.arch.make_warp_uniform( + thread // _WARP_GROUP_SIZE, + ) + warp_index = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + thread_in_group = thread % _WARP_GROUP_SIZE + allocator = cutlass.utils.SmemAllocator() + storage = allocator.allocate(self.shared_storage) + producer_value_work_by_warp = storage.producer_value_work_by_warp.get_tensor( + cute.make_layout( + (_PRODUCER_SIGNAL_WARPS,), + stride=(1,), + ), + ) + + value_tiles = cutlass.Int32(VALUE_SIZE // self.value_tile) + sequence_stride = num_v_heads * value_tiles + sequence_rank = work_index // sequence_stride + value_work = work_index - sequence_rank * sequence_stride + sequence = sequence_rank + if cutlass.const_expr(self.length_ranked_sequence_order): + if warp_index == cutlass.Int32(4): + scheduled_sequence = _stable_lpt32_sequence( + cu_seqlens, + sequence_rank, + num_sequences, + thread % cutlass.Int32(32), + ) + if thread % cutlass.Int32(32) == cutlass.Int32(0): + producer_value_work_by_warp[0] = scheduled_sequence + cute.arch.sync_threads() + sequence = cute.arch.make_warp_uniform( + producer_value_work_by_warp[0], + ) + elif cutlass.const_expr(self.sequence_wave_rotation > 0): + if num_sequences > cutlass.Int32(self.sequence_wave_rotation): + sequence = sequence + cutlass.Int32( + self.sequence_wave_rotation, + ) + if sequence >= num_sequences: + sequence = sequence - num_sequences + value_head = value_work // value_tiles + value_tile_index = value_work - value_head * value_tiles + value_start = value_tile_index * cutlass.Int32(self.value_tile) + group_size = num_v_heads // num_q_heads + q_head = value_head // group_size + + sequence_start_i64 = cutlass.Int64(cu_seqlens[sequence]) + sequence_end_i64 = cutlass.Int64( + cu_seqlens[sequence + cutlass.Int32(1)], + ) + if ( + sequence_start_i64 < cutlass.Int64(0) + or sequence_end_i64 <= sequence_start_i64 + or sequence_end_i64 > cutlass.Int64(total_tokens) + ): + _device_fail_closed() + if sequence == cutlass.Int32(0) and sequence_start_i64 != cutlass.Int64(0): + _device_fail_closed() + if sequence == num_sequences - cutlass.Int32(1) and sequence_end_i64 != cutlass.Int64(total_tokens): + _device_fail_closed() + raw_q = storage.raw_q.get_tensor( + raw_layout.outer, + swizzle=raw_layout.inner, + ) + raw_k = storage.raw_k.get_tensor( + raw_layout.outer, + swizzle=raw_layout.inner, + ) + raw_b = storage.raw_b.get_tensor( + raw_layout.outer, + swizzle=raw_layout.inner, + ) + raw_g = storage.raw_g.get_tensor(g_staging_layout) + raw_v = storage.raw_v.get_tensor( + raw_layout.outer, + swizzle=raw_layout.inner, + ) + raw_w = storage.raw_w.get_tensor( + raw_layout.outer, + swizzle=raw_layout.inner, + ) + shared_q = storage.q_bar.get_tensor( + key_operand_layout.outer, + swizzle=key_operand_layout.inner, + ) + shared_erase = storage.erase_bar.get_tensor( + key_operand_layout.outer, + swizzle=key_operand_layout.inner, + ) + shared_key_tail = storage.key_tail.get_tensor( + state_update_layout.outer, + swizzle=state_update_layout.inner, + ) + shared_aqk = storage.aqk_scaled.get_tensor( + token_operand_layout.outer, + swizzle=token_operand_layout.inner, + ) + shared_akk = storage.akk_inverse.get_tensor( + token_operand_layout.outer, + swizzle=token_operand_layout.inner, + ) + shared_write = storage.write_value.get_tensor(write_layout) + shared_output = storage.output.get_tensor( + output_layout.outer, + swizzle=output_layout.inner, + ) + shared_factor_k = storage.factor_workspace.get_tensor( + factor_workspace_layout.outer, + swizzle=factor_workspace_layout.inner, + ) + # The FP16 Gram/inverse aliases the raw-G staging arena. Raw G for + # chunk c is dead once both state warp groups finish preparation + # (which factor_ready orders before the factor stage), and WG0 issues + # the chunk c+1 G loads only after the chunk c factor stage returns, + # so the alias never overlaps a live read or an in-flight TMA write. + raw_g_address = storage.raw_g.data_ptr().toint() + shared_inverse = cute.make_tensor( + cute.make_ptr( + cutlass.Float16, + raw_g_address, + cute.AddressSpace.smem, + assumed_align=128, + ), + factor_inverse_layout, + ) + shared_gamma_end = storage.gamma_end.get_tensor(gamma_layout) + shared_gs_delta = storage.gs_delta.get_tensor(gs_delta_layout) + shared_gs_pair = storage.gs_pair.get_tensor(gs_pair_layout) + qkb0_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.qkb0_barriers.data_ptr(), + num_stages=1, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + _PRODUCER_SIGNAL_WARPS, + ), + tx_count=_QKBG_TRANSACTION_BYTES, + cta_layout_vmnk=cute.make_layout((1, 1, 1, 1)), + ) + qkb1_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.qkb1_barriers.data_ptr(), + num_stages=1, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + _PRODUCER_SIGNAL_WARPS, + ), + tx_count=_QKBG_TRANSACTION_BYTES, + cta_layout_vmnk=cute.make_layout((1, 1, 1, 1)), + ) + vw0_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.vw0_barriers.data_ptr(), + num_stages=_VW_PRIVATE_STAGES, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + _PRODUCER_SIGNAL_WARPS, + ), + tx_count=_VW_TRANSACTION_BYTES, + cta_layout_vmnk=cute.make_layout((1, 1, 1, 1)), + ) + vw1_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.vw1_barriers.data_ptr(), + num_stages=_VW_PRIVATE_STAGES, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + _PRODUCER_SIGNAL_WARPS, + ), + tx_count=_VW_TRANSACTION_BYTES, + cta_layout_vmnk=cute.make_layout((1, 1, 1, 1)), + ) + raw_handoff = pipeline.PipelineAsync.create( + barrier_storage=storage.raw_handoff_barriers.data_ptr(), + num_stages=1, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + _WARP_GROUP_SIZE, + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 2 * _WARP_GROUP_SIZE, + ), + ) + factor_ready_handoff = pipeline.PipelineAsync.create( + barrier_storage=storage.factor_ready_barriers.data_ptr(), + num_stages=_INPUT_STAGES, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 2 * _WARP_GROUP_SIZE, + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + _WARP_GROUP_SIZE, + ), + ) + factor_done_handoff = pipeline.PipelineAsync.create( + barrier_storage=storage.factor_done_barriers.data_ptr(), + num_stages=_INPUT_STAGES, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + _WARP_GROUP_SIZE, + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 2 * _WARP_GROUP_SIZE, + ), + ) + output_handoff = pipeline.PipelineAsync.create( + barrier_storage=storage.output_handoff_barriers.data_ptr(), + num_stages=_OUTPUT_STAGES, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + _WARP_GROUP_SIZE if self.single_state_owner else 2 * _WARP_GROUP_SIZE, + ), + consumer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + _WARP_GROUP_SIZE, + ), + ) + store_pipeline = pipeline.PipelineTmaStore.create( + num_stages=_OUTPUT_STAGES, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + _WARP_GROUP_SIZE, + ), + ) + + if warp_index == cutlass.Int32(0): + cpasync.prefetch_descriptor(q_atom) + cpasync.prefetch_descriptor(k_atom) + cpasync.prefetch_descriptor(b_atom) + cpasync.prefetch_descriptor(g_atom) + cpasync.prefetch_descriptor(v_atom) + cpasync.prefetch_descriptor(w_atom) + cpasync.prefetch_descriptor(output_atom) + + if warp_group == cutlass.Int32(0): + cute.arch.warpgroup_reg_dealloc( + _PRODUCER_REGISTER_TARGET, + ) + producer_work_index, _, _ = cute.arch.block_idx() + producer_sequence_rank = producer_work_index // sequence_stride + producer_value_work = producer_work_index - producer_sequence_rank * sequence_stride + producer_sequence = sequence + if cutlass.const_expr( + not self.length_ranked_sequence_order and self.sequence_wave_rotation > 0, + ): + producer_sequence = producer_sequence_rank + if num_sequences > cutlass.Int32( + self.sequence_wave_rotation, + ): + producer_sequence = producer_sequence + cutlass.Int32( + self.sequence_wave_rotation, + ) + if producer_sequence >= num_sequences: + producer_sequence = producer_sequence - num_sequences + if thread_in_group % cutlass.Int32(32) == cutlass.Int32(0): + producer_value_work_by_warp[warp_index] = producer_value_work + cute.arch.sync_warp() + producer_sequence_start = cutlass.Int32( + cu_seqlens[producer_sequence], + ) + producer_sequence_end = cutlass.Int32( + cu_seqlens[producer_sequence + cutlass.Int32(1)], + ) + producer_sequence_chunks = ( + producer_sequence_end - producer_sequence_start + cutlass.Int32(CHUNK_SIZE - 1) + ) // cutlass.Int32( + CHUNK_SIZE, + ) + output_wait = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, + _OUTPUT_STAGES, + ) + output_release = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, + _OUTPUT_STAGES, + ) + for pipeline_step in cutlass.range( + producer_sequence_chunks + cutlass.Int32(1), + unroll=1, + ): + factor_stage = cutlass.Int32(0) + factor_valid_tokens = cutlass.Int32(0) + if pipeline_step < producer_sequence_chunks: + qkb0_producer = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, + 1, + ) + qkb1_producer = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, + 1, + ) + local_chunk = pipeline_step + chunk_start = producer_sequence_start + local_chunk * cutlass.Int32(CHUNK_SIZE) + valid_tokens = cutlass.Int32(CHUNK_SIZE) + if local_chunk + cutlass.Int32(1) == producer_sequence_chunks: + valid_tokens = producer_sequence_end - chunk_start + factor_stage = pipeline_step % cutlass.Int32(_INPUT_STAGES) + factor_valid_tokens = valid_tokens + raw_handoff.producer_acquire( + pipeline.PipelineState( + 1, + pipeline_step, + cutlass.Int32(0), + cutlass.Int32(1) - pipeline_step % cutlass.Int32(2), + ), + ) + + q_use = cute.domain_offset( + (chunk_start, cutlass.Int32(0), cutlass.Int32(0)), + q_tma, + ) + k_use = cute.domain_offset( + (chunk_start, cutlass.Int32(0), cutlass.Int32(0)), + k_tma, + ) + b_use = cute.domain_offset( + (chunk_start, cutlass.Int32(0), cutlass.Int32(0)), + b_tma, + ) + g_use = cute.domain_offset( + (chunk_start, cutlass.Int32(0), cutlass.Int32(0)), + g_tma, + ) + q_tiles = cute.local_tile( + q_use[None, None, q_head], + (CHUNK_SIZE, _WGMMA_K), + (None, None), + ) + k_tiles = cute.local_tile( + k_use[None, None, q_head], + (CHUNK_SIZE, _WGMMA_K), + (None, None), + ) + b_tiles = cute.local_tile( + b_use[None, None, q_head], + (CHUNK_SIZE, _WGMMA_K), + (None, None), + ) + g_tiles = cute.local_tile( + g_use[None, None, q_head], + (CHUNK_SIZE, _WGMMA_K), + (None, None), + ) + q_smem, q_gmem = cpasync.tma_partition( + q_atom, + 0, + cute.make_layout(1), + cute.group_modes(raw_q, 0, 2), + cute.group_modes(q_tiles, 0, 2), + ) + k_smem, k_gmem = cpasync.tma_partition( + k_atom, + 0, + cute.make_layout(1), + cute.group_modes(raw_k, 0, 2), + cute.group_modes(k_tiles, 0, 2), + ) + b_smem, b_gmem = cpasync.tma_partition( + b_atom, + 0, + cute.make_layout(1), + cute.group_modes(raw_b, 0, 2), + cute.group_modes(b_tiles, 0, 2), + ) + g_smem, g_gmem = cpasync.tma_partition( + g_atom, + 0, + cute.make_layout(1), + cute.group_modes(raw_g, 0, 2), + cute.group_modes(g_tiles, 0, 2), + ) + + if warp_index == cutlass.Int32(0): + qkb0_pipeline.producer_acquire(qkb0_producer) + qkb0_barrier = qkb0_pipeline.producer_get_barrier( + qkb0_producer, + ) + cute.copy( + q_atom, + q_gmem[(None, 0, cutlass.Int32(0))], + q_smem[(None, cutlass.Int32(0))], + tma_bar_ptr=qkb0_barrier, + ) + cute.copy( + k_atom, + k_gmem[(None, 0, cutlass.Int32(0))], + k_smem[(None, cutlass.Int32(0))], + tma_bar_ptr=qkb0_barrier, + ) + cute.copy( + b_atom, + b_gmem[(None, 0, cutlass.Int32(0))], + b_smem[(None, cutlass.Int32(0))], + tma_bar_ptr=qkb0_barrier, + ) + cute.copy( + g_atom, + g_gmem[(None, 0, cutlass.Int32(0))], + g_smem[(None, cutlass.Int32(0))], + tma_bar_ptr=qkb0_barrier, + ) + qkb0_pipeline.producer_commit(qkb0_producer) + qkb0_producer.advance() + + qkb1_pipeline.producer_acquire(qkb1_producer) + qkb1_barrier = qkb1_pipeline.producer_get_barrier( + qkb1_producer, + ) + cute.copy( + q_atom, + q_gmem[ + ( + None, + 0, + cutlass.Int32(_QKB_STREAM_TILES), + ) + ], + q_smem[(None, cutlass.Int32(1))], + tma_bar_ptr=qkb1_barrier, + ) + cute.copy( + k_atom, + k_gmem[ + ( + None, + 0, + cutlass.Int32(_QKB_STREAM_TILES), + ) + ], + k_smem[(None, cutlass.Int32(1))], + tma_bar_ptr=qkb1_barrier, + ) + cute.copy( + b_atom, + b_gmem[ + ( + None, + 0, + cutlass.Int32(_QKB_STREAM_TILES), + ) + ], + b_smem[(None, cutlass.Int32(1))], + tma_bar_ptr=qkb1_barrier, + ) + cute.copy( + g_atom, + g_gmem[ + ( + None, + 0, + cutlass.Int32(_QKB_STREAM_TILES), + ) + ], + g_smem[(None, cutlass.Int32(1))], + tma_bar_ptr=qkb1_barrier, + ) + qkb1_pipeline.producer_commit(qkb1_producer) + qkb1_producer.advance() + + # State may begin consuming Q/K/B/G after each private + # stream has its first tile in flight. V/W is deliberately + # issued only after this generation's factor is complete, + # so its single shared write buffer cannot obstruct the + # next factor generation. + raw_handoff.producer_commit( + pipeline.PipelineState( + 1, + pipeline_step, + cutlass.Int32(0), + cutlass.Int32(1) - pipeline_step % cutlass.Int32(2), + ), + ) + + for local_key_tile in cutlass.range( + 1, + _QKB_STREAM_TILES, + unroll=1, + ): + if warp_index == cutlass.Int32(0): + qkb0_pipeline.producer_acquire(qkb0_producer) + qkb0_barrier = qkb0_pipeline.producer_get_barrier( + qkb0_producer, + ) + cute.copy( + q_atom, + q_gmem[(None, 0, local_key_tile)], + q_smem[(None, cutlass.Int32(0))], + tma_bar_ptr=qkb0_barrier, + ) + cute.copy( + k_atom, + k_gmem[(None, 0, local_key_tile)], + k_smem[(None, cutlass.Int32(0))], + tma_bar_ptr=qkb0_barrier, + ) + cute.copy( + b_atom, + b_gmem[(None, 0, local_key_tile)], + b_smem[(None, cutlass.Int32(0))], + tma_bar_ptr=qkb0_barrier, + ) + cute.copy( + g_atom, + g_gmem[(None, 0, local_key_tile)], + g_smem[(None, cutlass.Int32(0))], + tma_bar_ptr=qkb0_barrier, + ) + qkb0_pipeline.producer_commit(qkb0_producer) + qkb0_producer.advance() + + high_key_tile = local_key_tile + cutlass.Int32(_QKB_STREAM_TILES) + qkb1_pipeline.producer_acquire(qkb1_producer) + qkb1_barrier = qkb1_pipeline.producer_get_barrier( + qkb1_producer, + ) + cute.copy( + q_atom, + q_gmem[(None, 0, high_key_tile)], + q_smem[(None, cutlass.Int32(1))], + tma_bar_ptr=qkb1_barrier, + ) + cute.copy( + k_atom, + k_gmem[(None, 0, high_key_tile)], + k_smem[(None, cutlass.Int32(1))], + tma_bar_ptr=qkb1_barrier, + ) + cute.copy( + b_atom, + b_gmem[(None, 0, high_key_tile)], + b_smem[(None, cutlass.Int32(1))], + tma_bar_ptr=qkb1_barrier, + ) + cute.copy( + g_atom, + g_gmem[(None, 0, high_key_tile)], + g_smem[(None, cutlass.Int32(1))], + tma_bar_ptr=qkb1_barrier, + ) + qkb1_pipeline.producer_commit(qkb1_producer) + qkb1_producer.advance() + + if pipeline_step > cutlass.Int32(0): + # Produce next-generation QKB above before draining the + # current V/W generation. This matches State's + # prepare-next-before-materialize-current order and + # removes the single-stage V/W/QKB circular wait in r22. + vw_chunk = pipeline_step - cutlass.Int32(1) + vw0_producer = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, + _VW_PRIVATE_STAGES, + ) + vw1_producer = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, + _VW_PRIVATE_STAGES, + ) + vw_chunk_start = producer_sequence_start + vw_chunk * cutlass.Int32(CHUNK_SIZE) + v_use = cute.domain_offset( + ( + vw_chunk_start, + value_start, + cutlass.Int32(0), + ), + v_tma, + ) + w_use = cute.domain_offset( + ( + vw_chunk_start, + value_start, + cutlass.Int32(0), + ), + w_tma, + ) + v_tiles = cute.local_tile( + v_use[None, None, value_head], + (CHUNK_SIZE, _WGMMA_K), + (None, None), + ) + w_tiles = cute.local_tile( + w_use[None, None, value_head], + (CHUNK_SIZE, _WGMMA_K), + (None, None), + ) + v_smem, v_gmem = cpasync.tma_partition( + v_atom, + 0, + cute.make_layout(1), + cute.group_modes(raw_v, 0, 2), + cute.group_modes(v_tiles, 0, 2), + ) + w_smem, w_gmem = cpasync.tma_partition( + w_atom, + 0, + cute.make_layout(1), + cute.group_modes(raw_w, 0, 2), + cute.group_modes(w_tiles, 0, 2), + ) + for local_value_tile in cutlass.range( + _VALUE_TILES // 2, + unroll=1, + ): + if warp_index == cutlass.Int32(0): + vw0_pipeline.producer_acquire(vw0_producer) + vw0_barrier = vw0_pipeline.producer_get_barrier( + vw0_producer, + ) + cute.copy( + v_atom, + v_gmem[(None, 0, local_value_tile)], + v_smem[(None, 0)], + tma_bar_ptr=vw0_barrier, + ) + cute.copy( + w_atom, + w_gmem[(None, 0, local_value_tile)], + w_smem[(None, 0)], + tma_bar_ptr=vw0_barrier, + ) + vw0_pipeline.producer_commit(vw0_producer) + vw0_producer.advance() + + if cutlass.const_expr(not self.single_state_owner): + high_value_tile = local_value_tile + cutlass.Int32(_VALUE_TILES // 2) + vw1_pipeline.producer_acquire(vw1_producer) + vw1_barrier = vw1_pipeline.producer_get_barrier( + vw1_producer, + ) + cute.copy( + v_atom, + v_gmem[(None, 0, high_value_tile)], + v_smem[(None, 1)], + tma_bar_ptr=vw1_barrier, + ) + cute.copy( + w_atom, + w_gmem[(None, 0, high_value_tile)], + w_smem[(None, 1)], + tma_bar_ptr=vw1_barrier, + ) + vw1_pipeline.producer_commit(vw1_producer) + vw1_producer.advance() + + if pipeline_step < producer_sequence_chunks: + factor_consumer_state = pipeline.PipelineState( + _INPUT_STAGES, + pipeline_step, + pipeline_step % cutlass.Int32(_INPUT_STAGES), + (pipeline_step // cutlass.Int32(_INPUT_STAGES)) % cutlass.Int32(2), + ) + factor_producer_state = pipeline.PipelineState( + _INPUT_STAGES, + pipeline_step, + pipeline_step % cutlass.Int32(_INPUT_STAGES), + cutlass.Int32(1) - ((pipeline_step // cutlass.Int32(_INPUT_STAGES)) % cutlass.Int32(2)), + ) + factor_ready_handoff.consumer_wait( + factor_consumer_state, + ) + factor_done_handoff.producer_acquire( + factor_producer_state, + ) + + factor_q = shared_q[None, None, factor_stage] + factor_erase = shared_erase[None, None, factor_stage] + factor_k = shared_factor_k[ + None, + None, + cutlass.Int32(0), + ] + factor_gs_delta = shared_gs_delta[ + None, + None, + factor_stage, + ] + factor_gs_pair = shared_gs_pair[ + None, + None, + factor_stage, + ] + factor_aqk = shared_aqk[None, None, factor_stage] + factor_akk = shared_akk[None, None, factor_stage] + factor_inverse = shared_inverse[ + None, + None, + cutlass.Int32(0), + ] + fill_static_upper = pipeline_step < cutlass.Int32( + _INPUT_STAGES, + ) + _publish_factor_blocks( + thread_in_group, + factor_q, + factor_erase, + factor_k, + factor_gs_delta, + factor_gs_pair, + factor_aqk, + factor_inverse, + factor_valid_tokens, + scale, + fill_static_upper, + ) + if fill_static_upper: + # One-time zero publication of the akk upper blocks: + # the triangular convert below never rewrites them and + # the arena is private per stage. + for pair in cutlass.range_constexpr( + len(_FACTOR_UPPER_PAIRS), + ): + upper_query = _FACTOR_UPPER_PAIRS[pair][0] + upper_key = _FACTOR_UPPER_PAIRS[pair][1] + for linear in cutlass.range( + thread_in_group, + _FACTOR_BLOCK * _FACTOR_BLOCK, + _WARP_GROUP_SIZE, + unroll=1, + ): + local_row = linear // cutlass.Int32(_FACTOR_BLOCK) + local_column = linear % cutlass.Int32(_FACTOR_BLOCK) + factor_akk[ + upper_query * _FACTOR_BLOCK + local_row, + upper_key * _FACTOR_BLOCK + local_column, + ] = cutlass.BFloat16(0.0) + cute.arch.barrier( + barrier_id=_INVERSE_BARRIER, + number_of_threads=_WARP_GROUP_SIZE, + ) + CollectiveInverse().run( + factor_inverse, + _INVERSE_BARRIER, + ) + cute.arch.barrier( + barrier_id=_INVERSE_BARRIER, + number_of_threads=_WARP_GROUP_SIZE, + ) + for row_block in cutlass.range_constexpr( + _FACTOR_SUB_BLOCKS, + ): + # The inverse is unit-lower-triangular blockwise, so + # only the lower band through the diagonal block needs + # conversion; the upper zeros were published once. + band_columns = (row_block + 1) * _FACTOR_BLOCK + for linear in cutlass.range( + thread_in_group, + _FACTOR_BLOCK * band_columns, + _WARP_GROUP_SIZE, + unroll=1, + ): + row = row_block * _FACTOR_BLOCK + linear // band_columns + column = linear % band_columns + factor_akk[row, column] = cutlass.BFloat16( + factor_inverse[row, column], + ) + # Last read of the inverse, which aliases the raw-G arena. + # Rendezvous before any warp can leave this iteration and + # issue the next chunk's raw-G TMA over it: raw_handoff for + # chunk c+1 is already released by the state warp groups at + # the end of prep(c), so nothing else holds warp 0 back. + cute.arch.barrier( + barrier_id=_INVERSE_BARRIER, + number_of_threads=_WARP_GROUP_SIZE, + ) + cute.arch.fence_proxy("async.shared", space="cta") + factor_done_handoff.producer_commit( + factor_producer_state, + ) + factor_ready_handoff.consumer_release( + factor_consumer_state, + ) + if pipeline_step > cutlass.Int32(0): + local_chunk = pipeline_step - cutlass.Int32(1) + chunk_start = producer_sequence_start + local_chunk * cutlass.Int32(CHUNK_SIZE) + valid_tokens = cutlass.Int32(CHUNK_SIZE) + if local_chunk + cutlass.Int32(1) == producer_sequence_chunks: + valid_tokens = producer_sequence_end - chunk_start + + output_handoff.consumer_wait(output_wait) + if valid_tokens == cutlass.Int32(CHUNK_SIZE): + output_view = cute.domain_offset( + ( + chunk_start, + value_start, + cutlass.Int32(0), + ), + output_tma, + ) + output_tile = cute.zipped_divide( + output_view[None, None, value_head], + (CHUNK_SIZE, self.value_tile), + )[ + ( + (None, None), + (cutlass.Int32(0), cutlass.Int32(0)), + ) + ] + output_stage = shared_output[ + None, + None, + output_wait.index, + ] + output_smem, output_gmem = cpasync.tma_partition( + output_atom, + 0, + cute.make_layout(1), + cute.group_modes(output_stage, 0, 2), + cute.group_modes(output_tile, 0, 2), + ) + if warp_index == cutlass.Int32(0): + cute.arch.fence_view_async_shared() + cute.copy( + output_atom, + output_smem, + output_gmem, + ) + store_pipeline.producer_commit() + if local_chunk > cutlass.Int32(0): + store_pipeline.producer_acquire() + cute.arch.barrier( + barrier_id=_STORE_WG_BARRIER, + number_of_threads=_WARP_GROUP_SIZE, + ) + if local_chunk > cutlass.Int32(0): + output_handoff.consumer_release(output_release) + output_release.advance() + if local_chunk + cutlass.Int32(1) == producer_sequence_chunks: + if warp_index == cutlass.Int32(0): + store_pipeline.producer_tail() + cute.arch.barrier( + barrier_id=_STORE_WG_BARRIER, + number_of_threads=_WARP_GROUP_SIZE, + ) + output_handoff.consumer_release(output_release) + output_release.advance() + else: + if warp_index == cutlass.Int32(0): + store_pipeline.producer_tail() + cute.arch.barrier( + barrier_id=_STORE_WG_BARRIER, + number_of_threads=_WARP_GROUP_SIZE, + ) + if local_chunk > cutlass.Int32(0): + output_handoff.consumer_release(output_release) + output_release.advance() + tail_value_work = producer_value_work_by_warp[warp_index] + tail_value_head = tail_value_work // value_tiles + tail_value_tile_index = tail_value_work - tail_value_head * value_tiles + tail_value_start = tail_value_tile_index * cutlass.Int32(self.value_tile) + for linear in cutlass.range( + thread_in_group, + valid_tokens * cutlass.Int32(self.value_tile), + _WARP_GROUP_SIZE, + unroll=1, + ): + local_token = linear // self.value_tile + value_index = linear % self.value_tile + output[ + chunk_start + local_token, + tail_value_head, + tail_value_start + value_index, + ] = shared_output[ + local_token, + value_index, + output_wait.index, + ] + cute.arch.barrier( + barrier_id=_STORE_WG_BARRIER, + number_of_threads=_WARP_GROUP_SIZE, + ) + output_handoff.consumer_release(output_release) + output_release.advance() + output_wait.advance() + + else: + cute.arch.warpgroup_reg_alloc(_STATE_REGISTER_TARGET) + state_sequence_start = cutlass.Int32(cu_seqlens[sequence]) + state_sequence_end = cutlass.Int32( + cu_seqlens[sequence + cutlass.Int32(1)], + ) + state_sequence_chunks = ( + state_sequence_end - state_sequence_start + cutlass.Int32(CHUNK_SIZE - 1) + ) // cutlass.Int32( + CHUNK_SIZE, + ) + state_slab = warp_group - cutlass.Int32(1) + shared_value_start = state_slab * cutlass.Int32(self.state_value_tile) + if cutlass.const_expr(self.single_state_owner): + shared_value_start = cutlass.Int32(0) + state_value_start = value_start + shared_value_start + + state_thread = state_mma.get_slice(thread_in_group) + token_thread = token_mma.get_slice(thread_in_group) + state_coordinates = state_thread.partition_C( + cute.make_identity_tensor( + (self.state_value_tile, HEAD_SIZE), + ), + ) + token_coordinates = token_thread.partition_C( + cute.make_identity_tensor( + (self.state_value_tile, CHUNK_SIZE), + ), + ) + state_accumulator = state_thread.make_fragment_C( + state_thread.partition_shape_C( + (self.state_value_tile, HEAD_SIZE), + ), + ) + for element in cutlass.range_constexpr( + cute.size(state_accumulator), + ): + value_index, key_index = state_coordinates[element] + state_value = cutlass.Float32(0.0) + if cutlass.const_expr(self.has_initial_state): + state_value = cutlass.Float32( + initial_state[ + sequence, + value_head, + state_value_start + value_index, + key_index, + ], + ) + state_accumulator[element] = state_value + + output_producer = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, + _OUTPUT_STAGES, + ) + for state_step in cutlass.range( + state_sequence_chunks + cutlass.Int32(1), + unroll=1, + ): + current_stage = cutlass.Int32(0) + current_valid_tokens = cutlass.Int32(0) + current_chunk = state_step - cutlass.Int32(1) + factor_done_consumer_state = pipeline.PipelineState( + _INPUT_STAGES, + current_chunk, + current_chunk % cutlass.Int32(_INPUT_STAGES), + (current_chunk // cutlass.Int32(_INPUT_STAGES)) % cutlass.Int32(2), + ) + if state_step > cutlass.Int32(0): + current_stage = current_chunk % cutlass.Int32(_INPUT_STAGES) + current_chunk_start = state_sequence_start + current_chunk * cutlass.Int32(CHUNK_SIZE) + current_valid_tokens = cutlass.Int32(CHUNK_SIZE) + if current_chunk + cutlass.Int32(1) == state_sequence_chunks: + current_valid_tokens = state_sequence_end - current_chunk_start + + factor_done_handoff.consumer_wait( + factor_done_consumer_state, + ) + factor_done_handoff.consumer_release( + factor_done_consumer_state, + ) + + if state_step < state_sequence_chunks: + prepare_chunk = state_step + prepare_stage = prepare_chunk % cutlass.Int32(_INPUT_STAGES) + prepare_chunk_start = state_sequence_start + prepare_chunk * cutlass.Int32(CHUNK_SIZE) + prepare_valid_tokens = cutlass.Int32(CHUNK_SIZE) + if prepare_chunk + cutlass.Int32(1) == state_sequence_chunks: + prepare_valid_tokens = state_sequence_end - prepare_chunk_start + + qkb_consumer = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, + 1, + ) + raw_handoff.consumer_wait( + pipeline.PipelineState( + 1, + prepare_chunk, + cutlass.Int32(0), + prepare_chunk % cutlass.Int32(2), + ), + ) + factor_ready_handoff.producer_acquire( + pipeline.PipelineState( + _INPUT_STAGES, + prepare_chunk, + (prepare_chunk % cutlass.Int32(_INPUT_STAGES)), + cutlass.Int32(1) - ((prepare_chunk // cutlass.Int32(_INPUT_STAGES)) % cutlass.Int32(2)), + ), + ) + for local_key_tile in cutlass.range( + _QKB_STREAM_TILES, + unroll=1, + ): + if state_slab == cutlass.Int32(0): + qkb_ready = qkb0_pipeline.consumer_try_wait( + qkb_consumer, + ) + qkb0_pipeline.consumer_wait( + qkb_consumer, + qkb_ready, + ) + else: + qkb_ready = qkb1_pipeline.consumer_try_wait( + qkb_consumer, + ) + qkb1_pipeline.consumer_wait( + qkb_consumer, + qkb_ready, + ) + + raw_stage = state_slab + key_tile = state_slab * cutlass.Int32(_QKB_STREAM_TILES) + local_key_tile + if cutlass.const_expr(self.subgroup_prefix): + # Sixteen independent eight-lane subgroups cover + # one 16-channel raw-G tile. Each lane scans one + # consecutive eight-token segment, then a + # subgroup shuffle scan distributes the carry. + subgroup_lane = thread_in_group % cutlass.Int32(8) + tile_channel = thread_in_group // cutlass.Int32(8) + token_base = subgroup_lane * cutlass.Int32(8) + local_prefix = cute.make_rmem_tensor( + 8, + cutlass.Float32, + ) + segment_total = cutlass.Float32(0.0) + for local_index in cutlass.range_constexpr(8): + local_token = token_base + cutlass.Int32(local_index) + if local_token < prepare_valid_tokens: + segment_total = segment_total + cutlass.Float32( + raw_g[ + local_token, + tile_channel, + raw_stage, + ], + ) + local_prefix[local_index] = segment_total + + inclusive_segment_total = segment_total + for log_offset in cutlass.range_constexpr(3): + offset = 1 << log_offset + prior = cute.arch.shuffle_sync_up( + inclusive_segment_total, + offset, + mask_and_clamp=0, + ) + if subgroup_lane >= cutlass.Int32(offset): + inclusive_segment_total = inclusive_segment_total + prior + if cutlass.const_expr( + self.subgroup_exclusive_carry, + ): + carry = inclusive_segment_total - segment_total + else: + carry = cutlass.Float32(0.0) + prior_segment_total = cute.arch.shuffle_sync_up( + inclusive_segment_total, + 1, + mask_and_clamp=0, + ) + if subgroup_lane > cutlass.Int32(0): + carry = prior_segment_total + + for local_index in cutlass.range_constexpr(8): + local_token = token_base + cutlass.Int32(local_index) + if local_token < prepare_valid_tokens: + raw_g[ + local_token, + tile_channel, + raw_stage, + ] = local_prefix[local_index] + carry + else: + raw_g[ + local_token, + tile_channel, + raw_stage, + ] = cutlass.Float32(0.0) + elif thread_in_group < cutlass.Int32(_WGMMA_K): + prefix = cutlass.Float32(0.0) + for local_token in cutlass.range_constexpr( + CHUNK_SIZE, + ): + if cutlass.Int32(local_token) < prepare_valid_tokens: + prefix = prefix + cutlass.Float32( + raw_g[ + local_token, + thread_in_group, + raw_stage, + ], + ) + raw_g[ + local_token, + thread_in_group, + raw_stage, + ] = prefix + else: + raw_g[ + local_token, + thread_in_group, + raw_stage, + ] = cutlass.Float32(0.0) + if state_slab == cutlass.Int32(0): + cute.arch.barrier( + barrier_id=_STATE0_PREFIX_BARRIER, + number_of_threads=_WARP_GROUP_SIZE, + ) + else: + cute.arch.barrier( + barrier_id=_STATE1_PREFIX_BARRIER, + number_of_threads=_WARP_GROUP_SIZE, + ) + # Block-boundary decay ratios for this 16-channel + # tile: delta[0] = exp(Gs(0)) and + # delta[m] = exp(Gs(m) - Gs(m-1)), all <= 1. Fully + # invalid tail blocks store 1 so downstream running + # products stay exact no-ops. + if thread_in_group < cutlass.Int32(_WGMMA_K * _FACTOR_SUB_BLOCKS): + tile_channel = thread_in_group % cutlass.Int32(_WGMMA_K) + sub_block = thread_in_group // cutlass.Int32(_WGMMA_K) + key_channel = key_tile * cutlass.Int32(_WGMMA_K) + tile_channel + block_start = sub_block * cutlass.Int32(_FACTOR_BLOCK) + delta_value = cutlass.Float32(1.0) + if block_start < prepare_valid_tokens: + block_start_g = cutlass.Float32( + raw_g[ + block_start, + tile_channel, + raw_stage, + ], + ) + previous_g = cutlass.Float32(0.0) + if sub_block > cutlass.Int32(0): + previous_g = cutlass.Float32( + raw_g[ + block_start - cutlass.Int32(_FACTOR_BLOCK), + tile_channel, + raw_stage, + ], + ) + delta_value = cute.math.exp2( + (block_start_g - previous_g) * cutlass.Float32(_INV_LN2), + fastmath=True, + ) + shared_gs_delta[ + key_channel, + sub_block, + prepare_stage, + ] = delta_value + if sub_block == cutlass.Int32(0): + # Precompute the distance-two and distance-three + # pair products with the same left-to-right + # multiply order the factor fold used inline, so + # the published values stay bitwise identical. + previous_g = cutlass.Float32( + raw_g[ + cutlass.Int32(0), + tile_channel, + raw_stage, + ], + ) + block_deltas = cute.make_rmem_tensor( + _FACTOR_SUB_BLOCKS - 1, + cutlass.Float32, + ) + for later_block in cutlass.range_constexpr( + 1, + _FACTOR_SUB_BLOCKS, + ): + later_start = cutlass.Int32( + later_block * _FACTOR_BLOCK, + ) + later_delta = cutlass.Float32(1.0) + if later_start < prepare_valid_tokens: + later_g = cutlass.Float32( + raw_g[ + later_start, + tile_channel, + raw_stage, + ], + ) + later_delta = cute.math.exp2( + (later_g - previous_g) * cutlass.Float32(_INV_LN2), + fastmath=True, + ) + previous_g = later_g + block_deltas[later_block - 1] = later_delta + product_two_low = block_deltas[0] * block_deltas[1] + product_two_high = block_deltas[1] * block_deltas[2] + product_three = product_two_low * block_deltas[2] + shared_gs_pair[ + key_channel, + cutlass.Int32(0), + prepare_stage, + ] = cutlass.Float16(product_two_low) + shared_gs_pair[ + key_channel, + cutlass.Int32(1), + prepare_stage, + ] = cutlass.Float16(product_two_high) + shared_gs_pair[ + key_channel, + cutlass.Int32(2), + prepare_stage, + ] = cutlass.Float16(product_three) + # Unroll two token rows so the LDS -> FMUL -> EX2 -> + # convert chains of neighbouring iterations overlap; + # this prep latency sits on the steady-state critical + # ring ahead of every factor generation. + for linear in cutlass.range( + thread_in_group, + CHUNK_SIZE * _WGMMA_K, + _WARP_GROUP_SIZE, + unroll=2, + ): + local_token = linear // _WGMMA_K + tile_channel = linear % _WGMMA_K + key_channel = key_tile * cutlass.Int32(_WGMMA_K) + tile_channel + q_value = cutlass.BFloat16(0.0) + erase_value = cutlass.BFloat16(0.0) + if local_token < prepare_valid_tokens: + g_value = cutlass.Float32( + raw_g[ + local_token, + tile_channel, + raw_stage, + ], + ) + block_start_g = cutlass.Float32( + raw_g[ + (local_token // cutlass.Int32(_FACTOR_BLOCK)) * cutlass.Int32(_FACTOR_BLOCK), + tile_channel, + raw_stage, + ], + ) + gamma = cute.math.exp2( + (g_value - block_start_g) * cutlass.Float32(_INV_LN2), + fastmath=True, + ) + raw_k_value = cutlass.Float32( + raw_k[ + local_token, + tile_channel, + raw_stage, + ], + ) + q_value = cutlass.BFloat16( + cutlass.Float32( + raw_q[ + local_token, + tile_channel, + raw_stage, + ], + ) + * gamma, + ) + erase_value = cutlass.BFloat16( + cutlass.Float32( + raw_b[ + local_token, + tile_channel, + raw_stage, + ], + ) + * raw_k_value + * gamma, + ) + shared_q[ + local_token, + key_channel, + prepare_stage, + ] = q_value + shared_erase[ + local_token, + key_channel, + prepare_stage, + ] = erase_value + + if cutlass.const_expr(self.retain_final_tail) and ( + local_key_tile + cutlass.Int32(1) == cutlass.Int32(_QKB_STREAM_TILES) + ): + retained_key_tail = cute.make_rmem_tensor( + 8, + cutlass.BFloat16, + ) + retained_gamma_end = cutlass.Float32(0.0) + for local_index in cutlass.range_constexpr(8): + linear = thread_in_group + cutlass.Int32( + local_index * _WARP_GROUP_SIZE, + ) + local_token = linear // _WGMMA_K + tile_channel = linear % _WGMMA_K + key_channel = key_tile * cutlass.Int32(_WGMMA_K) + tile_channel + key_value = cutlass.BFloat16(0.0) + factor_key_value = cutlass.BFloat16(0.0) + if local_token < prepare_valid_tokens: + g_value = cutlass.Float32( + raw_g[ + local_token, + tile_channel, + raw_stage, + ], + ) + g_end = cutlass.Float32( + raw_g[ + prepare_valid_tokens - cutlass.Int32(1), + tile_channel, + raw_stage, + ], + ) + tail_gamma = cute.math.exp2( + (g_end - g_value) * cutlass.Float32(_INV_LN2), + fastmath=True, + ) + block_start_g = cutlass.Float32( + raw_g[ + (local_token // cutlass.Int32(_FACTOR_BLOCK)) * cutlass.Int32(_FACTOR_BLOCK), + tile_channel, + raw_stage, + ], + ) + key_value = cutlass.BFloat16( + cutlass.Float32( + raw_k[ + local_token, + tile_channel, + raw_stage, + ], + ) + * tail_gamma, + ) + factor_key_value = cutlass.BFloat16( + cutlass.Float32( + raw_k[ + local_token, + tile_channel, + raw_stage, + ], + ) + * cute.math.exp2( + (block_start_g - g_value) * cutlass.Float32(_INV_LN2), + fastmath=True, + ), + ) + retained_key_tail[local_index] = key_value + shared_factor_k[ + local_token, + key_channel, + cutlass.Int32(0), + ] = factor_key_value + if local_token == cutlass.Int32(0): + last_block_start = ( + (prepare_valid_tokens - cutlass.Int32(1)) // cutlass.Int32(_FACTOR_BLOCK) + ) * cutlass.Int32(_FACTOR_BLOCK) + retained_gamma_end = cute.math.exp2( + ( + cutlass.Float32( + raw_g[ + prepare_valid_tokens - cutlass.Int32(1), + tile_channel, + raw_stage, + ], + ) + - cutlass.Float32( + raw_g[ + last_block_start, + tile_channel, + raw_stage, + ], + ) + ) + * cutlass.Float32(_INV_LN2), + fastmath=True, + ) + + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + factor_ready_handoff.producer_commit( + pipeline.PipelineState( + _INPUT_STAGES, + prepare_chunk, + (prepare_chunk % cutlass.Int32(_INPUT_STAGES)), + cutlass.Int32(1) - ((prepare_chunk // cutlass.Int32(_INPUT_STAGES)) % cutlass.Int32(2)), + ), + ) + + for local_index in cutlass.range_constexpr(8): + linear = thread_in_group + cutlass.Int32( + local_index * _WARP_GROUP_SIZE, + ) + local_token = linear // _WGMMA_K + tile_channel = linear % _WGMMA_K + key_channel = key_tile * cutlass.Int32(_WGMMA_K) + tile_channel + shared_key_tail[ + key_channel, + local_token, + prepare_stage, + ] = retained_key_tail[local_index] + if thread_in_group < cutlass.Int32(_WGMMA_K): + key_channel = key_tile * cutlass.Int32(_WGMMA_K) + thread_in_group + shared_gamma_end[ + key_channel, + prepare_stage, + ] = retained_gamma_end + else: + # Same two-row unroll as the q~/e~ loop above: the + # k~ rebase chain is equally EX2-latency-bound. + for linear in cutlass.range( + thread_in_group, + CHUNK_SIZE * _WGMMA_K, + _WARP_GROUP_SIZE, + unroll=2, + ): + local_token = linear // _WGMMA_K + tile_channel = linear % _WGMMA_K + key_channel = key_tile * cutlass.Int32(_WGMMA_K) + tile_channel + key_value = cutlass.BFloat16(0.0) + factor_key_value = cutlass.BFloat16(0.0) + if local_token < prepare_valid_tokens: + g_value = cutlass.Float32( + raw_g[ + local_token, + tile_channel, + raw_stage, + ], + ) + g_end = cutlass.Float32( + raw_g[ + prepare_valid_tokens - cutlass.Int32(1), + tile_channel, + raw_stage, + ], + ) + tail_gamma = cute.math.exp2( + (g_end - g_value) * cutlass.Float32(_INV_LN2), + fastmath=True, + ) + block_start_g = cutlass.Float32( + raw_g[ + (local_token // cutlass.Int32(_FACTOR_BLOCK)) * cutlass.Int32(_FACTOR_BLOCK), + tile_channel, + raw_stage, + ], + ) + key_value = cutlass.BFloat16( + cutlass.Float32( + raw_k[ + local_token, + tile_channel, + raw_stage, + ], + ) + * tail_gamma, + ) + factor_key_value = cutlass.BFloat16( + cutlass.Float32( + raw_k[ + local_token, + tile_channel, + raw_stage, + ], + ) + * cute.math.exp2( + (block_start_g - g_value) * cutlass.Float32(_INV_LN2), + fastmath=True, + ), + ) + shared_key_tail[ + key_channel, + local_token, + prepare_stage, + ] = key_value + shared_factor_k[ + local_token, + key_channel, + cutlass.Int32(0), + ] = factor_key_value + if local_token == cutlass.Int32(0): + last_block_start = ( + (prepare_valid_tokens - cutlass.Int32(1)) // cutlass.Int32(_FACTOR_BLOCK) + ) * cutlass.Int32(_FACTOR_BLOCK) + shared_gamma_end[ + key_channel, + prepare_stage, + ] = cute.math.exp2( + ( + cutlass.Float32( + raw_g[ + prepare_valid_tokens - cutlass.Int32(1), + tile_channel, + raw_stage, + ], + ) + - cutlass.Float32( + raw_g[ + last_block_start, + tile_channel, + raw_stage, + ], + ) + ) + * cutlass.Float32(_INV_LN2), + fastmath=True, + ) + if state_slab == cutlass.Int32(0): + qkb0_pipeline.consumer_release(qkb_consumer) + else: + qkb1_pipeline.consumer_release(qkb_consumer) + qkb_consumer.advance() + + raw_handoff.consumer_release( + pipeline.PipelineState( + 1, + prepare_chunk, + cutlass.Int32(0), + prepare_chunk % cutlass.Int32(2), + ), + ) + + # Both State WGs publish disjoint halves of the common + # factor operands. The 256-producer factor-ready mbarrier + # is the only rendezvous: each producer fences its + # preceding stores before the release arrive, and Factor + # WG0 returns from the acquire wait only after all 256 + # arrivals. + if cutlass.const_expr( + not self.elide_state_common_barrier, + ): + cute.arch.barrier( + barrier_id=_STATE_COMMON_BARRIER, + number_of_threads=2 * _WARP_GROUP_SIZE, + ) + if cutlass.const_expr(not self.retain_final_tail): + cute.arch.fence_proxy("async.shared", space="cta") + factor_ready_handoff.producer_commit( + pipeline.PipelineState( + _INPUT_STAGES, + prepare_chunk, + (prepare_chunk % cutlass.Int32(_INPUT_STAGES)), + cutlass.Int32(1) - ((prepare_chunk // cutlass.Int32(_INPUT_STAGES)) % cutlass.Int32(2)), + ), + ) + + if state_step > cutlass.Int32(0) and ( + cutlass.const_expr(not self.single_state_owner) or state_slab == cutlass.Int32(0) + ): + # Prioritize the next factor-ready publication above. V/W + # for the current generation was issued after its factor + # completed and can now materialize while WG0 starts the + # next factor. + vw_consumer = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, + _VW_PRIVATE_STAGES, + ) + for local_value_tile in cutlass.range( + _VALUE_TILES // 2, + unroll=1, + ): + if state_slab == cutlass.Int32(0): + vw_ready = vw0_pipeline.consumer_try_wait( + vw_consumer, + ) + vw0_pipeline.consumer_wait( + vw_consumer, + vw_ready, + ) + else: + vw_ready = vw1_pipeline.consumer_try_wait( + vw_consumer, + ) + vw1_pipeline.consumer_wait( + vw_consumer, + vw_ready, + ) + + for linear in cutlass.range( + thread_in_group, + CHUNK_SIZE * _WGMMA_K, + _WARP_GROUP_SIZE, + unroll=1, + ): + local_token = linear // _WGMMA_K + tile_value = linear % _WGMMA_K + value_index = ( + cutlass.Int32( + local_value_tile * _WGMMA_K, + ) + + tile_value + ) + write_value = cutlass.BFloat16(0.0) + if local_token < current_valid_tokens: + write_value = cutlass.BFloat16( + cutlass.Float32( + raw_v[ + local_token, + tile_value, + state_slab, + ], + ) + * cutlass.Float32( + raw_w[ + local_token, + tile_value, + state_slab, + ], + ), + ) + shared_write[ + local_token, + shared_value_start + value_index, + cutlass.Int32(0), + ] = write_value + + if state_slab == cutlass.Int32(0): + vw0_pipeline.consumer_release(vw_consumer) + else: + vw1_pipeline.consumer_release(vw_consumer) + vw_consumer.advance() + + if state_slab == cutlass.Int32(0): + cute.arch.barrier( + barrier_id=_STATE0_WRITE_BARRIER, + number_of_threads=_WARP_GROUP_SIZE, + ) + else: + cute.arch.barrier( + barrier_id=_STATE1_WRITE_BARRIER, + number_of_threads=_WARP_GROUP_SIZE, + ) + + input_stage = current_stage + + q_stage = shared_q[None, None, input_stage] + erase_stage = shared_erase[None, None, input_stage] + key_stage = shared_key_tail[None, None, input_stage] + aqk_stage = shared_aqk[None, None, input_stage] + akk_stage = shared_akk[None, None, input_stage] + + q_token_blocks = cute.flat_divide( + q_stage, + (_FACTOR_BLOCK, HEAD_SIZE), + ) + erase_token_blocks = cute.flat_divide( + erase_stage, + (_FACTOR_BLOCK, HEAD_SIZE), + ) + aqk_operand = token_thread.make_fragment_B( + token_thread.partition_B(aqk_stage), + ) + aqk_stages = aqk_operand + akk_operand = token_thread.make_fragment_B( + token_thread.partition_B(akk_stage), + ) + akk_stages = akk_operand + key_operand = state_thread.make_fragment_B( + state_thread.partition_B(key_stage), + ) + key_stages = key_operand + + gs_delta_stage = shared_gs_delta[ + None, + None, + input_stage, + ] + + output_accumulator = token_thread.make_fragment_C( + token_thread.partition_shape_C( + (self.state_value_tile, CHUNK_SIZE), + ), + ) + erase_projection = token_thread.make_fragment_C( + token_thread.partition_shape_C( + (self.state_value_tile, CHUNK_SIZE), + ), + ) + # The projections consume blockwise-rebased q~/e~, so the + # state fragments advance by the bounded per-channel block + # delta before each 16-token slice. Each slice issues as + # its own WGMMA group against a private BF16 state copy — + # the FP32 state advances underneath without a hazard — so + # up to two slices stay in flight (depth-2 pipelining) and + # a completed slice is drained into the chunk accumulators + # while the next one runs. After the last slice the state + # carries exp(Gs(last)) and the end-of-chunk gamma + # completes the exact exp(G_end) recurrence scale. + slice_outputs = [] + slice_erases = [] + for token_block in cutlass.range_constexpr( + _FACTOR_SUB_BLOCKS, + ): + state_as_token_a = cute.make_rmem_tensor_like( + _convert_c_layout_to_a_layout( + state_accumulator.layout, + token_mma.tv_layout_A.shape[1], + ), + cutlass.BFloat16, + ) + operand_view = cute.make_tensor( + state_as_token_a.iterator, + state_accumulator.layout, + ) + for element in cutlass.range_constexpr( + cute.size(state_accumulator), + ): + _, key_index = state_coordinates[element] + advanced = state_accumulator[element] * cutlass.Float32( + gs_delta_stage[ + key_index, + token_block, + ], + ) + state_accumulator[element] = advanced + operand_view[element] = cutlass.BFloat16(advanced) + q_block_operand = token_thread.make_fragment_B( + token_thread.partition_B( + q_token_blocks[ + None, + None, + token_block, + 0, + ], + ), + ) + erase_block_operand = token_thread.make_fragment_B( + token_thread.partition_B( + erase_token_blocks[ + None, + None, + token_block, + 0, + ], + ), + ) + output_block = token_thread.make_fragment_C( + token_thread.partition_shape_C( + (self.state_value_tile, _FACTOR_BLOCK), + ), + ) + erase_block = token_thread.make_fragment_C( + token_thread.partition_shape_C( + (self.state_value_tile, _FACTOR_BLOCK), + ), + ) + slice_outputs.append(output_block) + slice_erases.append(erase_block) + _fence_register_fragment(state_as_token_a) + warpgroup.fence() + _wgmma_gemm( + token_mma, + output_block, + state_as_token_a, + q_block_operand, + False, + ) + _wgmma_gemm( + token_mma, + erase_block, + state_as_token_a, + erase_block_operand, + False, + ) + warpgroup.commit_group() + if token_block >= 1: + warpgroup.wait_group(1) + cute.autovec_copy( + slice_outputs[token_block - 1][(None, None, 0)], + output_accumulator[(None, None, token_block - 1)], + ) + cute.autovec_copy( + slice_erases[token_block - 1][(None, None, 0)], + erase_projection[(None, None, token_block - 1)], + ) + warpgroup.wait_group(0) + cute.autovec_copy( + slice_outputs[_FACTOR_SUB_BLOCKS - 1][(None, None, 0)], + output_accumulator[(None, None, _FACTOR_SUB_BLOCKS - 1)], + ) + cute.autovec_copy( + slice_erases[_FACTOR_SUB_BLOCKS - 1][(None, None, 0)], + erase_projection[(None, None, _FACTOR_SUB_BLOCKS - 1)], + ) + for element in cutlass.range_constexpr( + cute.size(output_accumulator), + ): + output_accumulator[element] = output_accumulator[element] * scale + + for element in cutlass.range_constexpr( + cute.size(erase_projection), + ): + value_index, token_index = token_coordinates[element] + erase_projection[element] = ( + cutlass.Float32( + shared_write[ + token_index, + shared_value_start + value_index, + cutlass.Int32(0), + ], + ) + - erase_projection[element] + ) + + residual_a = _make_acc_into_op( + erase_projection, + token_mma, + ) + value_new = token_thread.make_fragment_C( + token_thread.partition_shape_C( + (self.state_value_tile, CHUNK_SIZE), + ), + ) + _fence_register_fragment(residual_a) + _fence_register_fragment(value_new) + warpgroup.fence() + _wgmma_gemm( + token_mma, + value_new, + residual_a, + akk_stages, + False, + ) + warpgroup.commit_group() + warpgroup.wait_group(0) + + value_new_a = _make_acc_into_op( + value_new, + token_mma, + ) + _fence_register_fragment(value_new_a) + _fence_register_fragment(output_accumulator) + warpgroup.fence() + _wgmma_gemm( + token_mma, + output_accumulator, + value_new_a, + aqk_stages, + True, + ) + warpgroup.commit_group() + warpgroup.wait_group(0) + output_handoff.producer_acquire(output_producer) + for element in cutlass.range_constexpr( + cute.size(output_accumulator), + ): + value_index, token_index = token_coordinates[element] + shared_output[ + token_index, + shared_value_start + value_index, + output_producer.index, + ] = cutlass.BFloat16(output_accumulator[element]) + cute.arch.fence_proxy("async.shared", space="cta") + output_handoff.producer_commit(output_producer) + output_producer.advance() + + for element in cutlass.range_constexpr( + cute.size(state_accumulator), + ): + _, key_index = state_coordinates[element] + state_accumulator[element] = state_accumulator[element] * shared_gamma_end[key_index, input_stage] + + value_new_as_state_a = _make_acc_into_op( + value_new, + state_mma, + ) + _fence_register_fragment(value_new_as_state_a) + _fence_register_fragment(state_accumulator) + warpgroup.fence() + _wgmma_gemm( + state_mma, + state_accumulator, + value_new_as_state_a, + key_stages, + True, + ) + warpgroup.commit_group() + warpgroup.wait_group(0) + if cutlass.const_expr( + not self.elide_state_iteration_done_barrier, + ): + cute.arch.barrier( + barrier_id=_STATE_ITERATION_DONE_BARRIER, + number_of_threads=2 * _WARP_GROUP_SIZE, + ) + + if state_step > cutlass.Int32(0): + if cutlass.const_expr(self.single_state_owner): + # WG2 remains a common-factor preparation helper in + # the V64/N=1 route. It must not reuse gamma_end's + # two-stage arena until sole owner WG1 has consumed + # the current generation. + cute.arch.barrier( + barrier_id=_STATE_ITERATION_DONE_BARRIER, + number_of_threads=2 * _WARP_GROUP_SIZE, + ) + + if cutlass.const_expr(self.store_final_state): + if cutlass.const_expr(not self.single_state_owner) or state_slab == cutlass.Int32(0): + for element in cutlass.range_constexpr( + cute.size(state_accumulator), + ): + value_index, key_index = state_coordinates[element] + final_state[ + sequence, + value_head, + state_value_start + value_index, + key_index, + ] = state_accumulator[element] From 59f75f91732b47dab8a1fdc1d43b1a14676a3c5f Mon Sep 17 00:00:00 2001 From: Hongyi Wu <62729549+Aharrypotter@users.noreply.github.com> Date: Fri, 21 Aug 2026 11:15:45 +0800 Subject: [PATCH 2/5] feat(gdn2): add the public packed-varlen chunk_gdn2 API Direct SM90a dispatch with no fallback. Tensor metadata is rejected before compilation or launch; validate_inputs=True additionally checks device-side value preconditions, including the [-5, 0] decay bound. Backend availability and the dispatch error share one nvidia-cutlass-dsl range, >=4.5.1,<4.7. The installed version is read through cula.ops._mlir_compat, so this backend and the shared gateway cannot disagree about which toolchain is in use; the upper bound is the gateway's own, and GDN2 only raises the floor to 4.5.1 because 4.4.x lacks cutlass.cute.nvgpu.OperandMajorMode and cannot import the kernel at all. Membership uses standard version ordering, so a local or post release of a supported version stays supported, and pre-release handling is pinned explicitly rather than inherited from the installed packaging, whose default inference has changed between releases. --- cula/gdn2/__init__.py | 18 ++ cula/gdn2/prefill.py | 490 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 508 insertions(+) create mode 100644 cula/gdn2/__init__.py create mode 100644 cula/gdn2/prefill.py diff --git a/cula/gdn2/__init__.py b/cula/gdn2/__init__.py new file mode 100644 index 00000000..cf9dab35 --- /dev/null +++ b/cula/gdn2/__init__.py @@ -0,0 +1,18 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Public Gated DeltaNet-2 operators.""" + +from .prefill import ( + chunk_gdn2, + get_sm90_gdn2_backend, + get_sm90_gdn2_backend_identity, + is_sm90_gdn2_available, +) + +__all__ = [ + "chunk_gdn2", + "get_sm90_gdn2_backend", + "get_sm90_gdn2_backend_identity", + "is_sm90_gdn2_available", +] diff --git a/cula/gdn2/prefill.py b/cula/gdn2/prefill.py new file mode 100644 index 00000000..64c71227 --- /dev/null +++ b/cula/gdn2/prefill.py @@ -0,0 +1,490 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Public packed-varlen Gated DeltaNet-2 prefill API for Hopper SM90.""" + +from __future__ import annotations + +import functools +import math +from dataclasses import dataclass + +import torch +from packaging.requirements import Requirement +from packaging.specifiers import SpecifierSet +from packaging.version import InvalidVersion, Version + +from cula.ops import _mlir_compat +from cula.ops.gdn2.sm90.config import ( + CUTLASS_DSL_REQUIREMENT, + HEAD_SIZE, + MAX_SEQUENCES, + SM90_BACKEND_ID, + SUPPORTED_G_MIN, + SUPPORTED_Q_HEADS, + SUPPORTED_V_HEADS, + VALUE_SIZE, +) + +__all__ = [ + "chunk_gdn2", + "get_sm90_gdn2_backend", + "get_sm90_gdn2_backend_identity", + "is_sm90_gdn2_available", +] + +_INT32_MAX = 2**31 - 1 +_ALIGNMENT = 16 + + +@dataclass(frozen=True) +class _GDN2Inputs: + q: torch.Tensor + k: torch.Tensor + v: torch.Tensor + g: torch.Tensor + b: torch.Tensor + w: torch.Tensor + output: torch.Tensor + initial_state: torch.Tensor | None + output_state: torch.Tensor | None + cu_seqlens: torch.Tensor + total_tokens: int + num_sequences: int + num_q_heads: int + num_v_heads: int + output_final_state: bool + scale: float + + +@functools.cache +def _installed_cutlass_dsl_version() -> str | None: + """Version of the CuTeDSL that is actually imported. + + Read through the shared gateway rather than from package metadata, so + this backend and ``cula.ops._mlir_compat`` can never disagree about which + toolchain is in use. Reading the attribute does not trip the gateway's + own contract check; only private-dialect access does. + """ + + return _mlir_compat.cutlass_dsl_version() + + +@functools.cache +def _cutlass_dsl_specifier() -> SpecifierSet: + return Requirement(CUTLASS_DSL_REQUIREMENT).specifier + + +@functools.cache +def _supported_cutlass_dsl_version() -> str | None: + """Return the installed nvidia-cutlass-dsl version iff it is supported. + + Membership is decided by ``CUTLASS_DSL_REQUIREMENT`` under standard + version ordering, so a local or post release of a supported version + (``4.5.1+cu13``) is supported while anything outside the range is not. + Pre-releases are excluded because only released versions are exercised. + """ + + version = _installed_cutlass_dsl_version() + if version is None: + return None + try: + parsed = Version(version) + except InvalidVersion: + return None + # prereleases is passed explicitly: the default is inferred from the + # specifier set and that inference has differed across packaging + # releases, which would make the gate environment-dependent. + if not _cutlass_dsl_specifier().contains(parsed, prereleases=False): + return None + return version + + +def get_sm90_gdn2_backend() -> str: + """Return the only GDN2 v1 backend.""" + + return "dsl" + + +def get_sm90_gdn2_backend_identity() -> str: + """Return the stable SM90a CuTe DSL implementation identity.""" + + return SM90_BACKEND_ID + + +def is_sm90_gdn2_available( + device: torch.device | int | str | None = None, +) -> bool: + """Return whether the frozen SM90a GDN2 backend is available.""" + + if device is not None and not isinstance(device, int): + device = torch.device(device) + if device.type != "cuda": + return False + if not torch.cuda.is_available(): + return False + if device is None: + device = torch.cuda.current_device() + properties = torch.cuda.get_device_properties(device) + if (properties.major, properties.minor) != (9, 0): + return False + return _supported_cutlass_dsl_version() is not None + + +def _check_tensor( + name: str, + tensor: torch.Tensor, + *, + device: torch.device, + ndim: int, + dtype: torch.dtype, +) -> None: + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if not tensor.is_cuda or tensor.device != device: + raise ValueError(f"{name} must be a CUDA tensor on {device}") + if tensor.ndim != ndim: + raise ValueError( + f"{name} must be rank {ndim}, got shape {tuple(tensor.shape)}", + ) + if tensor.dtype != dtype: + raise TypeError(f"{name} must have dtype {dtype}, got {tensor.dtype}") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + if tensor.data_ptr() % _ALIGNMENT: + raise ValueError( + f"{name} data pointer must be {_ALIGNMENT}-byte aligned", + ) + + +def _storage_interval(tensor: torch.Tensor) -> tuple[int, int]: + start = tensor.data_ptr() + return start, start + tensor.numel() * tensor.element_size() + + +def _overlaps(left: torch.Tensor, right: torch.Tensor) -> bool: + left_start, left_end = _storage_interval(left) + right_start, right_end = _storage_interval(right) + return left_start < right_end and right_start < left_end + + +def _reject_writable_overlap( + name: str, + writable: torch.Tensor, + read_only: dict[str, torch.Tensor], +) -> None: + for other_name, other in read_only.items(): + if _overlaps(writable, other): + raise ValueError( + f"{name} must not overlap read-only tensor {other_name}", + ) + + +def _validate_device_contents( + cu_seqlens: torch.Tensor, + g: torch.Tensor, + b: torch.Tensor, + w: torch.Tensor, + *, + total_tokens: int, +) -> None: + """Synchronously validate value preconditions for diagnostics only.""" + + offsets = tuple(int(value) for value in cu_seqlens.detach().cpu().tolist()) + if offsets[0] != 0: + raise ValueError(f"cu_seqlens[0] must be 0, got {offsets[0]}") + if offsets[-1] != total_tokens: + raise ValueError( + f"cu_seqlens[-1] must equal total_tokens={total_tokens}, got {offsets[-1]}", + ) + if offsets[-1] > _INT32_MAX: + raise ValueError( + f"packed token count must not exceed {_INT32_MAX}", + ) + if any(end <= start for start, end in zip(offsets, offsets[1:])): + raise ValueError( + "zero-length or decreasing sequences are unsupported", + ) + if not bool(torch.isfinite(g).all()) or not bool((g <= 0).all()): + raise ValueError("g must contain finite non-positive log decays") + if not bool((g >= SUPPORTED_G_MIN).all()): + raise ValueError( + f"g must be elementwise >= {SUPPORTED_G_MIN} (see docs/gdn2_sm90_stable_factor.md)", + ) + for name, gate in (("b", b), ("w", w)): + if not bool(torch.isfinite(gate).all()) or not bool((gate >= 0).all()) or not bool((gate <= 1).all()): + raise ValueError( + f"{name} must contain finite values in [0,1]", + ) + + +def _prepare_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + b: torch.Tensor, + w: torch.Tensor, + *, + initial_state: torch.Tensor | None, + output_final_state: bool, + cu_seqlens: torch.Tensor, + scale: float | None, + output: torch.Tensor | None, + output_state: torch.Tensor | None, + validate_inputs: bool, +) -> _GDN2Inputs: + if not isinstance(q, torch.Tensor): + raise TypeError("q must be a torch.Tensor") + device = q.device + _check_tensor( + "q", + q, + device=device, + ndim=3, + dtype=torch.bfloat16, + ) + for name, tensor, dtype in ( + ("k", k, torch.bfloat16), + ("v", v, torch.bfloat16), + ("g", g, torch.float32), + ("b", b, torch.bfloat16), + ("w", w, torch.bfloat16), + ): + _check_tensor( + name, + tensor, + device=device, + ndim=3, + dtype=dtype, + ) + _check_tensor( + "cu_seqlens", + cu_seqlens, + device=device, + ndim=1, + dtype=torch.int64, + ) + + total_tokens, num_q_heads, key_size = q.shape + if not 1 <= total_tokens <= _INT32_MAX: + raise ValueError( + f"total_tokens must be in [1,{_INT32_MAX}]", + ) + if key_size != HEAD_SIZE: + raise ValueError(f"q key dimension must be {HEAD_SIZE}") + expected_q_shape = (total_tokens, num_q_heads, HEAD_SIZE) + for name, tensor in (("k", k), ("g", g), ("b", b)): + if tuple(tensor.shape) != expected_q_shape: + raise ValueError( + f"{name} must have shape {expected_q_shape}, got {tuple(tensor.shape)}", + ) + num_v_heads = v.shape[1] + expected_v_shape = (total_tokens, num_v_heads, VALUE_SIZE) + if tuple(v.shape) != expected_v_shape or tuple(w.shape) != expected_v_shape: + raise ValueError( + f"v and w must have shape {expected_v_shape}", + ) + if num_q_heads <= 0 or num_v_heads <= 0: + raise ValueError("head counts must be positive") + if num_q_heads != SUPPORTED_Q_HEADS: + raise NotImplementedError( + f"GDN2 SM90a prefill requires Hq={SUPPORTED_Q_HEADS}", + ) + if num_v_heads not in SUPPORTED_V_HEADS: + raise NotImplementedError( + f"GDN2 SM90a prefill requires Hv in {SUPPORTED_V_HEADS}", + ) + if num_v_heads < num_q_heads: + raise NotImplementedError( + "GQA is outside the GDN2 v1 contract", + ) + if num_v_heads % num_q_heads: + raise ValueError( + "GVA requires Hv to be an integer multiple of Hq", + ) + if cu_seqlens.numel() < 2: + raise ValueError( + "cu_seqlens must contain at least [0,total_tokens]", + ) + num_sequences = cu_seqlens.numel() - 1 + if num_sequences > MAX_SEQUENCES: + raise NotImplementedError( + f"GDN2 SM90a prefill requires 1 <= N <= {MAX_SEQUENCES}", + ) + + if not isinstance(output_final_state, bool): + raise TypeError("output_final_state must be a bool") + if not isinstance(validate_inputs, bool): + raise TypeError("validate_inputs must be a bool") + if validate_inputs: + _validate_device_contents( + cu_seqlens, + g, + b, + w, + total_tokens=total_tokens, + ) + + state_shape = ( + num_sequences, + num_v_heads, + VALUE_SIZE, + HEAD_SIZE, + ) + if initial_state is not None: + _check_tensor( + "initial_state", + initial_state, + device=device, + ndim=4, + dtype=torch.float32, + ) + if tuple(initial_state.shape) != state_shape: + raise ValueError( + f"initial_state must have public [N,Hv,V,K] shape {state_shape}, got {tuple(initial_state.shape)}", + ) + + output_shape = (total_tokens, num_v_heads, VALUE_SIZE) + with torch.cuda.device(device): + if output is None: + output = torch.empty( + output_shape, + dtype=torch.bfloat16, + device=device, + ) + else: + _check_tensor( + "output", + output, + device=device, + ndim=3, + dtype=torch.bfloat16, + ) + if tuple(output.shape) != output_shape: + raise ValueError( + f"output must have shape {output_shape}, got {tuple(output.shape)}", + ) + + if output_state is not None and not output_final_state: + raise ValueError( + "output_state requires output_final_state=True", + ) + if output_final_state and output_state is None: + output_state = torch.empty( + state_shape, + dtype=torch.float32, + device=device, + ) + elif output_state is not None: + _check_tensor( + "output_state", + output_state, + device=device, + ndim=4, + dtype=torch.float32, + ) + if tuple(output_state.shape) != state_shape: + raise ValueError( + f"output_state must have public [N,Hv,V,K] shape {state_shape}, got {tuple(output_state.shape)}", + ) + + read_only = { + "q": q, + "k": k, + "v": v, + "g": g, + "b": b, + "w": w, + "cu_seqlens": cu_seqlens, + } + if initial_state is not None: + read_only["initial_state"] = initial_state + _reject_writable_overlap("output", output, read_only) + if output_state is not None: + _reject_writable_overlap("output_state", output_state, read_only) + if _overlaps(output, output_state): + raise ValueError( + "output and output_state must not overlap", + ) + + scale_value = HEAD_SIZE**-0.5 if scale is None else float(scale) + if not math.isfinite(scale_value): + raise ValueError(f"scale must be finite, got {scale_value}") + return _GDN2Inputs( + q=q, + k=k, + v=v, + g=g, + b=b, + w=w, + output=output, + initial_state=initial_state, + output_state=output_state, + cu_seqlens=cu_seqlens, + total_tokens=total_tokens, + num_sequences=num_sequences, + num_q_heads=num_q_heads, + num_v_heads=num_v_heads, + output_final_state=output_final_state, + scale=scale_value, + ) + + +def chunk_gdn2( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + b: torch.Tensor, + w: torch.Tensor, + *, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.Tensor, + scale: float | None = None, + output: torch.Tensor | None = None, + output_state: torch.Tensor | None = None, + validate_inputs: bool = False, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Run packed MHA/GVA GDN2 forward prefill on Hopper SM90. + + The default path validates tensor metadata only. CUDA-resident offset and + gate values are caller preconditions unless ``validate_inputs=True`` is + explicitly selected; that diagnostic mode synchronizes. + """ + + inputs = _prepare_inputs( + q, + k, + v, + g, + b, + w, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + scale=scale, + output=output, + output_state=output_state, + validate_inputs=validate_inputs, + ) + properties = torch.cuda.get_device_properties(inputs.q.device) + if (properties.major, properties.minor) != (9, 0): + raise RuntimeError( + f"GDN2 SM90 requires compute capability 9.0, got {properties.major}.{properties.minor}", + ) + if _supported_cutlass_dsl_version() is None: + installed = _installed_cutlass_dsl_version() + raise RuntimeError( + f"GDN2 SM90 requires {CUTLASS_DSL_REQUIREMENT}; " + f"installed: {installed if installed is not None else 'not installed'}", + ) + + from cula.ops.gdn2.sm90.prefill import launch_sm90_gdn2 + + launch_sm90_gdn2(inputs) + if output_final_state: + assert inputs.output_state is not None + return inputs.output, inputs.output_state + return inputs.output From 038cecc3d5e442b6b21dd5a9238710943b7c017d Mon Sep 17 00:00:00 2001 From: Hongyi Wu <62729549+Aharrypotter@users.noreply.github.com> Date: Fri, 21 Aug 2026 11:15:45 +0800 Subject: [PATCH 3/5] test(gdn2): add correctness, adversarial-decay, stress, and sanitizer coverage An independent tokenwise PyTorch reference with no cuLA imports; production, boundary, and irregular packed shapes across MHA/GVA2/GVA4 and all state modes; adversarial decay at g = -1, -2, and the -5 contract boundary, plus mixed strong decay and both erase-gate endpoints; out-of-contract rejection; compile-cache route boundaries; the DSL version gate; a deterministic bitwise stress matrix; and a four-tool Compute Sanitizer runner. --- tests/gdn2/__init__.py | 1 + tests/gdn2/reference.py | 94 +++ tests/gdn2/run_compute_sanitizer_sm90.sh | 152 +++++ tests/gdn2/stress_gdn2_sm90.py | 706 +++++++++++++++++++++++ tests/gdn2/test_gdn2_prefill_sm90.py | 597 +++++++++++++++++++ 5 files changed, 1550 insertions(+) create mode 100644 tests/gdn2/__init__.py create mode 100644 tests/gdn2/reference.py create mode 100755 tests/gdn2/run_compute_sanitizer_sm90.sh create mode 100755 tests/gdn2/stress_gdn2_sm90.py create mode 100644 tests/gdn2/test_gdn2_prefill_sm90.py diff --git a/tests/gdn2/__init__.py b/tests/gdn2/__init__.py new file mode 100644 index 00000000..e6c451c9 --- /dev/null +++ b/tests/gdn2/__init__.py @@ -0,0 +1 @@ +"""GDN2 test package.""" diff --git a/tests/gdn2/reference.py b/tests/gdn2/reference.py new file mode 100644 index 00000000..e39787ca --- /dev/null +++ b/tests/gdn2/reference.py @@ -0,0 +1,94 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Independent PyTorch reference for packed Gated DeltaNet-2 prefill.""" + +from __future__ import annotations + +from itertools import pairwise + +import torch + + +def _expand_qk_heads( + tensor: torch.Tensor, + value_heads: int, +) -> torch.Tensor: + """Expand query-owned channels to their value-head owners in FP32.""" + + group_size = value_heads // tensor.shape[1] + owner = torch.arange( + value_heads, + device=tensor.device, + ).div(group_size, rounding_mode="floor") + return tensor.index_select(1, owner).float() + + +@torch.inference_mode() +def tokenwise_gdn2_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + b: torch.Tensor, + w: torch.Tensor, + *, + cu_seqlens: torch.Tensor, + initial_state: torch.Tensor | None, + output_final_state: bool, + scale: float, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Evaluate the GDN2 recurrence token by token using only PyTorch ops. + + The public recurrent state is accepted and returned in ``[N,Hv,V,K]`` + orientation. Accumulation and the returned reference output use FP32. + """ + + offsets = tuple(int(value) for value in cu_seqlens.detach().cpu().tolist()) + value_heads = v.shape[1] + key_size = q.shape[-1] + value_size = v.shape[-1] + qh = _expand_qk_heads(q, value_heads) + kh = _expand_qk_heads(k, value_heads) + gh = _expand_qk_heads(g, value_heads) + bh = _expand_qk_heads(b, value_heads) + vf = v.float() + wf = w.float() + output = torch.empty_like(v, dtype=torch.float32) + final_states: list[torch.Tensor] = [] + + for sequence, (start, end) in enumerate(pairwise(offsets)): + if initial_state is None: + state = torch.zeros( + value_heads, + key_size, + value_size, + device=v.device, + dtype=torch.float32, + ) + else: + state = initial_state[sequence].transpose(-1, -2).contiguous().clone() + for token in range(start, end): + decayed = state * gh[token].exp().unsqueeze(-1) + erase_key = bh[token] * kh[token] + erase_read = torch.einsum( + "hk,hkv->hv", + erase_key, + decayed, + ) + new_value = wf[token] * vf[token] - erase_read + state = decayed + kh[token].unsqueeze(-1) * new_value.unsqueeze(-2) + output[token] = scale * torch.einsum( + "hk,hkv->hv", + qh[token], + state, + ) + if output_final_state: + final_states.append( + state.transpose(-1, -2).contiguous(), + ) + + return ( + output, + (torch.stack(final_states) if output_final_state else None), + ) diff --git a/tests/gdn2/run_compute_sanitizer_sm90.sh b/tests/gdn2/run_compute_sanitizer_sm90.sh new file mode 100755 index 00000000..17b73eae --- /dev/null +++ b/tests/gdn2/run_compute_sanitizer_sm90.sh @@ -0,0 +1,152 @@ +#!/usr/bin/env bash +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +if [[ $# -lt 1 || $# -gt 2 ]]; then + echo "usage: $0 OUTPUT_ROOT [TOTAL_PRODUCT_LAUNCHES_PER_TOOL]" >&2 + exit 2 +fi + +output_root=$1 +iterations=${2:-120} +script_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +repo_root=$(cd "$script_dir/../.." && pwd) +python=${PYTHON:-python3} +source_manifest=${GDN2_SOURCE_MANIFEST:-} +required_gpu_uuid=${GDN2_REQUIRED_GPU_UUID:-} + +if [[ ! $iterations =~ ^[1-9][0-9]*$ ]]; then + echo "TOTAL_PRODUCT_LAUNCHES_PER_TOOL must be a positive integer" >&2 + exit 2 +fi +if (( iterations < 120 )); then + echo "formal sanitizer coverage requires at least 120 product launches per tool" >&2 + exit 2 +fi +if [[ -e $output_root ]]; then + echo "fresh OUTPUT_ROOT required: $output_root" >&2 + exit 2 +fi +if ! command -v compute-sanitizer >/dev/null 2>&1; then + echo "compute-sanitizer is required" >&2 + exit 2 +fi + +mkdir -p "$output_root" +export CUBLAS_WORKSPACE_CONFIG=:4096:8 +export CUTE_DSL_ARCH=sm_90a +export CUTE_DSL_KEEP=1 +export CUTE_DSL_NO_CACHE=1 +export PYTHONDONTWRITEBYTECODE=1 +export PYTHONHASHSEED=0 +export PYTHONPATH=$repo_root +export PYTORCH_NO_CUDA_MEMORY_CACHING=1 + +for tool in memcheck initcheck synccheck racecheck; do + tool_root=$output_root/$tool + cache_root=$output_root/cache/$tool + mkdir -p \ + "$tool_root" \ + "$cache_root/cuda" \ + "$cache_root/cute" \ + "$cache_root/torchinductor" \ + "$cache_root/triton" \ + "$cache_root/xdg" + + export CUDA_CACHE_PATH=$cache_root/cuda + export CUTE_DSL_CACHE_DIR=$cache_root/cute + export TORCHINDUCTOR_CACHE_DIR=$cache_root/torchinductor + export TRITON_CACHE_DIR=$cache_root/triton + export XDG_CACHE_HOME=$cache_root/xdg + + command=( + compute-sanitizer + --tool "$tool" + --target-processes all + --error-exitcode 86 + ) + case $tool in + memcheck) + command+=(--leak-check full) + ;; + synccheck) + command+=(--check-warpgroup-mma yes) + ;; + racecheck) + command+=( + --racecheck-report hazard + --racecheck-memcpy-async yes + --racecheck-trace-sync yes + ) + ;; + esac + command+=( + "$python" + "$script_dir/stress_gdn2_sm90.py" + --iterations "$iterations" + --warmup 1 + --device 0 + --progress-every 0 + --source-root "$repo_root" + --output "$tool_root/result.json" + ) + if [[ -n $source_manifest ]]; then + command+=(--source-manifest "$source_manifest") + fi + if [[ -n $required_gpu_uuid ]]; then + command+=(--required-gpu-uuid "$required_gpu_uuid") + fi + + printf '%q ' "${command[@]}" >"$tool_root/command.txt" + printf '\n' >>"$tool_root/command.txt" + set +e + "${command[@]}" \ + >"$tool_root/stdout.log" \ + 2>"$tool_root/stderr.log" + return_code=$? + set -e + printf '%s\n' "$return_code" >"$tool_root/returncode.txt" + if [[ $return_code -ne 0 ]]; then + echo "$tool failed with return code $return_code" >&2 + exit "$return_code" + fi + + if [[ $tool == racecheck ]]; then + zero_summary="RACECHECK SUMMARY: 0 hazards displayed (0 errors, 0 warnings)" + else + zero_summary="ERROR SUMMARY: 0 errors" + fi + if ! grep -Fq "$zero_summary" "$tool_root/stdout.log"; then + echo "$tool did not report the required zero-error summary" >&2 + exit 87 + fi + "$python" - "$tool_root/result.json" "$iterations" <<'PY' +import json +import pathlib +import sys + +path = pathlib.Path(sys.argv[1]) +minimum = int(sys.argv[2]) +payload = json.loads(path.read_text(encoding="utf-8")) +if payload["status"] != "PASS": + raise SystemExit(f"stress receipt did not pass: {path}") +if payload["product_launches"] < minimum: + raise SystemExit( + f"incomplete product launch count: " + f"{payload['product_launches']} < {minimum}", + ) +if payload["protocol"]["matrix_rows"] != 6: + raise SystemExit("sanitizer stress matrix must contain six rows") +PY +done + +printf '%s\n' PASS >"$output_root/DONE" +find "$output_root" -type f ! -name evidence-manifest.sha256 -print0 \ + | sort -z \ + | xargs -0 sha256sum \ + >"$output_root/evidence-manifest.sha256" +printf 'GDN2_SM90_SANITIZERS_PASS tools=4 launches_per_tool=%s output=%s\n' \ + "$iterations" \ + "$output_root" diff --git a/tests/gdn2/stress_gdn2_sm90.py b/tests/gdn2/stress_gdn2_sm90.py new file mode 100755 index 00000000..2e7f1006 --- /dev/null +++ b/tests/gdn2/stress_gdn2_sm90.py @@ -0,0 +1,706 @@ +#!/usr/bin/env python3 +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Run source-bound deterministic stress for the SM90 GDN2 product path.""" + +from __future__ import annotations + +import argparse +import datetime as dt +import hashlib +import importlib.metadata +import json +import pathlib +import platform +import subprocess +import sys +import time +from dataclasses import dataclass +from typing import Any + +import torch + +REPO_ROOT = pathlib.Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO_ROOT)) + +from cula.gdn2 import ( # noqa: E402 + chunk_gdn2, + get_sm90_gdn2_backend_identity, +) +from cula.ops.gdn2.sm90.config import ( # noqa: E402 + HEAD_SIZE, + SM90_BACKEND_ID, + SUPPORTED_Q_HEADS, + VALUE_SIZE, +) + +S3_LENGTHS = ( + 63, + 129, + 257, + 31, + 512, + 65, + 128, + 17, + 333, + 91, + 211, + 7, + 401, + 255, + 144, + 73, + 289, + 377, + 19, + 694, +) + + +@dataclass(frozen=True) +class StressSpec: + case_id: str + lengths: tuple[int, ...] + value_heads: int + initial_state: bool + output_final_state: bool + seed: int + + +STRESS_MATRIX = ( + StressSpec( + "S1-MHA-T64", + (64,), + 16, + False, + False, + 7801, + ), + StressSpec( + "S2-MHA-T1024", + (1024,), + 16, + True, + True, + 7802, + ), + StressSpec( + "S3-MHA-PACKED-T4096", + S3_LENGTHS, + 16, + False, + True, + 7803, + ), + StressSpec( + "N32-MHA-IRREGULAR", + (1, 63, 64, 65) * 8, + 16, + True, + False, + 7804, + ), + StressSpec( + "GVA2-PACKED", + (1, 63, 65, 2), + 32, + False, + True, + 7805, + ), + StressSpec( + "GVA4-PACKED", + (65, 1, 129, 63), + 64, + True, + True, + 7806, + ), +) + + +@dataclass +class StressCase: + spec: StressSpec + inputs: dict[str, torch.Tensor | None] + input_hashes_before: dict[str, str] + output_storage: torch.Tensor + output: torch.Tensor + output_redzones_before: tuple[str, str] + state_storage: torch.Tensor | None + output_state: torch.Tensor | None + state_redzones_before: tuple[str, str] | None + baseline_output: torch.Tensor + baseline_state: torch.Tensor | None + baseline_output_sha256: str + baseline_state_sha256: str | None + output_mismatches: torch.Tensor + state_mismatches: torch.Tensor + nonfinite_values: torch.Tensor + stress_launches: int = 0 + + +def _tensor_sha256(tensor: torch.Tensor) -> str: + byte_view = tensor.detach().contiguous().view(torch.uint8) + if byte_view.is_cuda: + byte_view = byte_view.cpu() + return hashlib.sha256( + memoryview(byte_view.numpy()), + ).hexdigest() + + +def _cpu_scalar(tensor: torch.Tensor) -> object: + """Read a scalar without PyTorch's pinned-host item() staging allocation.""" + return tensor.detach().cpu().item() + + +def _redzone_hashes(storage: torch.Tensor) -> tuple[str, str]: + return ( + _tensor_sha256(storage[:1]), + _tensor_sha256(storage[-1:]), + ) + + +def _make_inputs( + spec: StressSpec, + device: torch.device, +) -> dict[str, torch.Tensor | None]: + generator = torch.Generator(device="cpu").manual_seed(spec.seed) + total_tokens = sum(spec.lengths) + q_shape = (total_tokens, SUPPORTED_Q_HEADS, HEAD_SIZE) + v_shape = (total_tokens, spec.value_heads, VALUE_SIZE) + + def bf16_normal( + shape: tuple[int, ...], + standard_deviation: float, + ) -> torch.Tensor: + return (torch.randn(shape, generator=generator) * standard_deviation).to(torch.bfloat16) + + offsets = [0] + for length in spec.lengths: + offsets.append(offsets[-1] + length) + cpu_inputs: dict[str, torch.Tensor | None] = { + "q": bf16_normal(q_shape, 0.03), + "k": bf16_normal(q_shape, 0.01), + "v": bf16_normal(v_shape, 0.1), + "g": ( + -torch.rand( + q_shape, + generator=generator, + dtype=torch.float32, + ) + * 0.05 + ), + "b": torch.rand( + q_shape, + generator=generator, + ).to(torch.bfloat16), + "w": torch.rand( + v_shape, + generator=generator, + ).to(torch.bfloat16), + "cu_seqlens": torch.tensor( + offsets, + dtype=torch.int64, + ), + "initial_state": None, + } + if spec.initial_state: + cpu_inputs["initial_state"] = ( + torch.randn( + ( + len(spec.lengths), + spec.value_heads, + VALUE_SIZE, + HEAD_SIZE, + ), + generator=generator, + ) + * 0.005 + ) + return {name: (None if tensor is None else tensor.to(device=device)) for name, tensor in cpu_inputs.items()} + + +def _input_hashes( + inputs: dict[str, torch.Tensor | None], +) -> dict[str, str]: + return {name: _tensor_sha256(tensor) for name, tensor in sorted(inputs.items()) if tensor is not None} + + +def _launch(case: StressCase) -> None: + result = chunk_gdn2( + case.inputs["q"], + case.inputs["k"], + case.inputs["v"], + case.inputs["g"], + case.inputs["b"], + case.inputs["w"], + initial_state=case.inputs["initial_state"], + output_final_state=case.spec.output_final_state, + cu_seqlens=case.inputs["cu_seqlens"], + scale=HEAD_SIZE**-0.5, + output=case.output, + output_state=case.output_state, + validate_inputs=False, + ) + if case.spec.output_final_state: + output, state = result + if output is not case.output or state is not case.output_state: + raise RuntimeError( + f"preallocated result identity drift: {case.spec.case_id}", + ) + elif result is not case.output: + raise RuntimeError( + f"preallocated output identity drift: {case.spec.case_id}", + ) + + +def _accumulate_exactness(case: StressCase) -> None: + case.output_mismatches.add_( + torch.count_nonzero( + case.output.view(torch.int16) != case.baseline_output.view(torch.int16), + ), + ) + case.nonfinite_values.add_( + torch.count_nonzero(~torch.isfinite(case.output)), + ) + if case.output_state is not None: + if case.baseline_state is None: + raise RuntimeError( + f"missing state baseline: {case.spec.case_id}", + ) + case.state_mismatches.add_( + torch.count_nonzero( + case.output_state.view(torch.int32) != case.baseline_state.view(torch.int32), + ), + ) + case.nonfinite_values.add_( + torch.count_nonzero( + ~torch.isfinite(case.output_state), + ), + ) + + +def _build_case( + spec: StressSpec, + device: torch.device, + warmup: int, +) -> StressCase: + inputs = _make_inputs(spec, device) + input_hashes_before = _input_hashes(inputs) + total_tokens = sum(spec.lengths) + output_storage = torch.full( + ( + total_tokens + 2, + spec.value_heads, + VALUE_SIZE, + ), + float("nan"), + dtype=torch.bfloat16, + device=device, + ) + output_storage[0].fill_(-123.0) + output_storage[-1].fill_(-123.0) + output = output_storage[1:-1] + output_redzones_before = _redzone_hashes(output_storage) + state_storage = None + output_state = None + state_redzones_before = None + if spec.output_final_state: + state_storage = torch.full( + ( + len(spec.lengths) + 2, + spec.value_heads, + VALUE_SIZE, + HEAD_SIZE, + ), + float("nan"), + dtype=torch.float32, + device=device, + ) + state_storage[0].fill_(-987654.0) + state_storage[-1].fill_(-987654.0) + output_state = state_storage[1:-1] + state_redzones_before = _redzone_hashes( + state_storage, + ) + + placeholder = StressCase( + spec=spec, + inputs=inputs, + input_hashes_before=input_hashes_before, + output_storage=output_storage, + output=output, + output_redzones_before=output_redzones_before, + state_storage=state_storage, + output_state=output_state, + state_redzones_before=state_redzones_before, + baseline_output=output, + baseline_state=output_state, + baseline_output_sha256="", + baseline_state_sha256=None, + output_mismatches=torch.zeros( + (), + dtype=torch.int64, + device=device, + ), + state_mismatches=torch.zeros( + (), + dtype=torch.int64, + device=device, + ), + nonfinite_values=torch.zeros( + (), + dtype=torch.int64, + device=device, + ), + ) + for _ in range(warmup): + _launch(placeholder) + torch.cuda.synchronize(device) + if not bool(_cpu_scalar(torch.isfinite(output).all())): + raise RuntimeError( + f"non-finite baseline output: {spec.case_id}", + ) + baseline_output = output.detach().clone() + baseline_state = None if output_state is None else output_state.detach().clone() + if baseline_state is not None and not bool( + _cpu_scalar(torch.isfinite(baseline_state).all()), + ): + raise RuntimeError( + f"non-finite baseline state: {spec.case_id}", + ) + placeholder.baseline_output = baseline_output + placeholder.baseline_state = baseline_state + placeholder.baseline_output_sha256 = _tensor_sha256( + baseline_output, + ) + placeholder.baseline_state_sha256 = None if baseline_state is None else _tensor_sha256(baseline_state) + return placeholder + + +def _normalise_gpu_uuid(value: object) -> str: + uuid = str(value) + return uuid if uuid.startswith("GPU-") else f"GPU-{uuid}" + + +def _sha256(path: pathlib.Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _verify_source_manifest( + source_root: pathlib.Path, + manifest_path: pathlib.Path, +) -> dict[str, Any]: + manifest = json.loads( + manifest_path.read_text(encoding="utf-8"), + ) + aggregate = hashlib.sha256() + observed: list[tuple[str, str, int]] = [] + for entry in manifest["files"]: + relative = pathlib.Path(entry["path"]) + if relative.is_absolute() or ".." in relative.parts: + raise ValueError( + f"unsafe source-manifest path: {relative}", + ) + path = (source_root / relative).resolve(strict=True) + if not path.is_relative_to(source_root) or not path.is_file(): + raise ValueError( + f"source path escapes root: {relative}", + ) + sha256 = _sha256(path) + size_bytes = path.stat().st_size + if sha256 != entry["sha256"] or size_bytes != entry["size_bytes"]: + raise RuntimeError( + f"source-manifest mismatch: {relative}", + ) + observed.append( + (relative.as_posix(), sha256, size_bytes), + ) + for relative, sha256, size_bytes in sorted(observed): + aggregate.update( + f"{relative}\0{sha256}\0{size_bytes}\n".encode(), + ) + if aggregate.hexdigest() != manifest["aggregate_sha256"]: + raise RuntimeError( + "source-manifest aggregate mismatch", + ) + return { + "mode": "manifest", + "manifest_path": str(manifest_path), + "manifest_sha256": _sha256(manifest_path), + "aggregate_sha256": aggregate.hexdigest(), + "file_count": len(observed), + } + + +def _git_source_identity( + source_root: pathlib.Path, +) -> dict[str, Any]: + try: + commit = subprocess.check_output( + ["git", "rev-parse", "HEAD"], + cwd=source_root, + text=True, + stderr=subprocess.DEVNULL, + ).strip() + status = subprocess.check_output( + ["git", "status", "--short"], + cwd=source_root, + text=True, + stderr=subprocess.DEVNULL, + ) + except (FileNotFoundError, subprocess.CalledProcessError): + commit, status = None, None + return { + "mode": "git", + "repo_root": str(source_root), + "commit": commit, + "worktree_status": status, + } + + +def _environment( + device: torch.device, +) -> dict[str, Any]: + properties = torch.cuda.get_device_properties(device) + return { + "hostname": platform.node(), + "python": sys.version, + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "cutlass_dsl": importlib.metadata.version( + "nvidia-cutlass-dsl", + ), + "gpu_name": properties.name, + "gpu_uuid": _normalise_gpu_uuid(properties.uuid), + "compute_capability": [ + properties.major, + properties.minor, + ], + } + + +def _case_result(case: StressCase) -> dict[str, Any]: + output_mismatches = int(_cpu_scalar(case.output_mismatches)) + state_mismatches = int(_cpu_scalar(case.state_mismatches)) + nonfinite_values = int(_cpu_scalar(case.nonfinite_values)) + input_hashes_after = _input_hashes(case.inputs) + output_redzones_after = _redzone_hashes( + case.output_storage, + ) + state_redzones_after = None if case.state_storage is None else _redzone_hashes(case.state_storage) + final_output_sha256 = _tensor_sha256(case.output) + final_state_sha256 = None if case.output_state is None else _tensor_sha256(case.output_state) + result = { + "case_id": case.spec.case_id, + "lengths": list(case.spec.lengths), + "total_tokens": sum(case.spec.lengths), + "num_sequences": len(case.spec.lengths), + "value_heads": case.spec.value_heads, + "initial_state": case.spec.initial_state, + "output_final_state": (case.spec.output_final_state), + "seed": case.spec.seed, + "stress_launches": case.stress_launches, + "output_mismatches": output_mismatches, + "state_mismatches": state_mismatches, + "nonfinite_values": nonfinite_values, + "input_immutability": (input_hashes_after == case.input_hashes_before), + "output_redzones": (output_redzones_after == case.output_redzones_before), + "state_redzones": (state_redzones_after == case.state_redzones_before), + "baseline_output_sha256": (case.baseline_output_sha256), + "final_output_sha256": final_output_sha256, + "baseline_state_sha256": (case.baseline_state_sha256), + "final_state_sha256": final_state_sha256, + } + result["status"] = ( + "PASS" + if ( + case.stress_launches > 0 + and output_mismatches == 0 + and state_mismatches == 0 + and nonfinite_values == 0 + and result["input_immutability"] + and result["output_redzones"] + and result["state_redzones"] + and final_output_sha256 == case.baseline_output_sha256 + and final_state_sha256 == case.baseline_state_sha256 + ) + else "FAIL" + ) + return result + + +def _write_json( + path: pathlib.Path, + payload: dict[str, Any], +) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(payload, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--iterations", type=int, default=100000) + parser.add_argument("--warmup", type=int, default=1) + parser.add_argument("--device", type=int, default=0) + parser.add_argument( + "--output", + type=pathlib.Path, + default=pathlib.Path("gdn2-sm90-stress.json"), + ) + parser.add_argument( + "--source-root", + type=pathlib.Path, + default=REPO_ROOT, + ) + parser.add_argument( + "--source-manifest", + type=pathlib.Path, + ) + parser.add_argument("--required-gpu-uuid") + parser.add_argument("--progress-every", type=int, default=10000) + parser.add_argument("--list-matrix", action="store_true") + args = parser.parse_args() + if args.iterations <= 0: + parser.error("--iterations must be positive") + if args.warmup <= 0: + parser.error("--warmup must be positive") + if args.progress_every < 0: + parser.error("--progress-every must be non-negative") + return args + + +def _run(args: argparse.Namespace) -> dict[str, Any]: + started_at_utc = dt.datetime.now(dt.UTC).isoformat() + source_root = args.source_root.resolve(strict=True) + source = ( + _git_source_identity(source_root) + if args.source_manifest is None + else _verify_source_manifest( + source_root, + args.source_manifest.resolve(strict=True), + ) + ) + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + device = torch.device("cuda", args.device) + torch.cuda.set_device(device) + environment = _environment(device) + if environment["compute_capability"] != [9, 0]: + raise RuntimeError( + "GDN2 deterministic stress requires SM90", + ) + if args.required_gpu_uuid is not None and environment["gpu_uuid"] != args.required_gpu_uuid: + raise RuntimeError( + f"GPU UUID mismatch: {environment['gpu_uuid']} != {args.required_gpu_uuid}", + ) + if get_sm90_gdn2_backend_identity() != SM90_BACKEND_ID: + raise RuntimeError("product backend identity drift") + + cases = [_build_case(spec, device, args.warmup) for spec in STRESS_MATRIX] + torch.cuda.synchronize(device) + started = time.perf_counter() + for iteration in range(args.iterations): + case = cases[iteration % len(cases)] + _launch(case) + _accumulate_exactness(case) + case.stress_launches += 1 + completed = iteration + 1 + if args.progress_every > 0 and completed % args.progress_every == 0: + print( + f"GDN2_SM90_STRESS_PROGRESS completed={completed}/{args.iterations}", + flush=True, + ) + torch.cuda.synchronize(device) + duration_seconds = time.perf_counter() - started + if get_sm90_gdn2_backend_identity() != SM90_BACKEND_ID: + raise RuntimeError("product backend identity drift after stress") + case_results = [_case_result(case) for case in cases] + status = ( + "PASS" + if ( + sum(case["stress_launches"] for case in case_results) == args.iterations + and all(case["status"] == "PASS" for case in case_results) + ) + else "FAIL" + ) + return { + "schema": "cula.gdn2.sm90.deterministic-stress.v1", + "status": status, + "started_at_utc": started_at_utc, + "finished_at_utc": dt.datetime.now(dt.UTC).isoformat(), + "source": source, + "environment": environment, + "backend_identity": SM90_BACKEND_ID, + "fallback": False, + "protocol": { + "iterations": args.iterations, + "warmup_per_case": args.warmup, + "matrix_rows": len(STRESS_MATRIX), + "order": "round_robin", + "preallocated_output_and_state": True, + "validate_inputs": False, + "bitwise_output_and_state_check_each_launch": True, + "finite_check_each_launch": True, + "host_synchronization_inside_stress_loop": False, + }, + "duration_seconds": duration_seconds, + "product_launches": args.iterations, + "cases": case_results, + "claim_boundary": ( + "One process, one CUDA device, fixed per-case inputs and " + "initial states, round-robin public product launches, and " + "device-side exactness/finite accumulation for every launch." + ), + } + + +def main() -> None: + args = _parse_args() + if args.list_matrix: + print( + json.dumps( + [ + { + "case_id": spec.case_id, + "lengths": list(spec.lengths), + "total_tokens": sum(spec.lengths), + "value_heads": spec.value_heads, + "initial_state": spec.initial_state, + "output_final_state": (spec.output_final_state), + "seed": spec.seed, + } + for spec in STRESS_MATRIX + ], + indent=2, + ), + ) + return + output = args.output.resolve() + if output.exists(): + raise FileExistsError(f"fresh output required: {output}") + payload = _run(args) + _write_json(output, payload) + if payload["status"] != "PASS": + raise RuntimeError( + f"GDN2 deterministic stress failed: {output}", + ) + print( + "GDN2_SM90_STRESS_PASS " + f"launches={payload['product_launches']} " + f"duration_seconds={payload['duration_seconds']:.3f} " + f"output={output}", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/tests/gdn2/test_gdn2_prefill_sm90.py b/tests/gdn2/test_gdn2_prefill_sm90.py new file mode 100644 index 00000000..ace4049f --- /dev/null +++ b/tests/gdn2/test_gdn2_prefill_sm90.py @@ -0,0 +1,597 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import ast +import importlib.metadata +from dataclasses import dataclass +from pathlib import Path + +import pytest +import torch + +from cula.gdn2 import ( + chunk_gdn2, + get_sm90_gdn2_backend, + get_sm90_gdn2_backend_identity, + is_sm90_gdn2_available, +) +from cula.ops.gdn2.sm90.config import ( + HEAD_SIZE, + MAX_SEQUENCES, + SM90_BACKEND_ID, + SUPPORTED_G_MIN, + SUPPORTED_Q_HEADS, + SUPPORTED_V_HEADS, + VALUE_SIZE, +) +from cula.ops.gdn2.sm90.prefill import _compiled + +from .reference import tokenwise_gdn2_reference + +_OUTPUT_RTOL = 0.01 +_OUTPUT_ATOL = 0.01 +_STATE_RTOL = 0.001 +_STATE_ATOL = 0.005 + + +@dataclass(frozen=True) +class _Case: + case_id: str + lengths: tuple[int, ...] + value_heads: int + initial_state: bool + output_final_state: bool + + +_CASES = ( + _Case("mha-single-token", (1,), 16, False, False), + _Case("mha-tail-and-init", (65, 1), 16, True, True), + _Case("mha-initial-no-final", (65, 63), 16, True, False), + _Case("mha-production-t1024", (1024,), 16, True, True), + _Case("mha-t64-short-baseline", (64,), 16, True, True), + _Case("mha-t65-v64-boundary", (65,), 16, True, True), + _Case("mha-max-sequences", (1,) * 32, 16, False, True), + _Case("gva2-packed-tails", (1, 63, 65, 2), 32, False, True), + _Case("gva4-init", (4,), 64, True, True), +) + + +def _is_supported_sm90() -> bool: + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (9, 0): + return False + try: + importlib.metadata.version("nvidia-cutlass-dsl") + except importlib.metadata.PackageNotFoundError: + return False + return True + + +requires_sm90 = pytest.mark.skipif( + not _is_supported_sm90(), + reason="requires compute capability 9.0 and nvidia-cutlass-dsl installed", +) + + +def _assert_finite(name: str, tensor: torch.Tensor) -> None: + assert bool(torch.isfinite(tensor).all()), f"{name} contains NaN or Inf" + + +def _make_inputs(case: _Case) -> dict[str, torch.Tensor | None]: + generator = torch.Generator(device="cpu").manual_seed(20260727) + total_tokens = sum(case.lengths) + q_shape = (total_tokens, SUPPORTED_Q_HEADS, HEAD_SIZE) + v_shape = (total_tokens, case.value_heads, VALUE_SIZE) + + def _bf16_normal(shape: tuple[int, ...], scale: float) -> torch.Tensor: + return (torch.randn(shape, generator=generator) * scale).to(torch.bfloat16).cuda() + + q = _bf16_normal(q_shape, 0.03) + k = _bf16_normal(q_shape, 0.01) + v = _bf16_normal(v_shape, 0.1) + g = (-torch.rand(q_shape, generator=generator, dtype=torch.float32) * 0.05).cuda() + b = torch.rand(q_shape, generator=generator).to(torch.bfloat16).cuda() + w = torch.rand(v_shape, generator=generator).to(torch.bfloat16).cuda() + offsets = [0] + for length in case.lengths: + offsets.append(offsets[-1] + length) + cu_seqlens = torch.tensor(offsets, dtype=torch.int64, device="cuda") + initial_state = None + if case.initial_state: + initial_state = ( + torch.randn( + ( + len(case.lengths), + case.value_heads, + VALUE_SIZE, + HEAD_SIZE, + ), + generator=generator, + ) + * 0.005 + ).cuda() + return { + "q": q, + "k": k, + "v": v, + "g": g, + "b": b, + "w": w, + "cu_seqlens": cu_seqlens, + "initial_state": initial_state, + } + + +def test_product_identity_contract() -> None: + assert get_sm90_gdn2_backend() == "dsl" + assert get_sm90_gdn2_backend_identity() == SM90_BACKEND_ID + assert SM90_BACKEND_ID == "sm90a_cutedsl_gdn2_prefill_v1" + assert MAX_SEQUENCES == 32 + assert SUPPORTED_Q_HEADS == 16 + assert SUPPORTED_V_HEADS == (16, 32, 64) + + +def test_gdn2_has_no_gdn_namespace_dependency() -> None: + source_root = Path(__file__).resolve().parents[2] / "cula" + imported_modules: set[str] = set() + for path in sorted((source_root / "gdn2").rglob("*.py")) + sorted( + (source_root / "ops" / "gdn2").rglob("*.py"), + ): + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom) and node.module is not None: + imported_modules.add(node.module) + elif isinstance(node, ast.Import): + imported_modules.update(alias.name for alias in node.names) + + forbidden = sorted( + module + for module in imported_modules + if module == "cula.gdn" + or module.startswith("cula.gdn.") + or module == "cula.ops.gdn" + or module.startswith("cula.ops.gdn.") + ) + assert forbidden == [] + + +@requires_sm90 +@pytest.mark.parametrize("case", _CASES, ids=lambda case: case.case_id) +@torch.inference_mode() +def test_product_matches_tokenwise_reference(case: _Case) -> None: + inputs = _make_inputs(case) + read_only_before = {name: tensor.detach().clone() for name, tensor in inputs.items() if isinstance(tensor, torch.Tensor)} + scale = HEAD_SIZE**-0.5 + expected_output, expected_state = tokenwise_gdn2_reference( + inputs["q"], + inputs["k"], + inputs["v"], + inputs["g"], + inputs["b"], + inputs["w"], + cu_seqlens=inputs["cu_seqlens"], + initial_state=inputs["initial_state"], + output_final_state=case.output_final_state, + scale=scale, + ) + + output_storage = torch.full( + (sum(case.lengths) + 2, case.value_heads, VALUE_SIZE), + -123.0, + dtype=torch.bfloat16, + device="cuda", + ) + output = output_storage[1:-1] + state_storage = None + output_state = None + if case.output_final_state: + state_storage = torch.full( + ( + len(case.lengths) + 2, + case.value_heads, + VALUE_SIZE, + HEAD_SIZE, + ), + -987654.0, + dtype=torch.float32, + device="cuda", + ) + output_state = state_storage[1:-1] + + actual = chunk_gdn2( + inputs["q"], + inputs["k"], + inputs["v"], + inputs["g"], + inputs["b"], + inputs["w"], + initial_state=inputs["initial_state"], + output_final_state=case.output_final_state, + cu_seqlens=inputs["cu_seqlens"], + scale=scale, + output=output, + output_state=output_state, + ) + torch.cuda.synchronize() + + if case.output_final_state: + actual_output, actual_state = actual + assert actual_state is output_state + else: + actual_output = actual + actual_state = None + assert actual_output is output + _assert_finite("reference output", expected_output) + _assert_finite("product output", actual_output) + torch.testing.assert_close( + actual_output, + expected_output.to(dtype=actual_output.dtype), + rtol=_OUTPUT_RTOL, + atol=_OUTPUT_ATOL, + ) + if expected_state is not None: + assert actual_state is not None + _assert_finite("reference final state", expected_state) + _assert_finite("product final state", actual_state) + torch.testing.assert_close( + actual_state, + expected_state, + rtol=_STATE_RTOL, + atol=_STATE_ATOL, + ) + + first_output = actual_output.detach().clone() + first_state = None if actual_state is None else actual_state.detach().clone() + repeated = chunk_gdn2( + inputs["q"], + inputs["k"], + inputs["v"], + inputs["g"], + inputs["b"], + inputs["w"], + initial_state=inputs["initial_state"], + output_final_state=case.output_final_state, + cu_seqlens=inputs["cu_seqlens"], + scale=scale, + output=output, + output_state=output_state, + ) + torch.cuda.synchronize() + if case.output_final_state: + repeated_output, repeated_state = repeated + assert repeated_state is output_state + assert first_state is not None + assert torch.equal(repeated_state, first_state) + else: + repeated_output = repeated + assert repeated_output is output + assert torch.equal(repeated_output, first_output) + _assert_finite("repeated product output", repeated_output) + + assert bool((output_storage[0] == -123.0).all()) + assert bool((output_storage[-1] == -123.0).all()) + if state_storage is not None: + assert bool((state_storage[0] == -987654.0).all()) + assert bool((state_storage[-1] == -987654.0).all()) + for name, before in read_only_before.items(): + torch.testing.assert_close(inputs[name], before, rtol=0, atol=0) + + +@requires_sm90 +@torch.inference_mode() +def test_unsupported_metadata_fails_before_compile() -> None: + def _attempt(*, q_heads: int, value_heads: int, sequences: int) -> None: + q = torch.zeros(sequences, q_heads, HEAD_SIZE, dtype=torch.bfloat16, device="cuda") + k = torch.zeros_like(q) + v = torch.zeros( + sequences, + value_heads, + VALUE_SIZE, + dtype=torch.bfloat16, + device="cuda", + ) + g = torch.zeros_like(q, dtype=torch.float32) + b = torch.zeros_like(q) + w = torch.zeros_like(v) + cu_seqlens = torch.arange(sequences + 1, dtype=torch.int64, device="cuda") + chunk_gdn2(q, k, v, g, b, w, cu_seqlens=cu_seqlens) + + compile_count = len(_compiled) + with pytest.raises(NotImplementedError, match="Hq=16"): + _attempt(q_heads=8, value_heads=16, sequences=1) + with pytest.raises(NotImplementedError, match="Hv in"): + _attempt(q_heads=16, value_heads=48, sequences=1) + with pytest.raises(NotImplementedError, match="1 <= N <= 32"): + _attempt(q_heads=16, value_heads=16, sequences=33) + assert len(_compiled) == compile_count + + +@requires_sm90 +def test_availability_on_current_device() -> None: + assert is_sm90_gdn2_available() + assert is_sm90_gdn2_available(torch.cuda.current_device()) + assert not is_sm90_gdn2_available("cpu") + + +def _expected_compile_key( + *, + num_sequences: int, + total_tokens: int, + value_heads: int, + has_initial_state: bool, + store_final_state: bool, +) -> tuple[int, int, bool, bool, bool, bool]: + """Mirror the documented compile-cache key derivation. + + Keep in sync with ``_compile`` in ``cula.ops.gdn2.sm90.prefill`` and the + "State modes and dynamic compilation" section of + ``docs/gdn2_sm90_pipeline.md``; this test is the contract regression + guard for both. + """ + + use_n1_hv16_v64 = ( + num_sequences == 1 and value_heads == 16 and has_initial_state and store_final_state and total_tokens > 64 + ) + retain_final_tail = store_final_state and not (num_sequences == 1 and total_tokens <= 64) + return ( + torch.cuda.current_device(), + value_heads, + has_initial_state, + store_final_state, + use_n1_hv16_v64, + retain_final_tail, + ) + + +@requires_sm90 +@torch.inference_mode() +def test_compile_cache_boundaries() -> None: + """The cache key follows the documented route boundaries exactly. + + Moving across ``N=1,T<=64`` / ``N=1,T>64`` / ``N>1`` compiles one new + specialization each for the final-state mode, while ``T``/``N`` stay + dynamic within a route and while ``output_final_state=False`` collapses + every shape onto one specialization. + """ + + from cula.ops.gdn2.sm90 import prefill as sm90_prefill + + def _run(case: _Case) -> tuple[int, int, bool, bool, bool, bool]: + inputs = _make_inputs(case) + before_keys = set(sm90_prefill._compiled) + expected_key = _expected_compile_key( + num_sequences=len(case.lengths), + total_tokens=sum(case.lengths), + value_heads=case.value_heads, + has_initial_state=case.initial_state, + store_final_state=case.output_final_state, + ) + chunk_gdn2( + inputs["q"], + inputs["k"], + inputs["v"], + inputs["g"], + inputs["b"], + inputs["w"], + initial_state=inputs["initial_state"], + output_final_state=case.output_final_state, + cu_seqlens=inputs["cu_seqlens"], + ) + torch.cuda.synchronize() + after_keys = set(sm90_prefill._compiled) + assert expected_key in after_keys + expected_growth = 0 if expected_key in before_keys else 1 + assert len(after_keys) == len(before_keys) + expected_growth + return expected_key + + def _case(lengths: tuple[int, ...], final_state: bool) -> _Case: + return _Case( + case_id=f"cache-{len(lengths)}seq-{sum(lengths)}tok-{final_state}", + lengths=lengths, + value_heads=16, + initial_state=True, + output_final_state=final_state, + ) + + # Final-state mode: three distinct routes across the documented + # boundaries... + key_short = _run(_case((64,), True)) + key_n1_long = _run(_case((65,), True)) + key_packed = _run(_case((40, 25), True)) + assert len({key_short, key_n1_long, key_packed}) == 3 + + # ...and T/N stay dynamic inside each route: repeats and different + # shapes on the same route map to the same key (asserted inside _run + # via expected_growth == 0). + assert _run(_case((32,), True)) == key_short + assert _run(_case((1024,), True)) == key_n1_long + assert _run(_case((30, 20, 14), True)) == key_packed + + # No-final-state mode: every boundary collapses onto one key. + key_no_final = _run(_case((64,), False)) + assert _run(_case((65,), False)) == key_no_final + assert _run(_case((40, 25), False)) == key_no_final + + +def test_cutlass_dsl_version_gate(monkeypatch: pytest.MonkeyPatch) -> None: + """The availability gate enforces the one documented version range.""" + + from cula.gdn2 import prefill as gdn2_prefill + + def _probe(version: str | None) -> str | None: + monkeypatch.setattr( + gdn2_prefill, + "_installed_cutlass_dsl_version", + lambda: version, + ) + gdn2_prefill._supported_cutlass_dsl_version.cache_clear() + try: + return gdn2_prefill._supported_cutlass_dsl_version() + finally: + gdn2_prefill._supported_cutlass_dsl_version.cache_clear() + + # Endpoints are exercised on H20; interior versions follow the vendor's + # release ordering. Local/post releases of a supported version stay + # supported; pre-releases and unparseable strings do not. + supported = ("4.5.1", "4.6.0", "4.6.2", "4.5.1+cu13", "4.5.1.post1") + unsupported = ( + None, + "4.4.2", + "4.5.0", + "4.7.0", + "5.0.0", + "4.5", + "4.6.0rc1", + "not-a-version", + ) + for version in supported: + assert _probe(version) == version, version + for version in unsupported: + assert _probe(version) is None, version + + +@dataclass(frozen=True) +class _DecayCase: + case_id: str + lengths: tuple[int, ...] + decay: float | None # None -> mixed uniform in [SUPPORTED_G_MIN, 0] + gate_mode: str # "random" | "zeros" | "ones" + + +_DECAY_CASES = ( + _DecayCase("uniform-g1", (150,), -1.0, "random"), + _DecayCase("uniform-g2", (129,), -2.0, "random"), + _DecayCase("uniform-g5-bound", (150,), SUPPORTED_G_MIN, "random"), + _DecayCase("mixed-strong-decay", (65, 40), None, "random"), + _DecayCase("gate-endpoint-zeros", (100,), -1.0, "zeros"), + _DecayCase("gate-endpoint-ones", (100,), -1.0, "ones"), +) + + +def _make_decay_inputs(case: _DecayCase) -> dict[str, torch.Tensor | None]: + generator = torch.Generator(device="cpu").manual_seed(20260814) + total_tokens = sum(case.lengths) + q_shape = (total_tokens, SUPPORTED_Q_HEADS, HEAD_SIZE) + v_shape = (total_tokens, 16, VALUE_SIZE) + + def _bf16_normal(shape: tuple[int, ...], scale: float) -> torch.Tensor: + return (torch.randn(shape, generator=generator) * scale).to(torch.bfloat16).cuda() + + q = _bf16_normal(q_shape, 0.03) + k = _bf16_normal(q_shape, 0.01) + v = _bf16_normal(v_shape, 0.1) + if case.decay is None: + g = (torch.rand(q_shape, generator=generator, dtype=torch.float32) * SUPPORTED_G_MIN).cuda() + else: + g = torch.full(q_shape, case.decay, dtype=torch.float32).cuda() + if case.gate_mode == "zeros": + b = torch.zeros(q_shape, dtype=torch.bfloat16).cuda() + elif case.gate_mode == "ones": + b = torch.ones(q_shape, dtype=torch.bfloat16).cuda() + else: + b = torch.rand(q_shape, generator=generator).to(torch.bfloat16).cuda() + w = torch.rand(v_shape, generator=generator).to(torch.bfloat16).cuda() + offsets = [0] + for length in case.lengths: + offsets.append(offsets[-1] + length) + cu_seqlens = torch.tensor(offsets, dtype=torch.int64, device="cuda") + initial_state = ( + torch.randn( + (len(case.lengths), 16, VALUE_SIZE, HEAD_SIZE), + generator=generator, + ) + * 0.005 + ).cuda() + return { + "q": q, + "k": k, + "v": v, + "g": g, + "b": b, + "w": w, + "cu_seqlens": cu_seqlens, + "initial_state": initial_state, + } + + +@requires_sm90 +@pytest.mark.parametrize("case", _DECAY_CASES, ids=lambda case: case.case_id) +@torch.inference_mode() +def test_adversarial_decay_matches_tokenwise_reference(case: _DecayCase) -> None: + """Strong in-contract decays stay finite and match the exact recurrence. + + The released chunk-start factorization overflowed FP32 for uniform + ``g <= -1.5`` (64-token channel prefixes beyond ~88.7 nats) and poisoned + the rest of the sequence with NaN. These cases pin the blockwise-rebased + factorization across the documented ``[-5, 0]`` contract, the erase-gate + endpoints, and the ``g = -5`` boundary itself. + """ + + inputs = _make_decay_inputs(case) + scale = HEAD_SIZE**-0.5 + expected_output, expected_state = tokenwise_gdn2_reference( + inputs["q"], + inputs["k"], + inputs["v"], + inputs["g"], + inputs["b"], + inputs["w"], + cu_seqlens=inputs["cu_seqlens"], + initial_state=inputs["initial_state"], + output_final_state=True, + scale=scale, + ) + actual_output, actual_state = chunk_gdn2( + inputs["q"], + inputs["k"], + inputs["v"], + inputs["g"], + inputs["b"], + inputs["w"], + initial_state=inputs["initial_state"], + output_final_state=True, + cu_seqlens=inputs["cu_seqlens"], + scale=scale, + validate_inputs=True, + ) + torch.cuda.synchronize() + _assert_finite("reference output", expected_output) + _assert_finite("product output", actual_output) + _assert_finite("reference final state", expected_state) + _assert_finite("product final state", actual_state) + torch.testing.assert_close( + actual_output, + expected_output.to(dtype=actual_output.dtype), + rtol=_OUTPUT_RTOL, + atol=_OUTPUT_ATOL, + ) + torch.testing.assert_close( + actual_state, + expected_state, + rtol=_STATE_RTOL, + atol=_STATE_ATOL, + ) + + +@requires_sm90 +@torch.inference_mode() +def test_decay_below_bound_rejected() -> None: + """validate_inputs enforces the documented elementwise g >= -5 bound.""" + + case = _DecayCase("reject", (65,), SUPPORTED_G_MIN, "random") + inputs = _make_decay_inputs(case) + inputs["g"][3, 5, 7] = SUPPORTED_G_MIN - 0.5 + with pytest.raises(ValueError, match="elementwise >="): + chunk_gdn2( + inputs["q"], + inputs["k"], + inputs["v"], + inputs["g"], + inputs["b"], + inputs["w"], + initial_state=inputs["initial_state"], + output_final_state=True, + cu_seqlens=inputs["cu_seqlens"], + validate_inputs=True, + ) From fa7d7221bb9ac5ce16df7327a2f721b414531f11 Mon Sep 17 00:00:00 2001 From: Hongyi Wu <62729549+Aharrypotter@users.noreply.github.com> Date: Fri, 21 Aug 2026 11:15:46 +0800 Subject: [PATCH 4/5] bench(gdn2): add the canonical five-row SM90 product/FLA matrix Five immutable rows covering MHA, GVA2, GVA4, packed variable-length input, all four state modes, T={64,1024,4096}, and N={1,20}. Compilation is recorded separately and excluded from CUDA-event timing. --- benchmarks/bench_gdn2_prefill.py | 500 +++++++++++++++++++++++++++++++ 1 file changed, 500 insertions(+) create mode 100644 benchmarks/bench_gdn2_prefill.py diff --git a/benchmarks/bench_gdn2_prefill.py b/benchmarks/bench_gdn2_prefill.py new file mode 100644 index 00000000..b404bce4 --- /dev/null +++ b/benchmarks/bench_gdn2_prefill.py @@ -0,0 +1,500 @@ +#!/usr/bin/env python3 +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Benchmark the canonical GDN2 prefill matrix on Hopper SM90. + +The five rows cover MHA, GVA2, GVA4, packed variable-length input, all four +initial/final-state modes, and the shortest and longest release sentinels. +Compilation is recorded separately and excluded from CUDA-event timing. + +When ``--implementation both`` is selected, the FLA callable is measured from +the same public logical-input boundary. Its required GVA head expansion is +inside the timed call. +""" + +from __future__ import annotations + +import argparse +import datetime as dt +import hashlib +import importlib.metadata +import json +import math +import pathlib +import platform +import statistics +import subprocess +import sys +import time +from collections.abc import Callable +from dataclasses import asdict, dataclass +from typing import Any + +import torch + +REPO_ROOT = pathlib.Path(__file__).resolve().parents[1] +sys.path.insert(0, str(REPO_ROOT)) + +from cula.gdn2 import ( # noqa: E402 + chunk_gdn2, + get_sm90_gdn2_backend, + get_sm90_gdn2_backend_identity, +) + +HEAD_SIZE = 128 +QUERY_HEADS = 16 +VALUE_SIZE = 128 +DTYPE = torch.bfloat16 + + +@dataclass(frozen=True) +class BenchmarkRow: + row_id: str + lengths: tuple[int, ...] + value_heads: int + initial_state: bool + output_final_state: bool + seed: int + + @property + def total_tokens(self) -> int: + return sum(self.lengths) + + +CANONICAL_MATRIX = ( + BenchmarkRow( + "S1-MHA-T64-H16-H0NONE-HTOFF", + (64,), + 16, + False, + False, + 6651, + ), + BenchmarkRow( + "S2-MHA-T1024-H16-H0-HTON", + (1024,), + 16, + True, + True, + 6652, + ), + BenchmarkRow( + "S3-MHA-PACKED-T4096-H16-H0NONE-HTON", + ( + 63, + 129, + 257, + 31, + 512, + 65, + 128, + 17, + 333, + 91, + 211, + 7, + 401, + 255, + 144, + 73, + 289, + 377, + 19, + 694, + ), + 16, + False, + True, + 6653, + ), + BenchmarkRow( + "S4-GVA4-T1024-H16-H64-H0-HTON", + (1024,), + 64, + True, + True, + 6654, + ), + BenchmarkRow( + "S5-GVA2-T1024-H16-H32-H0NONE-HTON", + (1024,), + 32, + False, + True, + 6655, + ), +) + + +def _tensor_sha256(tensor: torch.Tensor) -> str: + bytes_tensor = tensor.detach().contiguous().view(torch.uint8) + if bytes_tensor.is_cuda: + bytes_tensor = bytes_tensor.cpu() + return hashlib.sha256(memoryview(bytes_tensor.numpy())).hexdigest() + + +def _make_inputs( + row: BenchmarkRow, + device: torch.device, +) -> dict[str, torch.Tensor | None]: + generator = torch.Generator(device="cpu").manual_seed(row.seed) + q_shape = (row.total_tokens, QUERY_HEADS, HEAD_SIZE) + v_shape = (row.total_tokens, row.value_heads, VALUE_SIZE) + + def bf16_normal( + shape: tuple[int, ...], + standard_deviation: float, + ) -> torch.Tensor: + return (torch.randn(shape, generator=generator) * standard_deviation).to(DTYPE) + + offsets = [0] + for length in row.lengths: + offsets.append(offsets[-1] + length) + cpu_inputs: dict[str, torch.Tensor | None] = { + "q": bf16_normal(q_shape, 0.03), + "k": bf16_normal(q_shape, 0.01), + "v": bf16_normal(v_shape, 0.1), + "g": -torch.rand(q_shape, generator=generator) * 0.05, + "b": torch.rand(q_shape, generator=generator).to(DTYPE), + "w": torch.rand(v_shape, generator=generator).to(DTYPE), + "cu_seqlens": torch.tensor(offsets, dtype=torch.int64), + "initial_state": None, + } + if row.initial_state: + cpu_inputs["initial_state"] = ( + torch.randn( + ( + len(row.lengths), + row.value_heads, + VALUE_SIZE, + HEAD_SIZE, + ), + generator=generator, + ) + * 0.005 + ) + cpu_hashes = {name: _tensor_sha256(tensor) for name, tensor in cpu_inputs.items() if tensor is not None} + inputs = {name: None if tensor is None else tensor.to(device=device) for name, tensor in cpu_inputs.items()} + torch.cuda.synchronize(device) + device_hashes = {name: _tensor_sha256(tensor) for name, tensor in inputs.items() if tensor is not None} + if device_hashes != cpu_hashes: + raise AssertionError(f"CPU/CUDA input bytes differ for {row.row_id}") + return inputs + + +def _product_call( + inputs: dict[str, torch.Tensor | None], + row: BenchmarkRow, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + return chunk_gdn2( + inputs["q"], + inputs["k"], + inputs["v"], + inputs["g"], + inputs["b"], + inputs["w"], + initial_state=inputs["initial_state"], + output_final_state=row.output_final_state, + cu_seqlens=inputs["cu_seqlens"], + scale=HEAD_SIZE**-0.5, + ) + + +def _fla_call( + inputs: dict[str, torch.Tensor | None], + row: BenchmarkRow, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + from fla.ops.gdn2.chunk import chunk_gdn2 as fla_chunk_gdn2 + + q = inputs["q"] + k = inputs["k"] + v = inputs["v"] + g = inputs["g"] + b = inputs["b"] + w = inputs["w"] + cu_seqlens = inputs["cu_seqlens"] + assert all(tensor is not None for tensor in (q, k, v, g, b, w, cu_seqlens)) + group_size = row.value_heads // QUERY_HEADS + if group_size > 1: + owner = torch.arange( + row.value_heads, + dtype=torch.int64, + device=q.device, + ).div(group_size, rounding_mode="floor") + q = q.index_select(1, owner) + k = k.index_select(1, owner) + g = g.index_select(1, owner) + b = b.index_select(1, owner) + result = fla_chunk_gdn2( + q=q.unsqueeze(0), + k=k.unsqueeze(0), + v=v.unsqueeze(0), + g=g.unsqueeze(0), + b=b.unsqueeze(0), + w=w.unsqueeze(0), + initial_state=inputs["initial_state"], + scale=HEAD_SIZE**-0.5, + output_final_state=row.output_final_state, + use_qk_l2norm_in_kernel=False, + use_gate_in_kernel=False, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=None, + safe_gate=False, + disable_recompute=False, + return_intermediate_states=False, + state_v_first=True, + ) + if not isinstance(result, tuple) or len(result) < 2: + raise TypeError("FLA chunk_gdn2 must return (output, final_state)") + output, final_state = result[:2] + output = output.squeeze(0) + if row.output_final_state: + if final_state is None: + raise RuntimeError("FLA did not return the requested final state") + return output, final_state + if final_state is not None: + raise RuntimeError("FLA returned a disabled final state") + return output + + +def _normalise_result( + result: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + *, + output_final_state: bool, +) -> tuple[torch.Tensor, torch.Tensor | None]: + if output_final_state: + if not isinstance(result, tuple) or len(result) != 2: + raise TypeError("expected (output, final_state)") + return result + if isinstance(result, tuple): + raise TypeError("expected an output tensor") + return result, None + + +def _measure( + call: Callable[ + [], + torch.Tensor | tuple[torch.Tensor, torch.Tensor], + ], + *, + output_final_state: bool, + warmup: int, + iterations: int, + device: torch.device, +) -> tuple[dict[str, Any], tuple[torch.Tensor, torch.Tensor | None]]: + started = time.perf_counter() + first = call() + torch.cuda.synchronize(device) + setup_ms = (time.perf_counter() - started) * 1000.0 + first_output, first_state = _normalise_result( + first, + output_final_state=output_final_state, + ) + + for _ in range(warmup): + call() + torch.cuda.synchronize(device) + starts = [torch.cuda.Event(enable_timing=True) for _ in range(iterations)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(iterations)] + for start, end in zip(starts, ends): + start.record() + call() + end.record() + torch.cuda.synchronize(device) + samples = [float(start.elapsed_time(end)) for start, end in zip(starts, ends)] + if not all(math.isfinite(value) and value > 0 for value in samples): + raise RuntimeError("invalid CUDA-event timing sample") + return ( + { + "setup_ms_excluded": setup_ms, + "warmup": warmup, + "iterations": iterations, + "timer": "cuda_event", + "raw_per_iteration_ms": samples, + "average_ms": statistics.fmean(samples), + "median_ms": statistics.median(samples), + "minimum_ms": min(samples), + "maximum_ms": max(samples), + "output_sha256": _tensor_sha256(first_output), + "final_state_sha256": (None if first_state is None else _tensor_sha256(first_state)), + }, + (first_output, first_state), + ) + + +def _source_identity() -> dict[str, Any]: + try: + commit = subprocess.check_output( + ["git", "rev-parse", "HEAD"], + cwd=REPO_ROOT, + text=True, + stderr=subprocess.DEVNULL, + ).strip() + status = subprocess.check_output( + ["git", "status", "--short"], + cwd=REPO_ROOT, + text=True, + ) + except (FileNotFoundError, subprocess.CalledProcessError): + commit, status = None, None + return { + "repo_root": str(REPO_ROOT), + "commit": commit, + "worktree_status": status, + "backend": get_sm90_gdn2_backend(), + "backend_identity": get_sm90_gdn2_backend_identity(), + } + + +def _environment(device: torch.device) -> dict[str, Any]: + properties = torch.cuda.get_device_properties(device) + result = { + "captured_at_utc": dt.datetime.now(dt.UTC).isoformat(), + "hostname": platform.node(), + "python": sys.version, + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "cutlass_dsl": importlib.metadata.version("nvidia-cutlass-dsl"), + "device": device.index, + "gpu_name": properties.name, + "compute_capability": [properties.major, properties.minor], + } + try: + result["fla"] = importlib.metadata.version( + "flash-linear-attention", + ) + except importlib.metadata.PackageNotFoundError: + result["fla"] = None + return result + + +def _write_json(path: pathlib.Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(payload, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--implementation", + choices=("product", "both"), + default="both", + ) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iterations", type=int, default=100) + parser.add_argument("--device", type=int, default=0) + parser.add_argument( + "--output", + type=pathlib.Path, + default=pathlib.Path("gdn2-sm90-benchmark.json"), + ) + parser.add_argument("--list-matrix", action="store_true") + args = parser.parse_args() + if args.warmup < 0 or args.iterations <= 0: + parser.error( + "--warmup must be non-negative and --iterations must be positive", + ) + return args + + +def main() -> None: + args = _parse_args() + if args.list_matrix: + print(json.dumps([asdict(row) for row in CANONICAL_MATRIX], indent=2)) + return + if not torch.cuda.is_available(): + raise SystemExit("CUDA is required") + device = torch.device("cuda", args.device) + capability = torch.cuda.get_device_capability(device) + if capability != (9, 0): + raise SystemExit(f"GDN2 prefill requires SM90, got {capability}") + + record: dict[str, Any] = { + "schema": "cula.gdn2.sm90.benchmark.v1", + "status": "RUNNING", + "protocol": { + "matrix_rows": len(CANONICAL_MATRIX), + "query_heads": QUERY_HEADS, + "supported_value_heads": [16, 32, 64], + "key_size": HEAD_SIZE, + "value_size": VALUE_SIZE, + "dtype": str(DTYPE), + "warmup": args.warmup, + "iterations": args.iterations, + "statistic": "arithmetic mean of raw CUDA-event samples", + "compile_and_setup_excluded": True, + "fla_gva_expansion_inside_timed_call": True, + }, + "source": _source_identity(), + "environment": _environment(device), + "rows": [], + } + for position, row in enumerate(CANONICAL_MATRIX, 1): + print( + f"[{position:02d}/{len(CANONICAL_MATRIX):02d}] {row.row_id}", + flush=True, + ) + inputs = _make_inputs(row, device) + immutable_before = {name: _tensor_sha256(tensor) for name, tensor in inputs.items() if tensor is not None} + product_timing, product_result = _measure( + lambda: _product_call(inputs, row), + output_final_state=row.output_final_state, + warmup=args.warmup, + iterations=args.iterations, + device=device, + ) + row_record: dict[str, Any] = { + **asdict(row), + "total_tokens": row.total_tokens, + "product": product_timing, + "status": "PASS", + } + if args.implementation == "both": + fla_timing, fla_result = _measure( + lambda: _fla_call(inputs, row), + output_final_state=row.output_final_state, + warmup=args.warmup, + iterations=args.iterations, + device=device, + ) + torch.testing.assert_close( + product_result[0].float(), + fla_result[0].float(), + rtol=0.01, + atol=0.01, + ) + if row.output_final_state: + assert product_result[1] is not None + assert fla_result[1] is not None + torch.testing.assert_close( + product_result[1], + fla_result[1], + rtol=0.001, + atol=0.005, + ) + row_record["fla"] = fla_timing + row_record["speedup_over_fla"] = fla_timing["average_ms"] / product_timing["average_ms"] + immutable_after = {name: _tensor_sha256(tensor) for name, tensor in inputs.items() if tensor is not None} + if immutable_after != immutable_before: + raise RuntimeError(f"input mutated for {row.row_id}") + record["rows"].append(row_record) + print( + f" product={product_timing['average_ms']:.6f} ms" + + ("" if "speedup_over_fla" not in row_record else f" speedup={row_record['speedup_over_fla']:.3f}x"), + flush=True, + ) + + record["status"] = "PASS" + record["coverage"] = f"{len(CANONICAL_MATRIX)}/{len(CANONICAL_MATRIX)}" + _write_json(args.output, record) + print(f"GDN2_SM90_BENCHMARK_PASS output={args.output}", flush=True) + + +if __name__ == "__main__": + main() From 36a494624922841f573a43ff2a0354791cfa2cca Mon Sep 17 00:00:00 2001 From: Hongyi Wu <62729549+Aharrypotter@users.noreply.github.com> Date: Fri, 21 Aug 2026 11:15:46 +0800 Subject: [PATCH 5/5] docs(gdn2): document the API, pipeline, and stable factorization Covers the public contract and its preconditions, the warp-group schedule, the real six-field compile-cache key and its route boundaries, the audited register and shared-memory figures, the relationship to the repository-wide CuTeDSL contract, and the blockwise-rebased factorization with the bound that produces the [-5, 0] decay contract. Registers cula/gdn2 and cula/ops/gdn2 in REPO_LAYOUT.md, and adds the README quick-start and reproduction commands. --- README.md | 49 +++++ REPO_LAYOUT.md | 9 + docs/gdn2_sm90_api.md | 215 ++++++++++++++++++++++ docs/gdn2_sm90_pipeline.md | 163 ++++++++++++++++ docs/gdn2_sm90_short_n1_specialization.md | 52 ++++++ docs/gdn2_sm90_stable_factor.md | 95 ++++++++++ 6 files changed, 583 insertions(+) create mode 100644 docs/gdn2_sm90_api.md create mode 100644 docs/gdn2_sm90_pipeline.md create mode 100644 docs/gdn2_sm90_short_n1_specialization.md create mode 100644 docs/gdn2_sm90_stable_factor.md diff --git a/README.md b/README.md index 91b25a08..36ab9914 100644 --- a/README.md +++ b/README.md @@ -140,6 +140,40 @@ print(f'Final state shape: {final_state.shape}') # [2, 32, 128, 128] - `beta` supports both `float32` and `bfloat16`; `initial_state` must be `float32`. - `cu_seqlens` (for variable-length sequences) must be `int32`. +### Gated DeltaNet-2 Prefill — Hopper (SM90) + +```python +import torch + +from cula.gdn2 import chunk_gdn2 + +q = torch.randn(66, 16, 128, device="cuda", dtype=torch.bfloat16) +k = torch.randn_like(q) +v = torch.randn(66, 32, 128, device="cuda", dtype=torch.bfloat16) +g = -torch.rand_like(q, dtype=torch.float32) * 0.05 +b = torch.rand_like(q) +w = torch.rand_like(v) +cu_seqlens = torch.tensor([0, 65, 66], device="cuda", dtype=torch.int64) + +output, final_state = chunk_gdn2( + q, + k, + v, + g, + b, + w, + cu_seqlens=cu_seqlens, + output_final_state=True, +) +``` + +The SM90a CuTe DSL backend supports packed MHA, GVA2, and GVA4 with +`Hq=16`, `Hv={16,32,64}`, `K=V=128`, and up to 32 non-empty sequences. +Q/K/V/B/W are BF16, G and recurrent state are FP32, and state uses public +`[N,Hv,V,K]` orientation. Unsupported inputs fail explicitly without a +fallback. See the [GDN2 SM90 API guide](docs/gdn2_sm90_api.md) and +[GDN2 SM90 pipeline](docs/gdn2_sm90_pipeline.md). + ## Usage See [USAGE.md](USAGE.md) for detailed usage examples and notes. @@ -176,6 +210,11 @@ python benchmarks/bench_la_decode_vs_fla.py --heads 64 --head-dim 128 # Hopper (SM90) python benchmarks/bench_kda_sm90_prefill.py --mode both python benchmarks/bench_kda_sm90_cp.py + +# GDN2 prefill — canonical five-row SM90 product/FLA matrix +python benchmarks/bench_gdn2_prefill.py \ + --implementation both \ + --output gdn2-sm90-benchmark.json ``` ## Tests @@ -193,6 +232,16 @@ python tests/test_lightning_sm100_prefill.py python -m pytest tests/test_lightning_attn_prefill_dispatch.py tests/test_lightning_attn_prefill_sm90.py -v # Tests for Lightning Attention decode python -m pytest tests/test_lightning_decode.py -v +# Tests for GDN2 SM90 product dispatch and tokenwise correctness +python -m pytest tests/gdn2/test_gdn2_prefill_sm90.py -v +# GDN2 SM90 deterministic stress: one process, 100,000 product launches +python tests/gdn2/stress_gdn2_sm90.py \ + --iterations 100000 \ + --output gdn2-sm90-stress.json +# GDN2 SM90 memcheck/initcheck/synccheck/racecheck, 120 launches per tool +tests/gdn2/run_compute_sanitizer_sm90.sh \ + gdn2-sm90-sanitizers \ + 120 # test_kda_sm100_chunk_vs_naive.py and test_kda_sm100_chunk_vs_fla.py support a fast/slow split. # Fast (default) — representative correctness paths for default CI and local iteration diff --git a/REPO_LAYOUT.md b/REPO_LAYOUT.md index 2ad40239..79e4400a 100644 --- a/REPO_LAYOUT.md +++ b/REPO_LAYOUT.md @@ -10,6 +10,8 @@ cuLA/ │ ├── cudac.py # Lazy proxy for the per-architecture CUDA extension │ ├── utils.py # Architecture, stream-buffer, and cu_seqlens helpers │ │ +│ ├── gdn2/ # [non-KDA] Gated DeltaNet-2 public API +│ │ │ ├── kda/ # KDA public API, wrappers, autograd, routing, and Triton support kernels │ │ ├── __init__.py # Lazy exports for chunk, prefill, and decode APIs │ │ ├── backends/ # kda_prefill runtime dispatch @@ -32,6 +34,12 @@ cuLA/ │ └── ops/ # CuTeDSL kernels and shared low-level helpers │ ├── inv.py / ptx.py # Shared low-level helpers │ ├── sm100/ptx.py # Shared SM100 PTX helpers +│ ├── gdn2/ # [non-KDA] Gated DeltaNet-2 prefill kernels +│ │ └── sm90/ # SM90a CuTe DSL implementation +│ │ ├── config.py # Host-side product contract and supported ranges +│ │ ├── prefill.py # Validation, compile cache, and TVM-FFI launch +│ │ ├── prefill_kernel.py # Fused packed recurrent prefill kernel +│ │ └── collective_inverse_hmma.py / inverse_helpers.py # Triangular inverse │ ├── kda/ │ │ ├── cp_mode.py # Shared intracard-CP mode vocabulary │ │ ├── sm100/ # Blackwell modular forward/backward kernels @@ -82,5 +90,6 @@ cuLA/ | `cula/ops/kda/` | Python (CuTeDSL) | CuTeDSL KDA kernels organized into SM100 modular kernels, SM90 FlashKDA K1+K2 and intracard CP, decode, and experimental code. The fully-fused SM90 implementation lives under `csrc/`, not here. | | `csrc/kda/{sm90,sm100}/` | CUDA C++ | Hopper fully-fused prefill and Blackwell modular chunk kernels. | | `csrc/api/` · `cula/cudac.py` | CUDA C++ / Python | Per-architecture `_cudac_sm90` and `_cudac_sm100` extensions, exposed lazily through the `cula.cudac` compatibility proxy. | +| `cula/gdn2/` · `cula/ops/gdn2/` | Python (CuTeDSL) | `[non-KDA]` Gated DeltaNet-2 packed-varlen prefill. `chunk_gdn2` dispatches directly to the SM90a backend with no fallback; unsupported inputs and out-of-range CuTeDSL versions fail closed. See [`docs/gdn2_sm90_api.md`](docs/gdn2_sm90_api.md). | | `cula/ops/lightning/` · `cula/ops/experimental/` | Python (CuTeDSL) | `[non-KDA]` Lightning/linear-attention kernels and prototypes. Lightning prefill dispatches to the SM90 or SM100 backend from Q's device capability. | | `cula/ops/{inv,ptx}.py` · `cula/ops/sm100/ptx.py` | Python | Shared low-level helpers used across operators. | diff --git a/docs/gdn2_sm90_api.md b/docs/gdn2_sm90_api.md new file mode 100644 index 00000000..bd68dfdd --- /dev/null +++ b/docs/gdn2_sm90_api.md @@ -0,0 +1,215 @@ +# Gated DeltaNet-2 Prefill API on Hopper SM90 + +cuLA provides a fully fused, packed variable-length Gated DeltaNet-2 (GDN2) +forward-prefill kernel for NVIDIA Hopper GPUs. The implementation is written in +CuTe DSL and has one product backend: +`sm90a_cutedsl_gdn2_prefill_v1`. Unsupported inputs fail explicitly; the public +entry point does not fall back to Triton or another cuLA kernel. + +## Requirements + +- NVIDIA compute capability 9.0 (Hopper SM90, including H20 and H100) +- Python 3.12 or newer +- CUDA and PyTorch versions supported by the surrounding cuLA installation +- `nvidia-cutlass-dsl>=4.5.1,<4.7`. `is_sm90_gdn2_available()` and the + dispatch error enforce exactly this range and report the backend + unavailable outside it. Both endpoints are exercised on H20: 4.5.1 and + 4.6.2 each pass the full product suite. This range is narrower than the + project-wide dependency in `pyproject.toml`, because the kernel needs + `cutlass.cute.nvgpu.OperandMajorMode`, which 4.4.x does not provide — the + backend cannot even be imported there. The upper bound matches the + repository-wide CuTeDSL contract in `cula/ops/_mlir_compat.py`, which is + enforced independently on every private-dialect access; GDN2 only raises + the floor, and reads the installed version through that same gateway. + +## API + +```python +from cula.gdn2 import chunk_gdn2 +``` + +The public tensors use packed-token layouts: + +| Argument | Dtype | Shape | Meaning | +|---|---|---|---| +| `q` | BF16 | `[T,16,128]` | queries | +| `k` | BF16 | `[T,16,128]` | keys | +| `v` | BF16 | `[T,Hv,128]` | values | +| `g` | FP32 | `[T,16,128]` | finite log decay in `[-5, 0]` | +| `b` | BF16 | `[T,16,128]` | erase gate in `[0,1]` | +| `w` | BF16 | `[T,Hv,128]` | write gate in `[0,1]` | +| `cu_seqlens` | INT64 | `[N+1]` | packed sequence prefixes | +| `initial_state` | FP32 | `[N,Hv,128,128]` | optional state in public `[V,K]` orientation | + +`Hv` may be 16, 32, or 64. This gives MHA (`Hv=16`), GVA2 (`Hv=32`), +or GVA4 (`Hv=64`). GQA and other head relationships are outside the first +product contract. `N` must be in `[1,32]`, and every sequence must contain at +least one token. + +All tensors must be contiguous, CUDA-resident on the same device, and at least +16-byte aligned. The default path validates tensor metadata without copying +device values to the host. The caller must ensure: + +- `cu_seqlens[0] == 0`; +- `cu_seqlens[-1] == T`; +- offsets are strictly increasing; +- `g` is finite and in `[-5, 0]` elementwise (the blockwise-rebased + factorization bound; see + [GDN2 SM90 stable factorization](gdn2_sm90_stable_factor.md)); +- `b` and `w` are finite and in `[0,1]`. + +Pass `validate_inputs=True` for a synchronous diagnostic check of those value +preconditions. Do not enable it on a latency-sensitive steady-state path. + +The output is BF16 `[T,Hv,128]`. With `output_final_state=True`, the function +returns `(output, final_state)`, where `final_state` is FP32 +`[N,Hv,128,128]` in the same public `[V,K]` orientation accepted by +`initial_state`. Callers may provide contiguous `output` and `output_state` +buffers. + +## Numerical validation + +The SM90 correctness suite uses an independent tokenwise PyTorch FP32 +recurrence and requires finite values from both implementations. The frozen +comparison policies are: + +| Result | Product dtype | Reference accumulation | `rtol` | `atol` | +|---|---|---|---:|---:| +| Output | BF16 | FP32, cast to BF16 for comparison | `0.01` | `0.01` | +| Final state | FP32 | FP32 | `0.001` | `0.005` | + +An unexpected `NaN` or `Inf` in either output or final state fails the test +before tolerance comparison. Repeated product launches with identical inputs, +initial state, seed, and configuration must be bitwise exact. + +## Example + +```python +import torch + +from cula.gdn2 import chunk_gdn2 + +device = "cuda" +lengths = (65, 1) +total_tokens = sum(lengths) +query_heads = 16 +value_heads = 32 # GVA2 + +q = torch.randn( + total_tokens, + query_heads, + 128, + device=device, + dtype=torch.bfloat16, +) +k = torch.randn_like(q) +v = torch.randn( + total_tokens, + value_heads, + 128, + device=device, + dtype=torch.bfloat16, +) +g = -torch.rand_like(q, dtype=torch.float32) * 0.05 +b = torch.rand_like(q) +w = torch.rand_like(v) +cu_seqlens = torch.tensor([0, 65, 66], device=device, dtype=torch.int64) + +output, final_state = chunk_gdn2( + q, + k, + v, + g, + b, + w, + cu_seqlens=cu_seqlens, + output_final_state=True, +) +``` + +Pass a returned `final_state` as `initial_state` in a later call to continue +the recurrence without changing orientation. + +## Unsupported behavior + +- compute capability other than 9.0; +- `Hq` other than 16 or `Hv` outside `{16,32,64}`; +- `N > 32`, zero-length sequences, or `T > 2^31-1`; +- key or value size other than 128; +- FP16, FP8, or FP32 Q/K/V; +- GQA, decode, backward, or intermediate state checkpoints; +- implicit fallback to FLA Triton or another backend. + +Unsupported metadata fails before compilation or launch with `ValueError`, +`TypeError`, `NotImplementedError`, or `RuntimeError`. + +## Canonical benchmark + +From the repository root, run: + +```bash +python benchmarks/bench_gdn2_prefill.py \ + --implementation both \ + --output gdn2-sm90-benchmark.json +``` + +The five immutable rows cover MHA, GVA2, GVA4, packed variable-length input, +all four initial/final-state modes, `T={64,1024,4096}`, and `N={1,20}`. +Compilation is recorded separately and excluded from CUDA-event timing. The +FLA comparison includes any GVA head expansion inside its timed public logical +call. Use `--list-matrix` to inspect the rows without running CUDA work. + +The standalone benchmark is a developer diagnostic. Release performance claims +must additionally use fresh processes, fresh caches, alternating implementation +order, exact source/input identities, and independently replayable raw timing +receipts. + +## Deterministic stress and Compute Sanitizer + +Run the deterministic SM90 stress matrix in one process: + +```bash +python tests/gdn2/stress_gdn2_sm90.py \ + --iterations 100000 \ + --warmup 1 \ + --device 0 \ + --output gdn2-sm90-stress.json +``` + +The six-row round-robin matrix covers all six product compile +specializations, MHA/GVA2/GVA4, all four state modes, `T={64,1024,4096}`, +`N=32`, and irregular packed tails. Every launch is compared bitwise with its +fixed-input output/final-state baseline using CUDA-stream-ordered checks. The +harness also accumulates non-finite counts and verifies input hashes, +redzones, backend identity, GPU identity, and the exact launch count. + +Run repeated stress under all four applicable NVIDIA Compute Sanitizer tools: + +```bash +tests/gdn2/run_compute_sanitizer_sm90.sh \ + gdn2-sm90-sanitizers \ + 120 +``` + +The runner uses separate empty compiler caches and records the exact command, +return code, stdout/stderr, JSON receipt, and SHA-256 evidence manifest for +each of: + +- `memcheck --leak-check full`; +- `initcheck`; +- `synccheck --check-warpgroup-mma yes`; +- `racecheck --racecheck-report hazard --racecheck-memcpy-async yes + --racecheck-trace-sync yes`. + +The runner takes its interpreter from `PYTHON` and falls back to `python3`. +Set it explicitly to the interpreter whose `nvidia-cutlass-dsl` is inside the +supported range — a bare `python3` may resolve to a system interpreter with a +different, unsupported version, in which case the backend now fails closed +rather than silently producing evidence on an unvalidated toolchain. + +For a source-bound release run, set `GDN2_SOURCE_MANIFEST` to the frozen +source manifest and `GDN2_REQUIRED_GPU_UUID` to the expected `GPU-...` +identity before invoking the sanitizer runner. + +For the kernel decomposition and scheduling policy, see +[GDN2 SM90 Pipeline](gdn2_sm90_pipeline.md). diff --git a/docs/gdn2_sm90_pipeline.md b/docs/gdn2_sm90_pipeline.md new file mode 100644 index 00000000..6eb42918 --- /dev/null +++ b/docs/gdn2_sm90_pipeline.md @@ -0,0 +1,163 @@ +# Fully Fused Gated DeltaNet-2 Prefill SM90 Pipeline + +> File: `cula/ops/gdn2/sm90/prefill_kernel.py` +> Class: `GDN2PrefillKernel` + +## Recurrence + +The public state uses `[sequence,Hv,V,K]`. For explanation, transpose one +head's state to the internal matrix `S_t` with shape `[K,V]`. At token `t`: + +$$ +\bar S_t = \operatorname{diag}(\exp(g_t)) S_{t-1} +$$ + +$$ +u_t = w_t \odot v_t - (b_t \odot k_t)^T \bar S_t +$$ + +$$ +S_t = \bar S_t + k_t u_t^T +$$ + +$$ +o_t = \mathrm{scale}\; q_t^T S_t +$$ + +`q`, `k`, `b`, and `g` have 16 query/key heads. A GVA output head maps to its +owner query/key head while retaining its own `v`, `w`, recurrent state, and +output. + +The kernel processes 64 tokens per chunk. It constructs the exact chunk-local +FP32 prefix from public raw `g`, builds the causal factor matrices for the +recurrence, and applies them without a global prefix workspace or a second +kernel launch. The intra-chunk factor matrices use the blockwise-rebased +factorization described in +[GDN2 SM90 stable factorization](gdn2_sm90_stable_factor.md): 16-token +sub-block operands with bounded exponents, warp-level MMA block pairs on the +producer warp group, and `n=16`-sliced state projections with per-block decay +deltas on the state warp groups. + +## Launch and thread roles + +Each CTA owns one `(sequence,value_head)` work unit and processes all of that +sequence's chunks in order. + +| Resource | Product configuration | +|---|---:| +| Grid | `(N * Hv, 1, 1)` | +| Threads | 384 (12 warps, 3 warp groups) | +| Chunk size | 64 tokens | +| Dynamic shared memory | 232,192 B (V128 routes), 207,616 B (V64 route) | +| Register allocation | 168 registers per thread, all specializations | +| Spill | none: zero local load/store traffic | +| Residency | one CTA per SM | + +The five V128 specializations sit 256 bytes under the 232,448-byte SM90a +per-CTA limit. That headroom is the binding constraint on the schedule: it +rules out a third input stage, a second factor-workspace stage, and any +de-aliasing of the FP16 Gram workspace on those routes. Only the V64 route +(`N=1`, `Hv=16`, initial and final state, `T>64`) has meaningful slack. + +Spill is measured as SASS local load/store traffic, which is zero for every +specialization. The 1,024-byte `launch__stack_size` reported by the profiler +is the driver's fixed ABI reserve, not per-kernel spill. + +The three warp groups have two coarse roles: + +| Warp group | Role | +|---|---| +| WG0 | TMA input production, raw-G/factor preparation, shared factor publication, output consumption, and global stores | +| WG1 | FP32 recurrent-state slab for value rows `[0,64)` plus state/output WGMMA | +| WG2 | FP32 recurrent-state slab for value rows `[64,128)` plus state/output WGMMA | + +WG1 and WG2 retain their FP32 state fragments in registers across chunks. They +exchange only staged operands and outputs through shared memory. CTA, named, +pipeline, and mbarrier synchronization provide explicit ownership transfers +between WG0 and the two state warp groups. + +## Chunk pipeline + +```text + WG0 producer / factor / store WG1 + WG2 state math + | | + |-- TMA raw Q,K,B,G,V,W ------------------------>| + | | + |-- build FP32 G prefix and factor matrices | + |-- publish Qbar, erase, A_QK, A_KK^-1 --------->| + | |-- decay state + | |-- read old state + | |-- form write value + | |-- update state + |<---------------- publish BF16 output tile -----| + |-- TMA/store valid output tokens | + | | + +------------------- next chunk ------------------+ +``` + +The input side is double-buffered by chunk stage. Value/write operands use +private producer stages, factor readiness and completion use explicit handoff +barriers, and output staging is pipelined independently. Invalid tail tokens +receive neutral values and are never written to global output. + +## Stable LPT32 sequence schedule + +The grid shape remains `N * Hv`, but the sequence dimension is remapped before +work starts: + +1. Each sequence's cost is `ceil(length / 64)`. +2. Sequences are ranked by descending cost. +3. Equal costs retain ascending original sequence index. +4. Every value head uses the same remapping. + +The rank is computed on-device from `cu_seqlens` for `N <= 32`. One warp +publishes the selected sequence through an existing shared slot, followed by a +single CTA synchronization. No host-side length sort, metadata copy, or extra +kernel is required. + +This schedule keeps the longest sequence waves at the front of the grid. It +reduces tail under-utilization for imbalanced packed batches while preserving +exact recurrence order within every sequence. + +## State modes and dynamic compilation + +The host adapter caches one specialization per compile key: + +```text +(device, Hv, has_initial_state, output_final_state, + use_n1_hv16_v64, retain_final_tail) +``` + +The last two fields are shape-derived dispatch booleans introduced by the +[short-sequence / N=1 specialization](gdn2_sm90_short_n1_specialization.md): + +- `retain_final_tail = output_final_state and not (N == 1 and T <= 64)`; +- `use_n1_hv16_v64 = N == 1 and Hv == 16 and has_initial_state and + output_final_state and T > 64`. + +`T` and `N` stay dynamic *within one dispatch route*, but crossing a route +boundary compiles a new specialization. Concretely: + +- with `output_final_state=False`, both derived fields are constant `False`, + so the key reduces to `(device, Hv, has_initial_state)` and `T`/`N` are + fully dynamic; +- with `output_final_state=True`, there are up to three route + specializations per `(device, Hv, has_initial_state)`: `N=1, T<=64`; + `N=1, T>64` (a distinct V64 route only for `Hv=16` with an initial state); + and every other supported shape. + +Across `Hv={16,32,64}` and the four state modes this gives at most 19 +specializations per device. Each first launch of a specialization pays a +one-time compilation on the order of tens of seconds; latency-sensitive +deployments should prewarm every route they will serve, including both sides +of the `N=1` / `T=64` boundaries when final states are requested. +Compilation and setup are excluded from steady-state benchmark timing. + +## Fail-closed behavior + +The public wrapper rejects unsupported tensor metadata before compilation. +The device kernel additionally checks sequence boundaries used by each CTA and +traps on invalid offsets. There is no alternate backend or silent fallback. + +For shapes, dtypes, value preconditions, and the canonical benchmark, see +[GDN2 SM90 API](gdn2_sm90_api.md). diff --git a/docs/gdn2_sm90_short_n1_specialization.md b/docs/gdn2_sm90_short_n1_specialization.md new file mode 100644 index 00000000..51ab25f6 --- /dev/null +++ b/docs/gdn2_sm90_short_n1_specialization.md @@ -0,0 +1,52 @@ +# GDN2 SM90 short-sequence / N=1 V64 specialization + +## Design + +A shape-driven dispatch extension to the production `GDN2PrefillKernel` that +improves the two highest-impact short/medium single-sequence routes without +changing the public API, numerical contract, or default behavior for any other +supported shape. + +### Dispatch rules(automatic, no environment switch) + +| Condition | Route | +|---|---| +| `N == 1 and T <= 64` | exact released preparation/commit schedule(V128, no retained tail) | +| `N == 1 and Hv == 16 and initial_state and final_state and T > 64` | V64 single State-WG with register-resident final-tail carry | +| all other supported shapes | unchanged V128 production path | + +The dispatch is derived purely from input metadata inside `_compile`; the +compile-cache key is extended by the two derived booleans +(`use_n1_hv16_v64`, `retain_final_tail`) so each route compiles exactly one +specialization. + +## Kernel changes(`GDN2PrefillKernel`) + +Three new constructor parameters, each validated at construction: + +- `value_tile: int = VALUE_SIZE` — `64` or `128`; +- `single_state_owner: bool = False` — must equal `(value_tile == 64)`; +- `retain_final_tail: bool = False` — register-resident final-tail carry for + final-state routes. + +All branches are compile-time (`cutlass.const_expr`), so the V128 production +path is instruction-identical to the released kernel when the new flags are +at their defaults. + +## Validation(H20, source-bound) + +- Correctness: `13 passed`(`9` matrix cases + `4` contract tests) including + explicit `mha-t64-short-baseline` and `mha-t65-v64-boundary` boundary cases. +- Determinism: `100,000` launches / 6 rows PASS. +- Compute Sanitizer: memcheck / initcheck / synccheck / racecheck, + `120` launches per tool, all PASS. +- Codegen/resources: no stack, local memory, or spill. +- Paired timing(S2/Q2 family): `0.9362x` vs released incumbent on the + targeted rows(≈ 6.4% latency reduction); full S1-S5 all rows faster than + pinned FLA. + +## Scope boundary + +This change is additive and does not alter the public `chunk_gdn2` signature, +the `[N,Hv,V,K]` state layout, supported-shape matrix, or any existing +behavior outside the two dispatch rules above. diff --git a/docs/gdn2_sm90_stable_factor.md b/docs/gdn2_sm90_stable_factor.md new file mode 100644 index 00000000..871939d0 --- /dev/null +++ b/docs/gdn2_sm90_stable_factor.md @@ -0,0 +1,95 @@ +# GDN2 SM90 blockwise-rebased intra-chunk factorization + +## Problem + +The released kernel formed the intra-chunk causal matrices from a +chunk-start-referenced split: + +```text +A[i,j] = sum_c ( q_i[c] * exp(G_i[c]) ) * ( k_j[c] * exp(-G_j[c]) ) +``` + +with `G` the chunk-local inclusive per-channel log-decay prefix. The +`exp(-G_j)` factor is an unbounded inverse decay: any 64-token chunk whose +per-channel accumulated decay exceeds `ln(FP32_MAX) ~ 88.72` overflows the +`k`-side operand even though every masked product `exp(G_i - G_j) <= 1` +stays finite. With a uniform per-element `g`, the cliff sits at +`|g| > 88.72 / 64 ~ 1.386`, inside the public contract `g <= 0`. +Because `G` is monotone, the overflow always pairs with an underflowed +`q`-side row, so masked-region entries become `0 * inf = NaN` and the NaN +propagates through the erase-Gram inverse and the recurrent state to every +later token of the sequence. + +## Factorization + +The chunk is split into four 16-token sub-blocks. `Gs(I)` denotes the +inclusive prefix at the first token of sub-block `I` (clamped to the last +valid token for partial tails). The prepared operands are: + +```text +q~[i] = q_i * exp(G_i - Gs(B(i))) # <= 1, span <= 15 tokens +e~[i] = b_i * k_i * exp(G_i - Gs(B(i))) # <= 1, span <= 15 tokens +k~'[j] = k_j * exp(Gs(B(j)) - G_j) # >= 1, span <= 15 tokens +``` + +where `B(t) = t // 16`. Per-pair block products then satisfy: + +- diagonal pairs `(I,I)`: `q~ k~'` is exactly `q k exp(G_i - G_j)`; +- off-diagonal pairs `(I,J), J < I`: the product must be corrected by the + per-channel factor `s'[c] = exp(Gs(I) - Gs(J))[c] <= 1`, folded into the + left fragment before the MMA. The folded left operand equals + `q exp(G_i - Gs(J)) <= 1`. + +The only intermediate with a positive exponent is `k~'`, whose span is 15 +token gaps instead of 63, moving the uniform-`g` overflow cliff from +`~1.386` to `88.72 / 15 ~ 5.91`. The documented contract is +`g in [-5, 0]`, matching the pinned FLA GDN2 `safe_gate` documented range; +at `g = -5` the largest `k~'` exponent is `75`, a factor `~9e5` below the +BF16/FP32 overflow boundary. + +The per-channel correction factors are served from a small FP32 SMEM +buffer of block-boundary decay ratios written during preparation: + +```text +delta[0][c] = exp(Gs(0)[c]) # = exp(g_0[c]) <= 1 +delta[m][c] = exp(Gs(m)[c] - Gs(m-1)[c]) # <= 1, m in {1,2,3} +``` + +`s'` for pair `(I,J)` is the running product `delta[J+1] * ... * delta[I]`, +computed in registers by the factor warp group; the recurrent-state warp +groups consume the same rows to advance the state scale block by block +(below). Fully-invalid tail blocks store `delta = 1`. + +## Factor stage + +`A_qk` (causal, scaled, BF16) and the strict-lower erase Gram (FP16, into +the collective-inverse workspace) are computed per sub-block pair by the +factor warp group with warp-level `m16n8k16` MMAs: ten lower/diagonal +16x16 blocks per matrix, four warps, the `k~'` right fragment shared +between both matrices of a pair. Upper blocks are zero-filled. The +collective inverse, the BF16 `A_kk` publication, and every downstream +consumer tile are unchanged. + +## Recurrent-state warp groups + +The inter-chunk output and erase projections previously consumed +`q_bar = q exp(G_i)` and `erase_bar = b k exp(G_i)` as single +`n=64` WGMMAs against the register-resident state. Those tiles now hold +`q~`/`e~`, so both projections are issued as four `n=16` slices with the +state fragments rescaled by `delta[I]` before slice `I`; after slice 3 the +state carries `exp(Gs(3))`, and the end-of-chunk update multiplies by the +retuned `gamma_end = exp(G_end - Gs(3)) <= 1` so the recurrence +`S <- S exp(G_end) + k_tail v_new^T` is unchanged. `key_tail` +(`k exp(G_end - G)`) and the final-tail carry path are untouched: they were +already in stable form. + +## Contract + +- `g` must be finite, non-positive, and elementwise `>= -5.0` + (`docs/gdn2_sm90_api.md`); `validate_inputs=True` enforces the bound + synchronously, and the default path documents it as a caller + precondition, exactly like the existing `g <= 0` and gate-range + preconditions. +- The bound is a per-15-token-channel-sum requirement + (`sum |g|` over any 15 consecutive in-block gaps `< 88.7`); the + elementwise `-5` bound is the documented sufficient condition.