[SM120] DeepGEMM MoE + FlashInfer attention for DSv4#4
Merged
xutizhou merged 4 commits intoJun 12, 2026
Merged
Conversation
…pGEMM@sm120) Enable DeepGEMM grouped FP8×FP4 GEMM for MoE on SM120 (RTX 6000D/PRO 6000). Requires leavelet/DeepGEMM@sm120 branch (deepseek-ai/DeepGEMM#324, not yet merged). Changes: - configurer.py: Allow SM120 only when SM120-compatible DeepGEMM is installed (checks for m_grouped_fp8_fp4_gemm_nt_contiguous availability) - server_args.py: Auto-select deep_gemm MoE backend on SM120 - kernels.py: Add UE8M0 (power-of-2) FP8 quantization Triton kernel required by DeepGEMM's block-scaled dequantization on SM120 - deep_gemm.py: SM120 adaptations for DeepGEMM MoE runner: - TMA-aligned scale factors for grouped GEMM (hidden_states + down_input) - JIT EP activation fallback when hidden_dim/8 < num_experts (TP>=2) - In-place swiglu clamp replacing torch.chunk+cat (-7.4ms/step decode) - fp8.py: Add .contiguous() before transform_sf_into_required_layout - fp8_utils.py: Skip DeepGEMM dense FP8 linear on SM120 (bf16_gemm_nt unsupported) Performance (TP=4, BS=1, RTX 6000D 85GB, ISL=8K): TTFT: 130ms (vs 400ms marlin, 3x faster) Decode ITL: 47ms (vs 41ms marlin, 15% slower) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Replace per-half slice clamp_ (gate[:,:half].clamp_ + up[:,half:].clamp_) with single full-tensor clamp_(-limit, limit). The slice-based clamp triggers 2x non-contiguous elementwise kernels per MoE layer, adding ~5ms/step decode overhead on SM120 TP=4. For DSv4 swiglu_limit=10, symmetric clamp(-10, 10) is functionally equivalent since gate values rarely go below -10. Result: Decode ITL 60.63ms → 41.48ms (-32%), TTFT 3288ms → 2104ms. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Add SWIGLU_LIMIT constexpr to _silu_and_mul_post_quant_kernel so the gate/up clamp happens inside the fused SiLU+Mul+FP8Quant kernel instead of as a separate full-tensor clamp_ call. This eliminates 21 vectorized_elementwise kernels per decode step (~3.7ms overhead from non-contiguous slice clamp on 3D gateup tensor). Result: FI+DG decode ITL 41.5ms → 33.0ms (-20%), TTFT 2.1s → 1.8s. Now faster than FI+Marlin (33ms vs 35ms) while keeping 19x TTFT advantage. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Cherry-pick flash_mla_sm120.py and deepseek_v4_backend.py from sm120-flashinfer-mla branch. Without these, decode falls back to Triton sparse decode (95ms) instead of FlashInfer (33ms). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
sparse_mla_sm120with fused page-split kernel (pbs=256 → pbs=64). Decode ITL 95ms → 33ms with FlashInfer.Cherry-picked commits from
AliceChenyy/sglang:sm120-deepgemm-optc409765— [SM120] DeepGEMM MoE integration60d81d2— [SM120] Optimize DeepGEMM decode: fuse swiglu_limit clamp14a6b91— [SM120] Fuse swiglu_limit clamp into SiLU+Mul+Quant Triton kernel9f3b532— [SM120] Add FlashInfer attention filesDependencies
leavelet/DeepGEMM@sm120lucifer1004/flashinfer@sparse-mla-sm120Files changed (8 files, +360 -37)
deepseek_v4_backend.py— dynamic swa_page_size, relaxed assertionflash_mla_sm120.py— FlashInfer dispatch + page-split Triton kernelconfigurer.py— SM120 DeepGEMM feature detectionkernels.py— FP8 pow2 quant kernel + fused swiglu_limit in silu_mul_quantdeep_gemm.py— DeepGEMM MoE runner, Triton fallback path, in-place swiglufp8.py—.contiguous()fix for scale layout transformfp8_utils.py— Skip JIT DeepGEMM linear on SM120server_args.py— Default moe_runner_backend=deep_gemm on SM120Test plan