From c284341f6de0e6f4113735face15126288276e5f Mon Sep 17 00:00:00 2001 From: huanxian Date: Mon, 9 Nov 2026 17:57:36 +0800 Subject: [PATCH 1/3] Port FlashKDA fp32 forward-substitution inverse --- cula/ops/kda/sm90/k1.py | 244 +++++++++++++++++++++------------------- 1 file changed, 131 insertions(+), 113 deletions(-) diff --git a/cula/ops/kda/sm90/k1.py b/cula/ops/kda/sm90/k1.py index 7edd60a9..6f1b5857 100644 --- a/cula/ops/kda/sm90/k1.py +++ b/cula/ops/kda/sm90/k1.py @@ -30,13 +30,136 @@ from cutlass.cute.nvgpu.warpgroup import SmemLayoutAtomKind, make_smem_layout_atom from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream -from cula.ops.kda.sm90._common import _stream_key, add_f16x2_u32, movm_t_b16 +from cula.ops.kda.sm90._common import _stream_key, movm_t_b16 CHUNK: int = 16 D: int = 128 THREADS_PER_CTA: int = 256 +@cute.jit +def _inv_fwd_subst_16(sL_fp32: cute.Tensor, sM_bf16: cute.Tensor, sINV_bf16: cute.Tensor, tidx): + """Invert ``I + L`` using FlashKDA's fp32-forward-substitution path. + + Two 8x8 diagonal blocks are inverted in fp32. Their block-diagonal + inverse ``P`` and the lower-left seed block ``M`` are then merged as + ``P + (-P @ M) @ P`` with bf16 operands and fp32 MMA accumulation. + """ + lane = tidx % 32 + row = lane % 8 + block_base = ((lane // 8) % 2) * 8 + group = lane // 8 + inv = cute.make_rmem_tensor(cute.make_layout((8,), stride=(1,)), cutlass.Float32) + + for p in cutlass.range_constexpr(8): + is_lower = cutlass.Float32(p < row) + is_diag = cutlass.Float32(p == row) + inv[p] = cutlass.Float32(sL_fp32[block_base + row, block_base + p]) * is_lower + is_diag + + for src_row in cutlass.range_constexpr(7): + row_scale = -inv[src_row] + target_lane = (lane // 8) * 8 + src_row + for p in cutlass.range_constexpr(src_row): + pivot = cute.arch.shuffle_sync_op( + value=inv[p], + offset=target_lane, + mask=0xFFFFFFFF, + mask_and_clamp=31, + ) + inv[p] = cutlass.select_(row > src_row, row_scale * pivot + inv[p], inv[p]) + inv[src_row] = cutlass.select_(row > src_row, row_scale, inv[src_row]) + + zero = cutlass.BFloat16(0.0) + if group == 0: + for col in cutlass.range_constexpr(16): + sINV_bf16[row, col] = cutlass.BFloat16(inv[col]) if col < 8 else zero + elif group == 1: + for col in cutlass.range_constexpr(16): + sINV_bf16[8 + row, col] = zero if col < 8 else cutlass.BFloat16(inv[col - 8]) + elif group == 2: + for col in cutlass.range_constexpr(16): + sM_bf16[row, col] = zero + else: + for col in cutlass.range_constexpr(16): + sM_bf16[8 + row, col] = cutlass.BFloat16(sL_fp32[8 + row, col]) if col < 8 else zero + cute.arch.sync_warp() + + mma_atom = warp.MmaF16BF16Op(cutlass.BFloat16, cutlass.Float32, (16, 8, 16)) + tiled_mma = cute.make_tiled_mma( + mma_atom, + atom_layout_mnk=(1, 1, 1), + permutation_mnk=(16, 16, 16), + ) + thr_mma = tiled_mma.get_slice(lane) + copy_atom_A = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), + cutlass.BFloat16, + ) + tiled_copy_A = cute.make_tiled_copy_A(copy_atom_A, tiled_mma) + thr_copy_A = tiled_copy_A.get_slice(lane) + copy_atom_stsm = cute.make_copy_atom( + warp.StMatrix8x8x16bOp(transpose=False, num_matrices=2), + cutlass.BFloat16, + ) + tiled_store_C = cute.make_tiled_copy_C_atom(copy_atom_stsm, tiled_mma) + thr_store_C = tiled_store_C.get_slice(lane) + + tCrP_A = thr_mma.make_fragment_A(thr_mma.partition_A(sINV_bf16)) + tCrP_A_cv = thr_copy_A.retile(tCrP_A) + cute.copy(tiled_copy_A, thr_copy_A.partition_S(sINV_bf16), tCrP_A_cv) + + tCrM_A = thr_mma.make_fragment_A(thr_mma.partition_A(sM_bf16)) + tCrM_A_cv = thr_copy_A.retile(tCrM_A) + cute.copy(tiled_copy_A, thr_copy_A.partition_S(sM_bf16), tCrM_A_cv) + tCrM_B = thr_mma.make_fragment_B(thr_mma.partition_B(sM_bf16)) + tCrM_A_u32 = cute.recast_tensor(tCrM_A, dtype=cutlass.Int32) + tCrM_B_u32 = cute.recast_tensor(tCrM_B, dtype=cutlass.Int32) + for ii in cutlass.range_constexpr(cute.size(tCrM_A_u32)): + tCrM_B_u32[ii] = movm_t_b16(cutlass.Int32(tCrM_A_u32[ii])) + + tCrDC = thr_mma.make_fragment_C(tiled_mma.partition_shape_C((16, 16))) + tCrDC.fill(0.0) + cute.gemm(tiled_mma, tCrDC, tCrP_A, tCrM_B, tCrDC) + tCrNegDC_bf16 = cute.make_fragment_like(tCrDC, cutlass.BFloat16) + for ii in cutlass.range_constexpr(cute.size(tCrDC)): + tCrNegDC_bf16[ii] = cutlass.BFloat16(-tCrDC[ii]) + cute.copy( + tiled_store_C, + thr_store_C.retile(tCrNegDC_bf16), + thr_store_C.partition_D(sM_bf16), + ) + cute.arch.sync_warp() + + tCrNegDC_A = thr_mma.make_fragment_A(thr_mma.partition_A(sM_bf16)) + tCrNegDC_A_cv = thr_copy_A.retile(tCrNegDC_A) + cute.copy(tiled_copy_A, thr_copy_A.partition_S(sM_bf16), tCrNegDC_A_cv) + tCrP_B = thr_mma.make_fragment_B(thr_mma.partition_B(sINV_bf16)) + tCrP_A_u32 = cute.recast_tensor(tCrP_A, dtype=cutlass.Int32) + tCrP_B_u32 = cute.recast_tensor(tCrP_B, dtype=cutlass.Int32) + for ii in cutlass.range_constexpr(cute.size(tCrP_A_u32)): + tCrP_B_u32[ii] = movm_t_b16(cutlass.Int32(tCrP_A_u32[ii])) + + tCrO = thr_mma.make_fragment_C(tiled_mma.partition_shape_C((16, 16))) + tCrO.fill(0.0) + cute.gemm(tiled_mma, tCrO, tCrNegDC_A, tCrP_B, tCrO) + tCrO_bf16 = cute.make_fragment_like(tCrO, cutlass.BFloat16) + for ii in cutlass.range_constexpr(cute.size(tCrO)): + tCrO_bf16[ii] = cutlass.BFloat16(tCrO[ii]) + cute.copy( + tiled_store_C, + thr_store_C.retile(tCrO_bf16), + thr_store_C.partition_D(sM_bf16), + ) + cute.arch.sync_warp() + + coord = cute.make_identity_tensor((16, 16)) + tCcC = thr_mma.partition_C(coord) + for ii in cutlass.range_constexpr(cute.size(tCcC)): + m, n = tCcC[ii] + sINV_bf16[m, n] = cutlass.BFloat16(cutlass.Float32(sINV_bf16[m, n]) + cutlass.Float32(sM_bf16[m, n])) + cute.arch.sync_warp() + + @cute.kernel def k1_kernel( tma_atom_q: cute.CopyAtom, @@ -87,15 +210,14 @@ def k1_kernel( s_q_decayed = cute.make_tensor(sQ.iterator, kinter_qk_layout) s_k_decayed = cute.make_tensor(sK.iterator, kinter_qk_layout) s_k_restored = cute.make_tensor(sG_raw.iterator, kinter_qk_layout) - sL_bf16 = smem.allocate_tensor(cutlass.BFloat16, cc_layout, 128) sMqk_bf16 = smem.allocate_tensor(cutlass.BFloat16, cc_layout, 128) sBetaSig = smem.allocate_tensor(cutlass.Float32, cute.make_layout((CHUNK,)), 128) - # Neumann buffers (alias sG_cumsum after decay_apply barrier) - sG_cumsum_fp16_ptr = cute.recast_ptr(sG_cumsum.iterator, dtype=cutlass.Float16) + # Inverse buffers alias sG_cumsum after decay_apply. L remains fp32; + # INV and the block-merge matrix use bf16 Tensor Core operands. + sL_fp32 = cute.make_tensor(sG_cumsum.iterator, cc_layout) sG_cumsum_bf16_ptr = cute.recast_ptr(sG_cumsum.iterator, dtype=cutlass.BFloat16) - sL_fp16 = cute.make_tensor(sG_cumsum_fp16_ptr, cc_layout) - sINV_fp16 = cute.make_tensor(sG_cumsum_fp16_ptr + (CHUNK * CHUNK), cc_layout) sINV_bf16 = cute.make_tensor(sG_cumsum_bf16_ptr + (2 * CHUNK * CHUNK), cc_layout) + sM_bf16 = cute.make_tensor(sG_cumsum_bf16_ptr + (3 * CHUNK * CHUNK), cc_layout) sMbar = smem.allocate_tensor(cutlass.Int64, cute.make_layout((1,)), 8) sMbar_ptr = sMbar.iterator @@ -368,14 +490,9 @@ def k1_kernel( keep = cutlass.Float32(1.0) if m > n else cutlass.Float32(0.0) tCrC_l[ii] = tCrC_l[ii] * keep * sBetaSig[m] - tCrC_l_bf16 = cute.make_fragment_like(tCrC_l, cutlass.BFloat16) for ii in cutlass.range_constexpr(cute.size(tCrC_l)): - tCrC_l_bf16[ii] = cutlass.BFloat16(tCrC_l[ii]) - cute.copy( - smem_tiled_store_mask_mma, - smem_thr_store_mask_mma.retile(tCrC_l_bf16), - smem_thr_store_mask_mma.partition_D(sL_bf16), - ) + crd = tCcC_l[ii] + sL_fp32[crd[0], crd[1]] = tCrC_l[ii] elif warp_idx == 1: sA_tile_m = cute.flat_divide(s_q_decayed, (CHUNK, 16)) sA_ref_m = sA_tile_m[None, None, 0, 0] @@ -413,107 +530,8 @@ def k1_kernel( ) cute.arch.barrier() - i = tidx // CHUNK - col = tidx % CHUNK - l_bf = cutlass.Float32(sL_bf16[i, col]) - sL_fp16[i, col] = cutlass.Float16(l_bf) - inv_init = cutlass.Float32(1.0 if i == col else 0.0) - l_bf - sINV_fp16[i, col] = cutlass.Float16(inv_init) - cute.arch.barrier() - if warp_idx == 0: - mma_atom_neumann = warp.MmaF16BF16Op(cutlass.Float16, cutlass.Float16, (16, 8, 16)) - tiled_mma_neumann = cute.make_tiled_mma( - mma_atom_neumann, - atom_layout_mnk=(1, 1, 1), - permutation_mnk=(16, 16, 16), - ) - thr_mma_neumann = tiled_mma_neumann.get_slice(tidx) - - copy_atom_A_neumann = cute.make_copy_atom( - warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), - cutlass.Float16, - ) - smem_tiled_copy_A_neumann = cute.make_tiled_copy_A(copy_atom_A_neumann, tiled_mma_neumann) - smem_thr_copy_A_neumann = smem_tiled_copy_A_neumann.get_slice(tidx) - - tCrL = thr_mma_neumann.make_fragment_A(thr_mma_neumann.partition_A(sL_fp16)) - tCrL_cv = smem_thr_copy_A_neumann.retile(tCrL) - cute.copy(smem_tiled_copy_A_neumann, smem_thr_copy_A_neumann.partition_S(sL_fp16), tCrL_cv) - - tCrInv = thr_mma_neumann.make_fragment_A(thr_mma_neumann.partition_A(sINV_fp16)) - tCrInv_cv = smem_thr_copy_A_neumann.retile(tCrInv) - cute.copy(smem_tiled_copy_A_neumann, smem_thr_copy_A_neumann.partition_S(sINV_fp16), tCrInv_cv) - - tCrLpowB = thr_mma_neumann.make_fragment_B(thr_mma_neumann.partition_B(sL_fp16)) - tCrLpow = thr_mma_neumann.make_fragment_C(tiled_mma_neumann.partition_shape_C((CHUNK, CHUNK))) - tCrDelta = thr_mma_neumann.make_fragment_C(tiled_mma_neumann.partition_shape_C((CHUNK, CHUNK))) - tCrLpowA = thr_mma_neumann.make_fragment_A(thr_mma_neumann.partition_A(sL_fp16)) - - tCrL_u32 = cute.recast_tensor(tCrL, dtype=cutlass.Int32) - tCrInv_u32 = cute.recast_tensor(tCrInv, dtype=cutlass.Int32) - tCrLpowB_u32 = cute.recast_tensor(tCrLpowB, dtype=cutlass.Int32) - tCrLpow_u32 = cute.recast_tensor(tCrLpow, dtype=cutlass.Int32) - tCrDelta_u32 = cute.recast_tensor(tCrDelta, dtype=cutlass.Int32) - tCrLpowA_u32 = cute.recast_tensor(tCrLpowA, dtype=cutlass.Int32) - - N_REGS_U32: cutlass.Constexpr[int] = 4 # 8 fp16 / thread = 4 u32 - - for ii in cutlass.range_constexpr(N_REGS_U32): - tCrLpowB_u32[ii] = movm_t_b16(cutlass.Int32(tCrL_u32[ii])) - tCrLpow.fill(0.0) - cute.gemm(tiled_mma_neumann, tCrLpow, tCrL, tCrLpowB, tCrLpow) - - for ii in cutlass.range_constexpr(N_REGS_U32): - tCrLpowB_u32[ii] = movm_t_b16(cutlass.Int32(tCrLpow_u32[ii])) - tCrDelta.fill(0.0) - cute.gemm(tiled_mma_neumann, tCrDelta, tCrInv, tCrLpowB, tCrDelta) - for ii in cutlass.range_constexpr(N_REGS_U32): - tCrInv_u32[ii] = add_f16x2_u32(cutlass.Int32(tCrInv_u32[ii]), cutlass.Int32(tCrDelta_u32[ii])) - - for ii in cutlass.range_constexpr(N_REGS_U32): - tCrLpowA_u32[ii] = tCrLpow_u32[ii] - tCrLpow.fill(0.0) - cute.gemm(tiled_mma_neumann, tCrLpow, tCrLpowA, tCrLpowB, tCrLpow) - - for ii in cutlass.range_constexpr(N_REGS_U32): - tCrLpowB_u32[ii] = movm_t_b16(cutlass.Int32(tCrLpow_u32[ii])) - tCrDelta.fill(0.0) - cute.gemm(tiled_mma_neumann, tCrDelta, tCrInv, tCrLpowB, tCrDelta) - for ii in cutlass.range_constexpr(N_REGS_U32): - tCrInv_u32[ii] = add_f16x2_u32(cutlass.Int32(tCrInv_u32[ii]), cutlass.Int32(tCrDelta_u32[ii])) - - for ii in cutlass.range_constexpr(N_REGS_U32): - tCrLpowA_u32[ii] = tCrLpow_u32[ii] - tCrLpow.fill(0.0) - cute.gemm(tiled_mma_neumann, tCrLpow, tCrLpowA, tCrLpowB, tCrLpow) - - for ii in cutlass.range_constexpr(N_REGS_U32): - tCrLpowB_u32[ii] = movm_t_b16(cutlass.Int32(tCrLpow_u32[ii])) - tCrDelta.fill(0.0) - cute.gemm(tiled_mma_neumann, tCrDelta, tCrInv, tCrLpowB, tCrDelta) - for ii in cutlass.range_constexpr(N_REGS_U32): - tCrInv_u32[ii] = add_f16x2_u32(cutlass.Int32(tCrInv_u32[ii]), cutlass.Int32(tCrDelta_u32[ii])) - - tCrInvC = thr_mma_neumann.make_fragment_C(tiled_mma_neumann.partition_shape_C((CHUNK, CHUNK))) - tCrInvC_u32 = cute.recast_tensor(tCrInvC, dtype=cutlass.Int32) - for ii in cutlass.range_constexpr(N_REGS_U32): - tCrInvC_u32[ii] = tCrInv_u32[ii] - tCrInvC_bf16 = cute.make_fragment_like(tCrInvC, cutlass.BFloat16) - for ii in cutlass.range_constexpr(cute.size(tCrInvC)): - tCrInvC_bf16[ii] = cutlass.BFloat16(cutlass.Float32(tCrInvC[ii])) - - copy_atom_stsm = cute.make_copy_atom( - warp.StMatrix8x8x16bOp(transpose=False, num_matrices=2), - cutlass.BFloat16, - ) - smem_tiled_store_C = cute.make_tiled_copy_C_atom(copy_atom_stsm, tiled_mma_neumann) - smem_thr_store_C = smem_tiled_store_C.get_slice(tidx) - cute.copy( - smem_tiled_store_C, - smem_thr_store_C.retile(tCrInvC_bf16), - smem_thr_store_C.partition_D(sINV_bf16), - ) + _inv_fwd_subst_16(sL_fp32, sM_bf16, sINV_bf16, tidx) # Generic-proxy smem writes -> visible to the async proxy. cute.arch.fence_view_async_shared() cute.arch.barrier() From a213bc208132ef742a264f968852d256acfc7ae8 Mon Sep 17 00:00:00 2001 From: huanxian Date: Mon, 9 Nov 2026 18:06:30 +0800 Subject: [PATCH 2/3] Test FlashKDA inverse on repeated tokens --- tests/test_kda_sm90_prefill_vs_fla.py | 55 ++++++++++++++++++++++++++- 1 file changed, 54 insertions(+), 1 deletion(-) diff --git a/tests/test_kda_sm90_prefill_vs_fla.py b/tests/test_kda_sm90_prefill_vs_fla.py index 7fcbae50..79e44c0f 100644 --- a/tests/test_kda_sm90_prefill_vs_fla.py +++ b/tests/test_kda_sm90_prefill_vs_fla.py @@ -27,7 +27,7 @@ import pytest import torch from fla.ops.kda.chunk import chunk_kda as fla_chunk_kda -from fla.utils import assert_close, device +from fla.utils import assert_close, device, get_err_ratio from cula.kda import flashkda_prefill as cula_kda_prefill @@ -195,3 +195,56 @@ def test_prefill_workspace_reuse_across_shapes(): lower_bound=-5.0, ) assert_close("o", ref_o, tri_o, 0.005) + + +@pytest.mark.kda_fast +def test_prefill_repeated_tokens_inverse_precision(): + """Near-collinear keys stress the old fp16 Neumann inverse.""" + B, T, H = 1, 16, 1 + torch.manual_seed(321) + token = torch.linspace(-1.0, 1.0, D, dtype=torch.float32, device=device) + token = torch.nn.functional.normalize(token, dim=0).to(torch.bfloat16) + q = token.view(1, 1, 1, D).expand(B, T, H, D).contiguous() + k = q.clone() + v = torch.randn(B, T, H, D, dtype=torch.bfloat16, device=device) + # sigmoid(-80) makes the intra-chunk decay nearly one, so consecutive + # normalized keys produce a strictly-lower L with values close to one. + g = torch.full_like(q, -80.0) + beta = torch.full((B, T, H), 0.999, dtype=torch.bfloat16, device=device) + A_log = torch.zeros(H, dtype=torch.float32, device=device) + dt_bias = torch.zeros(H, D, dtype=torch.float32, device=device) + + with torch.no_grad(): + ref_o, ref_ht = fla_chunk_kda( + q, + k, + v, + g, + beta, + A_log=A_log, + dt_bias=dt_bias, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + safe_gate=True, + lower_bound=-5.0, + ) + tri_o, tri_ht_vk = cula_kda_prefill( + q, + k, + v, + g, + beta, + A_log=A_log, + dt_bias=dt_bias, + output_final_state=True, + safe_gate=True, + lower_bound=-5.0, + use_intracard_cp=False, + ) + tri_ht = tri_ht_vk.transpose(-2, -1) + + o_err_ratio = get_err_ratio(ref_o, tri_o) + ht_err_ratio = get_err_ratio(ref_ht, tri_ht) + assert o_err_ratio < 0.005, f"repeated-token output precision regressed: {o_err_ratio=:g}" + assert ht_err_ratio < 0.01, f"repeated-token final-state precision regressed: {ht_err_ratio=:g}" From d78145f5123fd1b8ca6b90e46f169d4f56c5716d Mon Sep 17 00:00:00 2001 From: huanxian Date: Tue, 10 Nov 2026 16:16:37 +0800 Subject: [PATCH 3/3] Fix K1 inverse shared-memory overlap --- cula/ops/kda/sm90/k1.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/cula/ops/kda/sm90/k1.py b/cula/ops/kda/sm90/k1.py index 6f1b5857..fad94522 100644 --- a/cula/ops/kda/sm90/k1.py +++ b/cula/ops/kda/sm90/k1.py @@ -51,6 +51,8 @@ def _inv_fwd_subst_16(sL_fp32: cute.Tensor, sM_bf16: cute.Tensor, sINV_bf16: cut group = lane // 8 inv = cute.make_rmem_tensor(cute.make_layout((8,), stride=(1,)), cutlass.Float32) + # Groups 0/2 invert A and groups 1/3 invert B. The redundant copies keep + # every lane named by the full-warp shuffle mask converged. for p in cutlass.range_constexpr(8): is_lower = cutlass.Float32(p < row) is_diag = cutlass.Float32(p == row) @@ -69,6 +71,7 @@ def _inv_fwd_subst_16(sL_fp32: cute.Tensor, sM_bf16: cute.Tensor, sINV_bf16: cut inv[p] = cutlass.select_(row > src_row, row_scale * pivot + inv[p], inv[p]) inv[src_row] = cutlass.select_(row > src_row, row_scale, inv[src_row]) + # P = diag(A^-1, B^-1), M = [[0, 0], [C, 0]]. zero = cutlass.BFloat16(0.0) if group == 0: for col in cutlass.range_constexpr(16): @@ -117,6 +120,7 @@ def _inv_fwd_subst_16(sL_fp32: cute.Tensor, sM_bf16: cute.Tensor, sINV_bf16: cut for ii in cutlass.range_constexpr(cute.size(tCrM_A_u32)): tCrM_B_u32[ii] = movm_t_b16(cutlass.Int32(tCrM_A_u32[ii])) + # dc = P @ M with bf16 operands and fp32 accumulation. tCrDC = thr_mma.make_fragment_C(tiled_mma.partition_shape_C((16, 16))) tCrDC.fill(0.0) cute.gemm(tiled_mma, tCrDC, tCrP_A, tCrM_B, tCrDC) @@ -139,6 +143,7 @@ def _inv_fwd_subst_16(sL_fp32: cute.Tensor, sM_bf16: cute.Tensor, sINV_bf16: cut for ii in cutlass.range_constexpr(cute.size(tCrP_A_u32)): tCrP_B_u32[ii] = movm_t_b16(cutlass.Int32(tCrP_A_u32[ii])) + # O = bf16(-dc) @ P; only its lower-left block is nonzero. tCrO = thr_mma.make_fragment_C(tiled_mma.partition_shape_C((16, 16))) tCrO.fill(0.0) cute.gemm(tiled_mma, tCrO, tCrNegDC_A, tCrP_B, tCrO) @@ -152,6 +157,7 @@ def _inv_fwd_subst_16(sL_fp32: cute.Tensor, sM_bf16: cute.Tensor, sINV_bf16: cut ) cute.arch.sync_warp() + # P and O have disjoint nonzero blocks, so this assembles (I + L)^-1. coord = cute.make_identity_tensor((16, 16)) tCcC = thr_mma.partition_C(coord) for ii in cutlass.range_constexpr(cute.size(tCcC)): @@ -212,9 +218,10 @@ def k1_kernel( s_k_restored = cute.make_tensor(sG_raw.iterator, kinter_qk_layout) sMqk_bf16 = smem.allocate_tensor(cutlass.BFloat16, cc_layout, 128) sBetaSig = smem.allocate_tensor(cutlass.Float32, cute.make_layout((CHUNK,)), 128) - # Inverse buffers alias sG_cumsum after decay_apply. L remains fp32; - # INV and the block-merge matrix use bf16 Tensor Core operands. - sL_fp32 = cute.make_tensor(sG_cumsum.iterator, cc_layout) + # L is written by warp 0 while warp 1 still reads s_k_inv, so it cannot + # alias sG_cumsum. The block-merge buffers are written only after both + # warps finish and may safely reuse that dead storage. + sL_fp32 = smem.allocate_tensor(cutlass.Float32, cc_layout, 128) sG_cumsum_bf16_ptr = cute.recast_ptr(sG_cumsum.iterator, dtype=cutlass.BFloat16) sINV_bf16 = cute.make_tensor(sG_cumsum_bf16_ptr + (2 * CHUNK * CHUNK), cc_layout) sM_bf16 = cute.make_tensor(sG_cumsum_bf16_ptr + (3 * CHUNK * CHUNK), cc_layout)