[SM90] Use native block32 DeepGEMM for DeepSeek-V4.1 by vladnosiv · Pull Request #42604 · sgl-project/sglang

vladnosiv · 2026-10-05T14:58:27Z

Motivation

DeepSeek-V4.1-Flash stores its dense FP8 weights in [32,32] blocks, compared with [128,128] in DeepSeek-V4-Flash.
These GEMMs cover attention, the shared expert, indexer, Engram, and DSPARK linears.
On Hopper, the SGLang dispatcher routes K-block != 128 to Triton, while the DeepGEMM wrapper only accepts [128,128]. Larger prefill GEMMs can use the existing BF16 weight-cache fast path.
sgl-project/DeepGEMM#95 added native SM90 block32 support. This PR integrates that path into SGLang using sgl-deep-gemm==0.2.2

Modifications

Routes supported SM90 block32 linears through the existing DeepGEMM wrapper with:
fp8_gemm_nt((Aq, As), (Wq, Ws), D, recipe=(1, 32, 32))
The integration preserves checkpoint FP8 weights and UE8M0 activation rounding, passes the block dimensions through the GEMM wrapper, and prepares block32 scale buffers during kernel warmup.
Both prefill and decode use native DeepGEMM. The unused BF16 weight expansion is skipped when DeepGEMM is selected. Explicit Triton selection and unsupported shapes retain their fallback.

Server configuration for the tests listed below

SGLANG_ENABLE_DSV41_ENGRAM_HOST_TABLE=1 \
sglang serve \
  --model-path deepseek-ai/DeepSeek-V4.1-Flash \
  --trust-remote-code \
  --random-seed 42 \
  --tp 4 \
  --mem-fraction-static 0.8 \
  --attention-backend dsv4 \
  --moe-runner-backend flashinfer_mxfp4 \
  --ep-size 4 \
  --enable-decoder-swa-bounded-replay \
  --cuda-graph-max-bs-decode 64 \
  --reasoning-parser auto \
  --tool-call-parser auto \
  --speculative-algorithm DSPARK \
  --speculative-dspark-block-size 5 \
  --host 0.0.0.0 \
  --port 30000

Accuracy Tests

GSM8K, correct answers

Before	After
1279/1319	1285/1319

Speed Tests and Profiling

Fixed workload 16k inp / 1k out, simulated acceptance length 5.5

Metric	Before	After	Delta
Total time	4.618 s	3.408 s	-26.2%
TTFT	1.017 s	1.002 s	-1.4%
Decode throughput	285.87 tokens/s	426.87 tokens/s	+49.3%

Trace analysis

	Before	After
W8A8 Share	41.47%	12.86%

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

Review and Merge Process

Ping Merge Oncalls to start the process. See the PR Merge Process.
Get approvals from CODEOWNERS and other reviewers.
Trigger CI tests with comments or contact authorized users to do so.
- Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

CI States

Latest PR Test (Base): ❌ Run #38011736527
Latest PR Test (Extra): ❌ Run #38011736377
Latest PR Test (AMD ROCm 10): ❌ Run #38011736554

BBuf · 2026-10-06T12:15:05Z

