Fix non-deterministic results in k_grouped_fp8_gemm_nt_contiguous - #375
Fix non-deterministic results in k_grouped_fp8_gemm_nt_contiguous#375Functionhx wants to merge 1 commit into
Conversation
…iguous The SM90 FP8 1D1D kernel updates TMA tensor map descriptors in-place during k-grouped GEMM group switches. The original code issued tensormap.replace operations before draining in-flight TMA data loads, committing both into the same cp.async.bulk commit group. Per the PTX ISA, operations within a commit group may complete in any order, which allows the tensor map descriptor modifications to become visible before pending TMA data loads complete. This race manifests as non-deterministic numerical errors when one expert has exactly 128 tokens (k=BLOCK_K, a single k-block) and a column of the activation tensor is all zeros. The shorter pipeline depth with k=128 exposes the race, while the all-zero column with a non-negligible scale factor (e.g., 1.0) makes wrong reads visible. Fix: drain all pending TMA data operations before modifying tensor map descriptors, then fence the modifications in a separate commit group. This ensures clean separation between data operations and descriptor updates. Fixes deepseek-ai#365 Test Plan: Reproduced the non-determinism with the reproducer from deepseek-ai#365 on a GH200/H100 GPU and confirmed the fix eliminates all failures across 10,000 iterations for all scale factor values: 1, 1e-4, 1e-8, 1e-30. Signed-off-by: Yuchen Fan <functionhx@gmail.com>
| @@ -191,26 +191,29 @@ sm90_fp8_gemm_1d1d_impl(__nv_fp8_e4m3* gmem_a_ptr, __nv_fp8_e4m3* gmem_b_ptr, | |||
| if (kGemmType == GemmType::KGroupedContiguous and last_group_idx != scheduler.current_group_idx) { | |||
| last_group_idx = scheduler.current_group_idx; | |||
|
|
|||
There was a problem hiding this comment.
🔵 suggestion: Root-cause wording is slightly imprecise. tensormap.replace is not a cp.async.bulk operation and is not part of any commit group, so the description's claim that replace and TMA loads 'compete within the same commit group' is not technically accurate. The real hazard is that the #343 fix placed the drain AFTER the replace, so the replace could run while a prior TMA descriptor read was still in flight (a missing happens-before). The fix itself is correct; please tighten the commit-message/comment framing to avoid future confusion.
🤖 v3
| // compete within the same commit group, which the PTX ISA allows to resolve in any order. | ||
| cute::tma_desc_commit_group(); | ||
| cute::tma_desc_wait_group(); | ||
| __syncwarp(1U << lane_idx); |
There was a problem hiding this comment.
🔵 suggestion: __syncwarp(1U << lane_idx) uses a single-lane mask, which is intentional here since this block runs under elect_one_sync() (only one lane active) and the __syncwarp acts purely as a compiler-ordering barrier. This diff removed the previous explanatory comment ('Only used to prevent ptxas from moving ... Shouldn't be needed otherwise, since we only use one thread'). Consider restoring a one-line comment on the single-lane __syncwarp so future readers don't mistake the mask for a bug.
🤖 v3
🤖 ds-review-bot Code Reviewv4v3This PR restructures the k-grouped GEMM group-switch tensor-map update in deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh into three phases: (1) drain all pending TMA data loads via tma_desc_commit_group + tma_desc_wait_group BEFORE any tensormap.replace, (2) issue the tensormap.replace operations then fence them with a second commit/wait group, and (3) copy the updated descriptors to global memory with release_gpu/acquire_gpu. Compared to the previous fix (#343), the only functional addition is the drain that now precedes the replace calls, and it is correct: cp.async.bulk.wait_group.read 0 (emitted by tma_desc_wait_group) waits for prior in-flight TMA tensor loads, establishing a happens-before so those loads complete before descriptor bytes are modified. The three-phase structure mirrors CUTLASS's tensormaps_perform_update / tensormaps_cp_fence_release / tensormaps_fence_acquire pattern (deep_gemm adds an explicit pre-drain where CUTLASS relies on producer_acquire/load_tail). The second commit/wait group after the replace is retained from #343 and correctly fences the modifications before the smem->gmem copy; the first commit_group may produce a harmless empty bulk-async group. Performance impact is negligible since group switches occur once per expert, not in the k-loop. Verdict: LGTM, sound and minimal hardening of a real race, with only minor documentation suggestions below. Files reviewed: 1 |
2 participants
Summary
Fixes #365: non-deterministic numerical results in FP8 grouped GEMM when one expert has exactly 128 tokens.
Root Cause
Race condition in sm90_fp8_gemm_1d1d.cuh: tensormap.replace operations and in-flight TMA data loads were in the same cp.async.bulk commit group. Per PTX ISA, operations within a commit group may complete in any order, allowing descriptor modifications to become visible before pending TMA reads complete.
Fix
Restructure group-switch into three phases: