Skip to content

[Bug] DeepGEMM sm100_fp8_mqa_logits kernel fails to compile for models with < 52 attention heads on Blackwell GPUs #19529

Description

@nvelingker

Checklist

  • I searched related issues but found no solution.
  • The bug persists in the latest version.
  • Issues without environment info and a minimal reproducible demo are hard to resolve and may receive no feedback.
  • If this is not a bug report but a general question, please start a discussion at https://github.com/sgl-project/sglang/discussions. Otherwise, it will be closed.
  • Please use English. Otherwise, it will be closed.

Describe the bug

Hey all. I encountered a crash when running zai-org/GLM-5-FP8 on an NVIDIA DGX B200. I had Codex debug it, and I hope it can be useful:

When running FP8 models with Native Sparse Attention (NSA) on Blackwell GPUs (SM100a), the DeepGEMM JIT compilation of sm100_fp8_mqa_logits fails with a static assertion error if the model has fewer than 52 attention heads.

The root cause is a hardcoded constant in deep_gemm/include/deep_gemm/impls/sm100_fp8_mqa_logits.cuh
line 281:

constexpr uint32_t kNumWeightsInReg = 52;
float weights[BLOCK_Q][kNumWeightsInReg];
DG_STATIC_ASSERT(kNumWeightsInReg <= kNumHeads and kNumWeightsInReg % 4 == 0, "Invalid kNumWeightsInReg");

This assumes kNumHeads >= 52, which holds for DeepSeek-V3 (128 heads) but fails for models like zai-org/GLM-5-FP8 which has kNumHeads=32.

NVCC error output:

sm100_fp8_mqa_logits.cuh(283): error: static assertion failed with "Invalid kNumWeightsInReg"
          static_assert(kNumWeightsInReg <= kNumHeads and kNumWeightsInReg % 4 == 0, "Invalid kNumWeightsInReg");
          ^
          detected during instantiation of "void deep_gemm::sm100_fp8_mqa_logits<...>(...)"
          [with kNumHeads=32U, kHeadDim=128U, ...]

Steps to reproduce

python -m sglang.launch_server \
  --model zai-org/GLM-5-FP8 \
  --tp 8 --dp 8 \
  --enable-dp-attention \
  --mem-fraction-static 0.9

The server launches and generates successfully until the NSA indexer path is triggered, at which point the fp8_mqa_logits kernel is JIT-compiled and fails.

Root cause analysis
The kNumHeads template parameter in sm100_fp8_mqa_logits does not correspond to the model's main attention head count (num_attention_heads), but rather to the NSA (Native Sparse Attention) indexer head count index_n_heads, which is a separate, typically smaller value in the model config.

For zai-org/GLM-5-FP8, the config has:

{
  "num_attention_heads": 64,
  "num_key_value_heads": 64,
  "index_n_heads": 32
}

The flow is:

  1. sglang/srt/configs/model_config.py::get_nsa_index_n_heads() reads config.index_n_heads32
  2. This is passed to NSAIndexer.__init__() as index_n_headsself.n_heads = 32
  3. The indexer calls deep_gemm.fp8_mqa_logits(q_fp8, ...) where q_fp8 has shape [seq_len, 32, head_dim]
  4. DeepGEMM instantiates sm100_fp8_mqa_logits<kNumHeads=32, kHeadDim=128, ...>
  5. The hardcoded kNumWeightsInReg = 52 fails the assertion 52 <= 32

The constant kNumWeightsInReg = 52 appears to have been tuned for DeepSeek-V3, which has index_n_heads = 128. Any NSA-enabled model with index_n_heads < 52 will hit this assertion. Since kNumWeightsInReg is purely a register-vs-shared-memory optimization split (the remaining kNumHeads - kNumWeightsInReg weights are read from shared memory), clamping it to min(kNumHeads, 52) is both safe and slightly more efficient for small head counts.

Suggested fix:

Change line 281 in deep_gemm/include/deep_gemm/impls/sm100_fp8_mqa_logits.cuh
from:

constexpr uint32_t kNumWeightsInReg = 52;

to:

constexpr uint32_t kNumWeightsInReg = (kNumHeads < 52) ? (kNumHeads / 4 * 4) : 52;

This dynamically clamps the register preload count to kNumHeads (rounded down to a multiple of 4) when the model has fewer than 52 heads, preserving the existing behavior for models with ≥ 52 heads.

Reproduction

Steps to reproduce

python -m sglang.launch_server \
  --model zai-org/GLM-5-FP8 \
  --tp 8 --dp 8 \
  --enable-dp-attention \
  --mem-fraction-static 0.9

Environment

[redacted]$ python3 -m sglang.check_env
nvcc warning : incompatible redefinition for option 'compiler-bindir', the last value of this option was used
Python: 3.11.14 | packaged by conda-forge | (main, Jan 26 2026, 23:48:32) [GCC 14.3.0]
CUDA available: True
GPU 0,1,2,3,4,5,6,7: NVIDIA B200
GPU 0,1,2,3,4,5,6,7 Compute Capability: 10.0
CUDA_HOME: [redacted]/.local/share/mamba/envs/sglang
NVCC: Cuda compilation tools, release 12.8, V12.8.93
CUDA Driver Version: 580.95.05
PyTorch: 2.9.1+cu128
sglang: 0.5.9
sgl_kernel: 0.3.21
flashinfer_python: 0.6.3
flashinfer_cubin: 0.6.3
flashinfer_jit_cache: Module Not Found
triton: 3.5.1
transformers: 5.2.0
torchao: 0.9.0
numpy: 2.4.2
aiohttp: 3.13.3
fastapi: 0.133.1
hf_transfer: 0.1.9
huggingface_hub: 1.5.0
interegular: 0.3.3
modelscope: 1.34.0
orjson: 3.11.7
outlines: 0.1.11
packaging: 26.0
psutil: 7.2.2
pydantic: 2.12.5
python-multipart: 0.0.22
pyzmq: 27.1.0
uvicorn: 0.41.0
uvloop: 0.22.1
vllm: Module Not Found
xgrammar: 0.1.27
openai: 2.6.1
tiktoken: 0.12.0
anthropic: 0.84.0
litellm: Module Not Found
decord2: 3.0.0
NVIDIA Topology:
GPU0 GPU1 GPU2 GPU3 GPU4 GPU5 GPU6 GPU7 NIC0 NIC1 NIC2 NIC3 NIC4 NIC5 NIC6 NIC7 NIC8 NIC9 NIC10 NIC11 NIC12 NIC13 NIC14 NIC15 CPU Affinity NUMA Affinity GPU NUMA ID
GPU0 X NV18 NV18 NV18 NV18 NV18 NV18 NV18 NODE NODE NODE NODE PXB NODE NODE NODE NODE NODE SYS SYS SYS SYS SYS SYS 0-27,112-139 0 N/A
GPU1 NV18 X NV18 NV18 NV18 NV18 NV18 NV18 NODE NODE NODE NODE NODE NODE NODE PXB NODE NODE SYS SYS SYS SYS SYS SYS 0-27,112-139 0 N/A
GPU2 NV18 NV18 X NV18 NV18 NV18 NV18 NV18 NODE NODE NODE NODE NODE NODE NODE NODE PXB NODE SYS SYS SYS SYS SYS SYS 0-27,112-139 0 N/A
GPU3 NV18 NV18 NV18 X NV18 NV18 NV18 NV18 NODE NODE NODE NODE NODE NODE NODE NODE NODE PXB SYS SYS SYS SYS SYS SYS 0-27,112-139 0 N/A
GPU4 NV18 NV18 NV18 NV18 X NV18 NV18 NV18 SYS SYS SYS SYS SYS SYS SYS SYS SYS SYS PXB NODE NODE NODE NODE NODE 56-83,168-195 1 N/A
GPU5 NV18 NV18 NV18 NV18 NV18 X NV18 NV18 SYS SYS SYS SYS SYS SYS SYS SYS SYS SYS NODE NODE NODE PXB NODE NODE 56-83,168-195 1 N/A
GPU6 NV18 NV18 NV18 NV18 NV18 NV18 X NV18 SYS SYS SYS SYS SYS SYS SYS SYS SYS SYS NODE NODE NODE NODE PXB NODE 56-83,168-195 1 N/A
GPU7 NV18 NV18 NV18 NV18 NV18 NV18 NV18 X SYS SYS SYS SYS SYS SYS SYS SYS SYS SYS NODE NODE NODE NODE NODE PXB 56-83,168-195 1 N/A
NIC0 NODE NODE NODE NODE SYS SYS SYS SYS X PIX PIX PIX NODE NODE NODE NODE NODE NODE SYS SYS SYS SYS SYS SYS
NIC1 NODE NODE NODE NODE SYS SYS SYS SYS PIX X PIX PIX NODE NODE NODE NODE NODE NODE SYS SYS SYS SYS SYS SYS
NIC2 NODE NODE NODE NODE SYS SYS SYS SYS PIX PIX X PIX NODE NODE NODE NODE NODE NODE SYS SYS SYS SYS SYS SYS
NIC3 NODE NODE NODE NODE SYS SYS SYS SYS PIX PIX PIX X NODE NODE NODE NODE NODE NODE SYS SYS SYS SYS SYS SYS
NIC4 PXB NODE NODE NODE SYS SYS SYS SYS NODE NODE NODE NODE X NODE NODE NODE NODE NODE SYS SYS SYS SYS SYS SYS
NIC5 NODE NODE NODE NODE SYS SYS SYS SYS NODE NODE NODE NODE NODE X PIX NODE NODE NODE SYS SYS SYS SYS SYS SYS
NIC6 NODE NODE NODE NODE SYS SYS SYS SYS NODE NODE NODE NODE NODE PIX X NODE NODE NODE SYS SYS SYS SYS SYS SYS
NIC7 NODE PXB NODE NODE SYS SYS SYS SYS NODE NODE NODE NODE NODE NODE NODE X NODE NODE SYS SYS SYS SYS SYS SYS
NIC8 NODE NODE PXB NODE SYS SYS SYS SYS NODE NODE NODE NODE NODE NODE NODE NODE X NODE SYS SYS SYS SYS SYS SYS
NIC9 NODE NODE NODE PXB SYS SYS SYS SYS NODE NODE NODE NODE NODE NODE NODE NODE NODE X SYS SYS SYS SYS SYS SYS
NIC10 SYS SYS SYS SYS PXB NODE NODE NODE SYS SYS SYS SYS SYS SYS SYS SYS SYS SYS X NODE NODE NODE NODE NODE
NIC11 SYS SYS SYS SYS NODE NODE NODE NODE SYS SYS SYS SYS SYS SYS SYS SYS SYS SYS NODE X PIX NODE NODE NODE
NIC12 SYS SYS SYS SYS NODE NODE NODE NODE SYS SYS SYS SYS SYS SYS SYS SYS SYS SYS NODE PIX X NODE NODE NODE
NIC13 SYS SYS SYS SYS NODE PXB NODE NODE SYS SYS SYS SYS SYS SYS SYS SYS SYS SYS NODE NODE NODE X NODE NODE
NIC14 SYS SYS SYS SYS NODE NODE PXB NODE SYS SYS SYS SYS SYS SYS SYS SYS SYS SYS NODE NODE NODE NODE X NODE
NIC15 SYS SYS SYS SYS NODE NODE NODE PXB SYS SYS SYS SYS SYS SYS SYS SYS SYS SYS NODE NODE NODE NODE NODE X

Legend:

X = Self
SYS = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)
NODE = Connection traversing PCIe as well as the interconnect between PCIe Host Bridges within a NUMA node
PHB = Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU)
PXB = Connection traversing multiple PCIe bridges (without traversing the PCIe Host Bridge)
PIX = Connection traversing at most a single PCIe bridge
NV# = Connection traversing a bonded set of # NVLinks

NIC Legend:

NIC0: mlx5_0
NIC1: mlx5_1
NIC2: mlx5_2
NIC3: mlx5_3
NIC4: mlx5_4
NIC5: mlx5_5
NIC6: mlx5_6
NIC7: mlx5_7
NIC8: mlx5_8
NIC9: mlx5_9
NIC10: mlx5_10
NIC11: mlx5_11
NIC12: mlx5_12
NIC13: mlx5_13
NIC14: mlx5_14
NIC15: mlx5_15

ulimit soft: 131072

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions