From 98a5a833b0cd8e2c581c5e2cad833a7df68f092d Mon Sep 17 00:00:00 2001 From: zky <51477259+zkyue@users.noreply.github.com> Date: Thu, 23 Jul 2026 04:22:10 +0000 Subject: [PATCH 1/2] Fix indexer_backward_sm100 W_LOADED handoff: whole-warp mbarrier arrive In _load_warp all 32 lanes of the load warp store sW / sGradSignal to SMEM, but only an elected lane arrives on MBAR_W_LOADED. Per the PTX memory model, mbarrier.arrive (release, cta scope) orders only the executing thread's prior accesses, so the other 31 lanes' stores have no happens-before edge to the compute warpgroup's mbarrier_wait and subsequent reads: a formal data race. Latent in practice: no corruption observed on the tested B200 / CUDA 13.3 / cutlass-dsl 4.6.1 build, whose captured SASS (topk=128 specialization) carries an unpredicated MEMBAR.ALL.CTA before the arrive; that compensation is not contractual. Fix: initialize MBAR_W_LOADED with count WARP_SIZE and have all 32 lanes arrive, closing the happens-before chain per lane. W_LOADED is a one-shot handoff (single arrive site, single phase-0 wait), so the count change is self-contained. Verified: compute-sanitizer racecheck hazards on kernel_gemm drop to 0 across 1-CTA, 512-CTA and batch=3/topk=512 shapes (previously the only flagged site in those runs); d_index_q / d_weights byte-identical to unpatched on the tested shapes; upstream DSA pytest results unchanged. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> --- .../indexer_backward/indexer_backward_sm100.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) 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 8b657cafb..ca2fef95a 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 @@ -541,7 +541,7 @@ class SharedStorage: 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, 1) + cute.arch.mbarrier_init(mbar + MBAR_W_LOADED, self.WARP_SIZE) cute.arch.mbarrier_init(mbar + MBAR_DQ_DONE, 1) cute.arch.sync_threads() @@ -745,9 +745,13 @@ def _load_warp( sW[idx] = mW[seq_idx, idx, batch_idx] cute.arch.fence_view_async_shared() - # Signal W + grad_signal loaded for compute warpgroup - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive(mbar + MBAR_W_LOADED) + # Signal W + grad_signal loaded for compute warpgroup. + # All 32 lanes arrive (count = WARP_SIZE): mbarrier.arrive has release + # semantics for the *executing thread* only, so a single elected + # arrival would not order the other 31 lanes' sW/sGradSignal stores + # before the consumer's mbarrier_wait (racecheck flags exactly those + # 31 lanes). Whole-warp arrival closes the happens-before chain. + cute.arch.mbarrier_arrive(mbar + MBAR_W_LOADED) # --- TMA Q load (dsa-next pattern: cute.copy with pre-partitioned tensors) --- Q_producer.reset() From df6884573477f852be5715f73419c3268f219976 Mon Sep 17 00:00:00 2001 From: zky <51477259+zkyue@users.noreply.github.com> Date: Fri, 24 Jul 2026 02:58:45 +0000 Subject: [PATCH 2/2] Elect a single lane for mbarrier init in indexer_backward_sm100 Review follow-up: the barrier-initialization block under `if warp_idx == 0:` was executed by all 32 lanes of warp 0, i.e. each mbarrier_init ran 32 times on the same SMEM barrier object. Redundant re-initialization before the sync_threads is benign on current hardware, but a single initializing thread is the contract the PTX ISA documents for mbarrier.init, and every other kernel in this package already wraps barrier init in an election. Wrap the block in `with cute.arch.elect_one():` so exactly one lane performs the init; the trailing sync_threads() ordering is unchanged. No functional change intended or observed. Re-verified on B200 / CUDA 13.3 / cutlass-dsl 4.6.1: compute-sanitizer racecheck 0 hazards on the 1-CTA (hazard-level report), 512-CTA and batch=3/topk=512 shapes; fe_api/dsa pytest results identical to the parent commit (same 26 passed / 4 skipped / same 4 environment-specific failures); 30-replay d_index_q / d_weights SHA-256 byte-identical to the parent commit. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> --- .../indexer_backward_sm100.py | 35 ++++++++++--------- 1 file changed, 18 insertions(+), 17 deletions(-) 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 ca2fef95a..df88de296 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 @@ -525,24 +525,25 @@ class SharedStorage: cute.group_modes(gdQ, 0, 2), ) - # Init all barriers (warp 0) + # Init all barriers (single lane of warp 0) if warp_idx == 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) - 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) + 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) + 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.sync_threads() # Pre-load topk indices into SMEM cooperatively (all 512 threads).