[zyongye]

Merges upstream's public release 26/09/30 (deepseek-ai#462) into dev. It brings in locality-aware Mega MoE, the GEMM epilogue classes with BF16 stochastic rounding, the reworked Lightning and Sparse Indexer kernels, and upstream's MQA API changes.

This changes the MQA API, so vLLM must update its indexer before bumping its DeepGEMM pin to this commit (details below).

Upstream features on the TORCH_LIBRARY bindings

Upstream adds its new APIs through pybind11, which dev no longer links. They are ported as follows:

• Epilogue classes: deep_gemm.epilogue.{Identity, Alpha, FP8Quantization, BF16StochasticRounding}(...) build a deep_gemm::Epilogue custom class through factory ops. All GEMM and einsum ops take it as epilogue=None. The objects are torch.ScriptObjects, so they don't work with isinstance.
• Locality domains: the helpers are locality_domain_* ops, exposed as deep_gemm._C.locality_domain, so upstream's deep_gemm/locality_domain.py works unchanged.
• MLOPart: the C++ side calls deep_gemm.utils.mlopart through the CPython limited API (taking the GIL itself) instead of pybind11. Python errors from it now surface as RuntimeError rather than AssertionError, and test_locality_domain accepts either.
• shutdown: a config op; deep_gemm/__init__.py registers it with atexit as upstream does.
• Stub generator: scripts/generate_pyi.py maps custom-class schema types to Any.

MQA: upstream API adopted

• Removed: fp8_mqa_logits, fp8_paged_mqa_logits, fp8_gemm_nt_skip_head_mid, and the clean_logits and logits_dtype arguments.
• fp8_fp4_mqa_logits requires max_seqlen_k > 0 and returns compressed logits. Callers mask entries beyond each row's valid KV span.
• SM100 takes MXFP8 or MXFP4 inputs with BF16 weights and returns BF16 logits. Paged MQA on SM100 is varlen-only (indices, next_n=1).
• SM120 needs no kernel changes. It keeps FP8 with per-token float KV scales or MXFP4, FP32 weights and FP32 logits. Paged queries can be fixed-length or varlen. The SM120 kernels already had the compressed-logits path.
• SM90 keeps dev's 32-row pages, next_n=4 (two-CTA cluster), and the zero-context scheduler test.

dev changes kept on top of upstream

• SM120 GEMMs move to the renamed epilogue::operators. The head-split index remap goes away with fp8_gemm_nt_skip_head_mid, as does test_sm120_skip_head.py.
• SiTU Mega MoE is regenerated from upstream's new kernel plus dev's SiTU diff, so it also gets the epilogue prefetch and locality-aware scheduling.
• NVFP4 Mega MoE uses the new scheduler with a single locality domain (no localization).
• SwiGLU alpha/beta is re-applied on top of upstream's prefetching epilogue.
• Mega Gate keeps Remove Mega-Gate finite-score device trap #11's removal of the finite-score device trap, which upstream still has.
• Upstream behaviour now on dev:
  • get_bf16_mega_gate_config is removed, and bf16_mega_gate defaults to scoring_func='sqrtsoftplus'.
  • FP8×FP4 and BF16 Mega MoE assert that the configured SM count equals the device's.

Testing

Run on 4× GB300 (SM103), CUDA 13.4, torch 2.12.1:

Test	Result
test_bf16, test_fp8_fp4, test_einsum, test_layout, test_hyperconnection, test_mega_gate, test_mega_mhc, test_legacy	pass
test_attention (DG_MQA_NUM_CASES=40: dense, paged, sparse)	pass
test_lazy_init --num-processes 2, test_locality_domain (MLOPart available)	pass
test_mega_moe EP4: FP8×FP4 SwiGLU with alpha=1.702, beta=1.0, clamp=7	pass
test_mega_moe EP4: FP8×FP4 with --use-localization	pass
test_mega_moe EP4: FP8×FP4 SiTU	pass
test_mega_moe EP4: BF16	pass
test_nvfp4_mega_moe EP4	pass

Not run (no hardware): SM90 and SM120. Their kernel headers compile with nvcc -arch=sm_90a / sm_120a, and the host code builds. Someone with those GPUs should run test_attention, test_fp8_fp4, test_einsum and test_sm120_* before this merges.

vLLM follow-up

Before vLLM bumps its DeepGEMM pin to this commit, its indexer (sparse_attn_indexer.py) must:

• drop clean_logits and pass max_seqlen_k;
• handle compressed logits;
• on SM100, pass MXFP8 or MXFP4 with BF16 weights and use the varlen paged interface.

🤖 Generated with Claude Code