From 2284689a3b7f529fb53744a6745dcc5d8c247807 Mon Sep 17 00:00:00 2001 From: Yuanhang Sun Date: Tue, 14 Jul 2026 09:23:34 -0700 Subject: [PATCH] SM120: fix FP8 MQA logits swizzle mode for head_dim < 128 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). --- deep_gemm/include/deep_gemm/impls/sm120_fp8_mqa_logits.cuh | 3 ++- .../include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh | 3 ++- tests/test_attention.py | 4 ++++ 3 files changed, 8 insertions(+), 2 deletions(-) diff --git a/deep_gemm/include/deep_gemm/impls/sm120_fp8_mqa_logits.cuh b/deep_gemm/include/deep_gemm/impls/sm120_fp8_mqa_logits.cuh index 93fc69daba..3987b081da 100644 --- a/deep_gemm/include/deep_gemm/impls/sm120_fp8_mqa_logits.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm120_fp8_mqa_logits.cuh @@ -58,7 +58,8 @@ void sm120_fp8_mqa_logits(const uint32_t seq_len, const uint32_t seq_len_kv, DG_STATIC_ASSERT(BLOCK_KV == kNumMathWarps * MMA_M, "BLOCK_KV = warps × MMA_M"); DG_STATIC_ASSERT(kHeadDim % MMA_K == 0 and kNumHeads % MMA_N == 0, "Alignment"); - static constexpr uint32_t kSwizzleMode = 128; + // Must match the TMA descriptor's swizzle (host passes swizzle_mode = head_dim) + static constexpr uint32_t kSwizzleMode = kHeadDim; static constexpr uint32_t kSwizzleAlignment = kHeadDim * 8; static constexpr uint32_t kSMEMKBytes = kHeadDim; diff --git a/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh b/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh index f17c29ac7c..d35e90f1e4 100644 --- a/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh @@ -53,7 +53,8 @@ void sm120_fp8_paged_mqa_logits(const uint32_t batch_size, static constexpr uint32_t kNumMathWarps = kNumMathThreads / 32; static constexpr uint32_t kWarpsPerGroup = BLOCK_KV / MMA_M; static constexpr uint32_t kNumGroups = kNumMathWarps / kWarpsPerGroup; - static constexpr uint32_t kSwizzleMode = 128; + // Must match the TMA descriptor's swizzle (host passes swizzle_mode = head_dim) + static constexpr uint32_t kSwizzleMode = kHeadDim; static constexpr uint32_t kSMEMKBytes = kHeadDim; DG_STATIC_ASSERT(kNumTMAThreads == 128, "Expected 128 TMA threads"); diff --git a/tests/test_attention.py b/tests/test_attention.py index 1eb62e8b15..f5a8bd1f37 100644 --- a/tests/test_attention.py +++ b/tests/test_attention.py @@ -116,8 +116,12 @@ def enumerate_mqa_logits(): 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 + (512, 8192, [(64, 64), (64, 32)]), ): for num_heads, head_dim in head_cfgs: + if is_fp4 and head_dim != 128: + continue for disable_cp in (False, True): yield is_fp4, logits_dtype, compressed_logits, clean_logits, seq_len, seq_len_kv, num_heads, head_dim, disable_cp