diff --git a/docs/fe-oss-apis/dsa.md b/docs/fe-oss-apis/dsa.md index a288d836d..063525b38 100644 --- a/docs/fe-oss-apis/dsa.md +++ b/docs/fe-oss-apis/dsa.md @@ -305,8 +305,9 @@ and normalization semantics differ from the indexer path. Three-stage sparse top-K pipeline that produces the training gradients for the indexer tower: -1. `ScoreGradSm90` / `ScoreGradSm100` (kernel 1) — in-place score-grad precompute from - `attn_score` (target) and `index_score` (predict). +1. `ScoreGradSm90` / `ScoreGradSm100` (kernel 1) — in-place score-grad precompute + that overwrites `attn_score` (target) and reads `index_score` (predict) + without modifying it. 2. `IndexerBackwardSm90` / `IndexerBackwardSm100` (kernel 2) — three warp-specialised GEMMs produce `d_index_q`, `d_weights`, and a `dIndexK_f32` accumulator. @@ -331,11 +332,11 @@ d_index_q, d_weights, d_index_k = ( When compressed forward returns its fused `softmax`, backward can skip both indexer Q@K score recompute and the separate logits softmax. Pass `softmax` -directly as `index_score`; backward consumes and overwrites this buffer, so -pass `softmax.clone()` if it must be preserved. `attn_score` must use the same -valid-slot mask. Because compressed forward returns global indices by default, -also pass `topk_indices_global=True` unless forward used -`topk_indices_global=False`. The public sparse `indexer_backward_wrapper` has a +directly as `index_score`; backward treats this buffer as read-only. +`attn_score` must use the same valid-slot mask. Because compressed forward +returns global indices by default, also pass `topk_indices_global=True` unless +forward used `topk_indices_global=False`. The public sparse +`indexer_backward_wrapper` has a BSHD-shaped interface; BF16 THD tensors can use zero-copy `B=1` views (squeeze the singleton K head and add a batch dimension) together with global Top-K indices. FP8 and MXFP8 indexer backward are not currently supported because @@ -388,10 +389,10 @@ the default backend. full dQ/dK contribution (`g' * w`, not scaled by `S`) is what differs. The default backend's multiply preserves subnormals, so reaching it takes an exact `sm_scale * S` below `2^-150` (~7.0e-46). -- **Scratch / workspace behavior** — `attn_score`/`index_score` are consumed - in place exactly like the default backend (`attn_score` is left holding - kernel 1's `grad_signal`; `sm_scale` folds inside kernel 2 without touching - the buffer). The backend owns one piece of per-plan workspace — the +- **Scratch / workspace behavior** — `attn_score` is consumed in place and left + holding kernel 1's `grad_signal`, while `index_score` is read-only and + preserved. `sm_scale` folds inside kernel 2 without touching either buffer. + The backend owns one piece of per-plan workspace — the dynamic-ticket counter — allocated on first execute and reused by every later one. A BF16 `d_index_k` additionally needs a `B * S_k * D` fp32 accumulator (2 MiB = 2,097,152 bytes at B=1, S_k=4096, D=128, growing with @@ -403,7 +404,7 @@ the default backend. - **Concurrency** — executions sharing one plan must not overlap on the device (the ticket counter is per-plan workspace). One plan serves one device: the workspace is device-resident, and execution rejects tensors on - any other device before touching the score buffers. The wrapper keys its + any other device before overwriting `attn_score`. The wrapper keys its plan cache on the CUDA device and on the **resolved** stream (`stream` when given, otherwise `torch.cuda.current_stream()` at call time), so calls that differ in device or stream get a private plan and private workspace; calls diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py index 9debd8b30..2b2e2ca35 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py @@ -169,11 +169,11 @@ class IndexerBackward(APIBase): 1. Runtime signature re-validation of the tensors against the descriptors captured at plan-build time (dtype, shape, and stride/layout), - *before* any kernel launch — kernel 1 mutates the score buffers in + *before* any kernel launch — kernel 1 overwrites ``attn_score`` in place, so a mismatched reuse of an exported plan must raise here to avoid a fail-dirty. - 2. Kernel 1 — in-place score-grad precompute (overwrites ``attn_score`` and - ``index_score``). + 2. Kernel 1 — in-place score-grad precompute (overwrites ``attn_score`` + with ``grad_signal`` and reads ``index_score`` without modifying it). 3. (conditional) fp32 ``d_index_k`` zero-init — when the ``d_index_k`` output buffer is fp32, ``execute`` zeroes it on the selected stream before the GEMM, because the dK epilogue atomic-adds into it. This is a @@ -191,8 +191,9 @@ class IndexerBackward(APIBase): of the 24 fp32 significand bits — measured ~675x lower gradient-matrix representation error than the default single-bf16 rounding); ``d_weights`` is additionally reduced deterministically (bitwise - run-to-run, as is ``d_index_q``). Same wrapper contract and in-place - score consumption. Output dtype selects output precision: ``d_weights`` + run-to-run, as is ``d_index_q``). Same wrapper contract: ``attn_score`` is + the in-place score scratch and ``index_score`` remains read-only. Output + dtype selects output precision: ``d_weights`` and ``d_index_k`` accept fp32 buffers, which receive the fp32 accumulators directly — **the headline accuracy gains require fp32 output buffers**; bf16 output buffers round the result and keep only the @@ -270,13 +271,13 @@ def __init__( def _validate_plan_shapes_and_layout(self) -> None: """Semantic shape + layout validation of the plan's tensor descriptors. - Runs before kernel 1 mutates the score buffers in place. + Runs before kernel 1 overwrites ``attn_score`` in place. ``check_support`` otherwise only checks dtypes and ``execute``'s signature check only compares each runtime tensor against the recorded descriptor, so without this a directly-built / first-call plan whose ``d_weights`` / ``d_index_q`` / ``d_index_k`` / score / top-k shape is already inconsistent with ``index_q`` would pass validation, kernel 1 - would corrupt the scores, and only the GEMM would fault or corrupt + would overwrite ``attn_score``, and only the GEMM would fault or corrupt memory. The relationships enforced (per the API contract): index_q (B, S_q, H, D) weights (B, S_q, H) @@ -366,7 +367,7 @@ def check_support(self) -> bool: # full metadata matrix: cross-tensor shapes, output dtypes, # devices, contiguity — the backend binds true views and uses # 16B-aligned row pointers, so reject bad metadata before - # kernel 1 mutates the score buffers. + # kernel 1 overwrites ``attn_score``. self._check_tensor_shape(self.w_desc, (b, s_q, h), name="weights") self._check_tensor_shape(self.ik_desc, (b, s_k, d), name="index_k") self._check_tensor_shape(self.diq_desc, (b, s_q, h, d), name="d_index_q") @@ -409,9 +410,9 @@ def check_support(self) -> bool: self._check_dtype(self.idx_score_desc, torch.float32, name="index_score") self._check_dtype(self.topk_desc, torch.int32, name="topk_indices") # Output-dtype contract, validated here — before compile()/execute(), - # i.e. before kernel 1 mutates attn_score/index_score in place (no - # fail-dirty on a bad output dtype). ``d_index_q`` is bf16-only (TMA - # store of the input dtype). ``d_index_k`` accepts bf16 or fp32 on both + # i.e. before kernel 1 overwrites ``attn_score`` (``index_score`` is + # read-only; no fail-dirty on a bad output dtype). ``d_index_q`` is + # bf16-only (TMA store of the input dtype). ``d_index_k`` accepts bf16 or fp32 on both # arches (the GEMM accumulates dK in fp32; a bf16 buffer gets a # trailing cast). ``d_weights`` is bf16-only on the default backend: # its kernel's dW store rounds the fp32 accumulator to bf16, so an @@ -435,8 +436,8 @@ def check_support(self) -> bool: extra_error_msg="fp32 buffers receive the fp32 accumulator directly", ) # Semantic shape relationships + compact-layout requirement, validated - # here (before compile()/execute(), i.e. before kernel 1 mutates the - # score buffers). Guards a directly-built / first-call plan whose output + # here (before compile()/execute(), i.e. before kernel 1 overwrites + # ``attn_score``). Guards a directly-built / first-call plan whose output # or score shapes are inconsistent, or whose layout is non-compact. self._validate_plan_shapes_and_layout() self._is_supported = True @@ -496,7 +497,7 @@ def _check_execute_signature(self, *entries) -> None: Each entry is ``(tensor, descriptor, name)``. Guards against reusing a directly-built/exported plan with a tensor whose dtype, shape, or stride/layout differs from what it was compiled for — kernel 1 - mutates the score buffers in place, so a mismatch must raise *before* + overwrites ``attn_score`` in place, so a mismatch must raise *before* the pipeline starts (no fail-dirty). """ for tensor, desc, name in entries: @@ -537,10 +538,10 @@ def execute( self._logger.debug("Entering execute") # Stage 1: runtime signature re-validation against the descriptors # captured at plan-build time, BEFORE any kernel launch. The plan is - # specialized to the sample dtypes and shapes. Kernel 1 - # mutates ``attn_score`` / ``index_score`` in place before the GEMM - # runs, so reusing an exported plan with a mismatched signature would - # otherwise fail (or corrupt memory) only after the scores were + # specialized to the sample dtypes and shapes. Kernel 1 overwrites + # ``attn_score`` in place and treats ``index_score`` as read-only before + # the GEMM runs, so reusing an exported plan with a mismatched signature + # would otherwise fail (or corrupt memory) only after ``attn_score`` was # already mutated. Raise a clean ValueError here — no fail-dirty. The # signature-keyed wrapper cache never hits this, but a directly- # built/exported plan can. @@ -563,8 +564,8 @@ def execute( ) # One plan serves one device: the v2 backend's per-plan # workspace (the ticket counter) lives on the sample tensors' - # device, so reject cross-device execution before kernel 1 mutates - # the score buffers. The backend re-checks every + # device, so reject cross-device execution before kernel 1 overwrites + # ``attn_score``. The backend re-checks every # tensor against index_q.device, so validating index_q covers all. if self.use_v2 and index_q.device != self.iq_desc.device: raise ValueError( @@ -785,9 +786,9 @@ def indexer_backward_wrapper( ``index_score`` contains the predict softmax aligned slot-for-slot with ``topk_indices``; the fused softmax returned by compressed indexer forward - can be passed directly. ``attn_score`` and ``index_score`` are consumed - in-place: the kernel overwrites ``attn_score`` with ``grad_signal`` and - ``index_score`` with ``sum_grad`` during the score-grad precompute stage. + can be passed directly. The kernel overwrites ``attn_score`` with + ``grad_signal`` during the score-grad precompute stage and treats + ``index_score`` as read-only. Args: topk_indices_global: whether ``topk_indices`` already contains global @@ -801,7 +802,7 @@ def indexer_backward_wrapper( index_score: FP32 predict probabilities over the selected indexer logits. The caller must use the same valid-slot mask when constructing ``attn_score`` so target and predict describe - identical slots. This buffer is overwritten in-place. + identical slots. This buffer is preserved. grad_loss: single-element float32 tensor on the same CUDA device as ``index_q``. The kernel reads its value at runtime, including on CUDA Graph replay. @@ -817,9 +818,9 @@ def indexer_backward_wrapper( outside its envelope — SM100, H == 64, D == 128, block_I == 128, topk % 128 == 0 with 128 <= topk <= 2048, sm_scale > 0, contiguous same-device tensors — and never silently falls back to the default - backend. It keeps the wrapper contract (same inputs, outputs, and - in-place score consumption: ``attn_score`` is left holding - exactly kernel 1's ``grad_signal``; ``sm_scale`` folds in-kernel) + backend. It keeps the wrapper contract (same inputs and outputs: + ``attn_score`` is left holding exactly kernel 1's ``grad_signal``, + ``index_score`` remains read-only, and ``sm_scale`` folds in-kernel) and computes the GEMM stage with fp32 weights and a two-term bf16 expansion of the gradient matrix, plus deterministic (bitwise run-to-run) ``d_weights`` / ``d_index_q``. diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py index c14238d31..3675e53c4 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py @@ -62,6 +62,8 @@ import cutlass import cutlass.cute as cute from cutlass import Float32, Int32, const_expr +from cutlass._mlir.dialects import nvvm +from cutlass.cutlass_dsl import dsl_user_op from cutlass.cute.nvgpu import cpasync import cutlass.cute.nvgpu.tcgen05 as tcgen05 import cutlass.utils as utils @@ -87,6 +89,27 @@ mul_packed_f32x2 = partial(cute.arch.mul_packed_f32x2, rnd="rn") fma_packed_f32x2 = partial(cute.arch.fma_packed_f32x2, rnd="rn") + +@dsl_user_op +def _tcgen05_fence_after_thread_sync(*, loc=None, ip=None): + """Order subsequent tcgen05 operations after an inter-thread wait.""" + nvvm.tcgen05_fence( + nvvm.Tcgen05FenceKind.AFTER_THREAD_SYNC, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _tcgen05_fence_before_thread_sync(*, loc=None, ip=None): + """Order prior tcgen05 operations before an inter-thread signal.""" + nvvm.tcgen05_fence( + nvvm.Tcgen05FenceKind.BEFORE_THREAD_SYNC, + loc=loc, + ip=ip, + ) + + DENOM_EPS = 1e-10 CLIP_LOG_MIN = -100.0 CLIP_PROB_MIN = math.exp(CLIP_LOG_MIN) @@ -354,7 +377,9 @@ class SharedStorage: MBAR_2Q_GS_LOADED_0 = 4 # Load → Compute (2-stage K pipeline) MBAR_2Q_GS_LOADED_1 = 5 MBAR_2Q_W_LOADED = 6 -NUM_2Q_BARRIERS = 7 +MBAR_2Q_DQ_DONE = 7 # MMA → Compute (one-shot, final q0/q1 GEMM3 completion) +MBAR_2Q_REDUCE_DONE = 8 # Reduce → Compute warp 0 (one-shot, final dK T2R completion) +NUM_2Q_BARRIERS = 9 class DenseIndexerBackward2QGemmSm100: @@ -369,7 +394,7 @@ class DenseIndexerBackward2QGemmSm100: Offset 256: dQ_q0 (128 cols) — Q token 0 persistent accumulator Offset 384: dQ_q1 (128 cols) — Q token 1 persistent accumulator - Barriers: 7 custom. + Barriers: 9 custom. Grid: (ceil(seqlen_q/2), batch, 1). """ @@ -525,7 +550,7 @@ def __call__( sdS_layout = _make_smem_layout_a(tmma3, self.gemm3_tiler, self.q_dtype, 2) # 2-stage sdS_store_layout = _make_smem_layout_epi( self.q_dtype, - LayoutEnum.COL_MAJOR, + LayoutEnum.ROW_MAJOR, (self.heads_padded, self.block_I), 2, ) @@ -963,15 +988,20 @@ class SharedStorage: cute.group_modes(gdQ_q1, 0, 2), ) - # Init all custom barriers (warp 0) - if warp_idx == 0: + # Initialize each custom barrier exactly once. + if tidx == 0: cute.arch.mbarrier_init(mbar + MBAR_2Q_S_FULL, 1) cute.arch.mbarrier_init(mbar + MBAR_2Q_DS_READY, self.WARPGROUP_SIZE) cute.arch.mbarrier_init(mbar + MBAR_2Q_DK_FULL, 1) cute.arch.mbarrier_init(mbar + MBAR_2Q_DK_EMPTY, self.WARPGROUP_SIZE) - cute.arch.mbarrier_init(mbar + MBAR_2Q_GS_LOADED_0, 1) - cute.arch.mbarrier_init(mbar + MBAR_2Q_GS_LOADED_1, 1) - cute.arch.mbarrier_init(mbar + MBAR_2Q_W_LOADED, 1) + cute.arch.mbarrier_init(mbar + MBAR_2Q_GS_LOADED_0, self.WARP_SIZE) + cute.arch.mbarrier_init(mbar + MBAR_2Q_GS_LOADED_1, self.WARP_SIZE) + cute.arch.mbarrier_init(mbar + MBAR_2Q_W_LOADED, self.WARP_SIZE) + cute.arch.mbarrier_init(mbar + MBAR_2Q_DQ_DONE, 1) + cute.arch.mbarrier_init( + mbar + MBAR_2Q_REDUCE_DONE, + self.WARPGROUP_SIZE, + ) cute.arch.sync_threads() # Pre-compute accumulator shapes/layouts @@ -1072,6 +1102,7 @@ class SharedStorage: sGradSignal_q1_0, sGradSignal_q1_1, sW_full, + sdS_store, sdS, sdQ_epi_slice, s_acc_shape, @@ -1096,6 +1127,18 @@ class SharedStorage: num_kv_blocks, ) if warp_idx == self.compute_warp_id[0]: + # The reduce warpgroup issues asynchronous TMEM loads. + # All 128 reducer threads arrive only after their final + # tcgen05.wait::ld fence, making TMEM deallocation safe. + # A zero-K batch issues no reducer load (and therefore no + # REDUCE_DONE arrival); preserve the pre-existing empty + # path instead of waiting on an event that cannot occur. + if num_kv_blocks > 0: + cute.arch.mbarrier_wait( + mbar + MBAR_2Q_REDUCE_DONE, + Int32(0), + ) + _tcgen05_fence_after_thread_sync() cute.arch.dealloc_tmem(tmem_ptr_base, self.tmem_alloc_cols) elif warp_idx in self.reduce_warp_id: @@ -1201,8 +1244,8 @@ def _load_warp_2q( sW_full[self.heads + idx] = mW_b[q1_local, idx] cute.arch.fence_view_async_shared() - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive(mbar + MBAR_2Q_W_LOADED) + # Every load-warp lane publishes its own sW stores. + cute.arch.mbarrier_arrive(mbar + MBAR_2Q_W_LOADED) # --- TMA Q load (q0) --- Q_q0_producer.reset() @@ -1271,8 +1314,7 @@ def _load_warp_2q( bi, ) cute.arch.fence_view_async_shared() - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive(mbar + MBAR_2Q_GS_LOADED_0) + cute.arch.mbarrier_arrive(mbar + MBAR_2Q_GS_LOADED_0) else: self._load_grad_signal_to_buf( mGS_b, @@ -1292,8 +1334,7 @@ def _load_warp_2q( bi, ) cute.arch.fence_view_async_shared() - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive(mbar + MBAR_2Q_GS_LOADED_1) + cute.arch.mbarrier_arrive(mbar + MBAR_2Q_GS_LOADED_1) # ========================================================================= # MMA warp (2Q): Sequential loop, 6 GEMMs per K block @@ -1363,6 +1404,9 @@ def _mma_warp_2q( # Wait for K[bi] ready K_handle = K_consumer.wait_and_advance() k_stage = K_handle.index + # This also orders the next tcgen05 MMA after the earlier + # DK_EMPTY wait when the single dK TMEM buffer is being reused. + _tcgen05_fence_after_thread_sync() # ---- Phase A (Q token 0) ---- # GEMM1: S = Q_q0 @ K[bi] @@ -1382,6 +1426,7 @@ def _mma_warp_2q( # Wait for dS from Compute (Phase A) cute.arch.mbarrier_wait(mbar + MBAR_2Q_DS_READY, ds_ready_phase) ds_ready_phase ^= 1 + _tcgen05_fence_after_thread_sync() # GEMM2: dK = dS_q0^T @ Q_q0 (ACCUMULATE=False, clear) tmma2.set(tcgen05.Field.ACCUMULATE, False) @@ -1395,7 +1440,7 @@ def _mma_warp_2q( ) tmma2.set(tcgen05.Field.ACCUMULATE, True) - # GEMM3: dQ_q0 += dS @ Kt[bi] + # GEMM3: dQ_q0 += dS_q0 @ K^T tmma3.set(tcgen05.Field.ACCUMULATE, not is_first_dq_q0) is_first_dq_q0 = False for k_block in cutlass.range(0, cute.size(tDQrDS, mode=[2]), unroll=4): @@ -1427,6 +1472,7 @@ def _mma_warp_2q( # Wait for dS from Compute (Phase B) cute.arch.mbarrier_wait(mbar + MBAR_2Q_DS_READY, ds_ready_phase) ds_ready_phase ^= 1 + _tcgen05_fence_after_thread_sync() # GEMM2: dK += dS_q1^T @ Q_q1 (ACCUMULATE=True!) tmma2.set(tcgen05.Field.ACCUMULATE, True) @@ -1439,7 +1485,7 @@ def _mma_warp_2q( tDkDk, ) - # GEMM3: dQ_q1 += dS @ Kt[bi] + # GEMM3: dQ_q1 += dS_q1 @ K^T tmma3.set(tcgen05.Field.ACCUMULATE, not is_first_dq_q1) is_first_dq_q1 = False for k_block in cutlass.range(0, cute.size(tDQrDS, mode=[2]), unroll=4): @@ -1459,6 +1505,12 @@ def _mma_warp_2q( # Release K stage K_handle.release() + # GEMM3 accumulates dQ in TMEM asynchronously. Commit a dedicated + # one-shot completion only after the final q0/q1 GEMM3 has been + # issued, so Compute cannot race its TMEM epilogue readback. + with cute.arch.elect_one(): + tcgen05.commit(mbar + MBAR_2Q_DQ_DONE) + # ========================================================================= # Helper: dS/dW computation for one block (same as 1Q) # ========================================================================= @@ -1475,11 +1527,20 @@ def _compute_ds_dw_block( w_reg_h1: Float32, ): """Compute dS and accumulate dW for one block using grad_signal from sGS.""" + use_stmatrix_ds = const_expr( + self.heads_padded == 64 and self.block_I == 128, + ) for ei in cutlass.range(0, cute.size(tSrS), 2, unroll_full=True): - h0 = cute.get(tCcS[ei], mode=[0, 0]) - n0 = cute.get(tCcS[ei], mode=[0, 1]) - h1 = cute.get(tCcS[ei + 1], mode=[0, 0]) - n1 = cute.get(tCcS[ei + 1], mode=[0, 1]) + if const_expr(use_stmatrix_ds): + h0 = cute.get(tCcS[ei], mode=[0]) + n0 = cute.get(tCcS[ei], mode=[1]) + h1 = cute.get(tCcS[ei + 1], mode=[0]) + n1 = cute.get(tCcS[ei + 1], mode=[1]) + else: + h0 = cute.get(tCcS[ei], mode=[0, 0]) + n0 = cute.get(tCcS[ei], mode=[0, 1]) + h1 = cute.get(tCcS[ei + 1], mode=[0, 0]) + n1 = cute.get(tCcS[ei + 1], mode=[0, 1]) tSrS[ei], tSrS[ei + 1] = mul_packed_f32x2( (tSrS[ei], tSrS[ei + 1]), @@ -1488,8 +1549,15 @@ def _compute_ds_dw_block( s0 = tSrS[ei] s1 = tSrS[ei + 1] - w0 = w_reg_h0 if h0 == my_first_h else w_reg_h1 - w1 = w_reg_h0 if h1 == my_first_h else w_reg_h1 + if const_expr(use_stmatrix_ds): + # 16dp256b8x alternates pairs between the thread's low/high + # head; both elements in a pair use the same weight. + pair_weight = w_reg_h0 if (ei // 2) % 2 == 0 else w_reg_h1 + w0 = pair_weight + w1 = pair_weight + else: + w0 = w_reg_h0 if h0 == my_first_h else w_reg_h1 + w1 = w_reg_h0 if h1 == my_first_h else w_reg_h1 gs0 = sGS[n0] gs1 = sGS[n1] @@ -1523,6 +1591,7 @@ def _compute_warpgroup_2q( sGradSignal_q1_0, sGradSignal_q1_1, sW_full, + sdS_store, sdS, sdQ_epi_slice, s_acc_shape, @@ -1557,45 +1626,110 @@ def _compute_warpgroup_2q( Float32, ) - # --- TMEM readback (S, single-buffered) --- - tiled_tmem_load_s = tcgen05.make_tmem_copy(tmem_load_atom, tStS) - thr_tmem_load_s = tiled_tmem_load_s.get_slice(wg_tidx) - tStS_t2r = thr_tmem_load_s.partition_S(tStS) - - # Logical GEMM views for direct dS writes (stage 0/1). - # Writing through the same A-operand view that GEMM2/GEMM3 read from - # guarantees layout compatibility (stmatrix epilogue layout differs - # from the swizzled A-operand layout, so coord writes are required). - sdS_gemm_view_0 = cute.composition( - sdS[None, None, None, 0], - cute.make_layout((self.heads_padded, self.block_I)), - ) - sdS_gemm_view_1 = cute.composition( - sdS[None, None, None, 1], - cute.make_layout((self.heads_padded, self.block_I)), + # The production H64 x I128 tile has the same 16dp256b8x ownership as + # the sparse kernel. Derive STMatrix R2S ownership directly from the + # TMEM load; other factory shapes retain native coordinate stores. + use_stmatrix_ds = const_expr( + self.heads_padded == 64 and self.block_I == 128, ) + if const_expr(use_stmatrix_ds): + tStS_epi = tStS[((None, None), 0, 0)] + tiled_tmem_load_s = tcgen05.make_tmem_copy( + tmem_load_atom, + tStS_epi, + ) + thr_tmem_load_s = tiled_tmem_load_s.get_slice(wg_tidx) + tStS_t2r = thr_tmem_load_s.partition_S(tStS_epi) + + smem_store_atom = sm100_utils_basic.get_smem_store_op( + LayoutEnum.ROW_MAJOR, + self.q_dtype, + self.acc_dtype, + tiled_tmem_load_s, + ) + tiled_smem_store = cute.make_tiled_copy_D( + smem_store_atom, + tiled_tmem_load_s, + ) + thr_smem_store = tiled_smem_store.get_slice(wg_tidx) + tRS_sdS = thr_smem_store.partition_D(sdS_store) - cS = cute.make_identity_tensor(s_acc_shape) - tCcS = thr_tmem_load_s.partition_D(cS) + cS = cute.make_identity_tensor( + (self.heads_padded, self.block_I), + ) + tCcS = thr_tmem_load_s.partition_D(cS) + my_first_h = warp_id_in_wg * 16 + lane_id // 4 + my_second_h = my_first_h + 8 + else: + tiled_tmem_load_s = tcgen05.make_tmem_copy( + tmem_load_atom, + tStS, + ) + thr_tmem_load_s = tiled_tmem_load_s.get_slice(wg_tidx) + tStS_t2r = thr_tmem_load_s.partition_S(tStS) + + sdS_gemm_view_0 = cute.composition( + sdS[None, None, None, 0], + cute.make_layout((self.heads_padded, self.block_I)), + ) + sdS_gemm_view_1 = cute.composition( + sdS[None, None, None, 1], + cute.make_layout((self.heads_padded, self.block_I)), + ) + cS = cute.make_identity_tensor(s_acc_shape) + tCcS = thr_tmem_load_s.partition_D(cS) + my_first_h = cute.get(tCcS[0], mode=[0, 0]) + my_second_h = my_first_h + for ei in cutlass.range(cute.size(tCcS), unroll_full=True): + h_check = cute.get(tCcS[ei], mode=[0, 0]) + if h_check != my_first_h: + my_second_h = h_check tSrS_shape = tCcS.shape - my_first_h = cute.get(tCcS[0], mode=[0, 0]) - my_second_h = my_first_h - for ei in cutlass.range(cute.size(tCcS), unroll_full=True): - h_check = cute.get(tCcS[ei], mode=[0, 0]) - if h_check != my_first_h: - my_second_h = h_check - - # --- TMEM readback (dQ_q0 and dQ_q1) --- - tiled_tmem_load_dq0 = tcgen05.make_tmem_copy(tmem_load_atom, tDqDq_0) + + # dQ uses the same canonical 2-D ownership for H64 x D128. + use_stmatrix_dq = const_expr( + self.heads_padded == 64 and self.head_dim_padded == 128, + ) + if const_expr(use_stmatrix_dq): + tDqDq_0_load_view = tDqDq_0[((None, None), 0, 0)] + tDqDq_1_load_view = tDqDq_1[((None, None), 0, 0)] + cDQ = cute.make_identity_tensor( + (self.heads_padded, self.head_dim_padded), + ) + else: + tDqDq_0_load_view = tDqDq_0 + tDqDq_1_load_view = tDqDq_1 + cDQ = cute.make_identity_tensor(dq_acc_shape) + + tiled_tmem_load_dq0 = tcgen05.make_tmem_copy( + tmem_load_atom, + tDqDq_0_load_view, + ) thr_tmem_load_dq0 = tiled_tmem_load_dq0.get_slice(wg_tidx) - tDqDq_0_t2r = thr_tmem_load_dq0.partition_S(tDqDq_0) - cDQ = cute.make_identity_tensor(dq_acc_shape) + tDqDq_0_t2r = thr_tmem_load_dq0.partition_S(tDqDq_0_load_view) tCcDQ = thr_tmem_load_dq0.partition_D(cDQ) tDQrDQ_shape = tCcDQ.shape - tiled_tmem_load_dq1 = tcgen05.make_tmem_copy(tmem_load_atom, tDqDq_1) + tiled_tmem_load_dq1 = tcgen05.make_tmem_copy( + tmem_load_atom, + tDqDq_1_load_view, + ) thr_tmem_load_dq1 = tiled_tmem_load_dq1.get_slice(wg_tidx) - tDqDq_1_t2r = thr_tmem_load_dq1.partition_S(tDqDq_1) + tDqDq_1_t2r = thr_tmem_load_dq1.partition_S(tDqDq_1_load_view) + + if const_expr(use_stmatrix_dq): + smem_store_atom_dq = sm100_utils_basic.get_smem_store_op( + LayoutEnum.ROW_MAJOR, + self.q_dtype, + self.acc_dtype, + tiled_tmem_load_dq0, + ) + tiled_smem_store_dq = cute.make_tiled_copy_D( + smem_store_atom_dq, + tiled_tmem_load_dq0, + ) + thr_smem_store_dq = tiled_smem_store_dq.get_slice(wg_tidx) + tRDQ_sdQ = thr_smem_store_dq.partition_D(sdQ_epi_slice) # Wait for W loaded cute.arch.mbarrier_wait(mbar + MBAR_2Q_W_LOADED, Int32(0)) @@ -1636,6 +1770,7 @@ def _compute_warpgroup_2q( # Wait for S ready from MMA (Phase A) cute.arch.mbarrier_wait(mbar + MBAR_2Q_S_FULL, s_full_phase) s_full_phase ^= 1 + _tcgen05_fence_after_thread_sync() cute.copy(tiled_tmem_load_s, tStS_t2r, tSrS) # Compute dS_q0 + accumulate dW_q0 @@ -1669,7 +1804,21 @@ def _compute_warpgroup_2q( for ei in cutlass.range(cute.size(tSrS), unroll_full=True): tSrS_f16[ei] = self.q_dtype(tSrS[ei]) - if bi % 2 == 0: + if const_expr(use_stmatrix_ds): + tRS_rdS = tiled_smem_store.retile(tSrS_f16) + if bi % 2 == 0: + cute.copy( + tiled_smem_store, + tRS_rdS, + tRS_sdS[(None, None, None, 0)], + ) + else: + cute.copy( + tiled_smem_store, + tRS_rdS, + tRS_sdS[(None, None, None, 1)], + ) + elif bi % 2 == 0: for ei in cutlass.range(cute.size(tSrS_f16), unroll_full=True): h = cute.get(tCcS[ei], mode=[0, 0]) n = cute.get(tCcS[ei], mode=[0, 1]) @@ -1681,6 +1830,7 @@ def _compute_warpgroup_2q( sdS_gemm_view_1[h, n] = tSrS_f16[ei] cute.arch.fence_proxy("async.shared", space="cta") + _tcgen05_fence_before_thread_sync() cute.arch.mbarrier_arrive(mbar + MBAR_2Q_DS_READY) # ---- Phase B (Q token 1, guarded) ---- @@ -1688,6 +1838,7 @@ def _compute_warpgroup_2q( # Wait for S ready from MMA (Phase B) cute.arch.mbarrier_wait(mbar + MBAR_2Q_S_FULL, s_full_phase) s_full_phase ^= 1 + _tcgen05_fence_after_thread_sync() cute.copy(tiled_tmem_load_s, tStS_t2r, tSrS) # Compute dS_q1 + accumulate dW_q1, using sW_full offset by heads @@ -1721,7 +1872,21 @@ def _compute_warpgroup_2q( for ei in cutlass.range(cute.size(tSrS), unroll_full=True): tSrS_f16_b[ei] = self.q_dtype(tSrS[ei]) - if bi % 2 == 0: + if const_expr(use_stmatrix_ds): + tRS_rdS_b = tiled_smem_store.retile(tSrS_f16_b) + if bi % 2 == 0: + cute.copy( + tiled_smem_store, + tRS_rdS_b, + tRS_sdS[(None, None, None, 0)], + ) + else: + cute.copy( + tiled_smem_store, + tRS_rdS_b, + tRS_sdS[(None, None, None, 1)], + ) + elif bi % 2 == 0: for ei in cutlass.range(cute.size(tSrS_f16_b), unroll_full=True): h = cute.get(tCcS[ei], mode=[0, 0]) n = cute.get(tCcS[ei], mode=[0, 1]) @@ -1733,6 +1898,7 @@ def _compute_warpgroup_2q( sdS_gemm_view_1[h, n] = tSrS_f16_b[ei] cute.arch.fence_proxy("async.shared", space="cta") + _tcgen05_fence_before_thread_sync() cute.arch.mbarrier_arrive(mbar + MBAR_2Q_DS_READY) # DK_FULL is committed after both dQ GEMM3 phases for every K block. @@ -1749,20 +1915,39 @@ def _compute_warpgroup_2q( cute.make_layout((self.heads_padded, self.head_dim_padded)), ) + # The per-block S barriers only cover GEMM1. Wait for the MMA warp's + # final GEMM3 commit before reading either persistent dQ accumulator. + cute.arch.mbarrier_wait(mbar + MBAR_2Q_DQ_DONE, Int32(0)) + _tcgen05_fence_after_thread_sync() + # ---- Epilogue: dQ for q0 via TMA store ---- tDQrDQ_q0 = cute.make_rmem_tensor(tDQrDQ_shape, Float32) cute.copy(tiled_tmem_load_dq0, tDqDq_0_t2r, tDQrDQ_q0) tDQrDQ_q0_bf16 = cute.make_rmem_tensor(tDQrDQ_q0.shape, self.q_dtype) - for ei in cutlass.range(cute.size(tDQrDQ_q0), unroll_full=True): - tDQrDQ_q0_bf16[ei] = self.q_dtype(tDQrDQ_q0[ei] * Float32(sm_scale)) + for ei in cutlass.range(0, cute.size(tDQrDQ_q0), 2): + scaled0, scaled1 = mul_packed_f32x2( + (tDQrDQ_q0[ei], tDQrDQ_q0[ei + 1]), + (Float32(sm_scale), Float32(sm_scale)), + ) + tDQrDQ_q0_bf16[ei] = self.q_dtype(scaled0) + tDQrDQ_q0_bf16[ei + 1] = self.q_dtype(scaled1) cute.arch.fence_view_async_tmem_load() + _tcgen05_fence_before_thread_sync() - for ei in cutlass.range(cute.size(tDQrDQ_q0_bf16), unroll_full=True): - h = cute.get(tCcDQ[ei], mode=[0, 0]) - d = cute.get(tCcDQ[ei], mode=[0, 1]) - sdQ_gemm_view[h, d] = tDQrDQ_q0_bf16[ei] + if const_expr(use_stmatrix_dq): + tRDQ_rdQ_q0 = tiled_smem_store_dq.retile(tDQrDQ_q0_bf16) + cute.copy( + tiled_smem_store_dq, + tRDQ_rdQ_q0, + tRDQ_sdQ, + ) + else: + for ei in cutlass.range(cute.size(tDQrDQ_q0_bf16), unroll_full=True): + h = cute.get(tCcDQ[ei], mode=[0, 0]) + d = cute.get(tCcDQ[ei], mode=[0, 1]) + sdQ_gemm_view[h, d] = tDQrDQ_q0_bf16[ei] self.compute_sync_barrier.arrive_and_wait() cute.arch.fence_proxy("async.shared", space="cta") @@ -1786,15 +1971,29 @@ def _compute_warpgroup_2q( cute.copy(tiled_tmem_load_dq1, tDqDq_1_t2r, tDQrDQ_q1) tDQrDQ_q1_bf16 = cute.make_rmem_tensor(tDQrDQ_q1.shape, self.q_dtype) - for ei in cutlass.range(cute.size(tDQrDQ_q1), unroll_full=True): - tDQrDQ_q1_bf16[ei] = self.q_dtype(tDQrDQ_q1[ei] * Float32(sm_scale)) + for ei in cutlass.range(0, cute.size(tDQrDQ_q1), 2): + scaled0, scaled1 = mul_packed_f32x2( + (tDQrDQ_q1[ei], tDQrDQ_q1[ei + 1]), + (Float32(sm_scale), Float32(sm_scale)), + ) + tDQrDQ_q1_bf16[ei] = self.q_dtype(scaled0) + tDQrDQ_q1_bf16[ei + 1] = self.q_dtype(scaled1) cute.arch.fence_view_async_tmem_load() + _tcgen05_fence_before_thread_sync() - for ei in cutlass.range(cute.size(tDQrDQ_q1_bf16), unroll_full=True): - h = cute.get(tCcDQ[ei], mode=[0, 0]) - d = cute.get(tCcDQ[ei], mode=[0, 1]) - sdQ_gemm_view[h, d] = tDQrDQ_q1_bf16[ei] + if const_expr(use_stmatrix_dq): + tRDQ_rdQ_q1 = tiled_smem_store_dq.retile(tDQrDQ_q1_bf16) + cute.copy( + tiled_smem_store_dq, + tRDQ_rdQ_q1, + tRDQ_sdQ, + ) + else: + for ei in cutlass.range(cute.size(tDQrDQ_q1_bf16), unroll_full=True): + h = cute.get(tCcDQ[ei], mode=[0, 0]) + d = cute.get(tCcDQ[ei], mode=[0, 1]) + sdQ_gemm_view[h, d] = tDQrDQ_q1_bf16[ei] self.compute_sync_barrier.arrive_and_wait() cute.arch.fence_proxy("async.shared", space="cta") @@ -1804,32 +2003,78 @@ def _compute_warpgroup_2q( cute.copy(tma_atom_dQ_q1, tdQsdQ_q1, tdQgdQ_q1_mkl) dQ_store_pipeline.producer_commit() - # ---- Epilogue: dW for q0 via warp reduction ---- - # mdW_b is the per-batch view: (S_q, H) BSHD or (T_q, H) THD with T-offset - # already applied. q*_local indexes within the batch. - HEADS_PER_WARP = const_expr(self.heads_padded // 4) - warp_base_h = warp_id_in_wg * Int32(HEADS_PER_WARP) - for h_local in cutlass.range_constexpr(HEADS_PER_WARP): - h = warp_base_h + h_local - my_partial = Float32(0.0) - for ei in cutlass.range(cute.size(dw_accum_q0), unroll_full=True): - if cute.get(tCcS[ei], mode=[0, 0]) == h: - my_partial = my_partial + dw_accum_q0[ei] - total = cute.arch.warp_reduction_sum(my_partial) - if lane_id == 0: - mdW_b[q0_local, h] = self.q_dtype(total) - - # ---- Epilogue: dW for q1 via warp reduction (guarded) ---- - if has_q1: + if warp_idx == compute_warp0: + dQ_store_pipeline.producer_acquire() + + # ---- Epilogue: dW for q0/q1 ---- + # The production ownership gives each thread two heads; four adjacent + # lanes cover disjoint columns for those heads. Reduce only that + # 4-lane subgroup instead of scanning the fragment once per head and + # issuing 16 full-warp reductions for each Q token. + if const_expr(use_stmatrix_ds): + q0_sum_low = Float32(0.0) + q0_sum_high = Float32(0.0) + q1_sum_low = Float32(0.0) + q1_sum_high = Float32(0.0) + for ei in cutlass.range( + cute.size(dw_accum_q0), + unroll_full=True, + ): + if (ei // 2) % 2 == 0: + q0_sum_low = q0_sum_low + dw_accum_q0[ei] + q1_sum_low = q1_sum_low + dw_accum_q1[ei] + else: + q0_sum_high = q0_sum_high + dw_accum_q0[ei] + q1_sum_high = q1_sum_high + dw_accum_q1[ei] + + q0_sum_low = cute.arch.warp_reduction_sum( + q0_sum_low, + threads_in_group=4, + ) + q0_sum_high = cute.arch.warp_reduction_sum( + q0_sum_high, + threads_in_group=4, + ) + q1_sum_low = cute.arch.warp_reduction_sum( + q1_sum_low, + threads_in_group=4, + ) + q1_sum_high = cute.arch.warp_reduction_sum( + q1_sum_high, + threads_in_group=4, + ) + if lane_id % 4 == 0: + h0 = warp_id_in_wg * 16 + lane_id // 4 + mdW_b[q0_local, h0] = self.q_dtype(q0_sum_low) + mdW_b[q0_local, h0 + 8] = self.q_dtype(q0_sum_high) + if has_q1: + mdW_b[q1_local, h0] = self.q_dtype(q1_sum_low) + mdW_b[q1_local, h0 + 8] = self.q_dtype(q1_sum_high) + else: + # General-shape fallback: retain the complete native-coordinate + # scan and full-warp reduction. + HEADS_PER_WARP = const_expr(self.heads_padded // 4) + warp_base_h = warp_id_in_wg * Int32(HEADS_PER_WARP) for h_local in cutlass.range_constexpr(HEADS_PER_WARP): h = warp_base_h + h_local - my_partial = Float32(0.0) - for ei in cutlass.range(cute.size(dw_accum_q1), unroll_full=True): + q0_partial = Float32(0.0) + for ei in cutlass.range(cute.size(dw_accum_q0), unroll_full=True): if cute.get(tCcS[ei], mode=[0, 0]) == h: - my_partial = my_partial + dw_accum_q1[ei] - total = cute.arch.warp_reduction_sum(my_partial) + q0_partial = q0_partial + dw_accum_q0[ei] + q0_total = cute.arch.warp_reduction_sum(q0_partial) if lane_id == 0: - mdW_b[q1_local, h] = self.q_dtype(total) + mdW_b[q0_local, h] = self.q_dtype(q0_total) + + if has_q1: + for h_local in cutlass.range_constexpr(HEADS_PER_WARP): + h = warp_base_h + h_local + q1_partial = Float32(0.0) + for ei in cutlass.range(cute.size(dw_accum_q1), unroll_full=True): + if cute.get(tCcS[ei], mode=[0, 0]) == h: + q1_partial = q1_partial + dw_accum_q1[ei] + q1_total = cute.arch.warp_reduction_sum(q1_partial) + if lane_id == 0: + mdW_b[q1_local, h] = self.q_dtype(q1_total) # Ensure the final q0/q1 TMA store has completed before the CTA exits. dQ_store_pipeline.producer_tail() @@ -1873,11 +2118,20 @@ def _reduce_warpgroup_2q( # 1. Wait DK_FULL cute.arch.mbarrier_wait(mbar + MBAR_2Q_DK_FULL, dk_full_phase) dk_full_phase ^= 1 + _tcgen05_fence_after_thread_sync() # 2. T2R readback dK tDKrDK = cute.make_rmem_tensor(tDKrDK_shape, Float32) cute.copy(tiled_tmem_load_dk, tDkDk_t2r, tDKrDK) cute.arch.fence_view_async_tmem_load() + _tcgen05_fence_before_thread_sync() + + # Every reducer thread participates in this one-shot lifetime + # handshake. fence_view_async_tmem_load lowers to + # tcgen05.wait::ld.sync.aligned, so arrival means this thread has + # fully consumed its final TMEM fragment. + if bi == num_kv_blocks - 1: + cute.arch.mbarrier_arrive(mbar + MBAR_2Q_REDUCE_DONE) # 3. Signal DK_EMPTY immediately after T2R (single-buffered) cute.arch.mbarrier_arrive(mbar + MBAR_2Q_DK_EMPTY) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py index 49b46f8d8..2702cb5eb 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py @@ -4,24 +4,30 @@ """ Indexer Backward — SM100 CuTe-DSL, 3-kernel design. -Three kernels launched sequentially on the same stream: +Three stream-ordered kernels; the first two use programmatic dependent launch +to overlap independent prologue work while preserving their data dependencies: Kernel 1 (CuTe DSL): score_grad — compute sum_grad and grad_signal from - AttnScore and IdxScore, overwrite both Score tensors in-place. + AttnScore and IdxScore, overwrite AttnScore with grad_signal. + The BF16-output path also clears its FP32 dK accumulation scratch in the + same launch, avoiding a standalone memset kernel. Unsupported inputs trigger an exception before this stage launches. Kernel 2 (CuTe DSL): kernel_gemm — warp-specialized GEMM kernel (below). - dK is accumulated in float32 via atomicAdd for correctness/perf. + dK is accumulated in float32; the optimized path stages padded rows in + SMEM and issues 512-byte cp.reduce.async.bulk FP32 additions. Kernel 3 (PyTorch): dk_convert — cast dK from float32 to output dtype (same as dQ, dW). Kernel 2 — Warp specialization (16 warps, 512 threads): - Warp 0: Load (Q via TMA, weights) + Warp 0: Load (Q via TMA, weights, grad_signal) Warp 1: MMA (3-stage sK pipeline, 2-stage TMEM S/dK: GEMM1 runs 1 block ahead) Warps 2-3: Idle Warps 4-7: Compute warpgroup (per-block sGradSignal load, TMEM readback S → dS → dW, dQ TMA store) - Warps 8-11: K loading warpgroup (sparse cp.async gather, 3-stage sK) - Warps 12-15: Reduce warpgroup (TMEM readback dK → atomicAdd to f32 gmem, 2-stage) + Warps 8-11: K loading warpgroup (TMA Gather4 for global IDs; manual + cp.async fallback for local IDs, 3-stage sK) + Warps 12-15: Reduce warpgroup (wide TMEM readback → padded ping-pong SMEM + → cp.reduce.async.bulk to f32 gmem, 2-stage) TopkIdxs are pre-loaded into SMEM cooperatively by all 512 threads before warp dispatch. K/dK are flattened in ``__call__`` to a 2D ``(B*S_k, D)`` view so the kernel @@ -33,8 +39,8 @@ ``cu_seqlens_k[b] + local`` indexes the ``(T_k, D)`` packed buffer.) grad_signal (precomputed by kernel 1) is loaded per topk-block by the compute warpgroup. -SMEM (kernel 2): sGradSignal[block_I] replaces the former sAttnScore/sIdxScore/sScratch, - freeing ~3.5 KB for a larger sTopkIdxs buffer. +SMEM (kernel 2): full-row grad_signal/top-k staging plus a four-warp padded + FP32 ping-pong buffer for bulk dK reduction. TMEM: S0/dK0 @0, dQ @128, S1/dK1 @256 (384/512 cols). Barriers for kernel 2: @@ -45,22 +51,41 @@ mbar[8-10]: K_loaded_0/1/2 (K-load arrives → MMA waits, 3-stage) mbar[11-13]:K_consumed_0/1/2(MMA commits after GEMM3 → K-load waits, 3-stage) mbar[14]: W_loaded (Load arrives → Compute waits) + mbar[15]: dQ_done (MMA commits after GEMM3 → Compute waits) + mbar[16]: reduce_done (Reduce arrives after T2R → TMEM owner waits) + mbar[17-18]:dS_half_0/1 (TopK=512 first-half dS publication) Each warp/warpgroup has its own independent loop, communicating via barriers. -No CLC persistent scheduling (simple grid = batch × seqlen). +TopK=128/256/384/512 can use a grid-stride persistent CTA path capped at one +CTA per SM; it retains TMEM across rows and consolidates per-row barrier +initialization. The short-row variants select it only when the grid has more +rows than SMs. Stage phases are derived from the number of 128-wide blocks per +row so the same schedule is valid for one through four blocks. +TopK=1024/2048 retain the one-query-per-CTA 2-D grid. """ from __future__ import annotations import math from functools import partial +from typing import Any, cast import torch import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute +from cutlass.cute import atom as cute_atom +from cutlass.cute import core as cute_core from cutlass import Float32, Int32, const_expr +import cutlass._mlir.dialects.cute as _cute_ir +import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir +from cutlass._mlir.dialects import llvm, nvvm +from cutlass.cutlass_dsl import T, dsl_user_op from cutlass.cute.nvgpu import cpasync +from cutlass.cute.nvgpu.cpasync.copy import ( + CopyBulkTensorTileG2SNonExecTrait, +) +from cutlass.cute.nvgpu.cpasync.helpers import TmaInfo import cutlass.cute.nvgpu.tcgen05 as tcgen05 import cutlass.utils as utils import cutlass.pipeline as pipeline @@ -75,18 +100,50 @@ import cutlass.utils.blackwell_helpers as sm100_utils_basic from cudnn.deepseek_sparse_attention.utils.compiler import compile_options +from cudnn.deepseek_sparse_attention.utils.copy import cpasync_reduce_bulk_add_f32 from cudnn.deepseek_sparse_attention.utils.runtime import ( resolve_stream as _resolve_stream, torch_stream_context as _torch_stream_context, ) +_HAS_TMA_GATHER4 = all( + hasattr(_cute_nvgpu_ir, name) + for name in ( + "atom_make_non_exec_2d_gather4_tma_load", + "GatherScatterTmaLoadEnum", + "TmaDescriptorTiledType", + "atom_make_exec_tma", + "get_tma_desc_addr", + ) +) + mul_packed_f32x2 = partial(cute.arch.mul_packed_f32x2, rnd="rn") fma_packed_f32x2 = partial(cute.arch.fma_packed_f32x2, rnd="rn") +@dsl_user_op +def _tcgen05_fence_after_thread_sync(*, loc=None, ip=None): + """Order subsequent tcgen05 operations after an inter-thread wait.""" + nvvm.tcgen05_fence( + nvvm.Tcgen05FenceKind.AFTER_THREAD_SYNC, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _tcgen05_fence_before_thread_sync(*, loc=None, ip=None): + """Order prior tcgen05 operations before an inter-thread signal.""" + nvvm.tcgen05_fence( + nvvm.Tcgen05FenceKind.BEFORE_THREAD_SYNC, + loc=loc, + ip=ip, + ) + + # Barrier indices for kernel_gemm — per-stage barriers for S_FULL, DS_READY, K_LOADED # to avoid phase-wrap when producer runs 2 blocks ahead of consumer. -# sK uses 3-stage pipeline (Opt-7): K_LOADED and K_CONSUMED are per-stage (×3). +# sK uses a 3-stage pipeline: K_LOADED and K_CONSUMED are per-stage (×3). MBAR_S_FULL_0 = 0 MBAR_S_FULL_1 = 1 MBAR_DS_READY_0 = 2 @@ -103,7 +160,16 @@ MBAR_K_CONSUMED_2 = 13 MBAR_W_LOADED = 14 MBAR_DQ_DONE = 15 -NUM_BARRIERS = 16 +MBAR_REDUCE_DONE = 16 +MBAR_DS_HALF_0 = 17 +MBAR_DS_HALF_1 = 18 +MBAR_ROW_FREE_0 = 19 +MBAR_ROW_FREE_1 = 20 +MBAR_DQ_FREE_0 = 21 +MBAR_DQ_FREE_1 = 22 +MBAR_GW_READY_0 = 23 +MBAR_GW_READY_1 = 24 +NUM_BARRIERS = 25 CLIP_LOG_MIN = -100.0 CLIP_PROB_MIN = math.exp(CLIP_LOG_MIN) @@ -111,6 +177,196 @@ _score_grad_cute_cache: dict = {} +@dsl_user_op +def _load_global_i32x4(gmem_ptr, *, loc=None, ip=None): + """Load four contiguous int32 values with one 16-byte global load.""" + result = llvm.inline_asm( + llvm.StructType.get_literal( + [T.i32(), T.i32(), T.i32(), T.i32()], + ), + [gmem_ptr.toint(loc=loc, ip=ip).ir_value()], + "ld.global.v4.u32 {$0,$1,$2,$3}, [$4];", + "=r,=r,=r,=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return ( + Int32(llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)), + Int32(llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)), + Int32(llvm.extractvalue(T.i32(), result, [2], loc=loc, ip=ip)), + Int32(llvm.extractvalue(T.i32(), result, [3], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def _load_global_f32x4(gmem_ptr, *, loc=None, ip=None): + """Load four contiguous FP32 values with one aligned 16-byte load.""" + result = llvm.inline_asm( + llvm.StructType.get_literal( + [T.f32(), T.f32(), T.f32(), T.f32()], + ), + [gmem_ptr.toint(loc=loc, ip=ip).ir_value()], + "ld.global.v4.f32 {$0,$1,$2,$3}, [$4];", + "=f,=f,=f,=f,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return ( + Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def _store_global_f32x4(gmem_ptr, value0, value1, value2, value3, *, loc=None, ip=None): + """Store four contiguous FP32 values with one aligned 16-byte store.""" + llvm.inline_asm( + None, + [ + gmem_ptr.toint(loc=loc, ip=ip).ir_value(), + Float32(value0).ir_value(loc=loc, ip=ip), + Float32(value1).ir_value(loc=loc, ip=ip), + Float32(value2).ir_value(loc=loc, ip=ip), + Float32(value3).ir_value(loc=loc, ip=ip), + ], + "st.global.v4.f32 [$0], {$1,$2,$3,$4};", + "l,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _make_tiled_tma_gather4_atom( + gmem_tensor, + gmem_coord_tensor, + smem_layout, + mma_tiler_mnk, + tiled_mma, + *, + loc=None, + ip=None, +): + """Build the SM100 2-D ``tile::gather4`` atom exposed by DSL 4.6.1+. + + The 4.6.1 wheel already ships the Gather4 MLIR operation and lowering, but + its public Python helper/export is absent. This is the minimal equivalent + of ``make_tiled_tma_atom(..., gmem_coord_tensor=...)`` documented by that + same wheel; it deliberately reuses the standard executable TMA-load trait. + DSL 4.5.x lacks the underlying MLIR operation, so callers capability-gate + this helper and retain the equivalent manual ``cp.async`` K-load path. + """ + smem_rank = cute_core.rank(smem_layout) + assert smem_rank == 3 or smem_rank == 4 + + stored_smem_layout = smem_layout + if smem_rank == 4: + smem_layout = cute_core.select(smem_layout, mode=[0, 1, 2]) + + # Match make_tiled_tma_atom_B: the B operand's MMA N/K projection is + # generally hierarchical and is not equivalent to a plain (N, K) tile. + ident = cute_core.make_identity_layout(gmem_tensor.shape, loc=loc, ip=ip) + mma_tiler_nk = (mma_tiler_mnk[1], *mma_tiler_mnk[2:]) + g_tile = cute_core.composition( + ident, + mma_tiler_nk, + loc=loc, + ip=ip, + ) + cta_v_map = tiled_mma._thrfrg_B(g_tile) + cta_v_map = cute_core.get(cta_v_map, mode=[1]) + cta_v_map = cute_core.dice( + cta_v_map, + (1, (1,) * cute_core.rank(g_tile)), + ) + + smem_for_ir = smem_layout + if isinstance(smem_for_ir, cute_core._ComposedLayout): + smem_for_ir = smem_for_ir.value + + op = cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE) + res = _cute_nvgpu_ir.atom_make_non_exec_2d_gather4_tma_load( + cast(Any, gmem_tensor).value, + gmem_coord_tensor.layout, + smem_for_ir, + cta_v_map, + _cute_nvgpu_ir.GatherScatterTmaLoadEnum.sm_100, + num_multicast=1, + loc=loc, + ip=ip, + ) + return TmaInfo( + cute_atom.CopyAtom(op, CopyBulkTensorTileG2SNonExecTrait(res[0])), + res[1], + stored_smem_layout, + ) + + +@dsl_user_op +def _tma_gather4_k_rows( + tma_atom, + smem_ptr, + column, + row0, + row1, + row2, + row3, + transaction_barrier, + *, + loc=None, + ip=None, +): + """Issue one lane-level 4-row x 128-byte SM100 Gather4 transaction.""" + desc_ptr_type = _cute_ir.PtrType.get( + _cute_nvgpu_ir.TmaDescriptorTiledType.get(), + cute.AddressSpace.generic, + 64, + ) + exec_atom = _cute_nvgpu_ir.atom_make_exec_tma( + tma_atom._trait.value, + loc=loc, + ip=ip, + ) + desc_ptr = _cute_nvgpu_ir.get_tma_desc_addr( + desc_ptr_type, + exec_atom, + loc=loc, + ip=ip, + ) + desc_ptr_i64 = desc_ptr.toint(loc=loc, ip=ip).ir_value() + smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value() + barrier_ptr_i32 = transaction_barrier.toint(loc=loc, ip=ip).ir_value() + llvm.inline_asm( + None, + [ + smem_ptr_i32, + desc_ptr_i64, + Int32(column).ir_value(), + Int32(row0).ir_value(), + Int32(row1).ir_value(), + Int32(row2).ir_value(), + Int32(row3).ir_value(), + barrier_ptr_i32, + ], + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" ".mbarrier::complete_tx::bytes" " [$0], [$1, {$2, $3, $4, $5, $6}], [$7];", + "r,l,r,r,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + class IndexerBackwardSm100: arch = 100 WARP_SIZE = 32 @@ -118,6 +374,13 @@ class IndexerBackwardSm100: NUM_WARPS = 16 THREADS_PER_CTA = 512 + # dK bulk-reduce staging: four reducer warps, two ping-pong buffers per + # warp, eight 512-byte rows per buffer, and 32-byte row padding. + DK_STAGE_ROW_FLOATS = 136 + DK_STAGE_ROWS = 8 + DK_STAGE_BUFFERS = 2 + DK_STAGE_ELEMENTS = 4 * DK_STAGE_BUFFERS * DK_STAGE_ROWS * DK_STAGE_ROW_FLOATS + # Warp assignments load_warp_id = 0 mma_warp_id = 1 @@ -126,11 +389,30 @@ class IndexerBackwardSm100: k_load_warp_id = (8, 9, 10, 11) reduce_warp_id = (12, 13, 14, 15) - def __init__(self, head_dim, heads=64, block_I=128, topk=512, topk_indices_global: bool = True): + def __init__( + self, + head_dim, + heads=64, + block_I=128, + topk=512, + total_seqlen_k: int | None = None, + total_rows: int = 1, + persistent_grid_size: int = 1, + topk_indices_global: bool = True, + enable_score_pdl: bool = False, + ): self.head_dim = head_dim self.heads = heads self.block_I = block_I self.topk = topk + self.total_seqlen_k = total_seqlen_k + self.total_rows = total_rows + self.persistent_grid_size = persistent_grid_size + # Only the full wrapper guarantees that the immediate predecessor is + # ScoreGrad and that Q/K/top-k are independent of it. Direct + # ``gemm_only`` calls compile a non-PDL variant to preserve arbitrary + # stream-predecessor semantics. + self.enable_score_pdl = enable_score_pdl # When True (default, matches the public fwd convention), mTopkIdx # carries global KV ids (``b * seqlen_k + local``); the kernel uses # them as flat ids into the (B*S_k, D) K/dK view directly. When @@ -138,11 +420,37 @@ def __init__(self, head_dim, heads=64, block_I=128, topk=512, topk_indices_globa # ``batch_idx * S_k_per_batch`` to convert. Const_expr-branched. self.topk_indices_global = topk_indices_global assert heads >= 64 + assert topk > 0 assert topk % block_I == 0 self.num_topk_blocks = topk // block_I self.head_dim_padded = int(math.ceil(head_dim / 16) * 16) self.heads_padded = int(math.ceil(heads / 8) * 8) + # Half-dS publication is specialized for the production TopK=512, + # H64 x I128 fragment map. Larger TopK values retain the full-dS drain + # order because the additional barrier traffic regresses them. + self.use_ds_half = topk == 512 and self.heads_padded == 64 and block_I == 128 + self.use_persistent = ( + topk in (128, 256, 384, 512) + and self.heads_padded == 64 + and self.head_dim_padded == 128 + and block_I == 128 + # For the new short-row variants, at most one row per resident CTA + # leaves nothing to amortize and is 1-5% slower. Preserve the + # established TopK=512 dispatch policy unchanged. + and (topk == 512 or total_rows > persistent_grid_size) + ) + # Gather4 consumes explicit row coordinates, so short-row local IDs can + # be normalized to flat global IDs in registers before issue. Keep the + # established TopK=512 local-ID fallback unchanged; the new policy is + # deliberately scoped to the 128/256/384 specializations evaluated + # here. + # Public DSL 4.5.x wheels do not contain the private MLIR operation + # needed to construct a Gather4 descriptor. Keep those wheels on the + # existing manual cp.async loader; both paths feed identical BF16 K + # tiles into the same FP32 GEMMs. + self.use_tma_gather = _HAS_TMA_GATHER4 and (topk_indices_global or (self.use_persistent and topk in (128, 256, 384))) + self.use_cross_row_persistent = self.use_persistent and self.use_tma_gather # GEMM tilers (M, N, K) — cute.gemm, SMEM operands, TMEM acc # GEMM1: S[H,TileN] = Q[H,D] @ K[TileN,D]. A=Q K-major, B=K K-major @@ -162,14 +470,17 @@ def __init__(self, head_dim, heads=64, block_I=128, topk=512, topk_indices_globa self.tmem_s0_offset = 0 self.tmem_dq_offset = 128 self.tmem_s1_offset = 256 + self.tmem_dq_p1_offset = 384 self.tmem_alloc_cols = 512 # Register budgets — must sum to 512 per thread (65536 regs / 128 threads per WG) # Compute needs 128+ (tSrS=64 + dw_accum=64), Reduce needs 128+ (tDKrDK=128) self.num_regs_wg0 = 40 - self.num_regs_compute = 200 + # The dS compute warpgroup owns the largest register working set; + # all four allocations exactly consume the per-CTA SM100 budget. + self.num_regs_compute = 224 self.num_regs_reduce = 200 - self.num_regs_kload = 32 + self.num_regs_kload = 48 self.buffer_align_bytes = 1024 @@ -261,16 +572,18 @@ def __call__( ) # SMEM layouts — primary views - # sK/sKt: 3-stage pipeline (Opt-7) for hiding K-load scatter latency. + # sK/sKt: 3-stage pipeline for hiding K-load scatter latency. # sdS: 2-stage pipeline (tied to TMEM S/dK 2-stage). sQ_layout = _make_smem_layout_a(tmma1, self.gemm1_tiler, self.q_dtype, 1) sK_layout = _make_smem_layout_b(tmma1, self.gemm1_tiler, self.k_dtype, 3) sdS_layout = _make_smem_layout_a(tmma3, self.gemm3_tiler, self.q_dtype, 2) # Epilogue-style store layout for stmatrix writes to sdS (same physical SMEM). - # COL_MAJOR (M-major) + square tile → physically compatible with A-operand layout. + # dS is logical [H, I]. GEMM3 consumes it as a K-major A operand, so + # its physical storage is row-major in that logical view. GEMM2 sees + # the same bytes as the transposed [I, H] MN-major A operand. sdS_store_layout = _make_smem_layout_epi( self.q_dtype, - LayoutEnum.COL_MAJOR, + LayoutEnum.ROW_MAJOR, (self.heads_padded, self.block_I), 2, ) @@ -299,6 +612,36 @@ def __call__( ) self.tma_copy_Q_bytes = cute.size_in_bytes(self.q_dtype, Q_smem_layout_tma) + tma_atom_K_gather = None + if const_expr(self.use_tma_gather): + # Hardware sparse gather. The index-coordinate tensor has the + # same logical 2-D shape as K, but its D mode is broadcast + # (stride 0): each group of four row coordinates supplies one + # Gather4 instruction. The coordinate tensor is descriptor + # metadata; the issuing lanes pass row IDs explicitly. + gI_desc = cute.make_tensor( + mTopkIdx.iterator, + cute.make_layout( + (self.total_seqlen_k, self.head_dim_padded), + stride=(1, 0), + ), + ) + mK_gather = cute.make_tensor( + mK.iterator, + cute.make_layout( + (self.total_seqlen_k, self.head_dim_padded), + stride=(self.head_dim_padded, 1), + ), + ) + K_smem_layout_tma = cute.select(sK_layout, mode=[0, 1, 2]) + tma_atom_K_gather, _ = _make_tiled_tma_gather4_atom( + mK_gather, + gI_desc, + K_smem_layout_tma, + self.gemm1_tiler, + tmma1, + ) + # Epilogue SMEM layout for dQ store (bf16, row-major = D contiguous) sdQ_epi_layout = _make_smem_layout_epi( self.q_dtype, @@ -320,6 +663,8 @@ def __call__( seqlen = cute.size(mQ.shape[0]) batch_size = cute.size(mQ.shape[3]) if cute.rank(mQ.shape) > 3 else 1 + grid_rows = min(self.total_rows, self.persistent_grid_size) if self.use_persistent else self.total_rows + launch_grid = (grid_rows, 1, 1) if self.use_persistent else (seqlen, batch_size, 1) self.kernel_gemm( mQ_tma, mW, @@ -341,16 +686,18 @@ def __call__( sQ_g2b_layout, sdS_store_layout, tma_atom_Q, + tma_atom_K_gather, tma_atom_dQ, sdQ_epi_layout, seqlen, batch_size, ).launch( - grid=(seqlen, batch_size, 1), + grid=launch_grid, block=[self.THREADS_PER_CTA, 1, 1], cluster=[1, 1, 1], stream=stream, min_blocks_per_mp=1, + use_pdl=self.enable_score_pdl, ) @cute.kernel @@ -376,69 +723,72 @@ def kernel_gemm( sQ_g2b_layout, sdS_store_layout, tma_atom_Q, + tma_atom_K_gather, tma_atom_dQ, sdQ_epi_layout, seqlen: Int32, batch_size: Int32, ): + # ScoreGrad launches this grid programmatically. Serial rows wait at + # entry; the cross-row persistent path delays the wait until its load + # warp has issued the independent first-row Q TMA, allowing the Q/K + # prologue and barrier initialization to overlap score_grad safely. + if const_expr(self.enable_score_pdl and not self.use_cross_row_persistent): + cute.arch.griddepcontrol_wait() tidx = cute.arch.thread_idx()[0] warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - seq_idx = cute.arch.block_idx()[0] - batch_idx = cute.arch.block_idx()[1] + if const_expr(self.use_persistent): + flat_row_idx = cute.arch.block_idx()[0] + batch_idx = flat_row_idx // seqlen + seq_idx = flat_row_idx - batch_idx * seqlen + else: + seq_idx = cute.arch.block_idx()[0] + batch_idx = cute.arch.block_idx()[1] + flat_row_idx = batch_idx * seqlen + seq_idx seqlen_k = cute.size(mK.shape[0]) # TMA descriptor prefetch (load warp only) if warp_idx == self.load_warp_id: cpasync.prefetch_descriptor(tma_atom_Q) cpasync.prefetch_descriptor(tma_atom_dQ) + if const_expr(self.use_tma_gather): + if warp_idx == self.k_load_warp_id[0]: + cpasync.prefetch_descriptor(tma_atom_K_gather) # SMEM allocation sQ_size = cute.cosize(sQ_layout) sK_size = cute.cosize(sK_layout) sdS_size = cute.cosize(sdS_layout) - - # Compute sTopkIdxs capacity from remaining SMEM (dsa-next pattern) - def _align_up(x, a): - return (x + a - 1) // a * a - - _elem_bytes = self.q_dtype.width // 8 - _tma_align = self.buffer_align_bytes # 1024 - _non_tma_align = 128 - - _offset = 0 - _offset += self.Q_mbar_size * 8 # Q_mbar: Int64 × Q_mbar_size - _offset += NUM_BARRIERS * 8 # mbar: Int64 × NUM_BARRIERS - _offset += 4 # tmem_holding_buf: Int32 - _offset = _align_up(_offset, _tma_align) - _offset += int(sQ_size) * _elem_bytes # sQ - _offset = _align_up(_offset, _tma_align) - _offset += int(sK_size) * _elem_bytes # sK - _offset = _align_up(_offset, _tma_align) - _offset += int(sdS_size) * _elem_bytes # sdS - _offset = _align_up(_offset, _non_tma_align) - _offset += self.topk * 4 # sGradSignal: Float32 × topk - _offset = _align_up(_offset, _non_tma_align) - # sTopkIdxs goes here — compute remaining space - _topk_idx_offset = _offset - # Account for sW that comes after sTopkIdxs, - # with worst-case alignment padding. - _tail = 0 - _tail += _non_tma_align # worst-case align padding before sW - _tail += self.heads * _elem_bytes # sW + sdQ_epi_size = cute.cosize(sdQ_epi_layout) + _row_operand_stages = 2 if self.use_cross_row_persistent else 1 + _sQ_storage_size = int(sQ_size) * _row_operand_stages + _grad_storage_size = self.topk * _row_operand_stages + _weight_storage_size = self.heads * _row_operand_stages + _dq_epi_storage_size = int(sdQ_epi_size) if self.use_cross_row_persistent else 1 + + # Cross-row persistence double-buffers top-k IDs; all other paths need + # exactly one complete row. SharedStorage's size assertion below is the + # single source of truth for legal specializations. + smem_topk_capacity = 2 * self.topk if self.use_cross_row_persistent else self.topk + _dk_stage_elements = self.DK_STAGE_ELEMENTS _max_smem_bytes = 227 * 1024 - smem_topk_capacity = (_max_smem_bytes - _topk_idx_offset - _tail) // 4 @cute.struct class SharedStorage: Q_mbar: cute.struct.MemRange[cutlass.Int64, self.Q_mbar_size] mbar: cute.struct.MemRange[cutlass.Int64, NUM_BARRIERS] tmem_holding_buf: Int32 - sQ: cute.struct.Align[cute.struct.MemRange[self.q_dtype, sQ_size], self.buffer_align_bytes] + sQ: cute.struct.Align[cute.struct.MemRange[self.q_dtype, _sQ_storage_size], self.buffer_align_bytes] sK: cute.struct.Align[cute.struct.MemRange[self.k_dtype, sK_size], self.buffer_align_bytes] sdS: cute.struct.Align[cute.struct.MemRange[self.q_dtype, sdS_size], self.buffer_align_bytes] - sGradSignal: cute.struct.Align[cute.struct.MemRange[Float32, self.topk], 128] + sGradSignal: cute.struct.Align[cute.struct.MemRange[Float32, _grad_storage_size], 128] sTopkIdxs: cute.struct.Align[cute.struct.MemRange[Int32, smem_topk_capacity], 128] - sW: cute.struct.Align[cute.struct.MemRange[self.q_dtype, self.heads], 128] + sW: cute.struct.Align[cute.struct.MemRange[self.q_dtype, _weight_storage_size], 128] + sdKStage: cute.struct.Align[cute.struct.MemRange[Float32, _dk_stage_elements], 128] + sdQEpilogue: cute.struct.Align[ + cute.struct.MemRange[self.q_dtype, _dq_epi_storage_size], + self.buffer_align_bytes, + ] assert SharedStorage.size_in_bytes() <= _max_smem_bytes, ( f"SharedStorage ({SharedStorage.size_in_bytes()} bytes) exceeds {_max_smem_bytes} bytes (227KB), " f"smem_topk_capacity={smem_topk_capacity}" @@ -456,6 +806,7 @@ class SharedStorage: ) # Swizzled SMEM tensors + sK_raw_ptr = storage.sK.data_ptr() sQ = storage.sQ.get_tensor(sQ_layout.outer, swizzle=sQ_layout.inner) sK = storage.sK.get_tensor(sK_layout.outer, swizzle=sK_layout.inner) sdS = storage.sdS.get_tensor(sdS_layout.outer, swizzle=sdS_layout.inner) @@ -469,6 +820,58 @@ class SharedStorage: sGradSignal = storage.sGradSignal.get_tensor(cute.make_layout((self.topk,), stride=(1,))) sTopkIdxs = storage.sTopkIdxs.get_tensor(cute.make_layout((smem_topk_capacity,), stride=(1,))) sW = storage.sW.get_tensor(cute.make_layout((self.heads,), stride=(1,))) + sdKStage = storage.sdKStage.get_tensor(cute.make_layout((_dk_stage_elements,), stride=(1,))) + + if const_expr(self.use_cross_row_persistent): + # K is live across row boundaries, so unlike the serial path dQ + # cannot alias its storage. This dedicated tile is reused only by + # the compute warpgroup, with the previous TMA store drained just + # before the next overwrite. + sdQ_epi_cross = storage.sdQEpilogue.get_tensor( + sdQ_epi_layout.outer, + swizzle=sdQ_epi_layout.inner, + ) + self._run_persistent_cross_row( + mQ, + mW, + mK, + mdQ, + mdW, + mdK_f32, + mGradSignal, + mTopkIdx, + sm_scale, + tmma1, + tmma2, + tmma3, + sQ_layout, + sdS_g2a_layout, + sK, + sKt, + sdS, + sQ_g2b_layout, + sdS_store, + sK_raw_ptr, + storage.sQ.data_ptr(), + storage.sGradSignal.data_ptr(), + storage.sTopkIdxs.data_ptr(), + storage.sW.data_ptr(), + sdKStage, + sdQ_epi_cross, + Q_mbar_ptr, + mbar, + tmem, + tma_atom_Q, + tma_atom_K_gather, + tma_atom_dQ, + seqlen, + batch_size, + seqlen_k, + flat_row_idx, + tidx, + warp_idx, + ) + return # dQ epilogue SMEM — reuses sK physical memory (safe: dQ store happens after all iterations) sdQ_epi = cute.make_tensor( @@ -484,6 +887,7 @@ class SharedStorage: consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), tx_count=self.tma_copy_Q_bytes, cta_layout_vmnk=cute.make_layout(self.cluster_shape), + defer_sync=self.use_persistent, ) Q_producer, Q_consumer = Q_pipeline.make_participants() @@ -528,62 +932,990 @@ class SharedStorage: cute.group_modes(gdQ, 0, 2), ) - # Init all barriers (single lane of warp 0) - if warp_idx == 0: - with cute.arch.elect_one(): - cute.arch.mbarrier_init(mbar + MBAR_S_FULL_0, 1) - cute.arch.mbarrier_init(mbar + MBAR_S_FULL_1, 1) - cute.arch.mbarrier_init(mbar + MBAR_DS_READY_0, self.WARPGROUP_SIZE) - cute.arch.mbarrier_init(mbar + MBAR_DS_READY_1, self.WARPGROUP_SIZE) - cute.arch.mbarrier_init(mbar + MBAR_DK_FULL_0, 1) - cute.arch.mbarrier_init(mbar + MBAR_DK_FULL_1, 1) - cute.arch.mbarrier_init(mbar + MBAR_DK_EMPTY_0, self.WARPGROUP_SIZE) - cute.arch.mbarrier_init(mbar + MBAR_DK_EMPTY_1, self.WARPGROUP_SIZE) + # Initialize each barrier exactly once. + if tidx == 0: + cute.arch.mbarrier_init(mbar + MBAR_S_FULL_0, 1) + cute.arch.mbarrier_init(mbar + MBAR_S_FULL_1, 1) + cute.arch.mbarrier_init(mbar + MBAR_DS_READY_0, self.WARPGROUP_SIZE) + cute.arch.mbarrier_init(mbar + MBAR_DS_READY_1, self.WARPGROUP_SIZE) + cute.arch.mbarrier_init(mbar + MBAR_DK_FULL_0, 1) + cute.arch.mbarrier_init(mbar + MBAR_DK_FULL_1, 1) + cute.arch.mbarrier_init(mbar + MBAR_DK_EMPTY_0, self.WARPGROUP_SIZE) + cute.arch.mbarrier_init(mbar + MBAR_DK_EMPTY_1, self.WARPGROUP_SIZE) + if const_expr(self.use_tma_gather): + # Transaction barrier: one explicit arrive-and-expect plus the + # 32 KiB Gather4 completion for each K stage. + cute.arch.mbarrier_init(mbar + MBAR_K_LOADED_0, 1) + cute.arch.mbarrier_init(mbar + MBAR_K_LOADED_1, 1) + cute.arch.mbarrier_init(mbar + MBAR_K_LOADED_2, 1) + else: cute.arch.mbarrier_init(mbar + MBAR_K_LOADED_0, self.WARPGROUP_SIZE) cute.arch.mbarrier_init(mbar + MBAR_K_LOADED_1, self.WARPGROUP_SIZE) cute.arch.mbarrier_init(mbar + MBAR_K_LOADED_2, self.WARPGROUP_SIZE) - cute.arch.mbarrier_init(mbar + MBAR_K_CONSUMED_0, 1) - cute.arch.mbarrier_init(mbar + MBAR_K_CONSUMED_1, 1) - cute.arch.mbarrier_init(mbar + MBAR_K_CONSUMED_2, 1) - cute.arch.mbarrier_init(mbar + MBAR_W_LOADED, self.WARP_SIZE) - cute.arch.mbarrier_init(mbar + MBAR_DQ_DONE, 1) + cute.arch.mbarrier_init(mbar + MBAR_K_CONSUMED_0, 1) + cute.arch.mbarrier_init(mbar + MBAR_K_CONSUMED_1, 1) + cute.arch.mbarrier_init(mbar + MBAR_K_CONSUMED_2, 1) + cute.arch.mbarrier_init(mbar + MBAR_W_LOADED, self.WARP_SIZE) + cute.arch.mbarrier_init(mbar + MBAR_DQ_DONE, 1) + cute.arch.mbarrier_init( + mbar + MBAR_REDUCE_DONE, + self.WARPGROUP_SIZE, + ) + cute.arch.mbarrier_init( + mbar + MBAR_DS_HALF_0, + self.WARPGROUP_SIZE, + ) + cute.arch.mbarrier_init( + mbar + MBAR_DS_HALF_1, + self.WARPGROUP_SIZE, + ) + if const_expr(self.use_persistent): + cute.arch.mbarrier_init_fence() + cute.arch.sync_threads() + + # Pre-load the complete top-k row into SMEM cooperatively. + # K/dK are flattened to (B*S_k, D) above, so consumers index by global + # flat KV ids. ``topk_indices_global=True`` (default): ``mTopkIdx`` already + # carries global ids (``b * seqlen_k + local``); load directly. + # ``topk_indices_global=False``: ids are local-per-batch. Only ids in + # ``[0, S_k_per_batch)`` are converted to global flat ids; negative or + # positive-OOB entries are normalized to -1 so they cannot alias a row + # in a neighboring batch after adding the batch offset. + seqlen_k_per_batch = seqlen_k // batch_size + batch_offset_l2g = Int32(0) if const_expr(self.topk_indices_global) else batch_idx * seqlen_k_per_batch + TOPK_PER_THREAD = const_expr((self.topk + self.THREADS_PER_CTA - 1) // self.THREADS_PER_CTA) + for ii in cutlass.range_constexpr(TOPK_PER_THREAD): + pos = ii * self.THREADS_PER_CTA + tidx + if pos < self.topk: + raw_id = Int32(mTopkIdx[seq_idx, pos, batch_idx]) + if const_expr(self.topk_indices_global): + sTopkIdxs[pos] = raw_id + else: + sTopkIdxs[pos] = raw_id + batch_offset_l2g if raw_id >= 0 and raw_id < seqlen_k_per_batch else Int32(-1) + cute.arch.sync_threads() + + # Pre-compute accumulator shapes/layouts from tmma before dispatch, + # so branches that don't run _mma_warp never touch the tmma objects + # (avoids MLIR SSA domination issues from tmma.set() inside _mma_warp). + s_acc_shape = tmma1.partition_shape_C(self.gemm1_tiler[:2]) + s_acc_layout = tmma1.make_fragment_C(s_acc_shape).layout + dq_acc_shape = tmma3.partition_shape_C(self.gemm3_tiler[:2]) + dq_acc_layout = tmma3.make_fragment_C(dq_acc_shape).layout + dk_acc_shape = tmma2.partition_shape_C(self.gemm2_tiler[:2]) + dk_acc_layout = tmma2.make_fragment_C(dk_acc_shape).layout + + # tcgen05 ``set(ACCUMULATE, ...)`` mutates the Python Atom wrapper. + # Keep pristine wrappers for the staged persistent loop so its MMA + # operands are rooted in values that dominate the loop region. + if const_expr(self.use_persistent): + tmma1_persistent = tmma1.__new_from_mlir_values__( + tmma1.__extract_mlir_values__(), + ) + tmma2_persistent = tmma2.__new_from_mlir_values__( + tmma2.__extract_mlir_values__(), + ) + tmma3_persistent = tmma3.__new_from_mlir_values__( + tmma3.__extract_mlir_values__(), + ) + + # ============================================================= + # Warp dispatch — setmaxnreg rebalances registers across WGs. + # ============================================================= + if warp_idx == self.load_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_wg0) + self._load_warp( + mW, + mGradSignal, + sW, + sGradSignal, + tma_atom_Q, + tQsQ, + tQgQ_mkl, + Q_producer, + seq_idx, + batch_idx, + tidx, + mbar, + ) + + elif warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_wg0) + tmem.wait_for_alloc() + tmem_ptr_base = tmem.retrieve_ptr(self.acc_dtype) + tStS_0, tStS_1, tDqDq, tDkDk_0, tDkDk_1 = self.get_tmem_tensor( + s_acc_layout, + dq_acc_layout, + dk_acc_layout, + tmem_ptr_base, + ) + self._mma_warp( + sQ, + sdS_g2a, + sK, + sKt, + sdS, + sQ_g2b, + tmma1, + tmma2, + tmma3, + tStS_0, + tStS_1, + tDqDq, + tDqDq, + tDkDk_0, + tDkDk_1, + Q_consumer, + mbar, + Int32(0), + Int32(0), + Int32(0), + mbar, + Int32(0), + ) + + elif warp_idx in self.compute_warp_id: + cute.arch.setmaxregister_increase(self.num_regs_compute) + if warp_idx == self.compute_warp_id[0]: + tmem.allocate(self.tmem_alloc_cols) + tmem.wait_for_alloc() + tmem_ptr_base = tmem.retrieve_ptr(self.acc_dtype) + tStS_0, tStS_1, tDqDq, tDkDk_0, tDkDk_1 = self.get_tmem_tensor( + s_acc_layout, + dq_acc_layout, + dk_acc_layout, + tmem_ptr_base, + ) + self._compute_warpgroup( + mdW, + sGradSignal, + sW, + sdS_store, + sdS, + sdQ_epi_slice, + s_acc_shape, + dq_acc_shape, + tStS_0, + tStS_1, + tDqDq, + tDqDq, + tma_atom_dQ, + tdQsdQ, + tdQgdQ_mkl, + dQ_store_pipeline, + sm_scale, + seq_idx, + batch_idx, + tidx, + warp_idx, + mbar, + Int32(0), + mbar + MBAR_W_LOADED, + Int32(0), + mbar, + ) + if warp_idx == self.compute_warp_id[0]: + cute.arch.mbarrier_wait(mbar + MBAR_REDUCE_DONE, Int32(0)) + _tcgen05_fence_after_thread_sync() + if const_expr(not self.use_persistent): + cute.arch.dealloc_tmem(tmem_ptr_base, self.tmem_alloc_cols) + if const_expr(self.use_persistent): + dQ_store_pipeline.producer_tail() + + elif warp_idx in self.k_load_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_kload) + self._k_load_warpgroup( + mK, + sK, + sK_raw_ptr, + sTopkIdxs, + mTopkIdx, + tma_atom_K_gather, + seq_idx, + batch_idx, + seqlen_k, + tidx, + mbar, + Int32(0), + ) + + elif warp_idx in self.reduce_warp_id: + cute.arch.setmaxregister_increase(self.num_regs_reduce) + tmem.wait_for_alloc() + tmem_ptr_base = tmem.retrieve_ptr(self.acc_dtype) + tStS_0, tStS_1, tDqDq, tDkDk_0, tDkDk_1 = self.get_tmem_tensor( + s_acc_layout, + dq_acc_layout, + dk_acc_layout, + tmem_ptr_base, + ) + self._reduce_warpgroup( + mdK_f32, + sTopkIdxs, + sdKStage, + dk_acc_shape, + tDkDk_0, + tDkDk_1, + sm_scale, + seqlen_k, + tidx, + mbar, + Int32(0), + ) + + else: + cute.arch.setmaxregister_decrease(self.num_regs_wg0) + + if const_expr(self.use_persistent): + # The local-id fallback fully drains each row before reusing the + # single operand set; TMEM remains allocated across rows. + cute.arch.sync_threads() + row_stride = cute.arch.grid_dim()[0] + for next_flat_row in cutlass.range( + flat_row_idx + row_stride, + self.total_rows, + row_stride, + ): + next_batch_idx = next_flat_row // seqlen + next_seq_idx = next_flat_row - next_batch_idx * seqlen + self._run_persistent_row_serial( + mQ, + mW, + mK, + mdQ, + mdW, + mdK_f32, + mGradSignal, + mTopkIdx, + sm_scale, + tmma1_persistent, + tmma2_persistent, + tmma3_persistent, + sQ, + sdS_g2a, + sK, + sKt, + sdS, + sQ_g2b, + sdS_store, + sK_raw_ptr, + sGradSignal, + sTopkIdxs, + sW, + sdKStage, + sdQ_epi, + Q_mbar_ptr, + mbar, + tmem, + tma_atom_Q, + tma_atom_K_gather, + tma_atom_dQ, + s_acc_shape, + s_acc_layout, + dq_acc_shape, + dq_acc_layout, + dk_acc_shape, + dk_acc_layout, + seqlen_k, + batch_size, + next_seq_idx, + next_batch_idx, + tidx, + warp_idx, + ) + cute.arch.sync_threads() + if const_expr(self.use_persistent): + if warp_idx == self.compute_warp_id[0]: + persistent_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + cute.arch.dealloc_tmem( + persistent_tmem_ptr, + self.tmem_alloc_cols, + ) + + @cute.jit + def _issue_cross_row_q( + self, + tmma1, + gQ, + sQ, + tma_atom_Q, + q_ready_barrier, + ): + """Issue one Q TMA using statically aligned SMEM/barrier operands.""" + tAgQ = tmma1.get_slice(0).partition_A(gQ) + tQsQ, tQgQ = cpasync.tma_partition( + tma_atom_Q, + 0, + cute.make_layout(1), + cute.group_modes(sQ, 0, 3), + cute.group_modes(tAgQ, 0, 3), + ) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + q_ready_barrier, + self.tma_copy_Q_bytes, + ) + # CopyBulkTensorTileG2SOp inserts its own elect_one. Nesting it under + # an explicit election is documented by the DSL as a deadlock. + cute.copy( + tma_atom_Q, + tQgQ[None, 0, 0], + tQsQ[None, 0], + tma_bar_ptr=q_ready_barrier, + ) + + @cute.jit + def _run_persistent_cross_row( + self, + mQ, + mW, + mK, + mdQ, + mdW, + mdK_f32, + mGradSignal, + mTopkIdx, + sm_scale, + tmma1, + tmma2, + tmma3, + sQ_layout, + sdS_g2a_layout, + sK, + sKt, + sdS, + sQ_g2b_layout, + sdS_store, + sK_raw_ptr, + sQ_storage_ptr, + sGrad_storage_ptr, + sTopk_storage_ptr, + sW_storage_ptr, + sdKStage, + sdQ_epi, + Q_mbar_ptr, + mbar, + tmem, + tma_atom_Q, + tma_atom_K_gather, + tma_atom_dQ, + seqlen, + batch_size, + seqlen_k, + first_flat_row, + tidx, + warp_idx, + ): + """Persistent CTA with dominance-safe role-local row loops. + + Only scalar row/stage/phase values are loop carried. In particular, + the mutable tcgen05 MMA wrappers are cloned inside ``_mma_warp`` and + never escape warp 1's dynamic branch. Barrier phases are derived + from ``num_topk_blocks`` so one-, two-, three-, and four-block rows can + share the same role decomposition without restarting the pipeline. + """ + s_acc_shape = tmma1.partition_shape_C(self.gemm1_tiler[:2]) + s_acc_layout = tmma1.make_fragment_C(s_acc_shape).layout + dq_acc_shape = tmma3.partition_shape_C(self.gemm3_tiler[:2]) + dq_acc_layout = tmma3.make_fragment_C(dq_acc_shape).layout + dk_acc_shape = tmma2.partition_shape_C(self.gemm2_tiler[:2]) + dk_acc_layout = tmma2.make_fragment_C(dk_acc_shape).layout + sdS_g2a = cute.make_tensor( + cute.recast_ptr(sdS.iterator, sdS_g2a_layout.inner), + sdS_g2a_layout.outer, + ) + + if tidx == 0: + cute.arch.mbarrier_init(Q_mbar_ptr, 1) + cute.arch.mbarrier_init(Q_mbar_ptr + 1, 1) + cute.arch.mbarrier_init(mbar + MBAR_S_FULL_0, 1) + cute.arch.mbarrier_init(mbar + MBAR_S_FULL_1, 1) + cute.arch.mbarrier_init( + mbar + MBAR_DS_READY_0, + self.WARPGROUP_SIZE, + ) + cute.arch.mbarrier_init( + mbar + MBAR_DS_READY_1, + self.WARPGROUP_SIZE, + ) + cute.arch.mbarrier_init(mbar + MBAR_DK_FULL_0, 1) + cute.arch.mbarrier_init(mbar + MBAR_DK_FULL_1, 1) + cute.arch.mbarrier_init( + mbar + MBAR_DK_EMPTY_0, + self.WARPGROUP_SIZE, + ) + cute.arch.mbarrier_init( + mbar + MBAR_DK_EMPTY_1, + self.WARPGROUP_SIZE, + ) + cute.arch.mbarrier_init(mbar + MBAR_K_LOADED_0, 1) + cute.arch.mbarrier_init(mbar + MBAR_K_LOADED_1, 1) + cute.arch.mbarrier_init(mbar + MBAR_K_LOADED_2, 1) + cute.arch.mbarrier_init(mbar + MBAR_K_CONSUMED_0, 1) + cute.arch.mbarrier_init(mbar + MBAR_K_CONSUMED_1, 1) + cute.arch.mbarrier_init(mbar + MBAR_K_CONSUMED_2, 1) + cute.arch.mbarrier_init(mbar + MBAR_W_LOADED, self.WARP_SIZE) + cute.arch.mbarrier_init(mbar + MBAR_DQ_DONE, 1) + cute.arch.mbarrier_init( + mbar + MBAR_REDUCE_DONE, + self.WARPGROUP_SIZE, + ) + cute.arch.mbarrier_init( + mbar + MBAR_DS_HALF_0, + self.WARPGROUP_SIZE, + ) + cute.arch.mbarrier_init( + mbar + MBAR_DS_HALF_1, + self.WARPGROUP_SIZE, + ) + cute.arch.mbarrier_init( + mbar + MBAR_ROW_FREE_0, + 3 * self.WARPGROUP_SIZE, + ) + cute.arch.mbarrier_init( + mbar + MBAR_ROW_FREE_1, + 3 * self.WARPGROUP_SIZE, + ) + cute.arch.mbarrier_init( + mbar + MBAR_DQ_FREE_0, + self.WARPGROUP_SIZE, + ) + cute.arch.mbarrier_init( + mbar + MBAR_DQ_FREE_1, + self.WARPGROUP_SIZE, + ) + cute.arch.mbarrier_init(mbar + MBAR_GW_READY_0, self.WARP_SIZE) + cute.arch.mbarrier_init(mbar + MBAR_GW_READY_1, self.WARP_SIZE) + cute.arch.mbarrier_init_fence() + cute.arch.sync_threads() + + row_stride = cute.arch.grid_dim()[0] + sQ_size = const_expr(int(cute.cosize(sQ_layout))) + sQ_p0 = cute.make_tensor( + cute.recast_ptr(sQ_storage_ptr, sQ_layout.inner), + sQ_layout.outer, + ) + sQ_p1 = cute.make_tensor( + cute.recast_ptr( + sQ_storage_ptr + sQ_size, + sQ_layout.inner, + ), + sQ_layout.outer, + ) + + if warp_idx == self.load_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_wg0) + lane_id = tidx % self.WARP_SIZE + for row in cutlass.range( + first_flat_row, + self.total_rows, + row_stride, + ): + it = (row - first_flat_row) // row_stride + parity = it & 1 + batch_idx = row // seqlen + seq_idx = row - batch_idx * seqlen + epoch = it // 2 + gQ = cute.local_tile( + mQ, + cute.select(self.gemm1_tiler, mode=[0, 2]), + (None, None, seq_idx, batch_idx), + ) + # Q's last reader is the MMA warp. Its DQ_DONE commit is + # ordered after every Q-consuming GEMM for row it-2, so Q can + # be refilled before the slower compute/K-load/reduce roles + # collectively release the rest of this operand parity. + if it >= 2: + cute.arch.mbarrier_wait( + mbar + MBAR_DQ_DONE, + Int32((it - 2) & 1), + ) + if parity == 0: + self._issue_cross_row_q( + tmma1, + gQ, + sQ_p0, + tma_atom_Q, + Q_mbar_ptr, + ) + else: + self._issue_cross_row_q( + tmma1, + gQ, + sQ_p1, + tma_atom_Q, + Q_mbar_ptr + 1, + ) + # The remaining row operands share one parity buffer and are + # overwritten only after all three consumer warpgroups retire. + if it >= 2: + cute.arch.mbarrier_wait( + mbar + MBAR_ROW_FREE_0 + parity, + Int32((epoch - 1) & 1), + ) + if const_expr(self.enable_score_pdl) and it == 0: + # grad_signal and the FP32 dK zeroing are produced by the + # preceding score grid. Q/K/top-k setup above is + # independent; wait immediately before the first + # dependent operand load. GW_READY -> dS -> dK barriers + # transitively keep the reduce role behind this wait. + cute.arch.griddepcontrol_wait() + sGrad_row = cute.make_tensor( + sGrad_storage_ptr + parity * self.topk, + cute.make_layout((self.topk,), stride=(1,)), + ) + sTopk_row = cute.make_tensor( + sTopk_storage_ptr + parity * self.topk, + cute.make_layout((self.topk,), stride=(1,)), + ) + sW_row = cute.make_tensor( + sW_storage_ptr + parity * self.heads, + cute.make_layout((self.heads,), stride=(1,)), + ) + seqlen_k_per_batch = seqlen_k // batch_size + batch_offset_l2g = Int32(0) if const_expr(self.topk_indices_global) else batch_idx * seqlen_k_per_batch + for step in cutlass.range_constexpr( + (self.topk + self.WARP_SIZE - 1) // self.WARP_SIZE, + ): + pos = step * self.WARP_SIZE + lane_id + if pos < self.topk: + raw_id = Int32(mTopkIdx[seq_idx, pos, batch_idx]) + if const_expr(self.topk_indices_global): + sTopk_row[pos] = raw_id + else: + sTopk_row[pos] = raw_id + batch_offset_l2g if raw_id >= 0 and raw_id < seqlen_k_per_batch else Int32(-1) + sGrad_row[pos] = mGradSignal[ + seq_idx, + pos, + batch_idx, + ] + for step in cutlass.range_constexpr( + (self.heads + self.WARP_SIZE - 1) // self.WARP_SIZE, + ): + h = step * self.WARP_SIZE + lane_id + if h < self.heads: + sW_row[h] = mW[seq_idx, h, batch_idx] + cute.arch.fence_view_async_shared() + # Every load-warp lane publishes its own SMEM writes. + cute.arch.mbarrier_arrive( + mbar + MBAR_GW_READY_0 + parity, + ) + elif warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_wg0) + tmem.wait_for_alloc() + tmem_ptr_base = tmem.retrieve_ptr(self.acc_dtype) + tStS_0, tStS_1, tDqDq_0, tDkDk_0, tDkDk_1 = self.get_tmem_tensor( + s_acc_layout, + dq_acc_layout, + dk_acc_layout, + tmem_ptr_base, + ) + tDqDq_1 = cute.make_tensor( + tmem_ptr_base + self.tmem_dq_p1_offset, + dq_acc_layout, + ) + for row in cutlass.range( + first_flat_row, + self.total_rows, + row_stride, + ): + it = (row - first_flat_row) // row_stride + parity = it & 1 + epoch = it // 2 + if parity == 0: + sQ_g2b_row = cute.make_tensor( + cute.recast_ptr( + sQ_p0.iterator, + sQ_g2b_layout.inner, + ), + sQ_g2b_layout.outer, + ) + self._mma_warp( + sQ_p0, + sdS_g2a, + sK, + sKt, + sdS, + sQ_g2b_row, + tmma1, + tmma2, + tmma3, + tStS_0, + tStS_1, + tDqDq_0, + tDqDq_1, + tDkDk_0, + tDkDk_1, + Q_mbar_ptr, + mbar, + Int32(it & 1), + it, + Int32(epoch & 1), + mbar + MBAR_DQ_FREE_0, + Int32((epoch - 1) & 1), + ) + else: + sQ_g2b_row = cute.make_tensor( + cute.recast_ptr( + sQ_p1.iterator, + sQ_g2b_layout.inner, + ), + sQ_g2b_layout.outer, + ) + self._mma_warp( + sQ_p1, + sdS_g2a, + sK, + sKt, + sdS, + sQ_g2b_row, + tmma1, + tmma2, + tmma3, + tStS_0, + tStS_1, + tDqDq_0, + tDqDq_1, + tDkDk_0, + tDkDk_1, + Q_mbar_ptr + 1, + mbar, + Int32(it & 1), + it, + Int32(epoch & 1), + mbar + MBAR_DQ_FREE_1, + Int32((epoch - 1) & 1), + ) + elif warp_idx in self.compute_warp_id: + cute.arch.setmaxregister_increase(self.num_regs_compute) + if warp_idx == self.compute_warp_id[0]: + tmem.allocate(self.tmem_alloc_cols) + tmem.wait_for_alloc() + tmem_ptr_base = tmem.retrieve_ptr(self.acc_dtype) + tStS_0, tStS_1, tDqDq_0, _, _ = self.get_tmem_tensor( + s_acc_layout, + dq_acc_layout, + dk_acc_layout, + tmem_ptr_base, + ) + tDqDq_1 = cute.make_tensor( + tmem_ptr_base + self.tmem_dq_p1_offset, + dq_acc_layout, + ) + dQ_store_pipeline = pipeline.PipelineTmaStore.create( + num_stages=1, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.WARPGROUP_SIZE, + ), + ) + sdQ_epi_slice = sdQ_epi[None, None, 0] + for row in cutlass.range( + first_flat_row, + self.total_rows, + row_stride, + ): + it = (row - first_flat_row) // row_stride + parity = it & 1 + epoch = it // 2 + batch_idx = row // seqlen + seq_idx = row - batch_idx * seqlen + sGrad_row = cute.make_tensor( + sGrad_storage_ptr + parity * self.topk, + cute.make_layout((self.topk,), stride=(1,)), + ) + sW_row = cute.make_tensor( + sW_storage_ptr + parity * self.heads, + cute.make_layout((self.heads,), stride=(1,)), + ) + gdQ = cute.local_tile( + mdQ, + (self.heads_padded, self.head_dim_padded), + (0, 0, seq_idx, batch_idx), + ) + tdQsdQ, tdQgdQ = cpasync.tma_partition( + tma_atom_dQ, + 0, + cute.make_layout(1), + cute.group_modes(sdQ_epi_slice, 0, 2), + cute.group_modes(gdQ, 0, 2), + ) + self._compute_warpgroup( + mdW, + sGrad_row, + sW_row, + sdS_store, + sdS, + sdQ_epi_slice, + s_acc_shape, + dq_acc_shape, + tStS_0, + tStS_1, + tDqDq_0, + tDqDq_1, + tma_atom_dQ, + tdQsdQ, + tdQgdQ, + dQ_store_pipeline, + sm_scale, + seq_idx, + batch_idx, + tidx, + warp_idx, + mbar, + Int32(it & 1), + mbar + MBAR_GW_READY_0 + parity, + Int32(epoch & 1), + mbar + MBAR_DQ_FREE_0 + parity, + ) + cute.arch.mbarrier_arrive( + mbar + MBAR_ROW_FREE_0 + parity, + ) + + last_it = (self.total_rows - 1 - first_flat_row) // row_stride + last_parity = last_it & 1 + cute.arch.mbarrier_wait( + mbar + MBAR_ROW_FREE_0 + last_parity, + Int32((last_it // 2) & 1), + ) + if warp_idx == self.compute_warp_id[0]: + dQ_store_pipeline.producer_tail() + _tcgen05_fence_after_thread_sync() + cute.arch.dealloc_tmem(tmem_ptr_base, self.tmem_alloc_cols) + + elif warp_idx in self.k_load_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_kload) + for row in cutlass.range( + first_flat_row, + self.total_rows, + row_stride, + ): + it = (row - first_flat_row) // row_stride + parity = it & 1 + batch_idx = row // seqlen + seq_idx = row - batch_idx * seqlen + sTopk_row = cute.make_tensor( + sTopk_storage_ptr + parity * self.topk, + cute.make_layout((self.topk,), stride=(1,)), + ) + self._k_load_warpgroup( + mK, + sK, + sK_raw_ptr, + sTopk_row, + mTopkIdx, + tma_atom_K_gather, + seq_idx, + batch_idx, + seqlen_k, + tidx, + mbar, + Int32(it), + ) + # A short-row gather can reach the row-release barrier before + # compute/reduce have completed the preceding row that used + # the same operand parity. Observe that older completion + # before publishing this row's gather arrival, otherwise the + # arrival can be counted in the still-open prior phase. + if const_expr(self.num_topk_blocks < 4) and it >= 2: + cute.arch.mbarrier_wait( + mbar + MBAR_ROW_FREE_0 + parity, + Int32((it // 2 - 1) & 1), + ) + cute.arch.mbarrier_arrive( + mbar + MBAR_ROW_FREE_0 + parity, + ) + + elif warp_idx in self.reduce_warp_id: + cute.arch.setmaxregister_increase(self.num_regs_reduce) + tmem.wait_for_alloc() + tmem_ptr_base = tmem.retrieve_ptr(self.acc_dtype) + _, _, _, tDkDk_0, tDkDk_1 = self.get_tmem_tensor( + s_acc_layout, + dq_acc_layout, + dk_acc_layout, + tmem_ptr_base, + ) + for row in cutlass.range( + first_flat_row, + self.total_rows, + row_stride, + ): + it = (row - first_flat_row) // row_stride + parity = it & 1 + batch_idx = row // seqlen + seq_idx = row - batch_idx * seqlen + sTopk_row = cute.make_tensor( + sTopk_storage_ptr + parity * self.topk, + cute.make_layout((self.topk,), stride=(1,)), + ) + self._reduce_warpgroup( + mdK_f32, + sTopk_row, + sdKStage, + dk_acc_shape, + tDkDk_0, + tDkDk_1, + sm_scale, + seqlen_k, + tidx, + mbar, + Int32(it & 1), + ) + cute.arch.mbarrier_arrive( + mbar + MBAR_ROW_FREE_0 + parity, + ) + lane_id = tidx % self.WARP_SIZE + if lane_id < 8: + cute.arch.cp_async_bulk_wait_group(0) + + else: + cute.arch.setmaxregister_decrease(self.num_regs_wg0) + + @cute.jit + def _run_persistent_row_serial( + self, + mQ, + mW, + mK, + mdQ, + mdW, + mdK_f32, + mGradSignal, + mTopkIdx, + sm_scale, + tmma1, + tmma2, + tmma3, + sQ, + sdS_g2a, + sK, + sKt, + sdS, + sQ_g2b, + sdS_store, + sK_raw_ptr, + sGradSignal, + sTopkIdxs, + sW, + sdKStage, + sdQ_epi, + Q_mbar_ptr, + mbar, + tmem, + tma_atom_Q, + tma_atom_K_gather, + tma_atom_dQ, + s_acc_shape, + s_acc_layout, + dq_acc_shape, + dq_acc_layout, + dk_acc_shape, + dk_acc_layout, + seqlen_k, + batch_size, + seq_idx, + batch_idx, + tidx, + warp_idx, + ): + """Run one fully drained row in a persistent CTA.""" + Q_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=Q_mbar_ptr, + num_stages=1, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1), + tx_count=self.tma_copy_Q_bytes, + cta_layout_vmnk=cute.make_layout(self.cluster_shape), + defer_sync=self.use_persistent, + ) + Q_producer, Q_consumer = Q_pipeline.make_participants() + gQ = cute.local_tile( + mQ, + cute.select(self.gemm1_tiler, mode=[0, 2]), + (None, None, seq_idx, batch_idx), + ) + gemm1_thr_mma = tmma1.get_slice(0) + tAgQ = gemm1_thr_mma.partition_A(gQ) + tQsQ, tQgQ_mkl = cpasync.tma_partition( + tma_atom_Q, + 0, + cute.make_layout(1), + cute.group_modes(sQ, 0, 3), + cute.group_modes(tAgQ, 0, 3), + ) + + dQ_store_pipeline = pipeline.PipelineTmaStore.create( + num_stages=1, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.WARPGROUP_SIZE, + ), + ) + gdQ = cute.local_tile( + mdQ, + (self.heads_padded, self.head_dim_padded), + (0, 0, seq_idx, batch_idx), + ) + sdQ_epi_slice = sdQ_epi[None, None, 0] + tdQsdQ, tdQgdQ_mkl = cpasync.tma_partition( + tma_atom_dQ, + 0, + cute.make_layout(1), + cute.group_modes(sdQ_epi_slice, 0, 2), + cute.group_modes(gdQ, 0, 2), + ) + + if tidx == 0: + cute.arch.mbarrier_init(mbar + MBAR_S_FULL_0, 1) + cute.arch.mbarrier_init(mbar + MBAR_S_FULL_1, 1) + cute.arch.mbarrier_init(mbar + MBAR_DS_READY_0, self.WARPGROUP_SIZE) + cute.arch.mbarrier_init(mbar + MBAR_DS_READY_1, self.WARPGROUP_SIZE) + cute.arch.mbarrier_init(mbar + MBAR_DK_FULL_0, 1) + cute.arch.mbarrier_init(mbar + MBAR_DK_FULL_1, 1) + cute.arch.mbarrier_init(mbar + MBAR_DK_EMPTY_0, self.WARPGROUP_SIZE) + cute.arch.mbarrier_init(mbar + MBAR_DK_EMPTY_1, self.WARPGROUP_SIZE) + if const_expr(self.use_tma_gather): + cute.arch.mbarrier_init(mbar + MBAR_K_LOADED_0, 1) + cute.arch.mbarrier_init(mbar + MBAR_K_LOADED_1, 1) + cute.arch.mbarrier_init(mbar + MBAR_K_LOADED_2, 1) + else: + cute.arch.mbarrier_init(mbar + MBAR_K_LOADED_0, self.WARPGROUP_SIZE) + cute.arch.mbarrier_init(mbar + MBAR_K_LOADED_1, self.WARPGROUP_SIZE) + cute.arch.mbarrier_init(mbar + MBAR_K_LOADED_2, self.WARPGROUP_SIZE) + cute.arch.mbarrier_init(mbar + MBAR_K_CONSUMED_0, 1) + cute.arch.mbarrier_init(mbar + MBAR_K_CONSUMED_1, 1) + cute.arch.mbarrier_init(mbar + MBAR_K_CONSUMED_2, 1) + cute.arch.mbarrier_init(mbar + MBAR_W_LOADED, self.WARP_SIZE) + cute.arch.mbarrier_init(mbar + MBAR_DQ_DONE, 1) + cute.arch.mbarrier_init( + mbar + MBAR_REDUCE_DONE, + self.WARPGROUP_SIZE, + ) + cute.arch.mbarrier_init( + mbar + MBAR_DS_HALF_0, + self.WARPGROUP_SIZE, + ) + cute.arch.mbarrier_init( + mbar + MBAR_DS_HALF_1, + self.WARPGROUP_SIZE, + ) + if const_expr(self.use_persistent): + cute.arch.mbarrier_init_fence() cute.arch.sync_threads() - # Pre-load topk indices into SMEM cooperatively (all 512 threads). - # Load up to smem_topk_capacity; reads beyond that fall back to global memory. - # K/dK are flattened to (B*S_k, D) above, so consumers index by global - # flat KV ids. ``topk_indices_global=True`` (default): ``mTopkIdx`` already - # carries global ids (``b * seqlen_k + local``); load directly. - # ``topk_indices_global=False``: ids are local-per-batch; add - # ``batch_idx * S_k_per_batch`` to convert. Invalid (-1) entries stay - # negative (skipped in the local→global add) and are rejected by the - # ``>= 0`` bounds check at consumers. - batch_offset_l2g = Int32(0) if const_expr(self.topk_indices_global) else batch_idx * (seqlen_k // batch_size) - _load_bound = const_expr(min(self.topk, smem_topk_capacity)) - TOPK_PER_THREAD = const_expr((_load_bound + self.THREADS_PER_CTA - 1) // self.THREADS_PER_CTA) - for ii in cutlass.range_constexpr(TOPK_PER_THREAD): + seqlen_k_per_batch = seqlen_k // batch_size + batch_offset_l2g = Int32(0) if const_expr(self.topk_indices_global) else batch_idx * seqlen_k_per_batch + topk_per_thread = const_expr((self.topk + self.THREADS_PER_CTA - 1) // self.THREADS_PER_CTA) + for ii in cutlass.range_constexpr(topk_per_thread): pos = ii * self.THREADS_PER_CTA + tidx - if pos < _load_bound: + if pos < self.topk: raw_id = Int32(mTopkIdx[seq_idx, pos, batch_idx]) if const_expr(self.topk_indices_global): sTopkIdxs[pos] = raw_id else: - sTopkIdxs[pos] = raw_id + batch_offset_l2g if raw_id >= 0 else raw_id + sTopkIdxs[pos] = raw_id + batch_offset_l2g if raw_id >= 0 and raw_id < seqlen_k_per_batch else Int32(-1) cute.arch.sync_threads() - # Pre-compute accumulator shapes/layouts from tmma before dispatch, - # so branches that don't run _mma_warp never touch the tmma objects - # (avoids MLIR SSA domination issues from tmma.set() inside _mma_warp). - s_acc_shape = tmma1.partition_shape_C(self.gemm1_tiler[:2]) - s_acc_layout = tmma1.make_fragment_C(s_acc_shape).layout - dq_acc_shape = tmma3.partition_shape_C(self.gemm3_tiler[:2]) - dq_acc_layout = tmma3.make_fragment_C(dq_acc_shape).layout - dk_acc_shape = tmma2.partition_shape_C(self.gemm2_tiler[:2]) - dk_acc_layout = tmma2.make_fragment_C(dk_acc_shape).layout - - # ============================================================= - # Warp dispatch — setmaxnreg rebalances registers across WGs. - # ============================================================= if warp_idx == self.load_warp_id: cute.arch.setmaxregister_decrease(self.num_regs_wg0) self._load_warp( @@ -600,7 +1932,6 @@ class SharedStorage: tidx, mbar, ) - elif warp_idx == self.mma_warp_id: cute.arch.setmaxregister_decrease(self.num_regs_wg0) tmem.wait_for_alloc() @@ -624,19 +1955,25 @@ class SharedStorage: tStS_0, tStS_1, tDqDq, + tDqDq, tDkDk_0, tDkDk_1, Q_consumer, mbar, + Int32(0), + Int32(0), + Int32(0), + mbar, + Int32(0), ) - elif warp_idx in self.compute_warp_id: cute.arch.setmaxregister_increase(self.num_regs_compute) if warp_idx == self.compute_warp_id[0]: - tmem.allocate(self.tmem_alloc_cols) + if const_expr(not self.use_persistent): + tmem.allocate(self.tmem_alloc_cols) tmem.wait_for_alloc() tmem_ptr_base = tmem.retrieve_ptr(self.acc_dtype) - tStS_0, tStS_1, tDqDq, tDkDk_0, tDkDk_1 = self.get_tmem_tensor( + tStS_0, tStS_1, tDqDq, _, _ = self.get_tmem_tensor( s_acc_layout, dq_acc_layout, dk_acc_layout, @@ -654,6 +1991,7 @@ class SharedStorage: tStS_0, tStS_1, tDqDq, + tDqDq, tma_atom_dQ, tdQsdQ, tdQgdQ_mkl, @@ -664,30 +2002,38 @@ class SharedStorage: tidx, warp_idx, mbar, + Int32(0), + mbar + MBAR_W_LOADED, + Int32(0), + mbar, ) if warp_idx == self.compute_warp_id[0]: - cute.arch.dealloc_tmem(tmem_ptr_base, self.tmem_alloc_cols) - + cute.arch.mbarrier_wait(mbar + MBAR_REDUCE_DONE, Int32(0)) + _tcgen05_fence_after_thread_sync() + if const_expr(not self.use_persistent): + cute.arch.dealloc_tmem(tmem_ptr_base, self.tmem_alloc_cols) + dQ_store_pipeline.producer_tail() elif warp_idx in self.k_load_warp_id: cute.arch.setmaxregister_decrease(self.num_regs_kload) self._k_load_warpgroup( mK, sK, + sK_raw_ptr, sTopkIdxs, mTopkIdx, + tma_atom_K_gather, seq_idx, batch_idx, seqlen_k, - batch_size, tidx, mbar, + Int32(0), ) - elif warp_idx in self.reduce_warp_id: cute.arch.setmaxregister_increase(self.num_regs_reduce) tmem.wait_for_alloc() tmem_ptr_base = tmem.retrieve_ptr(self.acc_dtype) - tStS_0, tStS_1, tDqDq, tDkDk_0, tDkDk_1 = self.get_tmem_tensor( + _, _, _, tDkDk_0, tDkDk_1 = self.get_tmem_tensor( s_acc_layout, dq_acc_layout, dk_acc_layout, @@ -696,19 +2042,16 @@ class SharedStorage: self._reduce_warpgroup( mdK_f32, sTopkIdxs, - mTopkIdx, + sdKStage, dk_acc_shape, tDkDk_0, tDkDk_1, sm_scale, - seq_idx, - batch_idx, seqlen_k, - batch_size, tidx, mbar, + Int32(0), ) - else: cute.arch.setmaxregister_decrease(self.num_regs_wg0) @@ -770,6 +2113,40 @@ def _load_warp( # ========================================================================= # Warp 1: MMA warp (3-stage sK pipeline, 2-stage TMEM S/dK) # ========================================================================= + @cute.jit + def _gemm_dq_parity( + self, + tmma3, + tDqDq, + tDqDq_p1, + frag_a, + frag_b, + row_iteration, + ): + """Issue GEMM3 through one of two statically aligned TMEM views.""" + if const_expr(self.use_cross_row_persistent): + if (row_iteration & 1) == 0: + cute.gemm(tmma3, tDqDq, frag_a, frag_b, tDqDq) + else: + cute.gemm(tmma3, tDqDq_p1, frag_a, frag_b, tDqDq_p1) + else: + cute.gemm(tmma3, tDqDq, frag_a, frag_b, tDqDq) + + @cute.jit + def _wait_rotating_k_loaded(self, mbar, row_iteration, block_index): + global_block = row_iteration * self.num_topk_blocks + block_index + stage = global_block % 3 + cute.arch.mbarrier_wait( + mbar + MBAR_K_LOADED_0 + stage, + Int32((global_block // 3) & 1), + ) + _tcgen05_fence_after_thread_sync() + + @cute.jit + def _commit_k_consumed(self, mbar, stage): + with cute.arch.elect_one(): + tcgen05.commit(mbar + MBAR_K_CONSUMED_0 + stage) + @cute.jit def _mma_warp( self, @@ -785,12 +2162,18 @@ def _mma_warp( tStS_0, tStS_1, tDqDq, + tDqDq_p1, tDkDk_0, tDkDk_1, - Q_consumer, + Q_consumer_or_barrier, mbar, + persistent_row_phase, + row_iteration, + q_ready_phase, + dq_free_barrier, + dq_free_phase, ): - """MMA warp: 3-stage sK pipeline (Opt-7), 2-stage TMEM S/dK. + """MMA warp: 3-stage sK pipeline, 2-stage TMEM S/dK. Structure: Prologue(Fill[0]) → Main(Fill[bi]+Drain[bi-1]) → Epilogue(Drain[last]) GEMM1(S) runs 1 block ahead, hiding Compute latency behind the next GEMM1. @@ -799,8 +2182,25 @@ def _mma_warp( TMEM S/dK accumulators remain 2-stage (bi%2). K_CONSUMED is per-sK-stage (3 barriers) so K-load can run 3 blocks ahead. """ - Q_consumer.reset() - Q_consumer.wait_and_advance() + # ``Atom.set`` replaces the Python wrapper's internal MLIR SSA value. + # Never mutate the kernel-argument wrappers here: this helper is + # emitted inside the runtime warp-role branch, so leaking a value + # defined in that branch into a later branch/loop violates SSA + # dominance. Role-local wrappers keep every atom_set_value and use in + # the same control-flow region. This is also what lets the continuous + # persistent path call this helper from a dynamic row loop safely. + tmma1 = tmma1.__new_from_mlir_values__(tmma1.__extract_mlir_values__()) + tmma2 = tmma2.__new_from_mlir_values__(tmma2.__extract_mlir_values__()) + tmma3 = tmma3.__new_from_mlir_values__(tmma3.__extract_mlir_values__()) + + if const_expr(self.use_cross_row_persistent): + cute.arch.mbarrier_wait( + Q_consumer_or_barrier, + Int32(q_ready_phase), + ) + else: + Q_consumer_or_barrier.reset() + Q_consumer_or_barrier.wait_and_advance() # --- A/B fragments from SMEM --- # sK/sKt: 3-stage (stage dim = last dim), sdS: 2-stage, sQ: 1-stage @@ -813,22 +2213,53 @@ def _mma_warp( dk_empty_0_phase = Int32(0) dk_empty_1_phase = Int32(0) - ds_ready_0_phase = Int32(0) - ds_ready_1_phase = Int32(0) + ds_ready_0_phase = Int32(persistent_row_phase if const_expr(((self.num_topk_blocks + 1) // 2) & 1) else 0) + ds_ready_1_phase = Int32(persistent_row_phase if const_expr((self.num_topk_blocks // 2) & 1) else 0) + ds_half_0_phase = Int32(0) + ds_half_1_phase = Int32(0) k_loaded_0_phase = Int32(0) - k_loaded_1_phase = Int32(0) - k_loaded_2_phase = Int32(0) + k_loaded_1_phase = Int32(persistent_row_phase) + k_loaded_2_phase = Int32(persistent_row_phase) is_first_dq = True + # The serial path reinitializes these barriers every row. The + # continuous path derives the reuse wait from the number of prior uses + # of each TMEM stage. This is the key distinction between one through + # four 128-wide blocks: odd block counts leave one stage toggled at the + # row boundary, while even counts may leave both stages aligned. + if const_expr(self.use_cross_row_persistent): + if row_iteration >= 1: + stage_0_uses = const_expr((self.num_topk_blocks + 1) // 2) + cute.arch.mbarrier_wait( + mbar + MBAR_DK_EMPTY_0, + Int32((row_iteration * stage_0_uses - 1) & 1), + ) + _tcgen05_fence_after_thread_sync() + # ============================================================= # Prologue: Fill block 0 (sK stage 0, TMEM stage 0) # ============================================================= - cute.arch.mbarrier_wait(mbar + MBAR_K_LOADED_0, k_loaded_0_phase) - k_loaded_0_phase ^= 1 + if const_expr(self.use_cross_row_persistent): + k_stage_0 = (row_iteration * self.num_topk_blocks) % 3 + self._wait_rotating_k_loaded(mbar, row_iteration, 0) + else: + k_stage_0 = const_expr(0) + cute.arch.mbarrier_wait( + mbar + MBAR_K_LOADED_0, + k_loaded_0_phase, + ) + k_loaded_0_phase ^= 1 + _tcgen05_fence_after_thread_sync() tmma1.set(tcgen05.Field.ACCUMULATE, False) for k_block in cutlass.range(0, cute.size(tSrQ, mode=[2]), unroll=4): - cute.gemm(tmma1, tStS_0, tSrQ[None, None, k_block, 0], tSrK[None, None, k_block, 0], tStS_0) + cute.gemm( + tmma1, + tStS_0, + tSrQ[None, None, k_block, 0], + tSrK[None, None, k_block, k_stage_0], + tStS_0, + ) tmma1.set(tcgen05.Field.ACCUMULATE, True) with cute.arch.elect_one(): tcgen05.commit(mbar + MBAR_S_FULL_0) @@ -840,10 +2271,26 @@ def _mma_warp( # ============================================================= for bi_offset in cutlass.range_constexpr(self.num_topk_blocks - 1): bi = bi_offset + 1 + if const_expr(self.use_cross_row_persistent): + row_block_base = row_iteration * self.num_topk_blocks + fill_k_stage = (row_block_base + bi) % 3 + drain_k_stage = (row_block_base + bi - 1) % 3 + else: + fill_k_stage = const_expr(bi % 3) + drain_k_stage = const_expr((bi - 1) % 3) # ------ Fill[bi]: GEMM1 for current block ------ # DK_EMPTY: wait for TMEM slot reuse (2-stage, bi%2) - if bi >= 2: + if const_expr(self.use_cross_row_persistent): + stage_uses_per_row = const_expr((self.num_topk_blocks + 1) // 2 if bi % 2 == 0 else self.num_topk_blocks // 2) + prior_uses = row_iteration * stage_uses_per_row + bi // 2 + if prior_uses > 0: + cute.arch.mbarrier_wait( + mbar + MBAR_DK_EMPTY_0 + bi % 2, + Int32((prior_uses - 1) & 1), + ) + _tcgen05_fence_after_thread_sync() + elif bi >= 2: if bi % 2 == 0: cute.arch.mbarrier_wait(mbar + MBAR_DK_EMPTY_0, dk_empty_0_phase) dk_empty_0_phase ^= 1 @@ -851,37 +2298,86 @@ def _mma_warp( cute.arch.mbarrier_wait(mbar + MBAR_DK_EMPTY_1, dk_empty_1_phase) dk_empty_1_phase ^= 1 - # K_LOADED: wait for sK data (3-stage, bi%3) - if bi % 3 == 0: - cute.arch.mbarrier_wait(mbar + MBAR_K_LOADED_0, k_loaded_0_phase) - k_loaded_0_phase ^= 1 - elif bi % 3 == 1: - cute.arch.mbarrier_wait(mbar + MBAR_K_LOADED_1, k_loaded_1_phase) - k_loaded_1_phase ^= 1 + # K_LOADED rotates continuously across persistent rows. + if const_expr(self.use_cross_row_persistent): + self._wait_rotating_k_loaded(mbar, row_iteration, bi) else: - cute.arch.mbarrier_wait(mbar + MBAR_K_LOADED_2, k_loaded_2_phase) - k_loaded_2_phase ^= 1 + if bi % 3 == 0: + cute.arch.mbarrier_wait( + mbar + MBAR_K_LOADED_0, + k_loaded_0_phase, + ) + k_loaded_0_phase ^= 1 + elif bi % 3 == 1: + cute.arch.mbarrier_wait( + mbar + MBAR_K_LOADED_1, + k_loaded_1_phase, + ) + k_loaded_1_phase ^= 1 + else: + cute.arch.mbarrier_wait( + mbar + MBAR_K_LOADED_2, + k_loaded_2_phase, + ) + k_loaded_2_phase ^= 1 + # Order the next tcgen05 MMA after both the K-loaded wait and, + # when reusing a TMEM stage, the earlier DK_EMPTY wait. + _tcgen05_fence_after_thread_sync() # GEMM1: tStS[bi%2] = Q @ sK[bi%3] tmma1.set(tcgen05.Field.ACCUMULATE, False) if bi % 2 == 0: for k_block in cutlass.range(0, cute.size(tSrQ, mode=[2]), unroll=4): - cute.gemm(tmma1, tStS_0, tSrQ[None, None, k_block, 0], tSrK[None, None, k_block, bi % 3], tStS_0) + cute.gemm(tmma1, tStS_0, tSrQ[None, None, k_block, 0], tSrK[None, None, k_block, fill_k_stage], tStS_0) tmma1.set(tcgen05.Field.ACCUMULATE, True) with cute.arch.elect_one(): tcgen05.commit(mbar + MBAR_S_FULL_0) else: for k_block in cutlass.range(0, cute.size(tSrQ, mode=[2]), unroll=4): - cute.gemm(tmma1, tStS_1, tSrQ[None, None, k_block, 0], tSrK[None, None, k_block, bi % 3], tStS_1) + cute.gemm(tmma1, tStS_1, tSrQ[None, None, k_block, 0], tSrK[None, None, k_block, fill_k_stage], tStS_1) tmma1.set(tcgen05.Field.ACCUMULATE, True) with cute.arch.elect_one(): tcgen05.commit(mbar + MBAR_S_FULL_1) # ------ Drain[bi-1]: GEMM2(dK) + GEMM3(dQ) for previous block ------ + if const_expr(self.use_cross_row_persistent and bi == 1): + # dQ has two TMEM parities. Release happens immediately after + # the compute warpgroup's T2R, before its SMEM/TMA epilogue. + if row_iteration >= 2: + cute.arch.mbarrier_wait( + dq_free_barrier, + Int32(dq_free_phase), + ) + _tcgen05_fence_after_thread_sync() if (bi - 1) % 2 == 0: # Prev TMEM stage 0 + if const_expr(self.use_ds_half): + cute.arch.mbarrier_wait( + mbar + MBAR_DS_HALF_0, + ds_half_0_phase, + ) + ds_half_0_phase ^= 1 + _tcgen05_fence_after_thread_sync() + tmma3.set(tcgen05.Field.ACCUMULATE, not is_first_dq) + is_first_dq = False + for k_block in cutlass.range( + 0, + cute.size(tDQrDS, mode=[2]) // 2, + unroll=4, + ): + self._gemm_dq_parity( + tmma3, + tDqDq, + tDqDq_p1, + tDQrDS[None, None, k_block, 0], + tDQrKt[None, None, k_block, drain_k_stage], + row_iteration, + ) + tmma3.set(tcgen05.Field.ACCUMULATE, True) + cute.arch.mbarrier_wait(mbar + MBAR_DS_READY_0, ds_ready_0_phase) ds_ready_0_phase ^= 1 + _tcgen05_fence_after_thread_sync() tmma2.set(tcgen05.Field.ACCUMULATE, False) for k_block in cutlass.range(0, cute.size(tDKrA_g2, mode=[2]), unroll=4): @@ -890,25 +2386,55 @@ def _mma_warp( with cute.arch.elect_one(): tcgen05.commit(mbar + MBAR_DK_FULL_0) - tmma3.set(tcgen05.Field.ACCUMULATE, not is_first_dq) - is_first_dq = False - for k_block in cutlass.range(0, cute.size(tDQrDS, mode=[2]), unroll=4): - cute.gemm(tmma3, tDqDq, tDQrDS[None, None, k_block, 0], tDQrKt[None, None, k_block, (bi - 1) % 3], tDqDq) + if const_expr(not self.use_ds_half): + tmma3.set(tcgen05.Field.ACCUMULATE, not is_first_dq) + is_first_dq = False + DQ_K_BEGIN = const_expr(cute.size(tDQrDS, mode=[2]) // 2 if self.use_ds_half else 0) + for k_block in cutlass.range( + DQ_K_BEGIN, + cute.size(tDQrDS, mode=[2]), + unroll=4, + ): + self._gemm_dq_parity( + tmma3, + tDqDq, + tDqDq_p1, + tDQrDS[None, None, k_block, 0], + tDQrKt[None, None, k_block, drain_k_stage], + row_iteration, + ) tmma3.set(tcgen05.Field.ACCUMULATE, True) - # K_CONSUMED: release sK stage (bi-1)%3 - if (bi - 1) % 3 == 0: - with cute.arch.elect_one(): - tcgen05.commit(mbar + MBAR_K_CONSUMED_0) - elif (bi - 1) % 3 == 1: - with cute.arch.elect_one(): - tcgen05.commit(mbar + MBAR_K_CONSUMED_1) - else: - with cute.arch.elect_one(): - tcgen05.commit(mbar + MBAR_K_CONSUMED_2) + # Release the continuously rotating sparse-K stage. + self._commit_k_consumed(mbar, drain_k_stage) else: # Prev TMEM stage 1 + if const_expr(self.use_ds_half): + cute.arch.mbarrier_wait( + mbar + MBAR_DS_HALF_1, + ds_half_1_phase, + ) + ds_half_1_phase ^= 1 + _tcgen05_fence_after_thread_sync() + tmma3.set(tcgen05.Field.ACCUMULATE, not is_first_dq) + is_first_dq = False + for k_block in cutlass.range( + 0, + cute.size(tDQrDS, mode=[2]) // 2, + unroll=4, + ): + self._gemm_dq_parity( + tmma3, + tDqDq, + tDqDq_p1, + tDQrDS[None, None, k_block, 1], + tDQrKt[None, None, k_block, drain_k_stage], + row_iteration, + ) + tmma3.set(tcgen05.Field.ACCUMULATE, True) + cute.arch.mbarrier_wait(mbar + MBAR_DS_READY_1, ds_ready_1_phase) ds_ready_1_phase ^= 1 + _tcgen05_fence_after_thread_sync() tmma2.set(tcgen05.Field.ACCUMULATE, False) for k_block in cutlass.range(0, cute.size(tDKrA_g2, mode=[2]), unroll=4): @@ -917,72 +2443,145 @@ def _mma_warp( with cute.arch.elect_one(): tcgen05.commit(mbar + MBAR_DK_FULL_1) - tmma3.set(tcgen05.Field.ACCUMULATE, not is_first_dq) - is_first_dq = False - for k_block in cutlass.range(0, cute.size(tDQrDS, mode=[2]), unroll=4): - cute.gemm(tmma3, tDqDq, tDQrDS[None, None, k_block, 1], tDQrKt[None, None, k_block, (bi - 1) % 3], tDqDq) + if const_expr(not self.use_ds_half): + tmma3.set(tcgen05.Field.ACCUMULATE, not is_first_dq) + is_first_dq = False + DQ_K_BEGIN = const_expr(cute.size(tDQrDS, mode=[2]) // 2 if self.use_ds_half else 0) + for k_block in cutlass.range( + DQ_K_BEGIN, + cute.size(tDQrDS, mode=[2]), + unroll=4, + ): + self._gemm_dq_parity( + tmma3, + tDqDq, + tDqDq_p1, + tDQrDS[None, None, k_block, 1], + tDQrKt[None, None, k_block, drain_k_stage], + row_iteration, + ) tmma3.set(tcgen05.Field.ACCUMULATE, True) - # K_CONSUMED: release sK stage (bi-1)%3 - if (bi - 1) % 3 == 0: - with cute.arch.elect_one(): - tcgen05.commit(mbar + MBAR_K_CONSUMED_0) - elif (bi - 1) % 3 == 1: - with cute.arch.elect_one(): - tcgen05.commit(mbar + MBAR_K_CONSUMED_1) - else: - with cute.arch.elect_one(): - tcgen05.commit(mbar + MBAR_K_CONSUMED_2) + # Release the continuously rotating sparse-K stage. + self._commit_k_consumed(mbar, drain_k_stage) # ============================================================= # Epilogue: Drain last block # TMEM stage: (num_topk_blocks-1)%2, sK stage: (num_topk_blocks-1)%3 # ============================================================= LAST_TMEM_STAGE = const_expr((self.num_topk_blocks - 1) % 2) - LAST_SK_STAGE = const_expr((self.num_topk_blocks - 1) % 3) + if const_expr(self.use_cross_row_persistent): + last_sk_stage = (row_iteration * self.num_topk_blocks + self.num_topk_blocks - 1) % 3 + else: + last_sk_stage = const_expr((self.num_topk_blocks - 1) % 3) + if const_expr(self.use_cross_row_persistent and self.num_topk_blocks == 1): + # With one block there is no main-loop Drain[0], so acquire the dQ + # accumulator parity immediately before the epilogue drain. + if row_iteration >= 2: + cute.arch.mbarrier_wait( + dq_free_barrier, + Int32(dq_free_phase), + ) + _tcgen05_fence_after_thread_sync() if LAST_TMEM_STAGE == 0: + if const_expr(self.use_ds_half): + cute.arch.mbarrier_wait( + mbar + MBAR_DS_HALF_0, + ds_half_0_phase, + ) + _tcgen05_fence_after_thread_sync() + tmma3.set(tcgen05.Field.ACCUMULATE, not is_first_dq) + is_first_dq = False + for k_block in cutlass.range( + 0, + cute.size(tDQrDS, mode=[2]) // 2, + unroll=4, + ): + self._gemm_dq_parity( + tmma3, + tDqDq, + tDqDq_p1, + tDQrDS[None, None, k_block, 0], + tDQrKt[None, None, k_block, last_sk_stage], + row_iteration, + ) + tmma3.set(tcgen05.Field.ACCUMULATE, True) cute.arch.mbarrier_wait(mbar + MBAR_DS_READY_0, ds_ready_0_phase) + _tcgen05_fence_after_thread_sync() tmma2.set(tcgen05.Field.ACCUMULATE, False) for k_block in cutlass.range(0, cute.size(tDKrA_g2, mode=[2]), unroll=4): cute.gemm(tmma2, tDkDk_0, tDKrA_g2[None, None, k_block, 0], tDKrB_g2[None, None, k_block, 0], tDkDk_0) tmma2.set(tcgen05.Field.ACCUMULATE, True) with cute.arch.elect_one(): tcgen05.commit(mbar + MBAR_DK_FULL_0) - tmma3.set(tcgen05.Field.ACCUMULATE, not is_first_dq) - for k_block in cutlass.range(0, cute.size(tDQrDS, mode=[2]), unroll=4): - cute.gemm(tmma3, tDqDq, tDQrDS[None, None, k_block, 0], tDQrKt[None, None, k_block, LAST_SK_STAGE], tDqDq) + if const_expr(not self.use_ds_half): + tmma3.set(tcgen05.Field.ACCUMULATE, not is_first_dq) + DQ_K_BEGIN = const_expr(cute.size(tDQrDS, mode=[2]) // 2 if self.use_ds_half else 0) + for k_block in cutlass.range( + DQ_K_BEGIN, + cute.size(tDQrDS, mode=[2]), + unroll=4, + ): + self._gemm_dq_parity( + tmma3, + tDqDq, + tDqDq_p1, + tDQrDS[None, None, k_block, 0], + tDQrKt[None, None, k_block, last_sk_stage], + row_iteration, + ) tmma3.set(tcgen05.Field.ACCUMULATE, True) - if LAST_SK_STAGE == 0: - with cute.arch.elect_one(): - tcgen05.commit(mbar + MBAR_K_CONSUMED_0) - elif LAST_SK_STAGE == 1: - with cute.arch.elect_one(): - tcgen05.commit(mbar + MBAR_K_CONSUMED_1) - else: - with cute.arch.elect_one(): - tcgen05.commit(mbar + MBAR_K_CONSUMED_2) + self._commit_k_consumed(mbar, last_sk_stage) with cute.arch.elect_one(): tcgen05.commit(mbar + MBAR_DQ_DONE) else: + if const_expr(self.use_ds_half): + cute.arch.mbarrier_wait( + mbar + MBAR_DS_HALF_1, + ds_half_1_phase, + ) + _tcgen05_fence_after_thread_sync() + tmma3.set(tcgen05.Field.ACCUMULATE, not is_first_dq) + is_first_dq = False + for k_block in cutlass.range( + 0, + cute.size(tDQrDS, mode=[2]) // 2, + unroll=4, + ): + self._gemm_dq_parity( + tmma3, + tDqDq, + tDqDq_p1, + tDQrDS[None, None, k_block, 1], + tDQrKt[None, None, k_block, last_sk_stage], + row_iteration, + ) + tmma3.set(tcgen05.Field.ACCUMULATE, True) cute.arch.mbarrier_wait(mbar + MBAR_DS_READY_1, ds_ready_1_phase) + _tcgen05_fence_after_thread_sync() tmma2.set(tcgen05.Field.ACCUMULATE, False) for k_block in cutlass.range(0, cute.size(tDKrA_g2, mode=[2]), unroll=4): cute.gemm(tmma2, tDkDk_1, tDKrA_g2[None, None, k_block, 1], tDKrB_g2[None, None, k_block, 0], tDkDk_1) tmma2.set(tcgen05.Field.ACCUMULATE, True) with cute.arch.elect_one(): tcgen05.commit(mbar + MBAR_DK_FULL_1) - tmma3.set(tcgen05.Field.ACCUMULATE, not is_first_dq) - for k_block in cutlass.range(0, cute.size(tDQrDS, mode=[2]), unroll=4): - cute.gemm(tmma3, tDqDq, tDQrDS[None, None, k_block, 1], tDQrKt[None, None, k_block, LAST_SK_STAGE], tDqDq) + if const_expr(not self.use_ds_half): + tmma3.set(tcgen05.Field.ACCUMULATE, not is_first_dq) + DQ_K_BEGIN = const_expr(cute.size(tDQrDS, mode=[2]) // 2 if self.use_ds_half else 0) + for k_block in cutlass.range( + DQ_K_BEGIN, + cute.size(tDQrDS, mode=[2]), + unroll=4, + ): + self._gemm_dq_parity( + tmma3, + tDqDq, + tDqDq_p1, + tDQrDS[None, None, k_block, 1], + tDQrKt[None, None, k_block, last_sk_stage], + row_iteration, + ) tmma3.set(tcgen05.Field.ACCUMULATE, True) - if LAST_SK_STAGE == 0: - with cute.arch.elect_one(): - tcgen05.commit(mbar + MBAR_K_CONSUMED_0) - elif LAST_SK_STAGE == 1: - with cute.arch.elect_one(): - tcgen05.commit(mbar + MBAR_K_CONSUMED_1) - else: - with cute.arch.elect_one(): - tcgen05.commit(mbar + MBAR_K_CONSUMED_2) + self._commit_k_consumed(mbar, last_sk_stage) with cute.arch.elect_one(): tcgen05.commit(mbar + MBAR_DQ_DONE) @@ -1003,6 +2602,7 @@ def _compute_warpgroup( tStS_0, tStS_1, tDqDq, + tDqDq_p1, tma_atom_dQ, tdQsdQ, tdQgdQ_mkl, @@ -1013,6 +2613,10 @@ def _compute_warpgroup( tidx, warp_idx, mbar, + persistent_row_phase, + gw_ready_barrier, + gw_ready_phase, + dq_free_barrier, ): """Compute/Epilogue warpgroup: TMEM readback S → register dS → stmatrix sdS, dQ/dW output. @@ -1034,44 +2638,141 @@ def _compute_warpgroup( Float32, ) - # --- TMEM readback (keep original partitioning for coordinate fidelity) --- - tiled_tmem_load_s_0 = tcgen05.make_tmem_copy(tmem_load_atom, tStS_0) - thr_tmem_load_s_0 = tiled_tmem_load_s_0.get_slice(wg_tidx) - tStS_t2r_0 = thr_tmem_load_s_0.partition_S(tStS_0) + # The STMatrix epilogue below uses the logical [H, I] ROW_MAJOR view + # required by GEMM3's K-major A operand. The same bytes are GEMM2's + # transposed [I, H] MN-major A operand. Other legal factory shapes + # retain the original coordinate store because this ownership mapping + # is specialized for the production H64 x I128 tile. + use_stmatrix_ds = const_expr( + self.heads_padded == 64 and self.block_I == 128, + ) + if const_expr(use_stmatrix_ds): + # --- TMEM readback (canonical 2D epilogue view) --- + tStS_epi_0 = tStS_0[((None, None), 0, 0)] + tiled_tmem_load_s_0 = tcgen05.make_tmem_copy( + tmem_load_atom, + tStS_epi_0, + ) + thr_tmem_load_s_0 = tiled_tmem_load_s_0.get_slice(wg_tidx) + tStS_t2r_0 = thr_tmem_load_s_0.partition_S(tStS_epi_0) + + tStS_epi_1 = tStS_1[((None, None), 0, 0)] + tiled_tmem_load_s_1 = tcgen05.make_tmem_copy( + tmem_load_atom, + tStS_epi_1, + ) + thr_tmem_load_s_1 = tiled_tmem_load_s_1.get_slice(wg_tidx) + tStS_t2r_1 = thr_tmem_load_s_1.partition_S(tStS_epi_1) + + # Derive R2S ownership from the TMEM-load ownership. This selects + # STSM and preserves each register element's logical coordinate. + smem_store_atom = sm100_utils_basic.get_smem_store_op( + LayoutEnum.ROW_MAJOR, + self.q_dtype, + self.acc_dtype, + tiled_tmem_load_s_0, + ) + tiled_smem_store = cute.make_tiled_copy_D( + smem_store_atom, + tiled_tmem_load_s_0, + ) + thr_smem_store = tiled_smem_store.get_slice(wg_tidx) + tRS_sdS = thr_smem_store.partition_D(sdS_store) - tiled_tmem_load_s_1 = tcgen05.make_tmem_copy(tmem_load_atom, tStS_1) - thr_tmem_load_s_1 = tiled_tmem_load_s_1.get_slice(wg_tidx) - tStS_t2r_1 = thr_tmem_load_s_1.partition_S(tStS_1) + cS = cute.make_identity_tensor( + (self.heads_padded, self.block_I), + ) + tCcS = thr_tmem_load_s_0.partition_D(cS) + else: + # General-shape fallback: preserve the MMA accumulator's complete + # coordinate hierarchy and write through its native sdS view. + tiled_tmem_load_s_0 = tcgen05.make_tmem_copy( + tmem_load_atom, + tStS_0, + ) + thr_tmem_load_s_0 = tiled_tmem_load_s_0.get_slice(wg_tidx) + tStS_t2r_0 = thr_tmem_load_s_0.partition_S(tStS_0) - # Logical GEMM views for direct dS writes (stage 0/1). - sdS_gemm_view_0 = cute.composition( - sdS[None, None, None, 0], - cute.make_layout((self.heads_padded, self.block_I)), - ) - sdS_gemm_view_1 = cute.composition( - sdS[None, None, None, 1], - cute.make_layout((self.heads_padded, self.block_I)), - ) + tiled_tmem_load_s_1 = tcgen05.make_tmem_copy( + tmem_load_atom, + tStS_1, + ) + thr_tmem_load_s_1 = tiled_tmem_load_s_1.get_slice(wg_tidx) + tStS_t2r_1 = thr_tmem_load_s_1.partition_S(tStS_1) - # Coordinate map matched to the same TMEM load partition used by tSrS. - cS = cute.make_identity_tensor(s_acc_shape) - tCcS = thr_tmem_load_s_0.partition_D(cS) + sdS_gemm_view_0 = cute.composition( + sdS[None, None, None, 0], + cute.make_layout((self.heads_padded, self.block_I)), + ) + sdS_gemm_view_1 = cute.composition( + sdS[None, None, None, 1], + cute.make_layout((self.heads_padded, self.block_I)), + ) + cS = cute.make_identity_tensor(s_acc_shape) + tCcS = thr_tmem_load_s_0.partition_D(cS) tSrS_shape = tCcS.shape - # --- TMEM readback (dQ — original 3-mode, NOT reduced) --- - tiled_tmem_load_dq = tcgen05.make_tmem_copy(tmem_load_atom, tDqDq) + # Use the canonical 2-D accumulator view when CuTe can lower dQ + # staging to STMatrix; other legal shapes use coordinate stores. + use_stmatrix_dq = const_expr(self.heads_padded == 64 and self.head_dim_padded == 128) + if const_expr(use_stmatrix_dq): + tDqDq_load_view = tDqDq[((None, None), 0, 0)] + if const_expr(self.use_cross_row_persistent): + tDqDq_p1_load_view = tDqDq_p1[((None, None), 0, 0)] + else: + tDqDq_load_view = tDqDq + if const_expr(self.use_cross_row_persistent): + tDqDq_p1_load_view = tDqDq_p1 + tiled_tmem_load_dq = tcgen05.make_tmem_copy( + tmem_load_atom, + tDqDq_load_view, + ) thr_tmem_load_dq = tiled_tmem_load_dq.get_slice(wg_tidx) - tDqDq_t2r = thr_tmem_load_dq.partition_S(tDqDq) - cDQ = cute.make_identity_tensor(dq_acc_shape) + tDqDq_t2r = thr_tmem_load_dq.partition_S(tDqDq_load_view) + if const_expr(self.use_cross_row_persistent): + tiled_tmem_load_dq_p1 = tcgen05.make_tmem_copy( + tmem_load_atom, + tDqDq_p1_load_view, + ) + thr_tmem_load_dq_p1 = tiled_tmem_load_dq_p1.get_slice(wg_tidx) + tDqDq_p1_t2r = thr_tmem_load_dq_p1.partition_S( + tDqDq_p1_load_view, + ) + if const_expr(use_stmatrix_dq): + cDQ = cute.make_identity_tensor( + (self.heads_padded, self.head_dim_padded), + ) + else: + cDQ = cute.make_identity_tensor(dq_acc_shape) tCcDQ = thr_tmem_load_dq.partition_D(cDQ) tDQrDQ_shape = tCcDQ.shape + if const_expr(use_stmatrix_dq): + smem_store_atom_dq = sm100_utils_basic.get_smem_store_op( + LayoutEnum.ROW_MAJOR, + self.q_dtype, + self.acc_dtype, + tiled_tmem_load_dq, + ) + tiled_smem_store_dq = cute.make_tiled_copy_D( + smem_store_atom_dq, + tiled_tmem_load_dq, + ) + thr_smem_store_dq = tiled_smem_store_dq.get_slice(wg_tidx) + tRDQ_sdQ = thr_smem_store_dq.partition_D(sdQ_epi_slice) # ---- Wait for W loaded by load warp ---- - cute.arch.mbarrier_wait(mbar + MBAR_W_LOADED, Int32(0)) + cute.arch.mbarrier_wait( + gw_ready_barrier, + Int32(gw_ready_phase), + ) # ---- Per topk-block iteration (2-stage S/dS) ---- - s_full_0_phase = Int32(0) - s_full_1_phase = Int32(0) + # Persistent rows do not necessarily leave both stage barriers at + # phase zero: a row has ceil(N/2) stage-0 blocks and floor(N/2) + # stage-1 blocks. Fold that per-row completion count into the first + # wait, then keep the existing in-row toggle sequence. + s_full_0_phase = Int32(persistent_row_phase if const_expr(((self.num_topk_blocks + 1) // 2) & 1) else 0) + s_full_1_phase = Int32(persistent_row_phase if const_expr((self.num_topk_blocks // 2) & 1) else 0) dw_accum = cute.make_rmem_tensor(tSrS_shape, Float32) for ei in cutlass.range_constexpr(cute.size(dw_accum)): @@ -1079,6 +2780,14 @@ def _compute_warpgroup( tSrS = cute.make_rmem_tensor(tSrS_shape, Float32) + # For the production 64x128 accumulator and 16dp256b8x ownership, + # every thread visits exactly two heads. Hoist those two conversions + # instead of reloading BF16 weights for every score pair. + if const_expr(use_stmatrix_ds): + h_base = warp_id_in_wg * 16 + lane_id // 4 + weight_lo = Float32(sW[h_base]) + weight_hi = Float32(sW[h_base + 8]) + for bi in cutlass.range(0, self.num_topk_blocks): i_st = bi * self.block_I @@ -1086,18 +2795,27 @@ def _compute_warpgroup( if bi % 2 == 0: cute.arch.mbarrier_wait(mbar + MBAR_S_FULL_0, s_full_0_phase) s_full_0_phase ^= 1 - cute.copy(tiled_tmem_load_s_0, tStS_t2r_0, tSrS) else: cute.arch.mbarrier_wait(mbar + MBAR_S_FULL_1, s_full_1_phase) s_full_1_phase ^= 1 + _tcgen05_fence_after_thread_sync() + if bi % 2 == 0: + cute.copy(tiled_tmem_load_s_0, tStS_t2r_0, tSrS) + else: cute.copy(tiled_tmem_load_s_1, tStS_t2r_1, tSrS) # Phase 1: Compute dS (→ tSrS), accumulate dW — paired f32x2. for ei in cutlass.range(0, cute.size(tSrS), 2): - h0 = cute.get(tCcS[ei], mode=[0, 0]) - n0 = cute.get(tCcS[ei], mode=[0, 1]) - h1 = cute.get(tCcS[ei + 1], mode=[0, 0]) - n1 = cute.get(tCcS[ei + 1], mode=[0, 1]) + if const_expr(use_stmatrix_ds): + h0 = cute.get(tCcS[ei], mode=[0]) + n0 = cute.get(tCcS[ei], mode=[1]) + h1 = cute.get(tCcS[ei + 1], mode=[0]) + n1 = cute.get(tCcS[ei + 1], mode=[1]) + else: + h0 = cute.get(tCcS[ei], mode=[0, 0]) + n0 = cute.get(tCcS[ei], mode=[0, 1]) + h1 = cute.get(tCcS[ei + 1], mode=[0, 0]) + n1 = cute.get(tCcS[ei + 1], mode=[0, 1]) tSrS[ei], tSrS[ei + 1] = mul_packed_f32x2( (tSrS[ei], tSrS[ei + 1]), @@ -1106,8 +2824,15 @@ def _compute_warpgroup( s0 = tSrS[ei] s1 = tSrS[ei + 1] - w0 = Float32(sW[h0]) - w1 = Float32(sW[h1]) + if const_expr(use_stmatrix_ds): + # Pairs alternate between the thread's low/high head; + # both elements in one pair share that head. + pair_weight = weight_lo if (ei // 2) % 2 == 0 else weight_hi + w0 = pair_weight + w1 = pair_weight + else: + w0 = Float32(sW[h0]) + w1 = Float32(sW[h1]) gs0 = sGradSignal[i_st + n0] gs1 = sGradSignal[i_st + n1] @@ -1122,28 +2847,86 @@ def _compute_warpgroup( (dw_accum[ei], dw_accum[ei + 1]), ) - tSrS[ei] = gs0 * w0 if s_pos_0 else Float32(0.0) - tSrS[ei + 1] = gs1 * w1 if s_pos_1 else Float32(0.0) + ds0, ds1 = mul_packed_f32x2( + (gs0, gs1), + (w0, w1), + ) + tSrS[ei] = ds0 if s_pos_0 else Float32(0.0) + tSrS[ei + 1] = ds1 if s_pos_1 else Float32(0.0) cute.arch.fence_view_async_tmem_load() - # Phase 2: Convert dS f32→bf16, write to sdS via coordinate mapping. + # Phase 2: Convert dS f32→bf16, then use STSM on the production + # tile or the native-layout coordinate fallback on other shapes. tSrS_f16 = cute.make_rmem_tensor(tSrS.shape, self.q_dtype) for ei in cutlass.range_constexpr(cute.size(tSrS)): tSrS_f16[ei] = self.q_dtype(tSrS[ei]) - if bi % 2 == 0: - for ei in cutlass.range_constexpr(cute.size(tSrS_f16)): - h = cute.get(tCcS[ei], mode=[0, 0]) - n = cute.get(tCcS[ei], mode=[0, 1]) - sdS_gemm_view_0[h, n] = tSrS_f16[ei] + if const_expr(use_stmatrix_ds): + tRS_rdS = tiled_smem_store.retile(tSrS_f16) + if const_expr(self.use_ds_half): + # The final retiled-copy mode is the two 64-column + # repetitions. Publish columns 0..63 first so MMA can + # issue the first half of dQ while this warpgroup stores + # columns 64..127; DS_READY still covers the full tile. + if bi % 2 == 0: + cute.copy( + tiled_smem_store, + tRS_rdS[(None, None, 0)], + tRS_sdS[(None, None, 0, 0)], + ) + else: + cute.copy( + tiled_smem_store, + tRS_rdS[(None, None, 0)], + tRS_sdS[(None, None, 0, 1)], + ) + cute.arch.fence_proxy("async.shared", space="cta") + _tcgen05_fence_before_thread_sync() + if bi % 2 == 0: + cute.arch.mbarrier_arrive(mbar + MBAR_DS_HALF_0) + else: + cute.arch.mbarrier_arrive(mbar + MBAR_DS_HALF_1) + + if bi % 2 == 0: + cute.copy( + tiled_smem_store, + tRS_rdS[(None, None, 1)], + tRS_sdS[(None, None, 1, 0)], + ) + else: + cute.copy( + tiled_smem_store, + tRS_rdS[(None, None, 1)], + tRS_sdS[(None, None, 1, 1)], + ) + else: + if bi % 2 == 0: + cute.copy( + tiled_smem_store, + tRS_rdS, + tRS_sdS[(None, None, None, 0)], + ) + else: + cute.copy( + tiled_smem_store, + tRS_rdS, + tRS_sdS[(None, None, None, 1)], + ) else: - for ei in cutlass.range_constexpr(cute.size(tSrS_f16)): - h = cute.get(tCcS[ei], mode=[0, 0]) - n = cute.get(tCcS[ei], mode=[0, 1]) - sdS_gemm_view_1[h, n] = tSrS_f16[ei] + if bi % 2 == 0: + for ei in cutlass.range_constexpr(cute.size(tSrS_f16)): + h = cute.get(tCcS[ei], mode=[0, 0]) + n = cute.get(tCcS[ei], mode=[0, 1]) + sdS_gemm_view_0[h, n] = tSrS_f16[ei] + else: + for ei in cutlass.range_constexpr(cute.size(tSrS_f16)): + h = cute.get(tCcS[ei], mode=[0, 0]) + n = cute.get(tCcS[ei], mode=[0, 1]) + sdS_gemm_view_1[h, n] = tSrS_f16[ei] cute.arch.fence_proxy("async.shared", space="cta") + _tcgen05_fence_before_thread_sync() if bi % 2 == 0: cute.arch.mbarrier_arrive(mbar + MBAR_DS_READY_0) @@ -1153,73 +2936,139 @@ def _compute_warpgroup( # ---- Step 3: After all iterations — dQ via TMA store, dW via warp reduction ---- # Wait for MMA warp to finish the final GEMM3 (dQ accumulation). - cute.arch.mbarrier_wait(mbar + MBAR_DQ_DONE, Int32(0)) + cute.arch.mbarrier_wait( + mbar + MBAR_DQ_DONE, + Int32(persistent_row_phase), + ) + _tcgen05_fence_after_thread_sync() tDQrDQ = cute.make_rmem_tensor(tDQrDQ_shape, Float32) - cute.copy(tiled_tmem_load_dq, tDqDq_t2r, tDQrDQ) + if const_expr(self.use_cross_row_persistent): + if (persistent_row_phase & 1) == 0: + cute.copy(tiled_tmem_load_dq, tDqDq_t2r, tDQrDQ) + else: + cute.copy( + tiled_tmem_load_dq_p1, + tDqDq_p1_t2r, + tDQrDQ, + ) + else: + cute.copy(tiled_tmem_load_dq, tDqDq_t2r, tDQrDQ) tDQrDQ_bf16 = cute.make_rmem_tensor(tDQrDQ.shape, self.q_dtype) - for ei in cutlass.range_constexpr(cute.size(tDQrDQ)): - tDQrDQ_bf16[ei] = self.q_dtype(tDQrDQ[ei] * Float32(sm_scale)) + for ei in cutlass.range(0, cute.size(tDQrDQ), 2): + scaled0, scaled1 = mul_packed_f32x2( + (tDQrDQ[ei], tDQrDQ[ei + 1]), + (Float32(sm_scale), Float32(sm_scale)), + ) + tDQrDQ_bf16[ei] = self.q_dtype(scaled0) + tDQrDQ_bf16[ei + 1] = self.q_dtype(scaled1) cute.arch.fence_view_async_tmem_load() - - # dQ staging via coordinate writes. - sdQ_gemm_view = cute.composition( - sdQ_epi_slice, - cute.make_layout((self.heads_padded, self.head_dim_padded)), - ) - for ei in cutlass.range_constexpr(cute.size(tDQrDQ_bf16)): - h = cute.get(tCcDQ[ei], mode=[0, 0]) - d = cute.get(tCcDQ[ei], mode=[0, 1]) - sdQ_gemm_view[h, d] = tDQrDQ_bf16[ei] + _tcgen05_fence_before_thread_sync() + + if const_expr(self.use_cross_row_persistent): + # MMA may start accumulating row it+2 into this TMEM parity as + # soon as every compute thread has completed its dQ T2R. + cute.arch.mbarrier_arrive(dq_free_barrier) + + # The dedicated dQ SMEM tile is single-buffered. Drain row it-1's + # TMA store before any thread overwrites it, then rendezvous the + # whole compute warpgroup. + if warp_idx == compute_warp0: + dQ_store_pipeline.producer_acquire() + self.compute_sync_barrier.arrive_and_wait() + + if const_expr(use_stmatrix_dq): + tRDQ_rdQ = tiled_smem_store_dq.retile(tDQrDQ_bf16) + cute.copy(tiled_smem_store_dq, tRDQ_rdQ, tRDQ_sdQ) + else: + # General-shape dQ staging via coordinate writes. + sdQ_gemm_view = cute.composition( + sdQ_epi_slice, + cute.make_layout((self.heads_padded, self.head_dim_padded)), + ) + for ei in cutlass.range_constexpr(cute.size(tDQrDQ_bf16)): + h = cute.get(tCcDQ[ei], mode=[0, 0]) + d = cute.get(tCcDQ[ei], mode=[0, 1]) + sdQ_gemm_view[h, d] = tDQrDQ_bf16[ei] self.compute_sync_barrier.arrive_and_wait() cute.arch.fence_proxy("async.shared", space="cta") self.compute_sync_barrier.arrive_and_wait() if warp_idx == compute_warp0: - dQ_store_pipeline.producer_acquire() + if const_expr(not self.use_cross_row_persistent): + dQ_store_pipeline.producer_acquire() cute.copy(tma_atom_dQ, tdQsdQ, tdQgdQ_mkl) dQ_store_pipeline.producer_commit() - HEADS_PER_WARP = const_expr(self.heads_padded // 4) - warp_base_h = warp_id_in_wg * Int32(HEADS_PER_WARP) - for h_local in cutlass.range_constexpr(HEADS_PER_WARP): - h = warp_base_h + h_local - my_partial = Float32(0.0) + if const_expr(use_stmatrix_ds): + # 16dp256b8x gives each thread two heads, selected by pair bit 0; + # the four lanes sharing lane//4 cover disjoint columns for those + # same heads. Reduce only that 4-lane subgroup instead of doing + # 16 full-warp reductions and repeatedly scanning all 64 values. + sum_low = Float32(0.0) + sum_high = Float32(0.0) for ei in cutlass.range_constexpr(cute.size(dw_accum)): - if cute.get(tCcS[ei], mode=[0, 0]) == h: - my_partial = my_partial + dw_accum[ei] - total = cute.arch.warp_reduction_sum(my_partial) - if lane_id == 0: - mdW[seq_idx, h, batch_idx] = self.q_dtype(total) + if (ei // 2) % 2 == 0: + sum_low = sum_low + dw_accum[ei] + else: + sum_high = sum_high + dw_accum[ei] + sum_low = cute.arch.warp_reduction_sum( + sum_low, + threads_in_group=4, + ) + sum_high = cute.arch.warp_reduction_sum( + sum_high, + threads_in_group=4, + ) + if lane_id % 4 == 0: + h0 = warp_id_in_wg * 16 + lane_id // 4 + mdW[seq_idx, h0, batch_idx] = self.q_dtype(sum_low) + mdW[seq_idx, h0 + 8, batch_idx] = self.q_dtype(sum_high) + else: + HEADS_PER_WARP = const_expr(self.heads_padded // 4) + warp_base_h = warp_id_in_wg * Int32(HEADS_PER_WARP) + for h_local in cutlass.range_constexpr(HEADS_PER_WARP): + h = warp_base_h + h_local + my_partial = Float32(0.0) + for ei in cutlass.range_constexpr(cute.size(dw_accum)): + if const_expr(use_stmatrix_ds): + elem_h = cute.get(tCcS[ei], mode=[0]) + else: + elem_h = cute.get(tCcS[ei], mode=[0, 0]) + if elem_h == h: + my_partial = my_partial + dw_accum[ei] + total = cute.arch.warp_reduction_sum(my_partial) + if lane_id == 0: + mdW[seq_idx, h, batch_idx] = self.q_dtype(total) # ========================================================================= - # Warps 12-15: Reduce warpgroup (dK T2R readback + 2-wide atomic_add, 2-stage) + # Warps 12-15: Reduce warpgroup (wide dK T2R + bulk FP32 reduce, 2-stage) # ========================================================================= @cute.jit def _reduce_warpgroup( self, mdK_f32, sTopkIdxs, - mTopkIdx, + sdKStage, dk_acc_shape, tDkDk_0, tDkDk_1, sm_scale: Float32 | float, - seq_idx, - batch_idx, seqlen_k, - batch_size, tidx, mbar, + persistent_row_phase, ): - """Reduce warpgroup: TMEM readback dK → 2-wide atomic_add to global f32 memory.""" + """Reduce dK through padded SMEM and bulk FP32 global additions.""" wg_tidx = tidx % self.WARPGROUP_SIZE + lane_id = wg_tidx % self.WARP_SIZE + warp_in_wg = wg_tidx // self.WARP_SIZE tmem_load_atom_dk = cute.make_copy_atom( - tcgen05.copy.Ld16x256bOp(tcgen05.copy.Repetition(8)), + tcgen05.copy.Ld16x256bOp(tcgen05.copy.Repetition(16)), Float32, ) @@ -1234,70 +3083,96 @@ def _reduce_warpgroup( thr_tmem_load_dk_1 = tiled_tmem_load_dk_1.get_slice(wg_tidx) tDkDk_t2r_1 = thr_tmem_load_dk_1.partition_S(tDkDk_1) - dk_full_0_phase = Int32(0) - dk_full_1_phase = Int32(0) + # A persistent row contributes ceil(N/2) completions to stage 0 and + # floor(N/2) to stage 1. Start each wait at the parity left by all + # preceding rows; the in-row toggles below then remain unchanged. + dk_full_0_phase = Int32(persistent_row_phase if const_expr(((self.num_topk_blocks + 1) // 2) & 1) else 0) + dk_full_1_phase = Int32(persistent_row_phase if const_expr((self.num_topk_blocks // 2) & 1) else 0) for bi in cutlass.range(0, self.num_topk_blocks): tDKrDK = cute.make_rmem_tensor(tDKrDK_shape, Float32) if bi % 2 == 0: cute.arch.mbarrier_wait(mbar + MBAR_DK_FULL_0, dk_full_0_phase) dk_full_0_phase ^= 1 + _tcgen05_fence_after_thread_sync() cute.copy(tiled_tmem_load_dk_0, tDkDk_t2r_0, tDKrDK) cute.arch.fence_view_async_tmem_load() + _tcgen05_fence_before_thread_sync() + if bi == self.num_topk_blocks - 1: + cute.arch.mbarrier_arrive(mbar + MBAR_REDUCE_DONE) cute.arch.mbarrier_arrive(mbar + MBAR_DK_EMPTY_0) else: cute.arch.mbarrier_wait(mbar + MBAR_DK_FULL_1, dk_full_1_phase) dk_full_1_phase ^= 1 + _tcgen05_fence_after_thread_sync() cute.copy(tiled_tmem_load_dk_1, tDkDk_t2r_1, tDKrDK) cute.arch.fence_view_async_tmem_load() + _tcgen05_fence_before_thread_sync() + if bi == self.num_topk_blocks - 1: + cute.arch.mbarrier_arrive(mbar + MBAR_REDUCE_DONE) cute.arch.mbarrier_arrive(mbar + MBAR_DK_EMPTY_1) - # dK reduction: 2-wide atomic_add per consecutive (n, d/d+1) pair. - # mdK_f32 is the flat (B*S_k, D) view; topk_idx (global) indexes - # directly. SMEM-cached ids are already global (preload converted - # local→global when topk_indices_global=False); gmem fallback - # mirrors that conversion via const_expr branch. - batch_offset_l2g = Int32(0) if const_expr(self.topk_indices_global) else batch_idx * (seqlen_k // batch_size) - for pair in cutlass.range_constexpr(cute.size(tDKrDK) // 2): - ei = pair * 2 - n = cute.get(tCcDK[ei], mode=[0, 0]) - d = cute.get(tCcDK[ei], mode=[0, 1]) - idx_pos = bi * self.block_I + n - topk_idx = Int32(0) - if idx_pos < cute.size(sTopkIdxs.layout): - topk_idx = Int32(sTopkIdxs[idx_pos]) - else: - raw_id = Int32(mTopkIdx[seq_idx, idx_pos, batch_idx]) - if const_expr(self.topk_indices_global): - topk_idx = raw_id - else: - topk_idx = raw_id + batch_offset_l2g if raw_id >= 0 else raw_id - if topk_idx >= 0 and topk_idx < seqlen_k: - dk_row = mdK_f32[topk_idx, None] - dk_pairs = cute.flat_divide(dk_row, (2,)) - rdK_pair = cute.make_rmem_tensor((2,), Float32) - rdK_pair[0] = tDKrDK[ei] * Float32(sm_scale) - rdK_pair[1] = tDKrDK[ei + 1] * Float32(sm_scale) - cute.arch.atomic_add( - dk_pairs[None, d // 2].iterator.llvm_ptr, - rdK_pair.load(), - ) + # The 16dp256b16x map materializes eight rows per warp into one + # padded SMEM ping-pong buffer; lanes 0..7 issue one 512-byte FP32 + # bulk reduction per row. + stage_row = lane_id // 4 + WARP_STAGE_FLOATS = const_expr(self.DK_STAGE_BUFFERS * self.DK_STAGE_ROWS * self.DK_STAGE_ROW_FLOATS) + BUFFER_FLOATS = const_expr(self.DK_STAGE_ROWS * self.DK_STAGE_ROW_FLOATS) + for pass_idx in cutlass.range_constexpr(4): + if lane_id < 8: + cute.arch.cp_async_bulk_wait_group(1, read=True) + cute.arch.sync_warp() + + buffer_base = warp_in_wg * WARP_STAGE_FLOATS + (pass_idx % 2) * BUFFER_FLOATS + for b6 in cutlass.range_constexpr(2): + for col_group in cutlass.range_constexpr(8): + ei = (pass_idx // 2) * 64 + b6 * 32 + col_group * 4 + (pass_idx % 2) * 2 + col = (lane_id % 4) * 2 + col_group * 8 + b6 * 64 + dst = buffer_base + stage_row * self.DK_STAGE_ROW_FLOATS + col + scaled0, scaled1 = mul_packed_f32x2( + (tDKrDK[ei], tDKrDK[ei + 1]), + (Float32(sm_scale), Float32(sm_scale)), + ) + sdKStage[dst] = scaled0 + sdKStage[dst + 1] = scaled1 + + cute.arch.fence_view_async_shared() + cute.arch.sync_warp() + if lane_id < 8: + n = warp_in_wg * 32 + (pass_idx // 2) * 16 + (pass_idx % 2) * 8 + lane_id + topk_idx = Int32(sTopkIdxs[bi * self.block_I + n]) + if topk_idx >= 0 and topk_idx < seqlen_k: + cpasync_reduce_bulk_add_f32( + sdKStage.iterator + buffer_base + lane_id * self.DK_STAGE_ROW_FLOATS, + mdK_f32[topk_idx, None].iterator, + self.head_dim_padded * 4, + ) + # Advance every issuing lane's group sequence even for an + # invalid sparse row, matching the bulk pipeline contract. + cute.arch.cp_async_bulk_commit_group() + + if const_expr(not self.use_cross_row_persistent) and lane_id < 8: + # Full completion (not just shared-read completion) before the + # following stream-ordered cast can consume dK_f32. + cute.arch.cp_async_bulk_wait_group(0) # ========================================================================= - # Warps 8-11: K loading warpgroup (3-stage sK, Opt-7) + # Warps 8-11: K loading warpgroup (3-stage sK) # ========================================================================= @cute.jit def _k_load_warpgroup( self, mK, sK, + sK_raw_ptr, sTopkIdxs, mTopkIdx, + tma_atom_K_gather, seq_idx, batch_idx, seqlen_k, - batch_size, tidx, mbar, + persistent_row_phase, ): """K loading warpgroup: sparse cp.async gather into 3-stage sK. @@ -1306,6 +3181,125 @@ def _k_load_warpgroup( """ wg_tidx = tidx % self.WARPGROUP_SIZE + if const_expr(self.use_tma_gather): + if const_expr(self.use_cross_row_persistent): + # Treat every row's K blocks as one flattened producer stream. + # Rotating by global block number avoids restarting stage 0 at + # row boundaries and lets gather remain three blocks ahead. + topk_row_ptr = mTopkIdx[seq_idx, None, batch_idx].iterator + batch_count = cute.size(mTopkIdx.shape[2]) + seqlen_k_per_batch = seqlen_k // batch_count + batch_offset_l2g = batch_idx * seqlen_k_per_batch + K_STAGE_BYTES = const_expr(self.block_I * self.head_dim_padded * self.k_dtype.width // 8) + warp_in_wg = wg_tidx // self.WARP_SIZE + lane_in_warp = wg_tidx % self.WARP_SIZE + for bi in cutlass.range_constexpr(self.num_topk_blocks): + global_block = persistent_row_phase * self.num_topk_blocks + bi + stage = global_block % 3 + stage_occurrence = global_block // 3 + if global_block >= 3: + cute.arch.mbarrier_wait( + mbar + MBAR_K_CONSUMED_0 + stage, + Int32((stage_occurrence - 1) & 1), + ) + + if wg_tidx == 0: + cute.arch.mbarrier_arrive_and_expect_tx( + mbar + MBAR_K_LOADED_0 + stage, + K_STAGE_BYTES, + ) + if lane_in_warp < 8: + n = warp_in_wg * 32 + lane_in_warp * 4 + idx_pos = bi * self.block_I + n + row0, row1, row2, row3 = _load_global_i32x4( + topk_row_ptr + idx_pos, + ) + if const_expr(self.topk_indices_global): + row0 = row0 if row0 >= 0 and row0 < seqlen_k else Int32(-1) + row1 = row1 if row1 >= 0 and row1 < seqlen_k else Int32(-1) + row2 = row2 if row2 >= 0 and row2 < seqlen_k else Int32(-1) + row3 = row3 if row3 >= 0 and row3 < seqlen_k else Int32(-1) + else: + row0 = row0 + batch_offset_l2g if row0 >= 0 and row0 < seqlen_k_per_batch else Int32(-1) + row1 = row1 + batch_offset_l2g if row1 >= 0 and row1 < seqlen_k_per_batch else Int32(-1) + row2 = row2 + batch_offset_l2g if row2 >= 0 and row2 < seqlen_k_per_batch else Int32(-1) + row3 = row3 + batch_offset_l2g if row3 >= 0 and row3 < seqlen_k_per_batch else Int32(-1) + for stripe in cutlass.range_constexpr(2): + dst_offset = stage * self.block_I * self.head_dim_padded + stripe * self.block_I * 64 + n * 64 + _tma_gather4_k_rows( + tma_atom_K_gather, + sK_raw_ptr + dst_offset, + stripe * 64, + row0, + row1, + row2, + row3, + mbar + MBAR_K_LOADED_0 + stage, + ) + return + + k_consumed_0_phase_tma = Int32(0) + k_consumed_1_phase_tma = Int32(persistent_row_phase & 1) + k_consumed_2_phase_tma = Int32(persistent_row_phase & 1) + K_STAGE_BYTES = const_expr(self.block_I * self.head_dim_padded * self.k_dtype.width // 8) + for bi in cutlass.range_constexpr(self.num_topk_blocks): + if bi >= 3: + if bi % 3 == 0: + cute.arch.mbarrier_wait( + mbar + MBAR_K_CONSUMED_0, + k_consumed_0_phase_tma, + ) + k_consumed_0_phase_tma ^= 1 + elif bi % 3 == 1: + cute.arch.mbarrier_wait( + mbar + MBAR_K_CONSUMED_1, + k_consumed_1_phase_tma, + ) + k_consumed_1_phase_tma ^= 1 + else: + cute.arch.mbarrier_wait( + mbar + MBAR_K_CONSUMED_2, + k_consumed_2_phase_tma, + ) + k_consumed_2_phase_tma ^= 1 + + stage = const_expr(bi % 3) + if wg_tidx == 0: + cute.arch.mbarrier_arrive_and_expect_tx( + mbar + MBAR_K_LOADED_0 + stage, + K_STAGE_BYTES, + ) + # One Gather4 atom is 4 rows x 64 bf16 columns. Lanes 0..7 of + # each warp issue two stripes, matching the hardware layout. + warp_in_wg = wg_tidx // self.WARP_SIZE + lane_in_warp = wg_tidx % self.WARP_SIZE + if lane_in_warp < 8: + n = warp_in_wg * 32 + lane_in_warp * 4 + idx_pos = bi * self.block_I + n + row0 = Int32(sTopkIdxs[idx_pos]) + row1 = Int32(sTopkIdxs[idx_pos + 1]) + row2 = Int32(sTopkIdxs[idx_pos + 2]) + row3 = Int32(sTopkIdxs[idx_pos + 3]) + row0 = row0 if row0 >= 0 and row0 < seqlen_k else Int32(-1) + row1 = row1 if row1 >= 0 and row1 < seqlen_k else Int32(-1) + row2 = row2 if row2 >= 0 and row2 < seqlen_k else Int32(-1) + row3 = row3 if row3 >= 0 and row3 < seqlen_k else Int32(-1) + for stripe in cutlass.range_constexpr(2): + # Raw (unswizzled) stage address. The descriptor's + # 128B swizzle maps this base to sK's UMMA layout. + dst_offset = stage * self.block_I * self.head_dim_padded + stripe * self.block_I * 64 + n * 64 + _tma_gather4_k_rows( + tma_atom_K_gather, + sK_raw_ptr + dst_offset, + stripe * 64, + row0, + row1, + row2, + row3, + mbar + MBAR_K_LOADED_0 + stage, + ) + return + async_copy_atom = cute.make_copy_atom( cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), self.k_dtype, @@ -1336,14 +3330,11 @@ def _k_load_warpgroup( sK[None, None, None, 2], cute.make_layout((self.block_I, self.head_dim_padded)), ) - # mK is the flat (B*S_k, D) view; topk_idx (global) indexes directly. - # gmem fallback mirrors the SMEM preload's local→global conversion when - # topk_indices_global=False; const_expr-folded to a no-op for default. - batch_offset_l2g = Int32(0) if const_expr(self.topk_indices_global) else batch_idx * (seqlen_k // batch_size) + # mK is the flat (B*S_k, D) view; SMEM already holds global IDs. k_consumed_0_phase_kload = Int32(0) - k_consumed_1_phase_kload = Int32(0) - k_consumed_2_phase_kload = Int32(0) + k_consumed_1_phase_kload = Int32(persistent_row_phase) + k_consumed_2_phase_kload = Int32(persistent_row_phase) for bi in cutlass.range_constexpr(self.num_topk_blocks): # Back-pressure: wait for MMA to finish using sK[bi%3]. @@ -1365,15 +3356,7 @@ def _k_load_warpgroup( for r in cutlass.range_constexpr(ROWS_PER_GROUP): row = r * NUM_GROUPS + group_idx_local idx_pos = bi * self.block_I + row - topk_idx = Int32(0) - if idx_pos < cute.size(sTopkIdxs.layout): - topk_idx = Int32(sTopkIdxs[idx_pos]) - else: - raw_id = Int32(mTopkIdx[seq_idx, idx_pos, batch_idx]) - if const_expr(self.topk_indices_global): - topk_idx = raw_id - else: - topk_idx = raw_id + batch_offset_l2g if raw_id >= 0 else raw_id + topk_idx = Int32(sTopkIdxs[idx_pos]) if topk_idx >= 0 and topk_idx < seqlen_k: gK_raw = mK[topk_idx, None] gK = cute.make_tensor( @@ -1439,16 +3422,10 @@ def indexer_backward_sm100( block_I=128, topk_indices_global: bool = True, ): - # ``batch``/``seqlen``/``seqlen_k`` are kept in the signature for a stable, - # backend-agnostic factory API (mirrors ``indexer_backward_sm90`` and the - # ``api.py`` call site) but are deliberately **not** forwarded into the - # compile cache key: in the kernel they are runtime values only (dynamic - # tensor extents + ``Int32`` grid/args, never ``const_expr``), so one - # compiled kernel serves every shape. Keying them would trigger spurious - # recompiles under varlen / changing batch. - # - # ``grad_scale`` is likewise runtime-only (forwarded into ``score_grad`` as - # a ``Float32`` at call time), so it is neither an argument here nor keyed. + # ``batch``/``seqlen``/``seqlen_k`` specialize the persistent-row schedule + # and Gather4 descriptor extent, so they are part of the GEMM compile key. + # ``grad_scale`` remains runtime-only (forwarded into ``score_grad`` as a + # ``Float32`` at call time), so changing loss scaling does not recompile. # # ``topk_indices_global`` selects the topk-id contract: # True (default): mTopkIdx carries global flat ids — load directly. @@ -1460,18 +3437,49 @@ def indexer_backward_sm100( # packed tensors as a single B=1 BSHD batch (sparse path's topk indices # already encode per-batch validity, so no kernel-side cu_seqlens are # needed). See ``_indexer_backward_sparse_thd`` in csrc/bwd/__init__.py. - return _build_cute_dsl_kernel(heads, dim, topk, sm_scale, block_I, topk_indices_global=topk_indices_global) + return _build_cute_dsl_kernel( + batch, + seqlen, + seqlen_k, + heads, + dim, + topk, + sm_scale, + block_I, + topk_indices_global=topk_indices_global, + ) class ScoreGradSm100: - """CuTe DSL kernel for in-place score_grad precompute.""" + """CuTe DSL kernel for in-place score_grad and optional dK zeroing.""" - THREADS_PER_CTA = 128 WARP_SIZE = 32 - NUM_WARPS = THREADS_PER_CTA // WARP_SIZE - def __init__(self, topk: int): + def __init__(self, topk: int, zero_dk_f32: bool = False): self.topk = topk + self.zero_dk_f32 = zero_dk_f32 + # For short TopK, each thread owns one aligned float4. TopK=128 maps + # one row to a warp and packs 16 rows per CTA; TopK=256 packs four + # two-warp rows; TopK=384/512 pack two three-/four-warp rows. These + # choices keep at least a few CTAs per SM on production grids while + # amortizing block scheduling. Larger rows retain the generic loop. + if topk == 128: + self.rows_per_cta = 16 + self.threads_per_cta = 512 + elif topk == 256: + self.rows_per_cta = 4 + self.threads_per_cta = 256 + elif topk == 384: + self.rows_per_cta = 2 + self.threads_per_cta = 192 + elif topk == 512: + self.rows_per_cta = 2 + self.threads_per_cta = 256 + else: + self.rows_per_cta = 1 + self.threads_per_cta = 256 + self.num_warps = self.threads_per_cta // self.WARP_SIZE + self.vectorized_short_row = topk in (128, 256, 384, 512) @cute.jit def __call__( @@ -1479,6 +3487,7 @@ def __call__( mAttnScore: cute.Tensor, mIndexScore: cute.Tensor, mGradLoss: cute.Tensor, + mDkF32: cute.Tensor, grad_scale: Float32 | float, stream: cuda.CUstream, ): @@ -1488,62 +3497,205 @@ def __call__( seqlen = cute.size(mAttnScore.shape[0]) batch_size = cute.size(mAttnScore.shape[2]) if cute.rank(mAttnScore.shape) > 2 else 1 - self.kernel_score_grad(mAttnScore, mIndexScore, mGradLoss, grad_scale).launch( - grid=(seqlen, batch_size, 1), - block=[self.THREADS_PER_CTA, 1, 1], + self.kernel_score_grad( + mAttnScore, + mIndexScore, + mGradLoss, + mDkF32, + grad_scale, + ).launch( + grid=(cute.ceil_div(seqlen, self.rows_per_cta), batch_size, 1), + block=[self.threads_per_cta, 1, 1], cluster=[1, 1, 1], stream=stream, min_blocks_per_mp=1, + use_pdl=True, ) @cute.kernel - def kernel_score_grad(self, mAttnScore, mIndexScore, mGradLoss, grad_scale: Float32 | float): + def kernel_score_grad( + self, + mAttnScore, + mIndexScore, + mGradLoss, + mDkF32, + grad_scale: Float32 | float, + ): + # Hint the next PDL-enabled GEMM launch at CTA entry. The dependent + # kernel waits before touching grad_signal/dK. This score grid can + # itself launch over the preceding call's dK cast, so wait before any + # global access to keep reuse of score/dK buffers race-free. + cute.arch.griddepcontrol_launch_dependents() + cute.arch.griddepcontrol_wait() tidx = cute.arch.thread_idx()[0] - seq_idx = cute.arch.block_idx()[0] batch_idx = cute.arch.block_idx()[1] # grad_scale is a compile/runtime scalar (loss_coeff / (b*sq)); # grad_loss lives in a shape-(1,) f32 GPU tensor (from autograd). # Fold them together once per CTA — the compiler will hoist. grad_scale_f32 = Float32(grad_scale) * Float32(mGradLoss[0]) - @cute.struct - class SharedStorage: - thread_sums: cute.struct.Align[cute.struct.MemRange[Float32, self.THREADS_PER_CTA], 128] - - smem = cutlass.utils.SmemAllocator() - storage = smem.allocate(SharedStorage) - thread_sums = storage.thread_sums.get_tensor(cute.make_layout((self.THREADS_PER_CTA,), stride=(1,))) - - local_sum = Float32(0.0) - for pos in cutlass.range(tidx, self.topk, self.THREADS_PER_CTA): - target = Float32(mAttnScore[seq_idx, pos, batch_idx]) - predict = Float32(mIndexScore[seq_idx, pos, batch_idx]) - target_eff = cute.arch.fmax(target, Float32(CLIP_PROB_MIN)) - log_clip_mask = Float32(1.0) if predict >= Float32(CLIP_PROB_MIN) else Float32(0.0) - local_sum += -target_eff * log_clip_mask * grad_scale_f32 - - thread_sums[tidx] = local_sum - cute.arch.sync_threads() - - if tidx == 0: - block_sum = Float32(0.0) - for i in cutlass.range(self.THREADS_PER_CTA, unroll_full=True): - block_sum += thread_sums[i] - thread_sums[0] = block_sum - cute.arch.sync_threads() - - sum_grad = thread_sums[0] - for pos in cutlass.range(tidx, self.topk, self.THREADS_PER_CTA): - target = Float32(mAttnScore[seq_idx, pos, batch_idx]) - predict = Float32(mIndexScore[seq_idx, pos, batch_idx]) - target_eff = cute.arch.fmax(target, Float32(CLIP_PROB_MIN)) - log_clip_mask = Float32(1.0) if predict >= Float32(CLIP_PROB_MIN) else Float32(0.0) - g_i = -target_eff * log_clip_mask * grad_scale_f32 - mAttnScore[seq_idx, pos, batch_idx] = g_i - predict * sum_grad - mIndexScore[seq_idx, pos, batch_idx] = sum_grad + if const_expr(self.topk == 128): + lane_idx = tidx % self.WARP_SIZE + row_in_cta = tidx // self.WARP_SIZE + seq_idx = cute.arch.block_idx()[0] * self.rows_per_cta + row_in_cta + seqlen = cute.size(mAttnScore.shape[0]) + if seq_idx < seqlen: + attn_ptr = mAttnScore[seq_idx, None, batch_idx].iterator + lane_idx * 4 + index_ptr = mIndexScore[seq_idx, None, batch_idx].iterator + lane_idx * 4 + target0, target1, target2, target3 = _load_global_f32x4(attn_ptr) + predict0, predict1, predict2, predict3 = _load_global_f32x4(index_ptr) + target0 = cute.arch.fmax(target0, Float32(CLIP_PROB_MIN)) + target1 = cute.arch.fmax(target1, Float32(CLIP_PROB_MIN)) + target2 = cute.arch.fmax(target2, Float32(CLIP_PROB_MIN)) + target3 = cute.arch.fmax(target3, Float32(CLIP_PROB_MIN)) + mask0 = Float32(1.0) if predict0 >= Float32(CLIP_PROB_MIN) else Float32(0.0) + mask1 = Float32(1.0) if predict1 >= Float32(CLIP_PROB_MIN) else Float32(0.0) + mask2 = Float32(1.0) if predict2 >= Float32(CLIP_PROB_MIN) else Float32(0.0) + mask3 = Float32(1.0) if predict3 >= Float32(CLIP_PROB_MIN) else Float32(0.0) + g0 = -target0 * mask0 * grad_scale_f32 + g1 = -target1 * mask1 * grad_scale_f32 + g2 = -target2 * mask2 * grad_scale_f32 + g3 = -target3 * mask3 * grad_scale_f32 + sum_grad = cute.arch.warp_reduction_sum(g0 + g1 + g2 + g3) + _store_global_f32x4( + attn_ptr, + g0 - predict0 * sum_grad, + g1 - predict1 * sum_grad, + g2 - predict2 * sum_grad, + g3 - predict3 * sum_grad, + ) + elif const_expr(self.vectorized_short_row): + row_threads = const_expr(self.topk // 4) + warps_per_row = const_expr(row_threads // self.WARP_SIZE) + row_in_cta = tidx // row_threads + vector_in_row = tidx % row_threads + seq_idx = cute.arch.block_idx()[0] * self.rows_per_cta + row_in_cta + seqlen = cute.size(mAttnScore.shape[0]) + row_is_valid = seq_idx < seqlen + + # Initialize the tail-row values so every thread can participate + # in the CTA-wide barrier when TopK=256 has one leftover row. + predict0 = Float32(0.0) + predict1 = Float32(0.0) + predict2 = Float32(0.0) + predict3 = Float32(0.0) + g0 = Float32(0.0) + g1 = Float32(0.0) + g2 = Float32(0.0) + g3 = Float32(0.0) + if row_is_valid: + attn_ptr = mAttnScore[seq_idx, None, batch_idx].iterator + vector_in_row * 4 + index_ptr = mIndexScore[seq_idx, None, batch_idx].iterator + vector_in_row * 4 + target0, target1, target2, target3 = _load_global_f32x4(attn_ptr) + predict0, predict1, predict2, predict3 = _load_global_f32x4(index_ptr) + target0 = cute.arch.fmax(target0, Float32(CLIP_PROB_MIN)) + target1 = cute.arch.fmax(target1, Float32(CLIP_PROB_MIN)) + target2 = cute.arch.fmax(target2, Float32(CLIP_PROB_MIN)) + target3 = cute.arch.fmax(target3, Float32(CLIP_PROB_MIN)) + mask0 = Float32(1.0) if predict0 >= Float32(CLIP_PROB_MIN) else Float32(0.0) + mask1 = Float32(1.0) if predict1 >= Float32(CLIP_PROB_MIN) else Float32(0.0) + mask2 = Float32(1.0) if predict2 >= Float32(CLIP_PROB_MIN) else Float32(0.0) + mask3 = Float32(1.0) if predict3 >= Float32(CLIP_PROB_MIN) else Float32(0.0) + g0 = -target0 * mask0 * grad_scale_f32 + g1 = -target1 * mask1 * grad_scale_f32 + g2 = -target2 * mask2 * grad_scale_f32 + g3 = -target3 * mask3 * grad_scale_f32 + + @cute.struct + class ShortRowStorage: + warp_sums: cute.struct.Align[cute.struct.MemRange[Float32, self.num_warps], 128] + + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(ShortRowStorage) + warp_sums = storage.warp_sums.get_tensor(cute.make_layout((self.num_warps,), stride=(1,))) + warp_sum = cute.arch.warp_reduction_sum(g0 + g1 + g2 + g3) + warp_idx = tidx // self.WARP_SIZE + with cute.arch.elect_one(): + warp_sums[warp_idx] = warp_sum + cute.arch.sync_threads() + + if row_is_valid: + sum_grad = Float32(0.0) + for row_warp in cutlass.range_constexpr(warps_per_row): + sum_grad += warp_sums[row_in_cta * warps_per_row + row_warp] + attn_ptr = mAttnScore[seq_idx, None, batch_idx].iterator + vector_in_row * 4 + _store_global_f32x4( + attn_ptr, + g0 - predict0 * sum_grad, + g1 - predict1 * sum_grad, + g2 - predict2 * sum_grad, + g3 - predict3 * sum_grad, + ) + else: + seq_idx = cute.arch.block_idx()[0] + + @cute.struct + class SharedStorage: + # One partial per warp. The former implementation staged all 128 + # thread partials and reduced them serially in thread 0. + warp_sums: cute.struct.Align[cute.struct.MemRange[Float32, self.num_warps], 128] + + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + warp_sums = storage.warp_sums.get_tensor(cute.make_layout((self.num_warps,), stride=(1,))) + + local_sum = Float32(0.0) + for pos in cutlass.range(tidx, self.topk, self.threads_per_cta): + target = Float32(mAttnScore[seq_idx, pos, batch_idx]) + predict = Float32(mIndexScore[seq_idx, pos, batch_idx]) + target_eff = cute.arch.fmax(target, Float32(CLIP_PROB_MIN)) + log_clip_mask = Float32(1.0) if predict >= Float32(CLIP_PROB_MIN) else Float32(0.0) + local_sum += -target_eff * log_clip_mask * grad_scale_f32 + + warp_idx = tidx // self.WARP_SIZE + warp_sum = cute.arch.warp_reduction_sum(local_sum) + with cute.arch.elect_one(): + warp_sums[warp_idx] = warp_sum + cute.arch.sync_threads() + + sum_grad = Float32(0.0) + for warp in cutlass.range_constexpr(self.num_warps): + sum_grad += warp_sums[warp] + for pos in cutlass.range(tidx, self.topk, self.threads_per_cta): + target = Float32(mAttnScore[seq_idx, pos, batch_idx]) + predict = Float32(mIndexScore[seq_idx, pos, batch_idx]) + target_eff = cute.arch.fmax(target, Float32(CLIP_PROB_MIN)) + log_clip_mask = Float32(1.0) if predict >= Float32(CLIP_PROB_MIN) else Float32(0.0) + g_i = -target_eff * log_clip_mask * grad_scale_f32 + mAttnScore[seq_idx, pos, batch_idx] = g_i - predict * sum_grad + + if const_expr(self.zero_dk_f32): + # Fold the caller-owned FP32 dK scratch clear into score_grad. Each + # lane writes float4 vectors after its row work; distributing the + # clear over the score grid removes a standalone memset launch and + # lets memory traffic from short/tail row CTAs overlap naturally. + grid_x = cute.arch.grid_dim()[0] + grid_y = cute.arch.grid_dim()[1] + flat_block = cute.arch.block_idx()[1] * grid_x + cute.arch.block_idx()[0] + num_blocks = grid_x * grid_y + vector_count = cute.size(mDkF32) // 4 + for vector_idx in cutlass.range( + flat_block * self.threads_per_cta + tidx, + vector_count, + num_blocks * self.threads_per_cta, + ): + _store_global_f32x4( + mDkF32.iterator + vector_idx * 4, + Float32(0.0), + Float32(0.0), + Float32(0.0), + Float32(0.0), + ) -def _score_grad_inplace_cute(AttnScore, IndexScore, GradLoss, grad_scale, current_stream=None): +def _score_grad_inplace_cute( + AttnScore, + IndexScore, + GradLoss, + grad_scale, + dIndexK_f32=None, + current_stream=None, +): from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor # Kernel reads ``mGradLoss[0]`` so it must be at least 1-D. ``to_cute_tensor`` @@ -1554,35 +3706,52 @@ def _score_grad_inplace_cute(AttnScore, IndexScore, GradLoss, grad_scale, curren GradLoss = GradLoss.reshape(1) _, _, topk = AttnScore.shape - compile_key = (topk,) + zero_dk_f32 = dIndexK_f32 is not None + compile_key = (topk, zero_dk_f32) s = _resolve_stream(current_stream) if compile_key not in _score_grad_cute_cache: - kernel_obj = ScoreGradSm100(topk=topk) + kernel_obj = ScoreGradSm100( + topk=topk, + zero_dk_f32=zero_dk_f32, + ) + # The unused dummy keeps the two specializations on one kernel + # signature; const_expr removes every access in score-only mode. + dk_arg = dIndexK_f32 if zero_dk_f32 else GradLoss _score_grad_cute_cache[compile_key] = cute.compile( kernel_obj, to_cute_tensor(AttnScore), to_cute_tensor(IndexScore), to_cute_tensor(GradLoss), + to_cute_tensor(dk_arg), cutlass.Float32(float(grad_scale)), s, - options=compile_options("--opt-level 3"), + options=compile_options("--opt-level 2"), ) + dk_arg = dIndexK_f32 if zero_dk_f32 else GradLoss _score_grad_cute_cache[compile_key]( AttnScore, IndexScore, GradLoss, + dk_arg, cutlass.Float32(float(grad_scale)), s, ) -def _score_grad_inplace(AttnScore, IndexScore, GradLoss, grad_scale, block_I=128, current_stream=None): +def _score_grad_inplace( + AttnScore, + IndexScore, + GradLoss, + grad_scale, + block_I=128, + dIndexK_f32=None, + current_stream=None, +): """Kernel 1: Compute clipped-log KL grad_signal from target/predict. - Results overwrite the two Score tensors in-place: - AttnScore ← grad_signal (per topk element) - IndexScore ← sum_grad (broadcast scalar per (batch, seqlen)) + Results overwrite AttnScore in-place with grad_signal. IndexScore remains + unchanged after being read as the predict input. grad_scale: Python float (loss_coeff / (b*sq)), passed as a runtime ``Float32`` arg — not in the kernel cache key. @@ -1605,43 +3774,90 @@ def _score_grad_inplace(AttnScore, IndexScore, GradLoss, grad_scale, block_I=128 ) if not can_use_cute: raise NotImplementedError("score_grad_inplace requires contiguous fp32 CUDA tensors with matching " "3D shapes; the torch fallback was removed") - _score_grad_inplace_cute(AttnScore, IndexScore, GradLoss, grad_scale, current_stream=current_stream) + _score_grad_inplace_cute( + AttnScore, + IndexScore, + GradLoss, + grad_scale, + dIndexK_f32=dIndexK_f32, + current_stream=current_stream, + ) -def _build_cute_dsl_kernel(heads, dim, topk, sm_scale, block_I, topk_indices_global: bool = True): +def _build_cute_dsl_kernel( + batch, + seqlen, + seqlen_k, + heads, + dim, + topk, + sm_scale, + block_I, + topk_indices_global: bool = True, +): from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor if torch.cuda.get_device_capability()[0] < 10: raise RuntimeError("Requires SM100+") - kernel_obj = IndexerBackwardSm100( - head_dim=dim, - heads=heads, - block_I=block_I, - topk=topk, - topk_indices_global=topk_indices_global, - ) + persistent_grid_size = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + def _make_kernel(enable_score_pdl: bool): + return IndexerBackwardSm100( + head_dim=dim, + heads=heads, + block_I=block_I, + topk=topk, + total_seqlen_k=batch * seqlen_k, + total_rows=batch * seqlen, + persistent_grid_size=persistent_grid_size, + topk_indices_global=topk_indices_global, + enable_score_pdl=enable_score_pdl, + ) - # Only params that change the generated code (mirrors forward _interface.py). - # sm_scale is a runtime Float32 arg (passed fresh per call), so it's not keyed. - compile_key = (heads, dim, topk, block_I, topk_indices_global) + kernel_objects = { + False: _make_kernel(False), + True: _make_kernel(True), + } + compile_key_base = ( + batch, + seqlen, + seqlen_k, + heads, + dim, + topk, + block_I, + topk_indices_global, + ) - def _ensure_compiled(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, AttnScore, TopkIndices, current_stream=None): - """Lazy-compile the GEMM kernel (kernel 2) on first execute (needs real tensors).""" + def _ensure_compiled(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, AttnScore, TopkIndices, enable_score_pdl: bool, current_stream=None): + """Lazy-compile the GEMM kernel (kernel 2).""" + compile_key = (*compile_key_base, enable_score_pdl) if compile_key not in _compile_cache: s = _resolve_stream(current_stream) cute_args = [to_cute_tensor(t) for t in [IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, AttnScore, TopkIndices]] _compile_cache[compile_key] = cute.compile( - kernel_obj, + kernel_objects[enable_score_pdl], *cute_args, cutlass.Float32(sm_scale), s, - options=compile_options("--opt-level 3"), + options=compile_options("--opt-level 2"), ) - def _run_gemm_only(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, GradSignal, TopkIndices, current_stream=None): - """Run only kernel 2 (GEMM). Caller must have run kernel 1 and zeroed dIndexK_f32.""" + def _run_gemm( + IndexQ, + Weights, + IndexK, + dIndexQ, + dWeights, + dIndexK_f32, + GradSignal, + TopkIndices, + enable_score_pdl: bool, + current_stream=None, + ): s = _resolve_stream(current_stream) - _ensure_compiled(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, GradSignal, TopkIndices, current_stream=current_stream) + compile_key = (*compile_key_base, enable_score_pdl) + _ensure_compiled(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, GradSignal, TopkIndices, enable_score_pdl, current_stream=current_stream) with torch.cuda.nvtx.range("indexer_backward_dsl_gemm"): _compile_cache[compile_key]( IndexQ, @@ -1656,35 +3872,106 @@ def _run_gemm_only(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, Grad s, ) + def _run_gemm_only(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, GradSignal, TopkIndices, current_stream=None): + """Run kernel 2 after arbitrary stream work (PDL overlap disabled).""" + _run_gemm( + IndexQ, + Weights, + IndexK, + dIndexQ, + dWeights, + dIndexK_f32, + GradSignal, + TopkIndices, + False, + current_stream=current_stream, + ) + + def _run_gemm_after_score( + IndexQ, + Weights, + IndexK, + dIndexQ, + dWeights, + dIndexK_f32, + GradSignal, + TopkIndices, + current_stream=None, + ): + """Run kernel 2 as ScoreGrad's PDL-dependent consumer.""" + _run_gemm( + IndexQ, + Weights, + IndexK, + dIndexQ, + dWeights, + dIndexK_f32, + GradSignal, + TopkIndices, + True, + current_stream=current_stream, + ) + def _run(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK, AttnScore, IndexScore, TopkIndices, GradLoss, grad_scale, current_stream=None): # ``grad_scale`` is a host scalar (Python float / 0-D fp32 tensor) # multiplied into ``score_grad`` as a runtime ``Float32`` arg — # changing it across calls does **not** trigger recompilation. score_grad = partial(_score_grad_inplace, block_I=block_I) - # Kernel 1: Compute grad_signal from scores (CuTe DSL only). - score_grad(AttnScore, IndexScore, GradLoss, grad_scale, current_stream=current_stream) - if dIndexK.dtype == torch.float32: - # fp32 output: the dK epilogue atomic-adds into this buffer, so it - # must start zeroed. Zero it internally on the selected stream - # (cheap; removes the fragile caller pre-zero contract) rather than - # trusting the caller. This zero-init is a promised stage of the - # execute pipeline (see the IndexerBackward docstring) and mirrors - # the SM90 backend and the DenseIndexerBackward fp32 paths, which - # zero their fp32 dK buffer the same way. - with _torch_stream_context(current_stream): - dIndexK.zero_() - _run_gemm_only(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK, AttnScore, TopkIndices, current_stream=current_stream) + # The dK epilogue atomic-adds into this caller-owned fp32 buffer, + # so clear it internally as promised by the public API. Fold the + # clear into ScoreGrad instead of launching a standalone memset, + # then let the PDL-dependent GEMM consume both grad_signal and dK. + score_grad( + AttnScore, + IndexScore, + GradLoss, + grad_scale, + dIndexK_f32=dIndexK, + current_stream=current_stream, + ) + _run_gemm_after_score( + IndexQ, + Weights, + IndexK, + dIndexQ, + dWeights, + dIndexK, + AttnScore, + TopkIndices, + current_stream=current_stream, + ) else: - # Need a separate f32 buffer for atomicAdd, then cast back to output dtype. + # Need a separate f32 buffer for atomicAdd. ScoreGrad clears it in + # the same launch, then kernel 2 accumulates and the epilogue casts. with _torch_stream_context(current_stream): - dIndexK_f32 = torch.zeros_like(dIndexK, dtype=torch.float32) - _run_gemm_only(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, AttnScore, TopkIndices, current_stream=current_stream) + dIndexK_f32 = torch.empty_like(dIndexK, dtype=torch.float32) + score_grad( + AttnScore, + IndexScore, + GradLoss, + grad_scale, + dIndexK_f32=dIndexK_f32, + current_stream=current_stream, + ) + _run_gemm_after_score( + IndexQ, + Weights, + IndexK, + dIndexQ, + dWeights, + dIndexK_f32, + AttnScore, + TopkIndices, + current_stream=current_stream, + ) with _torch_stream_context(current_stream): dIndexK.copy_(dIndexK_f32) _run.score_grad = partial(_score_grad_inplace, block_I=block_I) + _run.score_grad_zero = partial(_score_grad_inplace, block_I=block_I) _run.gemm_only = _run_gemm_only + _run.gemm_after_score = _run_gemm_after_score return _run diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_v2_sm100.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_v2_sm100.py index 13d3267e7..c3dd048a8 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_v2_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_v2_sm100.py @@ -1722,8 +1722,8 @@ def _run( # One plan serves one device: the per-plan ticket counter is # device-resident, so bind the plan to the device of its first - # execute and reject any other device BEFORE kernel 1 mutates the - # score buffers. api.py keys its plan cache on the device; this + # execute and reject any other device BEFORE kernel 1 overwrites + # AttnScore. api.py keys its plan cache on the device; this # check keeps direct users of the factory safe independently of # that cache. plan_device = plan_ws.get("device") @@ -1735,7 +1735,7 @@ def _run( ) # Kernel 1: shared in-place score-grad precompute. - # AttnScore <- grad_signal, IndexScore <- sum_grad + # AttnScore <- grad_signal; IndexScore is read-only and preserved. _score_grad_inplace(AttnScore, IndexScore, GradLoss, grad_scale, block_I=block_I, current_stream=current_stream) # true views (validated contiguous above — cannot copy) diff --git a/test/python/fe_api/dsa/test_DSA_indexer_backward.py b/test/python/fe_api/dsa/test_DSA_indexer_backward.py index 837b8cad8..2ba6130b7 100644 --- a/test/python/fe_api/dsa/test_DSA_indexer_backward.py +++ b/test/python/fe_api/dsa/test_DSA_indexer_backward.py @@ -1,8 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import importlib +import math + import pytest import torch +import torch.nn.functional as F import inspect import threading @@ -124,9 +128,9 @@ def test_DSA_indexer_backward_wrapper( torch_stream = torch.cuda.Stream() stream = cuda.CUstream(torch_stream.cuda_stream) - # The kernel mutates attn_score + index_score in-place during its - # score-grad stage. Keep pre-call copies so the reference can consume the - # same inputs the kernel was given. + # The score-grad stage overwrites attn_score but treats index_score as + # read-only. Keep pre-call copies for both reference input and the + # preservation check. attn_score_ref = attn_score.clone() index_score_ref = index_score.clone() torch_stream.wait_stream(torch.cuda.current_stream()) @@ -148,6 +152,8 @@ def test_DSA_indexer_backward_wrapper( pytest.skip(f"Unsupported testcase: {e}") torch_stream.synchronize() + assert torch.equal(index_score, index_score_ref) + d_index_q = result["d_index_q"] d_weights = result["d_weights"] d_index_k = result["d_index_k"] @@ -178,7 +184,7 @@ def test_DSA_indexer_backward_wrapper( # =========================================================================== # Regression coverage for the output/plan-signature validation on the default # indexer backward (SM100/SM90): -# * illegal output dtypes raise BEFORE kernel 1 mutates the score buffers; +# * illegal output dtypes raise BEFORE kernel 1 overwrites attn_score; # * wrong-rank / wrong-shape / non-contiguous plans are rejected up front; # * a direct plan or a cached wrapper plan reused with a mismatched # shape/stride signature raises (no fail-dirty); @@ -227,7 +233,7 @@ def _idxbwd_inputs(): @torch_fork_set_rng(seed=0) def test_DSA_indexer_backward_illegal_output_dtype_raises_before_mutation(): """Each illegal output dtype raises ValueError, and the raise happens - before kernel 1 mutates attn_score / index_score in place (no fail-dirty).""" + before kernel 1 overwrites attn_score (index_score is always read-only).""" _require_sm100() DSA, _ = _import_dsa() @@ -308,7 +314,7 @@ def _noncontiguous_like(sample): def test_DSA_indexer_backward_direct_plan_shape_stride_mismatch_raises(): """A directly-built plan rejects a runtime tensor whose shape or stride/layout differs from the descriptor it was compiled for, raising a - clean ValueError BEFORE kernel 1 mutates the score buffers (no fail-dirty).""" + clean ValueError BEFORE kernel 1 overwrites attn_score (no fail-dirty).""" _require_sm100() DSA, _ = _import_dsa() @@ -404,7 +410,7 @@ def test_DSA_indexer_backward_wrapper_cache_hit_stride_mismatch_raises(): def test_DSA_indexer_backward_wrong_shape_plan_rejected_before_kernel1(): """A first-call / directly-built plan whose output/score shape is inconsistent with ``index_q`` is rejected by the semantic shape validation - in ``check_support`` BEFORE kernel 1 mutates the score buffers — not + in ``check_support`` BEFORE kernel 1 overwrites attn_score — not silently run on a mismatched signature that only faults in the GEMM. Covers both the wrapper (cache miss) and a directly-built plan. Fail-hard.""" _require_sm100() @@ -483,7 +489,7 @@ def test_DSA_indexer_backward_noncontiguous_index_k_plan_rejected(): plan.check_support() # (b) Wrapper first call (cache miss) with a non-contiguous index_k is - # rejected before kernel 1 mutates the score buffers. + # rejected before kernel 1 overwrites attn_score. aq = attn.clone() isc = idx.clone() aq_pre = aq.clone() @@ -546,12 +552,11 @@ def _v2_call( d_index_k=None, stream=None, ): - """Run the wrapper with backend="sm100_v2" on CLONED score buffers. + """Run the wrapper with backend="sm100_v2" on cloned score inputs. - Returns (result, grad_signal, predict) -- the two in-place score buffers - after the call: grad_signal is the attn_score scratch (kernel 1's output) - and predict is the consumed index_score buffer. Both are shared - bit-for-bit with the default backend. + Returns (result, grad_signal, predict): grad_signal is the overwritten + attn_score scratch (kernel 1's output), while predict is the preserved, + read-only index_score input. Both match the default backend bit-for-bit. """ from cudnn import DSA from cuda.bindings import driver as cuda @@ -974,7 +979,7 @@ def test_DSA_indexer_backward_wrapper_v2_envelope_rejection( request, ): """check_support must reject out-of-envelope requests with ValueError - BEFORE kernel 1 mutates the score buffers: topk bounds (2176 above the + BEFORE kernel 1 overwrites attn_score: topk bounds (2176 above the smem cap, non-multiple-of-128), non-positive sm_scale, unsupported output dtypes, non-contiguous metadata, and an index_k whose dim 0 / dim 2 disagree with index_q. (topk 128/256 are now inside the envelope and are @@ -1176,11 +1181,10 @@ def test_DSA_indexer_backward_wrapper_v2_sm_scale( folds the scale into the grad signal -- an autograd-reference check alone historically could not catch a mis-applied scale here (its noise floor has since dropped to rms_rel <= 0.005, but the fp64 recompute stays - the authoritative scale check); (b) both in-place score - buffers are left bitwise identical to the default backend's, i.e. the - scratch holds exactly kernel 1's grad_signal and ``index_score`` is - consumed the same way (the scale folds inside kernel 2, no host-side - buffer mutation).""" + the authoritative scale check); (b) score handling is bitwise identical + to the default backend, i.e. the attn_score scratch holds exactly kernel + 1's grad_signal and ``index_score`` remains unchanged (the scale folds + inside kernel 2, with no host-side buffer mutation).""" try: from cudnn import DSA from cuda.bindings import driver as cuda @@ -1384,7 +1388,7 @@ def test_DSA_indexer_backward_wrapper_v2_multi_stream( streams = [torch.cuda.Stream(), torch.cuda.Stream()] cu_streams = [cuda.CUstream(s.cuda_stream) for s in streams] - # pre-clone the in-place score buffers for every iteration up front so + # pre-clone the score inputs for every iteration up front so # the interleaved phase enqueues only wrapper work clones = [[(inp[3].clone(), inp[4].clone()) for _ in range(n_iters)] for inp in inputs] results = [[None] * n_iters for _ in range(2)] @@ -1696,7 +1700,7 @@ def test_DSA_indexer_backward_wrapper_v2_multi_device( Run the identical input bits on each device, then interleave the devices, checking dq/dw bitwise against each device's serial result and dk in the fp32-atomic class; finally, executing a plan with tensors on the wrong - device must raise ValueError before kernel 1 mutates the score buffers.""" + device must raise ValueError before kernel 1 overwrites attn_score.""" try: from cudnn import DSA from cuda.bindings import driver as cuda # noqa: F401 @@ -1773,7 +1777,7 @@ def call(dev): assert _rms_rel(r["d_index_k"], refs[dev]["d_index_k"].double()) < _DK_ATOMIC_BAND, f"device {dev} iter {it}: d_index_k diverged" # direct-object contract: a plan built on device 0 must reject device-1 - # tensors BEFORE kernel 1 mutates the score buffers + # tensors BEFORE kernel 1 overwrites attn_score iq0, w0, ik0, attn0, index0, tki0 = inputs[0] with torch.cuda.device(0): plan = DSA.IndexerBackward( @@ -2220,3 +2224,290 @@ def test_DSA_indexer_backward_wrapper_legacy_positional_call( assert torch.isfinite(d_index_q.float()).all() assert torch.isfinite(d_weights.float()).all() assert torch.isfinite(d_index_k.float()).all() + + +@pytest.mark.L1 +def test_DSA_indexer_backward_sm100_persistent_short_topk_dispatch_threshold(): + """Use persistence only once a CTA can amortize setup across rows.""" + if torch.cuda.get_device_capability()[0] < 10: + pytest.skip("SM100+ required") + + try: + from cudnn.deepseek_sparse_attention.indexer_backward.indexer_backward_sm100 import IndexerBackwardSm100, _HAS_TMA_GATHER4 + except ImportError: + pytest.skip("Environment not supported: cudnn[cutedsl] not installed") + + sm_count = torch.cuda.get_device_properties(0).multi_processor_count + for topk in (128, 256, 384): + short = IndexerBackwardSm100( + 128, + heads=64, + block_I=128, + topk=topk, + total_seqlen_k=512, + total_rows=sm_count, + persistent_grid_size=sm_count, + topk_indices_global=True, + ) + long = IndexerBackwardSm100( + 128, + heads=64, + block_I=128, + topk=topk, + total_seqlen_k=512, + total_rows=sm_count + 1, + persistent_grid_size=sm_count, + topk_indices_global=True, + ) + assert not short.use_persistent + assert long.use_persistent + assert short.use_tma_gather is _HAS_TMA_GATHER4 + assert long.use_tma_gather is _HAS_TMA_GATHER4 + assert long.use_cross_row_persistent is _HAS_TMA_GATHER4 + + +@pytest.mark.L1 +@pytest.mark.parametrize("topk", [128, 256, 384, 512]) +def test_DSA_indexer_backward_sm100_score_grad_packed_rows_and_fused_dk_zero(topk): + """Packed tail rows preserve score math and clear the FP32 dK scratch.""" + if torch.cuda.get_device_capability()[0] < 10: + pytest.skip("SM100+ required") + + try: + from cudnn.deepseek_sparse_attention.indexer_backward.indexer_backward_sm100 import indexer_backward_sm100 + except ImportError: + pytest.skip("Environment not supported: cudnn[cutedsl] not installed") + + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(20260826 + topk) + batch, seqlen_q, seqlen_k, heads, head_dim = 2, 151, max(512, topk), 64, 128 + target = torch.softmax( + torch.randn((batch, seqlen_q, topk), device=device, dtype=torch.float32, generator=generator), + dim=-1, + ) + predict = torch.softmax( + torch.randn((batch, seqlen_q, topk), device=device, dtype=torch.float32, generator=generator), + dim=-1, + ) + target[:, ::23, ::31] = 0.0 + predict[:, ::29, ::37] = 0.0 + grad_loss = torch.tensor([0.7], device=device, dtype=torch.float32) + grad_scale = 0.13 + actual = target.clone() + d_index_k_f32 = torch.full( + (batch, seqlen_k, head_dim), + 7.0, + device=device, + dtype=torch.float32, + ) + + kernel = indexer_backward_sm100( + batch, + seqlen_q, + seqlen_k, + heads, + head_dim, + topk, + sm_scale=head_dim**-0.5, + block_I=128, + topk_indices_global=True, + ) + kernel.score_grad_zero( + actual, + predict, + grad_loss, + grad_scale, + dIndexK_f32=d_index_k_f32, + ) + torch.cuda.synchronize() + + clip_probability_min = math.exp(-100.0) + grad = -target.clamp_min(clip_probability_min) * (predict >= clip_probability_min) * (grad_scale * grad_loss.item()) + expected = grad - predict * grad.sum(dim=-1, keepdim=True) + torch.testing.assert_close(actual, expected, rtol=2e-5, atol=2e-7) + torch.testing.assert_close(d_index_k_f32, torch.zeros_like(d_index_k_f32), rtol=0.0, atol=0.0) + + +@pytest.mark.L1 +@pytest.mark.parametrize("topk", [128, 256, 384]) +def test_DSA_indexer_backward_sm100_score_grad_pdl_full_pipeline_matches_serial(topk): + """PDL prologue overlap remains ordered across back-to-back calls.""" + if torch.cuda.get_device_capability()[0] < 10: + pytest.skip("SM100+ required") + + try: + from cudnn.deepseek_sparse_attention.indexer_backward.indexer_backward_sm100 import indexer_backward_sm100 + except ImportError: + pytest.skip("Environment not supported: cudnn[cutedsl] not installed") + + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(20260827 + topk) + batch, seqlen_q, seqlen_k, heads, head_dim = 1, 149, 512, 64, 128 + index_q = torch.randn((batch, seqlen_q, heads, head_dim), device=device, dtype=torch.bfloat16, generator=generator) + index_k = torch.randn((batch, seqlen_k, head_dim), device=device, dtype=torch.bfloat16, generator=generator) + weights = torch.randn((batch, seqlen_q, heads), device=device, dtype=torch.bfloat16, generator=generator).abs().mul_(0.1) + topk_indices = torch.randint(0, seqlen_k, (batch, seqlen_q, topk), device=device, dtype=torch.int32, generator=generator) + target = torch.softmax(torch.randn((batch, seqlen_q, topk), device=device, dtype=torch.float32, generator=generator), dim=-1) + predict = torch.softmax(torch.randn((batch, seqlen_q, topk), device=device, dtype=torch.float32, generator=generator), dim=-1) + grad_loss = torch.tensor([0.7], device=device, dtype=torch.float32) + grad_scale = 0.13 + clip_probability_min = math.exp(-100.0) + grad = -target.clamp_min(clip_probability_min) * (predict >= clip_probability_min) * (grad_scale * grad_loss.item()) + grad_signal = (grad - predict * grad.sum(dim=-1, keepdim=True)).contiguous() + + kernel = indexer_backward_sm100( + batch, + seqlen_q, + seqlen_k, + heads, + head_dim, + topk, + sm_scale=head_dim**-0.5, + block_I=128, + topk_indices_global=True, + ) + d_index_q_ref = torch.empty_like(index_q) + d_weights_ref = torch.empty_like(weights) + d_index_k_ref_f32 = torch.zeros_like(index_k, dtype=torch.float32) + kernel.gemm_only( + index_q, + weights, + index_k, + d_index_q_ref, + d_weights_ref, + d_index_k_ref_f32, + grad_signal, + topk_indices, + ) + d_index_k_ref = d_index_k_ref_f32.to(torch.bfloat16) + + d_index_q = torch.empty_like(index_q) + d_weights = torch.empty_like(weights) + d_index_k = torch.empty_like(index_k) + for _ in range(2): + kernel( + index_q, + weights, + index_k, + d_index_q, + d_weights, + d_index_k, + target.clone(), + predict, + topk_indices, + grad_loss, + grad_scale, + ) + torch.cuda.synchronize() + + for name, actual, expected in ( + ("dQ", d_index_q, d_index_q_ref), + ("dW", d_weights, d_weights_ref), + ("dK", d_index_k, d_index_k_ref), + ): + assert torch.isfinite(actual.float()).all(), f"{name} contains NaN/Inf" + torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2, msg=lambda msg, name=name: f"{name}: {msg}") + + +@pytest.mark.L1 +@pytest.mark.parametrize( + "topk,topk_indices_global,force_manual_k_load", + [ + pytest.param(topk, topk_indices_global, False, id=f"{'global' if topk_indices_global else 'local'}-{topk}") + for topk_indices_global in (True, False) + for topk in (128, 256, 384) + ] + + [pytest.param(512, True, True, id="global-512-manual-k-load")], +) +def test_DSA_indexer_backward_sm100_persistent_short_topk_matches_reference( + topk, + topk_indices_global, + force_manual_k_load, + monkeypatch, + request, +): + """Check persistent phases and the manual K-load fallback against the reference.""" + if torch.cuda.get_device_capability()[0] < 10: + pytest.skip("SM100+ required") + + try: + indexer_backward_sm100_module = importlib.import_module("cudnn.deepseek_sparse_attention.indexer_backward.indexer_backward_sm100") + except ImportError: + pytest.skip("Environment not supported: cudnn[cutedsl] not installed") + + if force_manual_k_load: + monkeypatch.setattr(indexer_backward_sm100_module, "_HAS_TMA_GATHER4", False) + indexer_backward_sm100_module._compile_cache.clear() + request.addfinalizer(indexer_backward_sm100_module._compile_cache.clear) + indexer_backward_sm100 = indexer_backward_sm100_module.indexer_backward_sm100 + + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(20260825 + topk + int(topk_indices_global)) + batch, seqlen_q, seqlen_k, heads, head_dim = 2, 151, 512, 64, 128 + sm_scale = head_dim**-0.5 + index_q = torch.randn((batch, seqlen_q, heads, head_dim), device=device, dtype=torch.bfloat16, generator=generator) + index_k = torch.randn((batch, seqlen_k, head_dim), device=device, dtype=torch.bfloat16, generator=generator) + weights = torch.randn((batch, seqlen_q, heads), device=device, dtype=torch.bfloat16, generator=generator).abs().mul_(0.1) + local_ids = torch.randint(0, seqlen_k, (batch, seqlen_q, topk), device=device, dtype=torch.int32, generator=generator) + local_ids[:, ::17, ::19] = -1 + valid = local_ids >= 0 + batch_offsets = torch.arange(batch, device=device, dtype=torch.int32)[:, None, None] * seqlen_k + global_ids = torch.where(valid, local_ids + batch_offsets, local_ids) + if topk_indices_global: + kernel_ids = global_ids.clone() + kernel_ids[0, 1, 3] = batch * seqlen_k + else: + kernel_ids = local_ids.clone() + kernel_ids[0, 1, 3] = seqlen_k + local_ids[0, 1, 3] = -1 + global_ids[0, 1, 3] = -1 + valid[0, 1, 3] = False + grad_signal = torch.randn((batch, seqlen_q, topk), device=device, dtype=torch.float32, generator=generator) + + d_index_q = torch.empty_like(index_q) + d_weights = torch.empty_like(weights) + d_index_k = torch.zeros_like(index_k, dtype=torch.float32) + kernel = indexer_backward_sm100( + batch, + seqlen_q, + seqlen_k, + heads, + head_dim, + topk, + sm_scale=sm_scale, + block_I=128, + topk_indices_global=topk_indices_global, + ) + kernel.gemm_only(index_q, weights, index_k, d_index_q, d_weights, d_index_k, grad_signal, kernel_ids) + torch.cuda.synchronize() + + safe_ids = local_ids.clamp(min=0).long() + selected_k = torch.gather( + index_k[:, None].expand(-1, seqlen_q, -1, -1), + 2, + safe_ids[..., None].expand(-1, -1, -1, head_dim), + ).float() + scores = torch.einsum("bqhd,bqtd->bqht", index_q.float(), selected_k) * sm_scale + d_scores = grad_signal[:, :, None, :] * weights.float()[:, :, :, None] * ((scores > 0) & valid[:, :, None, :]) + d_index_q_ref = torch.einsum("bqht,bqtd->bqhd", d_scores, selected_k) * sm_scale + d_weights_ref = (grad_signal[:, :, None, :] * torch.relu(scores) * valid[:, :, None, :]).sum(dim=-1) + d_index_k_ref = torch.zeros_like(index_k, dtype=torch.float32) + d_index_k_contrib = torch.einsum("bqht,bqhd->bqtd", d_scores, index_q.float()) * sm_scale + d_index_k_ref.view(-1, head_dim).index_add_( + 0, + global_ids.clamp(min=0).reshape(-1).long(), + d_index_k_contrib.reshape(-1, head_dim), + ) + + for name, actual, expected in ( + ("dQ", d_index_q.float(), d_index_q_ref), + ("dW", d_weights.float(), d_weights_ref), + ("dK", d_index_k, d_index_k_ref), + ): + actual_flat = actual.flatten() + expected_flat = expected.flatten() + cosine = F.cosine_similarity(actual_flat, expected_flat, dim=0).item() + rms_relative = ((actual_flat - expected_flat).square().mean().sqrt() / expected_flat.square().mean().sqrt().clamp_min(1e-12)).item() + assert torch.isfinite(actual_flat).all(), f"{name} contains NaN/Inf" + assert cosine >= 0.99, f"{name} cosine={cosine:.6f}" + assert rms_relative <= 0.02, f"{name} RMS-relative error={rms_relative:.6f}"