SM120: fix FP8 MQA logits swizzle mode for head_dim < 128 - #379
Conversation
kSwizzleMode was hardcoded to 128 while the TMA descriptor uses
swizzle_mode = head_dim (and kSwizzleAlignment = head_dim * 8 = the
swizzle span). For head_dim in {32, 64} the kernel read SMEM with a
128B swizzle over a 32/64B-swizzled layout, mis-addressing reads into
other (valid, already-written) pipeline stages: racecheck-clean but
timing-dependent, so logits were both wrong and nondeterministic at
deep KV pipelines. Set kSwizzleMode = kHeadDim (SwizzleContext
supports B32/B64/B128). Same fix applied to the paged kernel, where
head_dim is currently always 128 (no behavior change).
Repro on RTX PRO 6000 (head_dim 32/64, seq_len_kv 8192): diff vs
reference ~0.98 and nondeterministic across runs; after the fix all
head_dim in {32, 64, 128} are deterministic with diff ~1e-6.
Also adds FP8 head_dim 32/64 coverage to test_mqa_logits (previously
only head_dim=128 was exercised, which is why this never fired).
| // Must match the TMA descriptor's swizzle (host passes swizzle_mode = head_dim) | ||
| static constexpr uint32_t kSwizzleMode = kHeadDim; |
There was a problem hiding this comment.
🔴 critical: The Q/KV SMEM is written in the swizzle chosen by the TMA descriptor (swizzle_mode=head_dim: 32B/64B/128B) but the read SwizzleContext uses a constant kSwizzleMode=128 (B128). For head_dim 32/64 the read layout and the written layout differ, mis-addressing reads into other already-written pipeline stages => wrong + nondeterministic logits. Set kSwizzleMode=kHeadDim so the read layout matches the descriptor (B32/B64/B128). The paged kernel needs the very same change.
🤖 v4
| for seq_len, seq_len_kv, head_cfgs in ( | ||
| *[(s, k, [(64, 128)]) for s in (128, 512, 2048, 4096) for k in (4096, 8192)], | ||
| *[(s, k, [(64, 128), (32, 128)]) for s in (510, 512) for k in (130560,)], | ||
| # head_dim < 128 coverage (FP8-only): kernel swizzle must follow head_dim |
There was a problem hiding this comment.
🔵 suggestion: The new head_dim 32/64 coverage is a good regression guard. Since the original bug was nondeterministic across runs, consider confirming the test's reference-comparison tolerance is tight enough to reliably fail on the pre-fix ~0.98 diff, and ideally exercise a deep KV pipeline (large seq_len_kv) where the mis-addressing manifested. A single tight comparison should catch it, but making the deterministic-across-runs aspect explicit would strengthen the guard.
🤖 v3
🤖 ds-review-bot Code Reviewv6The kernel-side swizzle now matches the TMA descriptor for all supported FP8 head dimensions, and the added tests cover the previously broken cases. v4The MR fixes a wrong hardcoded kSwizzleMode=128 in the two SM120 FP8 MQA logits kernels (dense + paged) and extends test_mqa_logits to FP8-only head_dim 32/64. Root cause: the host encodes swizzle_mode=head_dim via mode_into_tensor_map_swizzle (32->32B, 64->64B, 128->128B) and writes Q/KV SMEM in that swizzled layout, while the read path used the constant 128 (B128). For head_dim 32/64 reads land in other already-written pipeline stages => wrong + nondeterministic (racecheck-clean, timing-dependent). Fix is kSwizzleMode=kHeadDim (SwizzleContext supports B32/B64/B128), leaving kSwizzleAlignment=head_dim*8 (the swizzle span) and kSMEMKBytes=head_dim unchanged; the DG_STATIC_ASSERT stage-size-alignment checks still hold. The paged kernel only runs head_dim=128 (=B128 in both old/new) so it is a no-op. Test coverage trigger: insert (512, 8192, [(64,64),(64,32)]) into head_cfgs, and skip FP4 for head_dim != 128. The exact code patch is: in both cuhs replace 'static constexpr uint32_t kSwizzleMode = 128;' with the comment '// Must match the TMA descriptor's swizzle (host passes swizzle_mode = head_dim)' followed by 'static constexpr uint32_t kSwizzleMode = kHeadDim;'. After the fix head_dim in {32,64,128} is deterministic with diff ~1e-6 versus the reference. v3This change fixes an SMEM swizzle-mode bug in the SM120 FP8 MQA logits kernels. Previously kSwizzleMode was hardcoded to 128 while the TMA descriptor writes SMEM with swizzle_mode = head_dim and kSwizzleAlignment = head_dim * 8. For head_dim in {32, 64} the kernel therefore read SMEM back with a 128B swizzle XOR pattern over a 32/64B-swizzled layout, mis-addressing reads into other (already-written) pipeline stages. This is racecheck-clean but timing-dependent, producing wrong and nondeterministic logits at deep KV pipelines. The fix sets kSwizzleMode = kHeadDim in both sm120_fp8_mqa_logits.cuh and sm120_fp8_paged_mqa_logits.cuh (the paged kernel always uses head_dim=128 today, so no behavior change there), and adds FP8 head_dim 32/64 coverage to test_mqa_logits. Correctness verification (all confirmed correct):
Overall the change is minimal, well-targeted, and consistent with the surrounding code. LGTM. Files reviewed: 3 |
|
@RayWang96 Please review, thanks |
2 participants
kSwizzleMode was hardcoded to 128 while the TMA descriptor uses swizzle_mode = head_dim (and kSwizzleAlignment = head_dim * 8 = the swizzle span). For head_dim in {32, 64} the kernel read SMEM with a 128B swizzle over a 32/64B-swizzled layout, mis-addressing reads into other (valid, already-written) pipeline stages: racecheck-clean but timing-dependent, so logits were both wrong and nondeterministic at deep KV pipelines. Set kSwizzleMode = kHeadDim (SwizzleContext supports B32/B64/B128). Same fix applied to the paged kernel, where head_dim is currently always 128 (no behavior change).
Repro on RTX PRO 6000 (head_dim 32/64, seq_len_kv 8192): diff vs reference ~0.98 and nondeterministic across runs; after the fix all head_dim in {32, 64, 128} are deterministic with diff ~1e-6.
Also adds FP8 head_dim 32/64 coverage to test_mqa_logits (previously only head_dim=128 was exercised, which is why this never fired).