Disable dsr1 prefill cudagraphs by default

[nvjullin]

Motivation

DSR1 + trtllm_mla attention backend fallbacks to FlashAttention under prefill cudagraphs, which regresses performance a lot.
Disable it by default.

Regression at 1k1k conc=64 is ~13% TPOT. Full results of PCG off

============ Serving Benchmark Result ============
Successful requests:                     640       
Benchmark duration (s):                  209.60    
Total input tokens:                      590227    
Total generated tokens:                  589961    
Request throughput (req/s):              3.05      
Output token throughput (tok/s):         2814.72   
Total Token throughput (tok/s):          5630.71   
---------------Time to First Token----------------
Mean TTFT (ms):                          654.26    
Median TTFT (ms):                        524.50    
P90 TTFT (ms):                           1030.20   
P99 TTFT (ms):                           2158.66   
P99.9 TTFT (ms):                         2160.71   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          21.46     
Median TPOT (ms):                        21.65     
P90 TPOT (ms):                           22.89     
P99 TPOT (ms):                           23.47     
P99.9 TPOT (ms):                         23.62     
---------------Inter-token Latency----------------
Mean ITL (ms):                           634.52    
Median ITL (ms):                         569.32    
P90 ITL (ms):                            788.06    
P99 ITL (ms):                            994.26    
P99.9 ITL (ms):                          1150.51   
----------------End-to-end Latency----------------
Mean E2EL (ms):                          20428.53  
Median E2EL (ms):                        20580.11  
P90 E2EL (ms):                           22634.90  
P99 E2EL (ms):                           23986.83  
P99.9 E2EL (ms):                         24347.97  
==================================================

PCG on

============ Serving Benchmark Result ============
Successful requests:                     640       
Benchmark duration (s):                  238.59    
Total input tokens:                      590227    
Total generated tokens:                  589961    
Request throughput (req/s):              2.68      
Output token throughput (tok/s):         2472.71   
Total Token throughput (tok/s):          4946.53   
---------------Time to First Token----------------
Mean TTFT (ms):                          766.11    
Median TTFT (ms):                        602.10    
P90 TTFT (ms):                           1356.78   
P99 TTFT (ms):                           2642.74   
P99.9 TTFT (ms):                         2644.19   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          24.49     
Median TPOT (ms):                        24.49     
P90 TPOT (ms):                           27.95     
P99 TPOT (ms):                           29.94     
P99.9 TPOT (ms):                         30.33     
---------------Inter-token Latency----------------
Mean ITL (ms):                           724.03    
Median ITL (ms):                         570.25    
P90 ITL (ms):                            1130.24   
P99 ITL (ms):                            1612.48   
P99.9 ITL (ms):                          1793.82   
----------------End-to-end Latency----------------
Mean E2EL (ms):                          23329.72  
Median E2EL (ms):                        23283.00  
P90 E2EL (ms):                           26920.38  
P99 E2EL (ms):                           30180.09  
P99.9 E2EL (ms):                         31263.80  
==================================================

Modifications

Accuracy Tests

Speed Tests and Profiling

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

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. 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
4. 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 #27413427468
Latest PR Test (Extra): ❌ Run #27413427384

[gemini-code-assist]

Code Review

This pull request introduces a check to disable the prefill CUDA graph by default when using the trtllm_mla attention backend with the DeepseekV3ForCausalLM architecture to avoid performance regressions. The reviewer suggests generalizing this logic to all models using the trtllm_mla backend, as the limitation is backend-specific rather than architecture-specific, which would also simplify the implementation and make it more robust.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The performance regression (falling back to FlashAttention under a captured prefill CUDA graph) is a limitation of the trtllm_mla attention backend itself, rather than being specific to the DeepseekV3ForCausalLM architecture.

