[PR #379] SM120: fix FP8 MQA logits swizzle mode for head_dim < 128 #400

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/379
Author: @leavelet
Created: 7/15/2026
Status: 🔄 Open

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

📝 Commits (1)

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

📊 Changes

3 files changed (+8 additions, -2 deletions)

View changed files

📝 deep_gemm/include/deep_gemm/impls/sm120_fp8_mqa_logits.cuh (+2 -1)
📝 deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh (+2 -1)
📝 tests/test_attention.py (+4 -0)

📄 Description

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).

_{🔄 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/379 **Author:** [@leavelet](https://github.com/leavelet) **Created:** 7/15/2026 **Status:** 🔄 Open **Base:** `nv_dev` ← **Head:** `sm120-mqa-swizzle-fix` --- ### 📝 Commits (1) - [`2284689`](https://github.com/deepseek-ai/DeepGEMM/commit/2284689a3b7f529fb53744a6745dcc5d8c247807) SM120: fix FP8 MQA logits swizzle mode for head_dim < 128 ### 📊 Changes **3 files changed** (+8 additions, -2 deletions) <details> <summary>View changed files</summary> 📝 `deep_gemm/include/deep_gemm/impls/sm120_fp8_mqa_logits.cuh` (+2 -1) 📝 `deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh` (+2 -1) 📝 `tests/test_attention.py` (+4 -0) </details> ### 📄 Description 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). --- <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:15 +00:00

bsakaguchi referenced this issue

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

[PR #400] [CLOSED] feat(native): add standalone CUDA core #420