[zcnrex]

Bring the SGL DeepGEMM fork up to upstream main, keeping the SGL wheel, its TORCH_LIBRARY binding and the downstream kernels. This is a single two-parent merge of upstream 1fc9c66 into dev b296ce3 (after #101 and #104), in the same shape as #86. The upstream side carries three commits, kept as-is:

• Fix Mega MoE task info slot release ordering deepseek-ai/DeepGEMM#441: Mega MoE task-info slot release ordering
• Public release 26/09/30 deepseek-ai/DeepGEMM#462: Public release 26/09/30
• Public release 10/14 deepseek-ai/DeepGEMM#471: FP8 epilogue FP8CastWithClamp and 128-byte paged sparse MQA strides

Breaking changes for sglang

Attention follows upstream's new API unchanged:

• fp8_mqa_logits, fp8_paged_mqa_logits and fp8_gemm_nt_skip_head_mid are removed.
• fp8_fp4_mqa_logits / fp8_fp4_paged_mqa_logits no longer take clean_logits or logits_dtype. Logits are compressed and not cleaned; the output is BF16 on SM100 and FP32 on SM90.
• SM100 paged MQA logits now require varlen indices with next_n == 1.
• get_bf16_mega_gate_config is removed upstream.

This is a format change, not only a rename: on SM100 the MQA logits now accept only MXFP8/MXFP4 inputs (packed UE8M0 scales), and plain E4M3 with an FP32 per-token scale is SM90-only. sglang's DSA index K cache stores E4M3 + FP32, so its indexers need porting before they can use a wheel built from this branch:

• kpool indexer (dsa_indexer_kpool.py): ported and tested on a B300 in zcnrex/sglang branch chunan/dg-mqa-api-260930. On SM100 the cache stores each power-of-two token scale as packed UE8M0, which is lossless and leaves the FP8 bytes unchanged.
• Main DSA indexer (dsa_indexer.py, DeepSeek-V3.2 / GLM-5) and DSV4.1 index_logits.py: not ported yet.

The SM120 MQA kernels (sm120_{fp4,fp8}_{,paged_}mqa_logits, smxx_clean_logits) only served the old attention API and are removed. Re-porting them to the new API needs SM120 hardware to validate.

What changed beyond the upstream merge

Mega MoE (SM100). Upstream adds locality-domain scheduling, 4D weight TMA, separate L1/L2 store blocks and a TMEM prefetch ring in the epilogue. Dev's features are carried over on top of that:

• NVFP4 / MXFP4: packed weights go through upstream's make_tma_weights_{2,3}d_desc on a byte view, using copy_nd<BLOCK_K_BYTES, …>.
• TRT-LLM weight layouts: kept on their own 2D/4D/5D descriptors. They are asserted to be non-localized.
• L1/L2 alphas, x-scales and SiTU: applied on the prefetched TMEM registers.
• FP8 combine: adopts upstream's shared-slot-first reduction order.
• The heuristic keeps dev's element-bit sizing, gran_k and global_num_tokens, and takes upstream's asymmetric FP8xFP4 store blocks.

GEMM epilogue classes. These are threaded through every SM100 dispatch. SM90 and SM120 assert that no custom epilogue is passed. The SM120 kernels drop their head-split epilogue template parameter, mirroring what upstream did for SM90 1d2d. Dev's use_padded_sf_layout k-grouped option is kept alongside upstream's packed SF strides.

TORCH_LIBRARY port of the new pybind-only features (csrc/torch_library_api.cpp):

• Epilogue classes: deep_gemm.epilogue.{Identity, Alpha, FP8CastWithClamp, BF16StochasticRounding} cross into every GEMM and einsum op as (kind, scalar, sfd, count). scalar is alpha or the clamp value. epilogue_sfd and epilogue_count are mutable schema arguments because FP8CastWithClamp writes them.
• Locality domains: get_num_locality_domains, localize, is_localized, destroy_localizer and the deep_gemm.locality_domain module are exported as ops. Allocations take a ScalarType, so FP8/FP4 work directly.
• MLOPart: _C stays CPython-free. deep_gemm/locality_domain.py installs ctypes callbacks into deep_gemm/utils/mlopart.py through set_mlopart_callbacks; a callback's Python exception is re-raised from the allocating call, so callers see the same errors as with upstream's pybind binding.
• shutdown is registered and hooked into atexit, as upstream does.
• deep_gemm/epilogue.py and deep_gemm/locality_domain.py are packaged into the wheel by build_sgl_deep_gemm.sh.

Tests.

• Upstream's test changes are 3-way merged into sgl_deep_gemm/tests, and every upstream test function is present.
• test_locality_domain.py is added and wired into run_tests.sh.
• test_attention.py is upstream's new file plus the SGL release-profile sampler, which takes the first and last case of each (format, dtype) pair.
• test_clean_logits_bounds.py is removed together with the API it tested.

Dev's per-expert post-norm (#104). Its normalization runs in upstream's restructured L2 epilogue and combine. #104's store block of 48 at block_m 192 becomes the L1 store block, and upstream's smaller FP8xFP4 L1/L2 store blocks at 240 are kept. DG_MEGA_MOE_TILES sets the L1 and L2 store blocks together.

Other. third-party/tilelang_ops/{ref_mhc,norm}.py are removed as upstream did, since nothing references them now. deep_jit stays at dev's 3732a3b, which already contains upstream's pin 2efdab4.

Validation

Built with build_sgl_deep_gemm.sh at the branch head 7db1164, installed into an isolated target dir, and tested on 8x B300 (SM103, CUDA 13.4, torch 2.13). sgl_deep_gemm/tests used the full profile, in parallel lanes:

• Pass: test_bf16, test_einsum (including Public release 10/14 deepseek-ai/DeepGEMM#471's clamp and clamped-count cases), test_fp8_fp4, test_hyperconnection, test_layout, test_legacy, test_attention, test_mega_gate, test_mega_mhc, test_lazy_init.
• Mega MoE pass: l1_fp4_accuracy, l1_sentinel, nvfp4_alphas, situ, situ_equal_transform, pre_dispatch, nvfp4_trtllm, bf16_trtllm, global_tokens.
• test_locality_domain: fails on this node because MLOPart is unavailable there (its MPS daemon reports one MLOPart device). It passed on another B300 at the pre-Public release 10/14 deepseek-ai/DeepGEMM#471 head, and Public release 10/14 deepseek-ai/DeepGEMM#471 does not touch the locality or MLOPart code.
• Multi-rank, 8x B300 with DeepEP at 1f88317: test_mega_moe.py passes its correctness check against the DeepEP baseline for FP8xFP4 (about 2.9 PFLOPS per rank, 1.7x the legacy-plus-shared baseline), FP8xFP8 (2.6 PFLOPS) and BF16 (1.3 PFLOPS). test_mega_moe_sgl.py also passes.
• [SM100] Support per-expert normalization for MegaMoE  #104's post-expert norm: test_mega_moe.py --post-expert-norm passes on 8 ranks for FP8xFP4, FP8xFP8 and BF16, which exercises the cross-rank partial-sum path. The no-gain and weight-first checks put the output RMS within 5e-4 of 1, and the standard form tracks rms(gain).
• Other mega MoE tests: l1_fp4_accuracy, l1_sentinel, nvfp4_alphas and global_tokens (4 ranks) pass at 1f88317. nvfp4_trtllm, bf16_trtllm, situ, situ_equal_transform and pre_dispatch pass on one rank.

Testing found three regressions, all fixed on this branch:

• The bf16_mega_gate wrapper defaulted scoring_func to 'identity', which upstream's gate no longer supports.
• The dev-only output-lifetime check in test_mega_gate.py assumed identity scoring.
• MLOPart allocation errors lost their Python exception when crossing the ctypes callback.

Still needs testing:

• test_mega_moe.py and test_mega_moe_sgl.py on a box where DeepEP's NCCL GIN works
• SM90 mega MoE and the SM120 GEMMs on their own hardware. Their host code builds, but their JIT kernels were never compiled or run on SM103

🤖 Generated with Claude Code