Restricting this check to the exact DeepseekV3ForCausalLM architecture makes the logic fragile. For example, it won't apply to draft models (whose architecture is renamed to DeepseekV3ForCausalLMNextN during speculative decoding setup) or other fine-tunes/architectures using trtllm_mla. Removing the architecture check simplifies the code, avoids unnecessary calls to self.get_model_config(), and makes the fallback disabling robust for all models using trtllm_mla.

    def _disable_prefill_cuda_graph_for_trtllm_mla(self):
        """Disable prefill CUDA graph by default when using the trtllm_mla
        attention backend. Under any captured prefill CUDA graph (tc_piecewise or
        breakable) trtllm_mla falls back to FlashAttention for prefill and regresses
        performance, so disable whichever prefill graph backend is in effect.
        """

        if (Phase.PREFILL, "backend") in self._cuda_graph_config_locked:
            return
        if self.cuda_graph_config.prefill.backend == Backend.DISABLED:
            return
        prefill_attention_backend, _ = self.get_attention_backends()
        if prefill_attention_backend != "trtllm_mla":
            return
        logger.warning(
            "Disabling prefill CUDA graph (%s) by default on "
            "the trtllm_mla attention backend (a captured prefill graph forces a "
            "FlashAttention fallback that regresses prefill). Set the prefill cuda graph "
            "backend explicitly (e.g. --cuda-graph-backend-prefill tc_piecewise) to override.",
            self.cuda_graph_config.prefill.backend,
        )
        self.cuda_graph_config.prefill.backend = Backend.DISABLED

[gemini-code-assist]

Rename the method to _disable_prefill_cuda_graph_for_trtllm_mla to reflect that this is a backend-specific limitation rather than being tied strictly to the DeepseekV3ForCausalLM architecture.

Suggested change

	# Must run after the attention backend is resolved so the trtllm_mla
	# default (auto-selected for DeepseekV3ForCausalLM on sm100) is visible.
	self._disable_prefill_cuda_graph_for_deepseek_trtllm_mla()
	# Must run after the attention backend is resolved so the trtllm_mla
	# default (auto-selected for DeepseekV3ForCausalLM on sm100) is visible.
	self._disable_prefill_cuda_graph_for_trtllm_mla()

[nvpohanh]

/tag-and-rerun-ci

[nvpohanh]

@nvjullin Please put some perf numbers in the PR description. Thanks!

cc @Fridge003 to assign someone to review.

[nvjullin]

Added benchmark results.

[nvpohanh]

@Fridge003 Could you review and merge this? This restores a DSR1 perf regression. Thanks!

[mmangkad]

I couldn’t reproduce this on my run. Do you have a more concrete command, or can you try rerunning on current main?

[nvpohanh]

@nvjullin could you provide detailed repro steps? thanks

[nvjullin]

Server

python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-R1-0528 --host 0.0.0.0 --port 8080 --tensor-parallel-size 8 --data-parallel-size 1 --ep-size 1 --cuda-graph-max-bs 128 --max-running-requests 128 --mem-fraction-static 0.82 --kv-cache-dtype fp8_e4m3 --chunked-prefill-size 32768 --max-prefill-tokens 32768 --enable-flashinfer-allreduce-fusion --scheduler-recv-interval 30 --disable-radix-cache --attention-backend trtllm_mla --stream-interval 30 --moe-runner-backend flashinfer_trtllm --quantization fp8 --weight-loader-prefetch-checkpoints

Client is 1k1k conc=64, roughly

python3 -m sglang.benchmark.serving --backend vllm --model deepseek-ai/DeepSeek-R1-0528 --base-url http://0.0.0.0:8080 --dataset-name random --random-input-len 1024 --random-output-len 1024 --random-range-ratio 0.8 --num-prompts 640 --max-concurrency 64 --request-rate inf --output-file concurrency_64.jsonl

[mmangkad]

I still couldn’t reproduce the claimed ~13% TPOT regression, but across several runs on both SM100 and SM103, PCG-on was consistently worse than disabled by up to ~4% TPOT, so the default disable looks justified.

[nvjullin]

Thanks for the thorough review.
I only realized now that I didn't provide an important context: PCG for DSR1 used to be disabled (almost accidentally) by https://github.com/sgl-project/sglang/pull/23351/changes#diff-700b5118b493d60d7b5994857f5f1e6a7e842ad702392b8ab199945764dfc8edL3196.
After #23351, PCG was enabled for DSR1 and we noticed a regression, hence this PR.