[PR #434] fix: derive the SM120 fp8 mqa_logits swizzle mode from head_dim #445

New issue

Open

opened 2026-10-26 21:17:57 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:17:57 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/434
Author: @aganhui
Created: 9/10/2026
Status: 🔄 Open

Base: nv_dev ← Head: fix/sm120-fp8-mqa-swizzle

📝 Commits (2)

184190b fix: match the SM120 fp8 mqa_logits read-side swizzle mode to the TMA descriptors (#433)
5825212 Address review: fix the fp8 paged variant too, add swizzle-mode static asserts (#433)

📊 Changes

2 files changed (+14 additions, -2 deletions)

View changed files

📝 deep_gemm/include/deep_gemm/impls/sm120_fp8_mqa_logits.cuh (+8 -1)
📝 deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh (+6 -1)

📄 Description

Fixes #433.

The SM120 fp8 MQA-logits readers un-swizzle shared memory with a hardcoded 128-byte pattern, while the host TMA descriptors and the tma::copy sites swizzle in head_dim-byte atoms (32B/64B/128B for D=32/64/128 — the fp8 row is head_dim bytes). For head_dim 32/64 the XOR bits don't match the write pattern, fragment loads address outside their row, and the kernel returns in-window logits that are silently wrong. D=128 is the only head_dim where the hardcoded value matches, which is why the main DSA path works.

This derives the mode from kHeadDim, matching the sibling kernels: SM90 (to_swizzle_cute_type<kHeadDim>()) and SM100 (kHeadDim / kPackFactor).

Both fp8 variants are fixed (second commit, addressing the review finding): the fp8 paged dispatch has no head_dim == 128 restriction — only its fp4 sibling does — so the paged kernel had the same mismatch. Notably, in the paged variant the resulting garbage is deterministic (misreads land on deterministically-stale shared memory), so self-consistency passes and only the accuracy comparison against the reference exposes it.

Validation on RTX 5090 D (CUDA 13.0, PyTorch 2.11.0+cu130, nv_dev @ 2642b32):

Dense, 384-config self-consistency scan: 189 failures → 0 (all failures were fp8 with D ∈ {32, 64}; zero at D=128, zero mxfp4)
Dense, full test_mqa_logits (20x bitwise + accuracy vs reference + bench): PASSED, all 384 configs, fp8 up to 411 TFLOPS
Paged, targeted fp8 repro (small configs; the full paged suite does not fit in 32 GB): accuracy diff vs reference D=32: 0.83 → 0.0008, D=64: 0.72 → 0.001, D=128 unchanged (~0.001)
Dense regression after adding the static asserts: fp8 D=32 accuracy diff 0.0006 (< 1e-3)

Also adds DG_STATIC_ASSERT(kHeadDim ∈ {32, 64, 128}) next to the derived constants (per review suggestion), so a future dispatch widening cannot silently instantiate an invalid CuTeSwizzle.

Unaffected kernels (verified): dense fp4 and both fp4 paged variants are head_dim-128-only (DG_HOST_ASSERT/DG_STATIC_ASSERT), where their constants already match. Dense fp8 and fp8 paged are the only instantiations whose dispatches allow multiple row widths, and both are now derived.

Scan script and raw logs: https://gist.github.com/aganhui/f783450b212217284c062c3b46300649

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepGEMM/pull/434 **Author:** [@aganhui](https://github.com/aganhui) **Created:** 9/10/2026 **Status:** 🔄 Open **Base:** `nv_dev` ← **Head:** `fix/sm120-fp8-mqa-swizzle` --- ### 📝 Commits (2) - [`184190b`](https://github.com/deepseek-ai/DeepGEMM/commit/184190b8a39287cd1e9d570ce05a0d5dee4bc224) fix: match the SM120 fp8 mqa_logits read-side swizzle mode to the TMA descriptors (#433) - [`5825212`](https://github.com/deepseek-ai/DeepGEMM/commit/58252126017ca9b5477760b998abdab6d74e4b23) Address review: fix the fp8 paged variant too, add swizzle-mode static asserts (#433) ### 📊 Changes **2 files changed** (+14 additions, -2 deletions) <details> <summary>View changed files</summary> 📝 `deep_gemm/include/deep_gemm/impls/sm120_fp8_mqa_logits.cuh` (+8 -1) 📝 `deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh` (+6 -1) </details> ### 📄 Description Fixes #433. The SM120 fp8 MQA-logits readers un-swizzle shared memory with a hardcoded 128-byte pattern, while the host TMA descriptors and the `tma::copy` sites swizzle in head_dim-byte atoms (32B/64B/128B for D=32/64/128 — the fp8 row is head_dim bytes). For head_dim 32/64 the XOR bits don't match the write pattern, fragment loads address outside their row, and the kernel returns in-window logits that are silently wrong. D=128 is the only head_dim where the hardcoded value matches, which is why the main DSA path works. This derives the mode from `kHeadDim`, matching the sibling kernels: SM90 (`to_swizzle_cute_type<kHeadDim>()`) and SM100 (`kHeadDim / kPackFactor`). **Both fp8 variants are fixed** (second commit, addressing the review finding): the fp8 paged dispatch has no `head_dim == 128` restriction — only its fp4 sibling does — so the paged kernel had the same mismatch. Notably, in the paged variant the resulting garbage is *deterministic* (misreads land on deterministically-stale shared memory), so self-consistency passes and only the accuracy comparison against the reference exposes it. Validation on RTX 5090 D (CUDA 13.0, PyTorch 2.11.0+cu130, nv_dev @ 2642b32): - Dense, 384-config self-consistency scan: **189 failures → 0** (all failures were fp8 with D ∈ {32, 64}; zero at D=128, zero mxfp4) - Dense, full `test_mqa_logits` (20x bitwise + accuracy vs reference + bench): **PASSED**, all 384 configs, fp8 up to 411 TFLOPS - Paged, targeted fp8 repro (small configs; the full paged suite does not fit in 32 GB): accuracy diff vs reference **D=32: 0.83 → 0.0008, D=64: 0.72 → 0.001**, D=128 unchanged (~0.001) - Dense regression after adding the static asserts: fp8 D=32 accuracy diff 0.0006 (< 1e-3) Also adds `DG_STATIC_ASSERT(kHeadDim ∈ {32, 64, 128})` next to the derived constants (per review suggestion), so a future dispatch widening cannot silently instantiate an invalid `CuTeSwizzle`. Unaffected kernels (verified): dense fp4 and both fp4 paged variants are head_dim-128-only (`DG_HOST_ASSERT`/`DG_STATIC_ASSERT`), where their constants already match. Dense fp8 and fp8 paged are the only instantiations whose dispatches allow multiple row widths, and both are now derived. Scan script and raw logs: https://gist.github.com/aganhui/f783450b212217284c062c3b46300649 --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-26 21:17:57 +00:00

bsakaguchi referenced this issue

2026-10-26 21:18:00 +00:00

[PR #445] fix(sm90): support parameterized SwiGLU in FP8 MegaMoE #451