[zcnrex]

Purpose

DeepGEMM 0.2.5 (#103) changed the MQA logits APIs and, on SM100, made the kernel accept only MX scales: packed UE8M0, four per 128-element token. sglang main is pinned to 0.2.3.post2 and calls the 0.2.3 APIs on SM90, SM100 and SM120, with FP32 index-K scales by default. With this PR, sglang main can move to the new wheel without code changes and run exactly the MQA logits code it runs today. The 0.2.5 APIs and kernels are unchanged.

What 0.2.3 callers get

• fp8_mqa_logits / fp8_paged_mqa_logits (removed in 0.2.5) are restored on top of the _legacy functions.
• fp8_fp4_mqa_logits / fp8_fp4_paged_mqa_logits take the _legacy path when the call uses the 0.2.3 signature: a bool clean_logits right after the fixed arguments, or a clean_logits= / logits_dtype= keyword. A dense call without max_seqlen_k also counts, since 0.2.5 requires it. The 0.2.5 signatures never match these.
• get_paged_mqa_logits_metadata without indices returns the 0.2.3 schedule on SM100, where the 0.2.5 metadata requires indices. On SM90 and SM120 it keeps returning the 0.2.5 metadata. For every SM with work, it is the same as 0.2.3's; it differs only in the end-of-schedule entries for idle SMs, which both kernels skip. The legacy test checks this bitwise.

Ambiguous by construction: a 0.2.3-style paged call with only the seven fixed arguments, or a dense call passing just max_seqlen_k=, uses the 0.2.5 path. sglang main never makes either call.

_legacy APIs (0.2.3.post2 signatures, SM90 / SM100 / SM120)

Function	0.2.3 equivalent
fp8_fp4_mqa_logits_legacy(q, kv, weights, ks, ke, clean_logits=False, max_seqlen_k=0, logits_dtype=torch.float, schedule_meta=None)	fp8_fp4_mqa_logits
fp8_fp4_paged_mqa_logits_legacy(q, kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, clean_logits=False, logits_dtype=torch.float, indices=None)	fp8_fp4_paged_mqa_logits
get_mqa_logits_metadata_legacy(ks, ke, num_kv_tokens, num_heads)	get_mqa_logits_metadata
get_paged_mqa_logits_metadata_legacy(context_lens, block_kv, num_sms, indices=None)	get_paged_mqa_logits_metadata

These stay available for callers that want to choose explicitly, e.g. FP32 index-K pools after sgl-project/sglang#44350.

Implementation

• deep_gemm/include/deep_gemm/legacy/ and csrc/jit_kernels/impls/legacy/ hold 0.2.3.post2's MQA logits code for all three architectures. That covers SM90 dense and paged, SM120 FP8/FP4 dense and paged, and SM100 including the BF16 contiguous kernel from release [SM100] mqa_logits: faster bf16-weight paths, including the dev contiguous kernel for MX scales (release/v0.2.3) #117. The only changes are include paths, the deep_gemm::legacy namespace and a legacy_ prefix on JIT kernel names. SM90 and SM120 files are byte-identical between 0.2.3.post2 and 3cca390^1; only SM100 carried release-only changes.
• csrc/apis/attention.hpp: the four _legacy entry points are 0.2.3.post2's host functions for all architectures, calling the legacy:: kernels.
• The legacy kernels compile against the current shared headers; dev kept tma::copy as a wrapper over copy_nd, and every other symbol they use is unchanged.
• To remove later: delete the two legacy/ directories, test_attention_legacy.py and the wiring.

Size: about +7,900 lines, almost all copied 0.2.3.post2 code.

Test

4x B300 (SM103), CUDA 13.0, torch 2.14.1+cu130, DG_TEST_PROFILE=release:

• test_attention_legacy.py: pass. Each case is also run through the 0.2.3 call shapes (fp8_mqa_logits, positional clean_logits, the shared metadata name) and must match the _legacy call bitwise. The JIT cache shows legacy_sm100_bf16_mqa_logits built, so the [SM100] mqa_logits: faster bf16-weight paths, including the dev contiguous kernel for MX scales (release/v0.2.3) #117 fast path runs.

• test_attention.py (0.2.5 APIs): pass, and it builds no legacy_* kernels.

• sglang main (38d9e4a3fe) DeepGEMM-related suites under 0.2.3.post2 and under this wheel: identical results in all 16 files. That covers 192 test_deepgemm_paged_mqa_logits, 192 CuteDSL paged, 95 FP4 indexer, DSA/DSV4 indexers, mega-MoE, FP8 linear and mHC. The 14 test_paged_mqa_metadata failures and 1 test_dsa failure fail identically on 0.2.3.post2.

• E2E, sglang main 38d9e4a3fe on 4x B300 (TP4), 0.2.3.post2 vs this wheel. Each run does GSM8K (few_shot_gsm8k, 1319 questions, run twice) and bench_serving (random 8192 in / 256 out, 256 prompts, concurrency 32):

  Model	Wheel	GSM8K	bench tok/s	TTFT ms	TPOT ms	MQA kernels compiled
  DeepSeek-V4-Flash	0.2.3.post2	0.951 / 0.957	645.6	1712	42.45	sm100_paged_mqa_logits
  DeepSeek-V4-Flash	this PR	0.955 / 0.954	670.5	1643	41.27	legacy_sm100_paged_mqa_logits
  DeepSeek-V3.2	0.2.3.post2	0.950 / 0.950	502.9	6486	38.41	sm100_mqa_logits, sm100_paged_mqa_logits_metadata
  DeepSeek-V3.2	this PR	0.951 / 0.950	503.5	6473	38.39	legacy_sm100_mqa_logits, legacy_sm100_paged_mqa_logits_metadata

  V4-Flash: flashinfer_mxfp4 MoE, EAGLE 3/1/4, chunked prefill 4096, reasoning effort pinned to low. No server errors in either run. With this wheel, sglang main ran only legacy_* MQA kernels.

Not run here: SM90 and SM120 (no box available; the release workflow's test job covers H200).

🤖 Generated with Claude Code