feat: add sm120 support for DeepGEMM - #324
Merged
Merged
Conversation
Open
leavelet
changed the title
This was referenced May 7, 2026
leavelet
marked this pull request as ready for review
May 9, 2026 04:37
|
@leavelet this is nice, after it is merged into the nv-dev branch, should vllm-project/vllm#41834 merge too in order for vllm to be able to work with the branch. |
Contributor
|
I add my benchmark at vllm-project/vllm#41834 as references |
jasl
added a commit
to jasl/tokenspeed
that referenced
this pull request
May 12, 2026
Replace the hand-written CUDA FP8 GEMV kernel (previously gated to tokens==1) with a port of the SM120 FP8 einsum kernel from upstream DeepGEMM's WIP SM120 support (deepseek-ai/DeepGEMM#324, file `deep_gemm/include/deep_gemm/impls/sm120_fp8_einsum.cuh`). The DeepGEMM kernel implements exactly the `bhr,hdr->bhd` einsum DeepSeek V4 needs, with per-thread-per-output-cell GEMV using fp8x4 vectorized loads and the same block-128 fp32 scale recipe. Removing the tokens==1 gate: the kernel handles all token counts that the SM12x dispatch predicate accepts (tokens <= 16 today; larger token batches will arrive once T1-α expands graph capture). Microbench (DSv4-Flash decode shape, groups=8, hidden=2048, out=1024, GPU idle): tokens=1 cuda 0.026ms triton 0.020ms speedup 0.72x (was 0.40x) tokens=2 cuda 0.027ms triton 0.020ms speedup 0.72x (was 0.73x*) tokens=8 cuda 0.075ms triton 0.020ms speedup 0.27x (was 0.21x*) * Triton-as-default after the previous tokens==1 hotfix. The kernel's grid is `tokens * groups * (out/128)`, one block per `(token, group, out_tile=128)` triple. Because each block reads its weight tile independently, total weight reads scale linearly with `num_tokens`. At graph bs=2 (today) this dominates: tokens<=2 is the production shape and the 0.72x is a real net win against the previous Triton-fallback default. At tokens=8 (future, post-T1-α) the kernel loses ~2x to Triton's m=16 cooperative tile; we will revisit with a multi-token tile design before T1-α exposes that shape to production. Earlier hand-written attempts (one-cell-per-block, per-thread B=16 accumulator tile, 1-warp m16n8 MMA, 4-warp m16n32 cooperative MMA, 4-warp m16n128 MMA) are documented in `docs/notes/2026-05-09-ds4-sm12x-rejected-experiments.md`. The MMA designs hit either occupancy collapse (80 regs/thread) or insufficient parallelism (64 blocks at decode shape vs Blackwell's 140 SMs), capping out at ~0.51x. The DeepGEMM design wins at the production shape by avoiding tensor cores entirely -- a per-thread GEMV with fp8x4 vectorization and L1/L2-friendly weight access fits the small-M decode profile better than the m=16 MMA tile. Attribution: kernel source ported under MIT license from upstream DeepGEMM (Copyright (c) 2025 DeepSeek). Tokenspeed adaptations are the tvm-ffi binding, stride/scale validation, and the SM12x dispatch integration; the dot-product math is unchanged. Signed-off-by: jasl <jasl9187@hotmail.com>
jasl
added a commit
to jasl/tokenspeed
that referenced
this pull request
May 12, 2026
Upstream PR lightseekorg#93 added a pre-flight DeepGEMM ``fp8_gemm_nt`` call to ``DeepseekV4Attention._compute_qr_kv``: on success it replaces the reference FP8 linear path, on failure it logs a WARNING per layer and falls back. DeepGEMM does not support SM120/SM121 yet (see PR ``deepseek-ai/DeepGEMM#324`` + ``reference_deepgemm_sm120`` memory), so on the RTX Pro 6000 workstation every layer fires: DeepSeek V4 DeepGEMM FP8 linear failed; falling back to reference FP8 linear. reason=RuntimeError: Assertion error (csrc/apis/layout.hpp:59): Unknown SF transformation The existing per-layer ``_deepseek_v4_deep_gemm_linear_disabled`` flag already catches this for steady-state replay, but it costs one failed call + one WARNING per layer at boot. Mirror the pattern used by ``_deepseek_v4_deepgemm_fp4_indexer_enabled_for_platform``: short- circuit ``_deepseek_v4_get_fp8_linear_deep_gemm`` to ``None`` on SM12x so the platform never tries the DeepGEMM path. Non-SM12x platforms keep the new fast path. Signed-off-by: jasl <jasl9187@hotmail.com>
jasl
added a commit
to jasl/vllm
that referenced
this pull request
May 18, 2026
Widen the direct FP8 MQA logits Triton fallback from BLOCK_M=8 to BLOCK_M=16 while keeping BLOCK_N=128 and the existing 4-warp launch. This reduces CTA count for late-context prefill without introducing a runtime switch. The direction was motivated by the tile-shape discussion in deepseek-ai/DeepGEMM#324, but this is a vLLM-owned Triton fallback adjustment and does not copy DeepGEMM code. On the SM120 long-context gate with prefix cache disabled, the 128K synthetic mean TTFT improved from 36.541s to 33.264s at C=1, 56.902s to 49.199s at C=2, and 96.317s to 82.181s at C=4. GSM8K exact_match_flexible stayed at 0.95. Signed-off-by: jasl <jasl9187@hotmail.com>
jasl
added a commit
to jasl/vllm
that referenced
this pull request
May 18, 2026
Widen the direct FP8 MQA logits Triton fallback from BLOCK_M=8 to BLOCK_M=16 while keeping BLOCK_N=128 and the existing 4-warp launch. This reduces CTA count for late-context prefill without introducing a runtime switch. The direction was motivated by the tile-shape discussion in deepseek-ai/DeepGEMM#324, but this is a vLLM-owned Triton fallback adjustment and does not copy DeepGEMM code. On the SM120 long-context gate with prefix cache disabled, the 128K synthetic mean TTFT improved from 36.541s to 33.264s at C=1, 56.902s to 49.199s at C=2, and 96.317s to 82.181s at C=4. GSM8K exact_match_flexible stayed at 0.95. Signed-off-by: jasl <jasl9187@hotmail.com>
jasl
added a commit
to jasl/vllm
that referenced
this pull request
May 19, 2026
Widen the direct FP8 MQA logits Triton fallback from BLOCK_M=8 to BLOCK_M=16 while keeping BLOCK_N=128 and the existing 4-warp launch. This reduces CTA count for late-context prefill without introducing a runtime switch. The direction was motivated by the tile-shape discussion in deepseek-ai/DeepGEMM#324, but this is a vLLM-owned Triton fallback adjustment and does not copy DeepGEMM code. On the SM120 long-context gate with prefix cache disabled, the 128K synthetic mean TTFT improved from 36.541s to 33.264s at C=1, 56.902s to 49.199s at C=2, and 96.317s to 82.181s at C=4. GSM8K exact_match_flexible stayed at 0.95. Signed-off-by: jasl <jasl9187@hotmail.com>
jasl
added a commit
to jasl/vllm
that referenced
this pull request
May 19, 2026
Widen the direct FP8 MQA logits Triton fallback from BLOCK_M=8 to BLOCK_M=16 while keeping BLOCK_N=128 and the existing 4-warp launch. This reduces CTA count for late-context prefill without introducing a runtime switch. The direction was motivated by the tile-shape discussion in deepseek-ai/DeepGEMM#324, but this is a vLLM-owned Triton fallback adjustment and does not copy DeepGEMM code. On the SM120 long-context gate with prefix cache disabled, the 128K synthetic mean TTFT improved from 36.541s to 33.264s at C=1, 56.902s to 49.199s at C=2, and 96.317s to 82.181s at C=4. GSM8K exact_match_flexible stayed at 0.95. Signed-off-by: jasl <jasl9187@hotmail.com>
jasl
added a commit
to jasl/vllm
that referenced
this pull request
May 20, 2026
Widen the direct FP8 MQA logits Triton fallback from BLOCK_M=8 to BLOCK_M=16 while keeping BLOCK_N=128 and the existing 4-warp launch. This reduces CTA count for late-context prefill without introducing a runtime switch. The direction was motivated by the tile-shape discussion in deepseek-ai/DeepGEMM#324, but this is a vLLM-owned Triton fallback adjustment and does not copy DeepGEMM code. On the SM120 long-context gate with prefix cache disabled, the 128K synthetic mean TTFT improved from 36.541s to 33.264s at C=1, 56.902s to 49.199s at C=2, and 96.317s to 82.181s at C=4. GSM8K exact_match_flexible stayed at 0.95. Signed-off-by: jasl <jasl9187@hotmail.com>
|
nice work! may I ask the hardware for testing here is either 5090 or RTX6000pro? |
jasl
added a commit
to jasl/vllm
that referenced
this pull request
May 22, 2026
Widen the direct FP8 MQA logits Triton fallback from BLOCK_M=8 to BLOCK_M=16 while keeping BLOCK_N=128 and the existing 4-warp launch. This reduces CTA count for late-context prefill without introducing a runtime switch. The direction was motivated by the tile-shape discussion in deepseek-ai/DeepGEMM#324, but this is a vLLM-owned Triton fallback adjustment and does not copy DeepGEMM code. On the SM120 long-context gate with prefix cache disabled, the 128K synthetic mean TTFT improved from 36.541s to 33.264s at C=1, 56.902s to 49.199s at C=2, and 96.317s to 82.181s at C=4. GSM8K exact_match_flexible stayed at 0.95. Signed-off-by: jasl <jasl9187@hotmail.com>
jasl
added a commit
to jasl/vllm
that referenced
this pull request
May 22, 2026
Widen the direct FP8 MQA logits Triton fallback from BLOCK_M=8 to BLOCK_M=16 while keeping BLOCK_N=128 and the existing 4-warp launch. This reduces CTA count for late-context prefill without introducing a runtime switch. The direction was motivated by the tile-shape discussion in deepseek-ai/DeepGEMM#324, but this is a vLLM-owned Triton fallback adjustment and does not copy DeepGEMM code. On the SM120 long-context gate with prefix cache disabled, the 128K synthetic mean TTFT improved from 36.541s to 33.264s at C=1, 56.902s to 49.199s at C=2, and 96.317s to 82.181s at C=4. GSM8K exact_match_flexible stayed at 0.95. Signed-off-by: jasl <jasl9187@hotmail.com>
DoradusResearch
pushed a commit
to DoradusResearch/vllm
that referenced
this pull request
May 23, 2026
Widen the direct FP8 MQA logits Triton fallback from BLOCK_M=8 to BLOCK_M=16 while keeping BLOCK_N=128 and the existing 4-warp launch. This reduces CTA count for late-context prefill without introducing a runtime switch. The direction was motivated by the tile-shape discussion in deepseek-ai/DeepGEMM#324, but this is a vLLM-owned Triton fallback adjustment and does not copy DeepGEMM code. On the SM120 long-context gate with prefix cache disabled, the 128K synthetic mean TTFT improved from 36.541s to 33.264s at C=1, 56.902s to 49.199s at C=2, and 96.317s to 82.181s at C=4. GSM8K exact_match_flexible stayed at 0.95. Signed-off-by: jasl <jasl9187@hotmail.com>
Phase 1a: infrastructure + dense FP8 GEMM kernel for SM120a (CC 12.0). Architecture: warp-level mma.sync with block-scaled UE8M0 scale factors, B128 XOR swizzle, persistent scheduling, register-based epilogue. New files: - SM120 heuristics, JIT codegen, MMA PTX wrappers, ldmatrix/swizzle utils - CUDA kernel with warp-specialized TMA/math pipeline (3-9 stages) Modified files: - Arch detection, compiler flags (-gencode for SM120a) - API dispatch (arch_major == 12), SF layout transform - Default recipe for SM120 Correctness: 8/8 shapes pass (diff < 0.001 cosine distance) Performance: ~73 TFLOPS (baseline, optimization pending) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
… only Drop the non-warp-specialized kernel path for SM120a (matching SM90/SM100 architecture), merging the warp-specialized implementation into the main sm120_fp8_fp4_gemm_1d1d kernel. Add FP4 GEMM support using packed SMEM with the mxf4nvf4 m16n8k64 MMA instruction. Key changes: - Consolidate: remove non-spec path, always BM=128/BK=128/384 threads - FP4: packed 4-bit SMEM (CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B), standard ldmatrix, uint16_t scale factors (scale_vec::2X), kKSteps=2 vs FP8's 4 - Heuristic: simplified to warp-spec only, correct SMEM sizing for FP4 - API: enable FP4 on SM120a (arch_major==12), add fp8_fp4_gemm_nt binding - Fix SF hoist bug: hoist SFA/SFB independently for mixed gran_k configs Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- Kernel: Add TMA descriptor runtime update (tensormap.replace) in producer loop for K-grouped group transitions, fix SF_K_ALIGNMENT to kGranKA*4, fix SMEM layout (pipeline data at offset 0 for B128 swizzle alignment, tensor map descriptors at end), fix epilogue bounds for multi-group output. - MMA wrappers: Replace CUTLASS mma_sm120.hpp dependency with custom inline asm using "+f" read-write constraints for accumulator registers. Eliminates CUTLASS header dependency and gives explicit control over MMA operand encoding for both FP8 (m16n8k32) and FP4 (m16n8k64) block-scaled MMA. - JIT launcher: Add sm120_k_grouped_fp8_fp4_gemm_1d1d() with proper TMA descriptor creation (first_k base, FP4-aware stride), SF TMA covering concatenated groups, CD TMA with num_groups outer dimension. - API dispatch: Add arch_major==12 path in k_grouped_fp8_gemm_nt_contiguous, relax recipe assertion to support gran_k=32/128, add SM120 SF layout transform with auto-detection of transposed K-major scale factors. - Tests: Add dedicated SM120 K-grouped test (7 configs including zero-K edge case), fix K-major selection for SM120 in generators, fix test dispatch for SM120 in test_fp8_fp4.py, update FP4 test with perf comparison. Tested: Dense FP8 8/8, Dense FP4 10/10, K-grouped FP8 7/7 — all PASS. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Einsum support: - Add GemmType::Batched to FP8/FP4 and BF16 kernels with 3D TMA load/store - Add IndexType::SF_K for batched SF coordinate computation - Add MN-major B support to BF16 kernel (scalar SMEM loads, single-atom constraint) - BF16 bhr,hdr->bhd: 384 TFLOPS, FP8: 681 TFLOPS (batch=8, b=8192) M-grouped BF16: - Add contiguous and masked M-grouped BF16 GEMM launchers HC prenorm TF32: - New fused GEMM + sqr_sum kernel using mma.sync.m16n8k8 TF32 (226T peak) - BF16 A -> FP32 cast with fused sqr_sum accumulation - Atom-aware FP32 B fragment loading from K-major SMEM - Split-K support for large K / small M shapes - 24/24 test shapes PASS, ~1.1 TB/s bandwidth on large shapes Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Dense (ragged): 651 TFLOPS peak (80% of FP8 MMA peak 814T), 40/40 tests pass. Paged (KV cache): 320 TFLOPS peak, 1.36 TB/s DRAM (91% of HBM BW), 8/8 tests pass. Kernel design: warp-specialized mma.sync m16n8k32 FP8 (no block_scale). 8 math warps × 16 KV rows each = 128 BLOCK_KV. In-warp 2-shfl reduction across 4 threads (lane%4) — only ~10 cycles, negligible vs MMA time. Global stores are fire-and-forget on SM120a, so no epilogue warps needed. Key parameters: block_qh=128, num_heads=64, head_dim=128, 2 Q stages, 3 KV stages, 84KB SMEM (83% of 101KB capacity). Paged variant: 2 groups of 4 warps, SPLIT_KV=128, per-group KV pipeline. Fixed metadata split_kv mismatch and register budget overflow (TMA regs 64→40 to stay within 65536 register limit). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- skip_head_mid: Add SM120 dispatch in attention.hpp for EpilogueHeadSplits. Fix three issues: TMA CD descriptor uses d.size(-1)/d.stride(-2) instead of n; kernel uses stride_d parameter for D row stride and bounds checks; TMA store coordinates apply epilogue N-index remapping. - MN-major B: Fix kernel TMA coordinate for M-grouped BF16 with MN-major B. Group offset moves to outer=K coordinate (not inner=N) when kBKMajor=false. - FP8 kernel stride_d: Add stride_d parameter to decouple D tensor stride from computation dimension n, enabling epilogue transforms that expand N.
Replace per-element scalar SMEM loads in the MN-major B path with ldmatrix.sync.aligned.x2.m8n8.trans.shared.b16 which natively loads column-major 8x8 BF16 matrices — directly producing MMA B fragments. Performance: MN-major B improves from ~220T to ~290T (dense) and ~330T (M-grouped), a 30-50% gain. Remaining gap vs K-major (400T) is due to heuristic selecting BLOCK_N=32/64 vs 128 (single swizzle atom constraint). Verified by micro benchmark b_bf16_4: fragment layout 32x32 lanes PASS, MMA pipeline 4 K-steps accumulation PASS.
Remove single-atom BLOCK_N constraint for MN-major B. ldmatrix.trans correctly handles multi-atom SMEM (verified by micro benchmark b_bf16_5). MN-major B now achieves 99-102% of K-major performance (was 53% with scalar loads, 80% with single-atom ldmatrix.trans).
…bhr->hdr New BF16 bmk,bnk->mn reduction kernel with split-S atomicAdd to FP32 output (188T peak, HBM BW limited). FP8 einsum dispatch for bhd,hdr->bhr and bhd,bhr->hdr via .contiguous() to K-major. Fix batched epilogue stride formula: replace single stride_d with stride_cd_m + stride_cd_batch to support arbitrary D layouts ([batch,M,N] vs [M,batch,N]). Add kBKMajor template parameter to FP8 kernel with verified scalar-load MN-major B path (correct but 3x slower than K-major ldmatrix, kept for future optimization).
K-grouped TN: single .t().contiguous() transpose with constant-stride TMA (kKGroupedConstantStride) — per-group only replaces addr+dim, not stride. TN achieves 99-101% of NT performance. New PTX tensor_map_replace_global_dim_in_smem. Paged MQA varlen: fix 4 kernel bugs in sm120_fp8_paged_mqa_logits.cuh: - TMA Q coordinate: use atom_to_token_idx() instead of hardcoded *kNextNAtom - Prefetch advance: use get_atom_advance() instead of hardcoded +1 - Math loop: conditional iteration count via is_paired_atom for unpaired atoms - KV block idx: reset kv_block_idx_ptr=32 on q_atom change
New kernels using mma.sync m16n8k64 block-scaled FP4 (mxf4nvf4, scale_vec::2X). Architecture: 8 math warps + 4 TMA warps, B64 swizzle, kKSteps=2. Block-scaled MMA folds UE8M0 SF into computation — no post-MMA scale. Dense FP4 MQA: 1022 TFLOPS peak (63% FP4 peak), 1.6x vs FP8. Paged FP4 MQA: 707 TFLOPS peak (varlen), 566 TFLOPS (non-varlen next_n=2).
|
@ormandj - you can revert to old behavior by setting |
|
Tested the merged Build/runtime:
What fails:
The engine fails during weight post-processing before server startup:
This selects the intended split: but still fails in MXFP4 MoE scale packing:
Working workaround for this stack is still to disable DeepGEMM globally: That boots cleanly with: So this is not just a vLLM linear-backend selection issue. On GB10, the |
Author
@arthur-drozdov vllm mainline has already merged sm120 support—please use mainline vllm. jasl's PR is out of support scope. Please see vllm-project/vllm#43477 for reference |
liz-badada
added a commit
to liz-badada/DeepGEMM
that referenced
this pull request
Jun 30, 2026
Extract only the paged FP4 MQA path required by sgl-project/sglang#27059 from deepseek-ai#324. The upstream feature commit 9160119 is not independently cherry-pickable because it relies on earlier SM120 JIT, MMA, and scheduler refactors, so this commit carries only the leaf support missing from sgl/dev.\n\nThe supported contract is intentionally limited to next_n=1, 64 query heads, head_dim=128, block_kv=64, FP32 logits, 2-D context lengths, and no varlen indices. Dense MQA, FP8 MQA, BF16 logits, wider shape support, and SM121 family reuse are deliberately excluded.
liz-badada
added a commit
to liz-badada/DeepGEMM
that referenced
this pull request
Jun 30, 2026
Extract only the paged FP4 MQA path required by sgl-project/sglang#27059 from deepseek-ai#324. The upstream feature commit 9160119 is not independently cherry-pickable because it relies on earlier SM120 JIT, MMA, and scheduler refactors, so this commit carries only the leaf support missing from sgl/dev. The supported contract is intentionally limited to next_n=1, 64 query heads, head_dim=128, block_kv=64, FP32 logits, 2-D context lengths, and no varlen indices. Dense MQA, FP8 MQA, BF16 logits, wider shape support, and SM121 family reuse are deliberately excluded.
liz-badada
added a commit
to liz-badada/DeepGEMM
that referenced
this pull request
Jul 1, 2026
Extract only the paged FP4 MQA path required by sgl-project/sglang#27059 from deepseek-ai#324. The upstream feature commit 9160119 is not independently cherry-pickable because it relies on earlier SM120 JIT, MMA, and scheduler refactors, so this commit carries only the leaf support missing from sgl/dev. The supported contract is intentionally limited to next_n=1, 64 query heads, head_dim=128, block_kv=64, FP32 logits, 2-D context lengths, and no varlen indices. Dense MQA, FP8 MQA, BF16 logits, wider shape support, and SM121 family reuse are deliberately excluded.
Merged
4 tasks
guqiong96
pushed a commit
to guqiong96/Lvllmds4
that referenced
this pull request
Jul 6, 2026
Widen the direct FP8 MQA logits Triton fallback from BLOCK_M=8 to BLOCK_M=16 while keeping BLOCK_N=128 and the existing 4-warp launch. This reduces CTA count for late-context prefill without introducing a runtime switch. The direction was motivated by the tile-shape discussion in deepseek-ai/DeepGEMM#324, but this is a vLLM-owned Triton fallback adjustment and does not copy DeepGEMM code. On the SM120 long-context gate with prefix cache disabled, the 128K synthetic mean TTFT improved from 36.541s to 33.264s at C=1, 56.902s to 49.199s at C=2, and 96.317s to 82.181s at C=4. GSM8K exact_match_flexible stayed at 0.95. Signed-off-by: jasl <jasl9187@hotmail.com>
leavelet
pushed a commit
to leavelet/DeepGEMM
that referenced
this pull request
Jul 14, 2026
…ped GEMM, einsum, prenorm) Cherry-pick of deepseek-ai/DeepGEMM nv_dev commit 2073ddb (PR deepseek-ai#324) onto sgl/dev. Stage 1 scope: dense + m-grouped FP8/FP4/BF16 GEMM, einsum, TF32 HC prenorm. Conflict resolutions: - scheduler/gemm.cuh: merged Scheduler template (dev kEnsureZeroPadding/kKAlignment/ kSFKSpan + sm120 kSplitKFactor). - heuristics/runtime.hpp: per-arch MK-alignment policy; preserved dev SM100 tuning {224,32,32}. - gemm.hpp: kept dev m-grouped calls (ensure_zero_padding) + added sm120 arch-12 arms. - Deferred to follow-ups: k-grouped GEMM (arch-12 -> UNREACHABLE, needs dev psum-SF API), MQA-logits (attention.hpp kept dev; smxx_*_mqa_logits.hpp removed; sm120 mqa .cuh orphaned). - Tests: took dev's harness; sm120 arch-12 test gating to be re-added.
AliceChenyy
added a commit
to xutizhou/sglang
that referenced
this pull request
Jul 21, 2026
…LA prefill On SM120, DeepSeek-V4 forced the sparse-MLA indexer to compute paged-MQA logits in eager torch ops (SGLANG_FP8_PAGED_MQA_LOGITS_TORCH=True), which scales with batch x context and dominates decode. Route to the DeepGEMM paged-MQA-logits CUDA kernel instead: - server_args: on SM120 stop forcing the torch fallback; also default SGLANG_OPT_FLASHMLA_SPARSE_PREFILL=False so sparse-MLA prefill uses the FlashInfer SM120 path (the unified sgl_kernel flash_mla_sparse_fwd prefill is SM90a/SM100f-only and errors on SM120). - metadata: on SM120 always use the DeepGEMM metadata kernel, never the JIT one (its dynamic-smem request exceeds SM120 limits, incl. large prefill). Requires DeepGEMM with SM120 support (deepseek-ai/DeepGEMM#324). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
AliceChenyy
added a commit
to xutizhou/sglang
that referenced
this pull request
Jul 21, 2026
Enable the DeepGEMM grouped-GEMM FP4 MoE runner on SM120 (--moe-runner-backend deep_gemm); marlin stays the SM120 default (faster decode), this is opt-in for prefill throughput / FP4 experts. Requires DeepGEMM with SM120 support (deepseek-ai/DeepGEMM#324): - configurer: enable JIT DeepGEMM on SM120 when the SM120 grouped-GEMM symbol is present; enable UE8M0 block scales on SM120. - moe_runner/deep_gemm + ep_moe/kernels: SM120 handling (pow2 UE8M0 FP8 quant, TMA-aligned e8m0 scales, JIT-EP-activation fallback for TP>=2, in-place swiglu_limit clamp on the Triton fallback path). - fp8/fp8_utils: contiguous scale for the layout transform; keep dense w8a8 fp8 linear off DeepGEMM on SM120. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
allenh1
pushed a commit
to allenh1/vllm
that referenced
this pull request
Aug 5, 2026
Widen the direct FP8 MQA logits Triton fallback from BLOCK_M=8 to BLOCK_M=16 while keeping BLOCK_N=128 and the existing 4-warp launch. This reduces CTA count for late-context prefill without introducing a runtime switch. The direction was motivated by the tile-shape discussion in deepseek-ai/DeepGEMM#324, but this is a vLLM-owned Triton fallback adjustment and does not copy DeepGEMM code. On the SM120 long-context gate with prefix cache disabled, the 128K synthetic mean TTFT improved from 36.541s to 33.264s at C=1, 56.902s to 49.199s at C=2, and 96.317s to 82.181s at C=4. GSM8K exact_match_flexible stayed at 0.95. Signed-off-by: jasl <jasl9187@hotmail.com>
allenh1
pushed a commit
to allenh1/vllm
that referenced
this pull request
Aug 6, 2026
Widen the direct FP8 MQA logits Triton fallback from BLOCK_M=8 to BLOCK_M=16 while keeping BLOCK_N=128 and the existing 4-warp launch. This reduces CTA count for late-context prefill without introducing a runtime switch. The direction was motivated by the tile-shape discussion in deepseek-ai/DeepGEMM#324, but this is a vLLM-owned Triton fallback adjustment and does not copy DeepGEMM code. On the SM120 long-context gate with prefix cache disabled, the 128K synthetic mean TTFT improved from 36.541s to 33.264s at C=1, 56.902s to 49.199s at C=2, and 96.317s to 82.181s at C=4. GSM8K exact_match_flexible stayed at 0.95. Signed-off-by: jasl <jasl9187@hotmail.com>
allenh1
pushed a commit
to allenh1/vllm
that referenced
this pull request
Aug 7, 2026
Widen the direct FP8 MQA logits Triton fallback from BLOCK_M=8 to BLOCK_M=16 while keeping BLOCK_N=128 and the existing 4-warp launch. This reduces CTA count for late-context prefill without introducing a runtime switch. The direction was motivated by the tile-shape discussion in deepseek-ai/DeepGEMM#324, but this is a vLLM-owned Triton fallback adjustment and does not copy DeepGEMM code. On the SM120 long-context gate with prefix cache disabled, the 128K synthetic mean TTFT improved from 36.541s to 33.264s at C=1, 56.902s to 49.199s at C=2, and 96.317s to 82.181s at C=4. GSM8K exact_match_flexible stayed at 0.95. Signed-off-by: jasl <jasl9187@hotmail.com>
allenh1
pushed a commit
to allenh1/vllm
that referenced
this pull request
Aug 8, 2026
Widen the direct FP8 MQA logits Triton fallback from BLOCK_M=8 to BLOCK_M=16 while keeping BLOCK_N=128 and the existing 4-warp launch. This reduces CTA count for late-context prefill without introducing a runtime switch. The direction was motivated by the tile-shape discussion in deepseek-ai/DeepGEMM#324, but this is a vLLM-owned Triton fallback adjustment and does not copy DeepGEMM code. On the SM120 long-context gate with prefix cache disabled, the 128K synthetic mean TTFT improved from 36.541s to 33.264s at C=1, 56.902s to 49.199s at C=2, and 96.317s to 82.181s at C=4. GSM8K exact_match_flexible stayed at 0.95. Signed-off-by: jasl <jasl9187@hotmail.com>
allenh1
pushed a commit
to allenh1/vllm
that referenced
this pull request
Aug 12, 2026
Widen the direct FP8 MQA logits Triton fallback from BLOCK_M=8 to BLOCK_M=16 while keeping BLOCK_N=128 and the existing 4-warp launch. This reduces CTA count for late-context prefill without introducing a runtime switch. The direction was motivated by the tile-shape discussion in deepseek-ai/DeepGEMM#324, but this is a vLLM-owned Triton fallback adjustment and does not copy DeepGEMM code. On the SM120 long-context gate with prefix cache disabled, the 128K synthetic mean TTFT improved from 36.541s to 33.264s at C=1, 56.902s to 49.199s at C=2, and 96.317s to 82.181s at C=4. GSM8K exact_match_flexible stayed at 0.95. Signed-off-by: jasl <jasl9187@hotmail.com>
allenh1
pushed a commit
to allenh1/vllm
that referenced
this pull request
Aug 12, 2026
Widen the direct FP8 MQA logits Triton fallback from BLOCK_M=8 to BLOCK_M=16 while keeping BLOCK_N=128 and the existing 4-warp launch. This reduces CTA count for late-context prefill without introducing a runtime switch. The direction was motivated by the tile-shape discussion in deepseek-ai/DeepGEMM#324, but this is a vLLM-owned Triton fallback adjustment and does not copy DeepGEMM code. On the SM120 long-context gate with prefix cache disabled, the 128K synthetic mean TTFT improved from 36.541s to 33.264s at C=1, 56.902s to 49.199s at C=2, and 96.317s to 82.181s at C=4. GSM8K exact_match_flexible stayed at 0.95. Signed-off-by: jasl <jasl9187@hotmail.com>
1 task
allenh1
pushed a commit
to allenh1/vllm
that referenced
this pull request
Aug 13, 2026
Widen the direct FP8 MQA logits Triton fallback from BLOCK_M=8 to BLOCK_M=16 while keeping BLOCK_N=128 and the existing 4-warp launch. This reduces CTA count for late-context prefill without introducing a runtime switch. The direction was motivated by the tile-shape discussion in deepseek-ai/DeepGEMM#324, but this is a vLLM-owned Triton fallback adjustment and does not copy DeepGEMM code. On the SM120 long-context gate with prefix cache disabled, the 128K synthetic mean TTFT improved from 36.541s to 33.264s at C=1, 56.902s to 49.199s at C=2, and 96.317s to 82.181s at C=4. GSM8K exact_match_flexible stayed at 0.95. Signed-off-by: jasl <jasl9187@hotmail.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
10 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.
This PR adds first-class
sm_120support to DeepGEMM, maintained on thenv_devbranch by the NVIDIA DevTech APAC team. It brings the full DeepGEMM kernel surface — dense, grouped/MoE, einsum, hyper-connection, and MQA-logits — to the sm120 and sm121 devices like RTX Pro 6000 and DGX Spark.co-authored with @lucifer1004.
Feature surface
Dense GEMM (
sm120_fp8_fp4_gemm_1d1d,sm120_bf16_gemm)mxf4nvf4, gran_k=32)mxf8f6f4 .e4m3.e2m1) and FP4_A × FP8_B(
.e2m1.e4m3) — both directionsGrouped GEMM / MoE
Batched GEMM / Einsum (
sm120_bmk_bnk_mn)bhr,hdr->bhd) and MN-major B (bhd,hdr->bhr)bmk,bnk->mn); FP8 / BF16; small-M AB-swapHC Prenorm (
sm120_tf32_hc_prenorm_gemm)**MQA Logits ** (
sm120_fp8_mqa_logits,sm120_fp4_mqa_logits+ paged)paged (next_n / kPadOddN support)
Performance (RTX PRO 6000 Blackwell, single GPU,
tests/test_*.py)Roofline used: FP8 ≈ 814 TFLOPS, FP4 ≈ 1628 TFLOPS (2× FP8), BF16 ≈ 445 TFLOPS. GB/s is effective bandwidth (cache-inclusive for small/cached shapes).
Dense GEMM
Grouped GEMM (MoE)
Einsum
Peak 742 TFLOPS (
b=8192,h=8,r=4096,d=1024); up to 2893 GB/s (small-batch, cache-bound); Split-S path scales down for small problems.HC Prenorm
m=8192, n=24, k=28672, splits=16: 29 TFLOPS / 1248 GB/s (bandwidth-bound).MQA Logits