Fix SM120 BF16 GEMM shared-memory reuse ordering

[Sunt-ing]

Summary

Fixes intermittent silent data corruption in SM120 BF16 GEMM caused by a generic-to-async proxy WAR (Write-After-Read) hazard.

Problem

Consumer warps read smem tiles via the generic proxy and signal the empty barrier without an async proxy fence. Producer TMA loads (async proxy) can subsequently overwrite the smem stage before generic reads complete across all consumer threads.

Fix

Add cutlass::arch::fence_view_async_shared() before the empty-barrier arrival. Placed outside the lane-0 check so all consumer threads issue the fence for their own smem reads.

Targets nv_dev (based on 572557e7ae9ad5331b81a1c250f141fba2c57962).

Reproduction and validation

Environment: RTX 5090 (SM120), Driver 590.48.01, PyTorch 2.13.0+cu130.

Added test_bf16_repeatability in tests/test_bf16.py ([4096, 1024] × [1024, 896], seed 233, 6,000 iterations against initial output).

To run the standalone regression:

cd tests
python -c "from test_bf16 import test_bf16_repeatability; test_bf16_repeatability()"

Note: The corruption is an intermittent race condition; a single baseline run passing does not guarantee correctness.

Repeatability (A/B/A Test)

• Original: 2 / 6,000 mismatches
• Patched: 0 / 6,000 mismatches
• Reverted: 2 / 6,000 mismatches

Performance Overhead

CUDA graph timing on [4096, 1024] × [1024, 896] (median per-call time across 9 batches of 200 calls):

Version	Median Latency
Original	50.67 μs
Patched	52.01 μs
Reverted	50.68 μs

Adds ~2.6% latency overhead on this shape from the fence.

Supplemental End-to-End Inference Validation

Internal TP2 inference test (weights/prompts omitted; public verification relies on the standalone GEMM test above).

• Setup: 32,640-token prompt, 89 decode steps -> re-prefill -> 16 decode steps vs. continuous generation. Unified prefill/decode attention was held constant; CUDA graphs were disabled.
• Original / Reverted: 6 / 12 runs with divergent logprobs.
• Patched: 0 / 12 runs diverged (token IDs and logprobs match strictly).

Related work

• fix: fence combine buffer reuse in mega moe #389: Addressed the same class of generic-to-async proxy hazard in SM100 MegaMoE combine.
• Integrate main APIs and SM120 kernels while preserving nv_dev features #447: The consumer release path (reviewed at 1f78b85a1fb974cc69ea63e864d817c603688cab) does not yet include this fence.

[ds-review-bot]

🤖 ds-review-bot Code Review

v6

未发现本次变更新引入的可确认缺陷。fence 位于共享内存读取之后、empty barrier 到达之前，且由所有消费者线程执行；新增重复性测试也会调用目标内核。静态检查通过，但当前环境缺少 PyTorch 和 CUDA 工具链，未执行 GPU 验证。

v5

本 MR 在 SM120 BF16 GEMM 的消费者释放路径中，于 empty barrier arrive 之前加入 cutlass::arch::fence_view_async_shared()（fence.proxy.async.shared::cta），修复 generic proxy（ldmatrix 读 smem）与 async proxy（TMA 写 smem）之间的 WAR 竞争。修复位置正确：fence 放在 lane_idx == 0 判断之外，保证每个消费者线程都对自身的 smem 读做排序后再由 lane 0 发出 arrive（fence.proxy.async 仅对执行线程自身访存生效，若只由 lane 0 执行则其余 lane 的读仍会乱序）；cutlass/arch/barrier.h 已在文件头包含，符号可正常解析。新增 test_bf16_repeatability 回归测试（seed 233，[4096,1024]×[1024,896]，6000 次 torch.equal 对比）设计合理，能在无修复时以约 2/6000 概率复现问题，且对非 SM120 架构直接返回、不影响其他架构 CI。A/B/A 复现数据、端到端 TP2 验证与 ~2.6% 的性能开销说明充分，改动最小且聚焦。建议合并。遗留问题：同类 lane-0 arrive 无 fence 的模式仍存在于其他 SM120 内核，建议在后续 MR（#447）中统一处理，详见评论。

v4

The change is correct and well-targeted. It fixes a real generic-to-async proxy WAR hazard in the SM120 BF16 GEMM consumer path: after the consumer warps finish reading the smem A/B tiles, they signal stage reuse via empty_barriers[stage]->arrive() without ordering those generic-proxy reads against the producer's subsequent async-proxy TMA write, so a fast producer can overwrite a stage while consumer reads are still in flight. Adding cutlass::arch::fence_view_async_shared() (fence.proxy.async.shared::cta) immediately before the release is the right primitive, and placing it outside the lane_idx == 0 guard is correct because the fence is a per-thread ordering primitive and every lane that issued smem loads must drain them before lane 0's release publishes the stage as reusable. <cutlass/arch/barrier.h> is already included, so no new include is needed. The added repeatability test is a reasonable stress test for the race (though probabilistic). The main concerns are: (1) the identical unfenced release pattern remains in the sibling SM120 GEMM consumers (FP8/FP4 1D1D, BMK/BNK MN, TF32 prenorm), so this class of silent corruption is only partially closed; and (2) the new soak test may pass on the buggy build a non-trivial fraction of the time and synchronizes on every iteration.

