NVIDIA · hwanseoc · Jun 3, 2026 · Jun 3, 2026
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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