[DSv4] Integrate TRT-LLM DSv4 Attention for SM100/103 by akhilg-nv · Pull Request #30805 · sgl-project/sglang

akhilg-nv · 2026-07-10T19:01:43Z

Motivation

Integrates TRT-LLM attention kernel for DSv4 style attention (CSA, HCA).

Modifications

Accuracy Tests

SGLANG_DSV4_ATTN_DECODE_BACKEND=flashmla/trtllm_gen \
python -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V4-Pro \
  --trust-remote-code --tp 8 --moe-runner-backend flashinfer_mxfp4 \
  --chunked-prefill-size 4096 --disable-flashinfer-autotune \
  --mem-fraction-static 0.88 --max-running-requests 32

GSM8K Accuracy

GSM8K V4-Pro flashmla: 0.965
GSM8K V4-Pro trtllm_gen: 0.985

GPQA

(add --chat-template to above server command)

python3 -m sglang.test.run_eval --port 21000 --eval-name gpqa \
  --num-examples 198 --max-tokens 120000 --repeat 8 \
  --thinking-mode deepseek-v3

FlashMLA: Repeat: 8, mean: 0.885
Scores: ['0.889', '0.879', '0.889', '0.879', '0.869', '0.899', '0.904', '0.874']

TRTLLM: Repeat: 8, mean: 0.907
Scores: ['0.919', '0.904', '0.894', '0.904', '0.909', '0.919', '0.904', '0.904']

Speed Tests and Profiling

Unit Kernel Benchmarking (ran with private benchmarking repo):

# prefill (varlen), long contexts
python benchmark.py --backends flashmla_dsv4,trtllm_gen_dsv4 --model deepseek-v4 \
  --hw b200 --dtype fp8 --mode prefill --tp 1 \
  --batch-sizes 1,2 --seq-lens 4096,16384,65536 --graph
# decode, low/mid concurrency x long contexts
python benchmark.py --backends flashmla_dsv4,trtllm_gen_dsv4 --model deepseek-v4 \
  --hw b200 --dtype fp8 --mode decode --tp 1 \
  --batch-sizes 1,4,16,64 --seq-lens 1024,4096,16384,65536,131072 --graph
# decode, high concurrency
python benchmark.py --backends flashmla_dsv4,trtllm_gen_dsv4 --model deepseek-v4 \
  --hw b200 --dtype fp8 --mode decode --tp 1 \
  --batch-sizes 128,256,512,1024 --seq-lens 256,1024,4096,16384 --graph

Prefill: ~1.2x faster
Decode: ~1.45x faster

The performance for TP > 1 scales better on trtllm-gen, since it seems the flashmla variant requires padding the head_dim to 64 or 128. I've attached some sample graphs showing the kernel perf comparison.

E2E bench-serving

python -m sglang.bench_serving --backend sglang --host 127.0.0.1 --port 21000 \
  --dataset-name random --random-input-len 131072 --random-output-len 128 \
  --random-range-ratio 1 --num-prompts 8 --max-concurrency 4

flashmla:

Benchmark duration (s):                  59.92
Total input tokens:                      1048576
Input token throughput (tok/s):          17500.55
Mean E2E Latency (ms):                   29950.67
Mean TTFT (ms):                          19009.87   (median 19636.34)
Mean TPOT (ms):                          86.15
Median ITL (ms):                         12.12

trtllm_gen:

Benchmark duration (s):                  60.09
Total input tokens:                      1048576
Input token throughput (tok/s):          17448.78
Mean E2E Latency (ms):                   30039.27
Mean TTFT (ms):                          16808.71   (median 15734.66)
Mean TPOT (ms):                          104.18
Median ITL (ms):                         13.11

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 #34310903207
Latest PR Test (Extra): ✅ Run #34310902955
Latest PR Test (AMD ROCm 7.2): ❌ Run #34310903174

gemini-code-assist · 2026-07-10T19:01:46Z

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

gemini-code-assist · 2026-07-25T00:45:22Z

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

nvpohanh · 2026-07-28T08:59:48Z

@akhilg-nv could you fix the conflicts?

Adds a uniform FP8 (e4m3, 448 nope + 64 rope = 512 bytes) KV pool.

b8zhong · 2026-08-04T22:42:34Z

@akhilg-nv Please fix the lint and fix conflicts

b8zhong

Can we have some more benchmarks in TP {2, 4, 8} with DSV4 head dims?

b8zhong · 2026-09-03T11:58:05Z

/rerun-test test/registered/models_e2e/test_deepseek_v4_flash_fp4_b200_trtllm.py

github-actions · 2026-09-03T11:58:28Z

Results for /rerun-test test/registered/models_e2e/test_deepseek_v4_flash_fp4_b200_trtllm.py:

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

cd test/ && python3 registered/models_e2e/test_deepseek_v4_flash_fp4_b200_trtllm.py

b8zhong · 2026-09-03T15:23:24Z

/rerun-test test/registered/backends/test_dsv4_fp8_trtllm_backend.py

b8zhong · 2026-09-03T15:23:54Z

/rerun-test test/registered/models_e2e/test_deepseek_v4_flash_fp4_b200_trtllm.py

github-actions · 2026-09-03T15:23:57Z

Results for /rerun-test test/registered/backends/test_dsv4_fp8_trtllm_backend.py:

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

cd test/ && python3 registered/backends/test_dsv4_fp8_trtllm_backend.py

github-actions · 2026-09-03T15:24:24Z

Results for /rerun-test test/registered/models_e2e/test_deepseek_v4_flash_fp4_b200_trtllm.py:

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

cd test/ && python3 registered/models_e2e/test_deepseek_v4_flash_fp4_b200_trtllm.py

akhilg-nv · 2026-09-03T18:23:17Z

            )
        )

+    def test_trtllm_semaphore_capacity_covers_configured_query_rows(self):


I think this test is a bit overkill. The semaphore issue is a known WAR that TRT-LLM also uses, and I've raised the issue with flashinfer for them to allocate a dedicated, appropriately sized workspace for it. We should pull in that change when it's made, but for now the monkeypatch approach is fine I'm not sure it warrants its own test.

Sure, can trim it in the next pr or we can deal with this if we need to address other comments together

…issue sgl-project#37838

…r the registered-test taxonomy

…allel guard

…ag removed on main)

nvpohanh · 2026-09-09T09:21:48Z

All NV pipelines have passed.

b8zhong · 2026-09-09T14:05:19Z

Do we need #36652 @YAMY1234 @akhilg-nv ?

b8zhong · 2026-09-09T14:06:20Z

Seperately: flashinfer-ai/flashinfer#4918 might be easier than #32975

akhilg-nv · 2026-09-09T21:02:45Z

Do we need #36652 @YAMY1234 @akhilg-nv ?

It is based off of this branch, we will need it as a follow-up to showcase the performance improvements once this is merged.

akhilg-nv · 2026-09-09T21:04:44Z

Seperately: flashinfer-ai/flashinfer#4918 might be easier than #32975

I will close #32975 since I have pulled those changes into #36652, and I will investigate the flashinfer change to see if we can use that instead.