+        and deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
+        and (backend.is_auto() or backend.is_deep_gemm())
+    ):
+        return partial(


DeepGEMM #95 reports 0.951× speedup for M=1–4, but this switches those shapes to DeepGEMM in auto mode too. Could you add single-request decode results without speculation, plus DSPARK results with real acceptance? That would help decide whether small M should stay on Triton.

Workload Before, tokens/s After, tokens/s Delta

16k input / 1024 output, bs=1, no speculation 126.62 125.66 −0.76%

16k input / 1024 output, bs=1, DSPARK simulated 5.5 acc len 291.71 420.31 +44.09%

Single GSM8K prompt, bs=1, DSPARK ~3.8 acc len 201.51 281.14 +39.52%

yuan-luo · 2026-10-06T13:22:20Z

        assert not deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
        assert input.dtype == torch.float8_e4m3fn
-        assert weight.shape[0] % 64 == 0 and weight.shape[1] % 128 == 0, (
+        assert shape_supported and block_supported, (


On the pre-quantized-activation path, an unsupported block32 shape now raises an AssertionError instead of falling back to Triton. This may convert a previously working request into a crash.

yuan-luo · 2026-10-06T13:23:15Z

            _is_cuda
            and get_platform().is_sm90
            and envs.SGLANG_OPT_HOPPER_BLOCK_FP8_BF16.get()
+            and not use_deepgemm_runner


This disables the Hopper BF16 weight cache for every SM90 block32 layer, which makes the new weight_bf16 parameter dead code and strips the BF16 fast path from exactly the shapes DeepGEMM cannot serve.

yuan-luo · 2026-10-06T13:24:35Z

    weight_scale: torch.Tensor,
    input_scale: Optional[torch.Tensor] = None,
    bias: Optional[torch.Tensor] = None,
+    act_scale_ue8m0: bool = False,


This parameter is added to deepgemm_w8a8_block_fp8_linear_with_fallback's signature but only the new block32 partial ever supplies it; the [128,128] DeepGEMM dispatch routes still return the bare function, so its Triton fallback silently drops ue8m0 activation rounding.

agree, btw seems like exactly the same behavior actually already exists in the main branch for the current block128

yuan-luo

LGTM

BBuf · 2026-10-08T02:30:33Z

@vladnosiv Can you add aime dataset's acc too?

BBuf · 2026-10-09T11:07:28Z

/rerun-test test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py test/registered/kernels/ops/gemm/test_fp8_blockwise_gemm.py

Targeted validation for b093f33.

github-actions · 2026-10-09T11:07:59Z

⚠️ Rebase Required Before Re-run

A major update has landed on main. Your PR is diverged relative to required base commit 60878c3fba50.

Re-run was not dispatched. What to do:

Rebase your branch onto the latest main and push again
Follow issue #21065 for context
CI-fix PRs may request the bypass-maintenance label to skip this check

BBuf · 2026-10-09T11:59:37Z

/rerun-test test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py test/registered/kernels/ops/gemm/test_fp8_blockwise_gemm.py

Targeted validation for a0fac1e after updating the branch from main.

github-actions · 2026-10-09T12:00:02Z

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py

🚀 1-gpu-5090 (2 tests): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py
cd test/ && python3 registered/kernels/ops/gemm/test_fp8_blockwise_gemm.py

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py

BBuf

LGTM. The current-head targeted tests pass on H100, B200 and RTX 5090, including the native SM90 block32 and UE8M0 activation-scale cases.

BBuf · 2026-10-10T01:05:05Z

/rerun-test test/registered/unit/layers/deep_gemm_wrapper/test_warmup_memory.py test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py test/registered/kernels/ops/gemm/test_fp8_blockwise_gemm.py

Targeted validation for 1d8096a after merging main and reconciling block32 warmup memory accounting with the scale-workspace fix.

github-actions · 2026-10-10T01:05:29Z

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/layers/deep_gemm_wrapper/test_warmup_memory.py

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py

🚀 1-gpu-5090 (2 tests): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py
cd test/ && python3 registered/kernels/ops/gemm/test_fp8_blockwise_gemm.py

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py

vladnosiv added 3 commits October 5, 2026 12:42

Use native DeepGEMM block32 FP8 GEMM on SM90

2ab93e8

Avoid unused BF16 caches for native block32 DeepGEMM

7818b24

Format native block32 DeepGEMM integration

b9aa964

vladnosiv requested review from Alisehen, AniZpZ, BBuf, DarkSharpness, Edwardf0t1, FlamingoPg, Fridge003, HaiShaw, HydraQYH, OrangeRedeng, Ying1123, b8zhong, celve, ch-wan, ispobock, merrymercy, mmangkad and yuan-luo as code owners October 5, 2026 14:58

github-actions Bot added dependencies Pull requests that update a dependency file jit-kernel labels Oct 5, 2026

BBuf reviewed Oct 6, 2026

View reviewed changes

yuan-luo reviewed Oct 6, 2026

View reviewed changes

vladnosiv added 2 commits October 6, 2026 16:49

Fix native block32 DeepGEMM fallbacks

848b1f3

Preserve UE8M0 activation policy in DeepGEMM

b093f33

vladnosiv requested a review from fzyzcjy as a code owner October 6, 2026 14:44

github-actions Bot added the deepseek label Oct 6, 2026

yuan-luo approved these changes Oct 7, 2026

View reviewed changes

Merge branch 'main' into native-block32-sm90

a0fac1e

BBuf approved these changes Oct 9, 2026

View reviewed changes

Merge main and reconcile block32 DeepGEMM warmup memory estimates

1d8096a

BBuf merged commit 2ff3e6d into sgl-project:main Oct 10, 2026
113 of 129 checks passed

Workload	Before, tokens/s	After, tokens/s	Delta
16k input / 1024 output, bs=1, no speculation	126.62	125.66	−0.76%
16k input / 1024 output, bs=1, DSPARK simulated 5.5 acc len	291.71	420.31	+44.09%
Single GSM8K prompt, bs=1, DSPARK ~3.8 acc len	201.51	281.14	+39.52%

Conversation

vladnosiv commented Oct 5, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

Review and Merge Process

CI States

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vladnosiv Oct 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yuan-luo left a comment

Choose a reason for hiding this comment

Uh oh!

BBuf commented Oct 8, 2026

Uh oh!

BBuf commented Oct 9, 2026

Uh oh!

github-actions Bot commented Oct 9, 2026

⚠️ Rebase Required Before Re-run

Uh oh!

BBuf commented Oct 9, 2026

Uh oh!

github-actions Bot commented Oct 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

BBuf left a comment

Choose a reason for hiding this comment

Uh oh!

BBuf commented Oct 10, 2026

Uh oh!

github-actions Bot commented Oct 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

vladnosiv commented Oct 5, 2026 •

edited by github-actions Bot

Loading

vladnosiv Oct 6, 2026 •

edited

Loading

github-actions Bot commented Oct 9, 2026 •

edited

Loading

github-actions Bot commented Oct 10, 2026 •

edited

Loading