diff --git a/include/cudnn_frontend/experimental/attention_utils.h b/include/cudnn_frontend/experimental/attention_utils.h index 34192aa7e..b44a55e15 100644 --- a/include/cudnn_frontend/experimental/attention_utils.h +++ b/include/cudnn_frontend/experimental/attention_utils.h @@ -64,7 +64,7 @@ find_log_2_floor(uint32_t x) { // Compute FastDivisor_t for the kernel's fastDivMod which uses: // div = __umulhi(2 * val, mul) >> shr -// This matches cuDNN's find_divisor_v2 (xmma/fast_math.h:118-125). +// This matches cuDNN backend's find_divisor_v2 fast-math helper. inline FastDivisor_t make_fast_divisor(uint32_t divisor) { FastDivisor_t d; 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 29e33d266..eb9501d02 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 @@ -1010,7 +1010,7 @@ def _compute_warpgroup( - 2D identity for scalar (h, n) coordinates: mode=[0] → h, mode=[1] → n - sW pre-loaded into bf16 register array (scalar h enables rW[h]) - Paired f32x2 packed ops (fma_packed_f32x2 for dW accumulation) - - stmatrix bulk store (8 warp-level instructions vs 32+ scalar STS per thread) + - stmatrix bulk store (8 warp-level instructions vs 32+ scalar shared-mem stores per thread) - Compute and store phases fully separated """ wg_tidx = tidx % self.WARPGROUP_SIZE diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm90.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm90.py index f7ec529fb..6e628d3d0 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm90.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm90.py @@ -730,18 +730,18 @@ def _compute_warpgroup( n_offset = compute_wg_idx * self.half_block_I - # STS dK staging: partition sdK_staging using GEMM2's MMA layout + # shared-mem store dK staging: partition sdK_staging using GEMM2's MMA layout sdK_staging_half = cute.local_tile(sdK_staging, (64, self.block_I), (compute_wg_idx, 0)) tCsDK_staging = thr_mma2.partition_C(sdK_staging_half) - # Fused pass sdS write via stmatrix (r2s bulk copy, replaces scalar STS) + # Fused pass sdS write via stmatrix (r2s bulk copy, replaces scalar shared-mem stores) sdS_half = cute.local_tile(sdS_view, (self.heads_padded, self.half_block_I), (0, compute_wg_idx)) stmatrix_atom_ds = cute.make_copy_atom(warp.StMatrix8x8x16bOp(), self.q_dtype) tiled_r2s_ds = cute.make_tiled_copy_C(stmatrix_atom_ds, tmma1) thr_r2s_ds = tiled_r2s_ds.get_slice(wg_tidx) tRdDS = thr_r2s_ds.partition_D(sdS_half) - # P4: 2D block view of sGradSignal for partition-based access (sparse only) + # Pass 4: 2D block view of sGradSignal for partition-based access (sparse only) if const_expr(not self.is_dense): sGS_per_block = cute.make_tensor( sGradSignal.iterator, @@ -813,7 +813,7 @@ def _compute_warpgroup( warpgroup.wait_group(0) # ----- Fused pass: compute dS → registers, then stmatrix → sdS ----- - # P4: broadcast partition of sGradSignal — eliminates n-coordinate cute.get + # Pass 4: broadcast partition of sGradSignal — eliminates n-coordinate cute.get if const_expr(self.is_dense): sGS_cur_block = sGradSignal else: @@ -846,7 +846,7 @@ def _compute_warpgroup( else: dw_h1 = dw_h1 + dw_val - # Bulk write dS to SMEM via stmatrix (replaces per-element scalar STS) + # Bulk write dS to SMEM via stmatrix (replaces per-element scalar shared-mem stores) tRsDS = tiled_r2s_ds.retile(acc_dS) cute.copy(tiled_r2s_ds, tRsDS, tRdDS) cute.arch.fence_view_async_shared() @@ -872,13 +872,13 @@ def _compute_warpgroup( else: gemm(tmma3, acc_dQ, tDQrDS, tDQrKt_s2, zero_init=False, wg_wait=-1) - # Signal other WG: it can start its GEMM2+3 while we do STS/memory + # Signal other WG: it can start its GEMM2+3 while we do shared-mem stores / memory if compute_wg_idx == 0: cute.arch.barrier_arrive(barrier_id=self.SCHED_BARRIER_WG1, number_of_threads=self.TOTAL_COMPUTE_THREADS) else: cute.arch.barrier_arrive(barrier_id=self.SCHED_BARRIER_WG0, number_of_threads=self.TOTAL_COMPUTE_THREADS) - warpgroup.wait_group(1) # GEMM2 done (acc_dK ready for STS) + warpgroup.wait_group(1) # GEMM2 done (acc_dK ready to write to shared mem) # Deferred DMA wait: previous iteration's bulk reduce must finish # reading sdK_staging before we overwrite it (both WGs do their own reduce) diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/dense_score_recompute_sm90.py b/python/cudnn/deepseek_sparse_attention/score_recompute/dense_score_recompute_sm90.py index a50bbdb55..62df05bfe 100644 --- a/python/cudnn/deepseek_sparse_attention/score_recompute/dense_score_recompute_sm90.py +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/dense_score_recompute_sm90.py @@ -536,7 +536,7 @@ def producer( for h_tile in cutlass.range_constexpr(self.num_head_tiles): eff_m_block = m_block * self.num_head_tiles + h_tile - # ---- Q TMA + Weights/LSE LDG ---- + # ---- Q TMA + Weights/LSE load-from-global ---- gQ = cute.local_tile(mQ_cur, (self.tile_m, self.tile_hdim), (eff_m_block, 0)) load_Q, _, _ = copy_ops.tma_get_copy_fn(tma_atom_Q, 0, cute.make_layout(1), gQ, sQ, single_stage=True) if warp_idx_in_wg == 0: @@ -663,7 +663,7 @@ def producer( # the unified `producer` has warp 0 doing: # # wait KVE0 → TMA stage0 → wait KVE1 → TMA stage1 → ... # # so stage-1's TMA issue is serialized behind stage-0's wait+issue, # - # even when both consumer WGs have released their KVEmpty. Perfsim # + # even when both consumer WGs have released their KVEmpty. Profiling # # shows WG1's wait KV is ~20-50c longer than WG0's under unified. # # # # With warp split, stage-0 and stage-1 pipelines are independent — # @@ -724,7 +724,7 @@ def producer_warp_split( for h_tile in cutlass.range_constexpr(self.num_head_tiles): eff_m_block = m_block * self.num_head_tiles + h_tile - # ---- Q TMA (warp-0) + Weights/LSE LDG (all 128 threads) ---- + # ---- Q TMA (warp-0) + Weights/LSE load-from-global (all 128 threads) ---- gQ = cute.local_tile(mQ_cur, (self.tile_m, self.tile_hdim), (eff_m_block, 0)) load_Q, _, _ = copy_ops.tma_get_copy_fn(tma_atom_Q, 0, cute.make_layout(1), gQ, sQ, single_stage=True) if warp_idx_in_wg == 0: diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py index 8194edb92..87024aeb3 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py @@ -938,7 +938,7 @@ def bwd( tdKVrQT = QdS_tiled_mma.make_fragment_A(sQT) tdKVrdS = QdS_tiled_mma.make_fragment_B(sdS) - # Ugly, but I don't know how to handle the correct modes for cute.gemm + # Awkward, but I don't know a cleaner way to handle the modes for cute.gemm tdKVrQT_shape = (tdKVrQT.shape[0], 1, tdKVrQT.shape[1], tdKVrQT.shape[2], tdKVrQT.shape[3]) tdKVrQT_stride = (tdKVrQT.stride[0], 0, tdKVrQT.stride[1], tdKVrQT.stride[2], tdKVrQT.stride[3]) tdKVrQT = cute.make_tensor(tdKVrQT.iterator, cute.make_layout(tdKVrQT_shape, stride=tdKVrQT_stride)) diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm90.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm90.py index cbe9d9d98..fa1cc5323 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm90.py +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm90.py @@ -383,7 +383,7 @@ def kernel( else: mdQaccum_cur = cute.domain_offset((padded_offset_q * self.head_dim_padded,), mdQaccum[head_idx, None]) - # HACK: Compiler doesn't seem to recognize that padding + # Note: compiler doesn't seem to recognize that padding # by padded_offset_q * self.head_dim_padded keeps alignment # since statically divisible by 4 @@ -1551,7 +1551,7 @@ def _wg0_one_n_block( ) # (6.5) G4_half_0: dQ[0:128] += dS_scaled(reg) @ sKVt_q0 (RS, immediate) - # Fire before STS(dS) — dS_scaled is still in registers, TC reads at issue time. + # Fire before storing dS to shared mem — dS_scaled is still in registers, TC reads at issue time. mma_dsk_fn_0(tCrA=tdQrdS_scaled, B_idx=None, zero_init=not dQ_accumulate, wg_wait=-1) # (7) R2S dS_scaled -> sdS (after G4_half_0 issued, regs still valid) diff --git a/python/cudnn/gemm_swiglu/dense_gemm_persistent_swiglu.py b/python/cudnn/gemm_swiglu/dense_gemm_persistent_swiglu.py index 8b6c9f6f6..28f6aff86 100644 --- a/python/cudnn/gemm_swiglu/dense_gemm_persistent_swiglu.py +++ b/python/cudnn/gemm_swiglu/dense_gemm_persistent_swiglu.py @@ -1046,7 +1046,7 @@ def kernel( # # Wait for accumulator buffer full # - acc_pipeline.consumer_wait(acc_consumer_state) # Get for the single CGA + acc_pipeline.consumer_wait(acc_consumer_state) # Get for the single cluster tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) bSG_gAB12 = cute.group_modes(bSG_gAB12, 1, cute.rank(bSG_gAB12)) bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) diff --git a/test/python/fe_api/test_sdpa_bwd.py b/test/python/fe_api/test_sdpa_bwd.py index 6685cc096..88f8950a1 100644 --- a/test/python/fe_api/test_sdpa_bwd.py +++ b/test/python/fe_api/test_sdpa_bwd.py @@ -5,7 +5,6 @@ uses Blackwell TMEM which produces NaN when GPU kernels from other CUDA contexts (i.e. other pytest-xdist workers) execute concurrently on the same device. Run with ``-m gpu_exclusive -n 0`` in CI. -See: https://gitlab-master.nvidia.com/cudnn/cudnn_frontend/-/jobs/299938128 """ import pytest