Skip to content

SM120: fix FP8 MQA logits swizzle mode for head_dim < 128 #379

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
leavelet wants to merge 1 commit into
base: nv_dev
Choose a base branch
from
+8 −2
Open

SM120: fix FP8 MQA logits swizzle mode for head_dim < 128 #379

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion deep_gemm/include/deep_gemm/impls/sm120_fp8_mqa_logits.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Comment on lines +61 to +62

@ds-review-bot ds-review-bot Jul 15, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔴 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

static constexpr uint32_t kSwizzleAlignment = kHeadDim * 8;
static constexpr uint32_t kSMEMKBytes = kHeadDim;

Expand Down
3 changes: 2 additions & 1 deletion deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand Down
4 changes: 4 additions & 0 deletions tests/test_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

@ds-review-bot ds-review-bot Jul 15, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔵 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

(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

Expand Down