Files reviewed: 2
Issues found: 🔴 3 critical | 🟡 1 warning | 🔵 4 suggestion
Inline comments posted: 5
General comments (无法定位到 diff): 3

---

📍 未定位到 diff 的评论

🔴 critical deep_gemm/include/deep_gemm/impls/sm120_fp8_fp4_gemm_1d1d.cuh:L666: Same unfenced release pattern remains here (also at lines 814 and 1143). The producer issues TMA loads via tma::copy + arrive_and_expect_tx (lines 321-331), the consumer reads the tiles through the generic proxy (load_a_fragment/load_b_fragment/ldmatrix) and then releases with if (lane_idx == 0) empty_barriers[stage]->arrive(); with no async proxy fence. This is the same generic->async WAR hazard fixed for BF16 and can produce the same intermittent silent corruption. Please apply cutlass::arch::fence_view_async_shared() before each of these arrivals (or factor the release into a helper that always fences). _{🤖 v4}

🔴 critical deep_gemm/include/deep_gemm/impls/sm120_bmk_bnk_mn.cuh:L202: This consumer has the identical structure to the BF16 kernel: the producer TMA-copies A/B into stages (lines 126-136), the consumer ldmatrix-reads them (lines 178-198), then releases the stage with a bare if (lane_idx == 0) empty_barriers[stage_idx]->arrive();. No proxy fence is present, so the producer can overwrite the stage before all generic reads complete. It should receive the same fix; otherwise the corruption class this MR addresses is only partially closed. _{🤖 v4}

🔴 critical deep_gemm/include/deep_gemm/impls/sm120_tf32_hc_prenorm_gemm.cuh:L231: Same pattern as the other SM120 consumers: TMA producer (lines 129-136) plus generic smem reads followed by a lane-0 empty-barrier arrival with no async proxy fence. Please audit and apply fence_view_async_shared() here as well. Given the references to prior similar fixes (#389, #447), centralizing the release so this fence cannot be forgotten again would prevent recurrence. _{🤖 v4}

[Sunt-ing]

@LyricZhao @zheanxu Could you please take a look at this fix? We found this while investigating reproducibility in LLM evaluation with SGLang on RTX 5090s. With a 32K prompt, continuing decode and re-prefilling the same token prefix sometimes produced different logprobs. Tracing the difference through the model led to a BF16 projection where repeated GEMM calls with identical inputs occasionally returned different values. We then reproduced it with the public random-matrix test in this PR.

[lucifer1004]

Have you checked if this is covered by #447?

[Sunt-ing]

I checked #447, but I couldn’t find the corresponding fix in the SM120 BF16 kernel. At its current head (1f78b85), the consumer still calls empty_barriers[stage]->arrive() without a preceding async proxy fence.

[lucifer1004]

@Sunt-ing Thanks for the careful A/B/A reproduction — confirmed the hazard class. Your fix has now been ported into #447 at 381e2d7, via the canonical SM120 device-layer repo (lucifer1004/DeepGEMM-sm120 v0.1.5, b31a688), which both SM120 forks vendor byte-for-byte; the same fix is also in the vllm fork's SM120 layer (vllm-project#10).

While merging it we audited all SM120 kernels for the same pattern — TMA-produced smem consumed via ldmatrix/ld_shared (generic proxy), with the empty-barrier arrive handing the stage back to the async proxy — and added the missing fence_view_async_shared() at 14 sites across 8 headers: bf16_gemm (your report), bmk_bnk_mn, tf32_hc_prenorm_gemm, fp8_fp4_gemm_1d1d (×3), fp8/fp4 dense MQA logits (empty_kv + empty_q each), and fp8/fp4 paged MQA logits (empty_q stage release + empty_kv each). The sparse MQA kernel is exempt (its producer is legacy cp.async, generic proxy). Upstream's SM100 kernels carry the same fence for the analogous UMMA case ("Release KV smem only after UMMA commits TMEM"), which corroborates the pattern.

For transparency on evidence: on our RTX PRO 6000 a 30k-iteration soak of the affected shapes did not reproduce the corruption either way, so the fix rests on the memory model plus your 5090 reproduction (2/6000 → 0/6000). Perf A/B (bf16 GEMM, fp8 dense/paged MQA, 3 runs each) shows all deltas within run-to-run noise — consistent with your +2.6% report.

[Sunt-ing]

Thanks for porting the fix and checking the other kernels!