* origin/main: (27 commits) [Simulator] Give the OFFLINE/BLOCKING comparison tolerances real headroom (sgl-project#38732) [Config] msgspec.Struct for the config tier (sgl-project#38753) [AMD] ci: move the miles nightlies from rocm700 to rocm10 (sgl-project#37495) [Config] One writer for the declaration stash; no exception to the write seal (sgl-project#38752) docker(xpu): drop redundant setvars.sh from torch_memory_saver RUN (sgl-project#38665) [XPU][Fix] Pack device-pointer tables as uint64 to avoid 64-bit address overflow (sgl-project#35051) [CI] Temporarily disable GB300 tests (sgl-project#38770) [diffusion] feat: spill large tensors over shared memory like numpy arrays (sgl-project#38656) [diffusion] refactor: refactor utility ownership and document helper placement (sgl-project#38699) [NPU]Support GLM5.2 and FP8 DSA&Indexer kvcache for 950 (sgl-project#38250) [CI] Answer unrecognized slash commands instead of skipping silently (sgl-project#38736) [AMD] Parallelize aiter spec-decode KV index building over token blocks (sgl-project#37659) [DSv4] Integrate TRT-LLM DSv4 Attention for SM100/103 (sgl-project#30805) Add Opt-In for GLM-5.3 Flash breakable prefill CUDA graphs (sgl-project#38522) [CI] Install helion 1.4.0 for the KDA Helion kernel tests (sgl-project#38688) [Rust] Gate health on startup warmup completion (sgl-project#37994) [HiCache] Replace skip_lock_node_ids with a segment lock protocol (sgl-project#36848) feat: add optimized Domino rollout to DFlash V2 (sgl-project#36899) [CI] Add /run-full-ci and /run-extra-ci slash commands (sgl-project#38734) [Model] Support GLM-5.3 Flash NVFP4 loading (sgl-project#38621) ...

) Co-authored-by: Yangmin Li <yangminl@nvidia.com> Co-authored-by: Po-Han Huang (NVIDIA) <53919306+nvpohanh@users.noreply.github.com>

github-actions Bot added the deepseek label Jul 10, 2026

akhilg-nv marked this pull request as ready for review July 25, 2026 00:45

akhilg-nv requested review from Ying1123 and merrymercy as code owners July 25, 2026 00:45

akhilg-nv requested review from Fridge003, HaiShaw, Qiaolin-Yu, alphabetc1, hanming-lu, hebiao064, hnyls2002, hzh0425, ispobock, xiezhq-hermann and yizhang2077 as code owners July 25, 2026 00:45

nvpohanh reviewed Jul 28, 2026

View reviewed changes

Comment thread python/sglang/srt/arg_groups/deepseek_v4_hook.py

nvpohanh reviewed Jul 28, 2026

View reviewed changes

Comment thread python/sglang/srt/arg_groups/deepseek_v4_hook.py Outdated

akhilg-nv force-pushed the dsv4_fp8_trtllm_gen branch from ddde870 to 7f14cea Compare July 30, 2026 23:37

akhilg-nv mentioned this pull request Jul 30, 2026

feat: enable MTP with DSv4 TRT-LLM Sparse Attention #32973

Closed

5 tasks

akhilg-nv force-pushed the dsv4_fp8_trtllm_gen branch from 7f14cea to 9c135ea Compare July 30, 2026 23:53

Integrate TRT-LLM DSv4 sparse attention kernels.

7991319

Adds a uniform FP8 (e4m3, 448 nope + 64 rope = 512 bytes) KV pool.

b8zhong self-assigned this Aug 4, 2026

b8zhong reviewed Aug 4, 2026

View reviewed changes

b8zhong mentioned this pull request Aug 5, 2026

[NVIDIA] DeepSeek V4 Perf Tracking #33636

Open

41 tasks

akhilg-nv mentioned this pull request Aug 6, 2026

perf: add fused norm + RoPE + uniform fp8 store for TRT-LLM DSv4 sparse attention #32975

Closed

5 tasks

Fridge003 added the blackwell SM100/SM120 label Aug 6, 2026

b8zhong reviewed Sep 3, 2026

View reviewed changes

Comment thread python/sglang/srt/arg_groups/deepseek_v4_hook.py

b8zhong reviewed Sep 3, 2026

View reviewed changes

Comment thread test/registered/e2e/dsv4/test_dsv4_fp8_trtllm_backend.py

YAMY1234 mentioned this pull request Sep 3, 2026

[Feature] DeepSeek-V4 trtllm attention backend: support PD disaggregation (uniform-FP8 KV layout in the PD handshake) #37838

Open

akhilg-nv commented Sep 3, 2026

View reviewed changes

nvpohanh and others added 2 commits September 4, 2026 14:25

Merge branch 'main' into dsv4_fp8_trtllm_gen

6c0337d

[DSv4] trtllm backend: point the PD disaggregation guard at tracking …

4708f64

…issue sgl-project#37838

YAMY1234 mentioned this pull request Sep 4, 2026

[Bug] test_unified_radix_cache_kl_cp.py (Qwen3-32B, prefill CP=2 + HiCache) livelocks on KV-pool-full retractions since main 2026-09-04 #38019

Closed

YAMY1234 added 3 commits September 8, 2026 18:06

Merge branch 'main' into dsv4_fp8_trtllm_gen

44e9376

[DSv4] Move the trtllm backend test under test/registered/e2e/dsv4 pe…

f432ba0

…r the registered-test taxonomy

[DSv4] trtllm hook: drop the removed CP V1 flags from the context-par…

1ed93d9

…allel guard

YAMY1234 mentioned this pull request Sep 9, 2026

[Bug] mock_model/test_e2e_pd.py TestPdTransferChecksumFullRealData: kv_canary SWEEP_K_FULL chain-hash violation (slot 1 all-zero after PD transfer) since main 2026-09-08 #38586

Closed

nvpohanh and others added 2 commits September 9, 2026 10:49

Merge branch 'main' into dsv4_fp8_trtllm_gen

4021365

[DSv4] trtllm BCG test: use --cuda-graph-max-bs-prefill (piecewise fl…

1d82b50

…ag removed on main)

Fridge003 approved these changes Sep 10, 2026

View reviewed changes

Fridge003 merged commit 880d6fa into sgl-project:main Sep 10, 2026
245 of 266 checks passed

bingxche mentioned this pull request Sep 10, 2026

[CI Monitor] Daily Report - 2026-09-10 bingxche/sglang-ci-bot#191

Open

mmangkad mentioned this pull request Sep 11, 2026

[CI] Trim DSV4 trtllm B200 tests #39013

Merged

Uh oh!

Conversation

akhilg-nv commented Jul 10, 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

GSM8K Accuracy

GPQA

Speed Tests and Profiling

Unit Kernel Benchmarking (ran with private benchmarking repo):

E2E bench-serving

Checklist

Review and Merge Process

CI States

Uh oh!

gemini-code-assist Bot commented Jul 10, 2026

Uh oh!

gemini-code-assist Bot commented Jul 25, 2026

Uh oh!

nvpohanh commented Jul 28, 2026

Uh oh!

Uh oh!

Uh oh!

b8zhong commented Aug 4, 2026

Uh oh!

b8zhong left a comment

Choose a reason for hiding this comment

Uh oh!

b8zhong commented Sep 3, 2026

Uh oh!

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

Uh oh!

Uh oh!

Uh oh!

Uh oh!

b8zhong commented Sep 3, 2026

Uh oh!

b8zhong commented Sep 3, 2026

Uh oh!

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

Uh oh!

Uh oh!

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

Uh oh!

Uh oh!

akhilg-nv Sep 3, 2026

Choose a reason for hiding this comment

Uh oh!

YAMY1234 Sep 3, 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!

nvpohanh commented Sep 9, 2026

Uh oh!

b8zhong commented Sep 9, 2026

Uh oh!

b8zhong commented Sep 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

akhilg-nv commented Sep 9, 2026

Uh oh!

akhilg-nv commented Sep 9, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

akhilg-nv commented Jul 10, 2026 •

edited by github-actions Bot

Loading

github-actions Bot commented Sep 3, 2026 •

edited

Loading

github-actions Bot commented Sep 3, 2026 •

edited

Loading

github-actions Bot commented Sep 3, 2026 •

edited

Loading

YAMY1234 Sep 3, 2026 •

edited

Loading

b8zhong commented Sep 9, 2026 •

edited

Loading