Support TBO prefill-only mode for non-PD deployments

[Qin-sx]

Motivation

This PR adds a prefill-only mode for two-batch-overlap (TBO).

On some single-node and multi-node non-PD-disaggregated deployments, it is useful to keep TBO enabled for prefill/extend while disabling it for decode and target-verify batches. This allows deployments to benefit from TBO on prefill-heavy workloads without forcing decode CUDA graph capture/replay and attention metadata paths to use TBO splitting.

By default, the existing behavior is unchanged. Decode TBO remains enabled unless SGLANG_DISABLE_TBO_FOR_DECODE=true is set.

Modifications

• Added a new environment variable:

  • SGLANG_DISABLE_TBO_FOR_DECODE
  • Default: false
  • When set to true, TBO is disabled for decode and target-verify batches while remaining enabled for prefill/extend.

• Added is_decode_tbo_enabled() in the MoE utility layer and exported it through sglang.srt.layers.moe.

• Updated TBO scheduling/metadata logic to skip decode-side TBO when the new env var is enabled:

  • In TboDPAttentionPreparer, decode/target-verify batches no longer set can_run_tbo when decode TBO is disabled.
  • In the decode CUDA graph runner, TBO capture and replay preparation are skipped for decode/target-verify when decode TBO is disabled.
  • Decode CUDA graph support checks still allow non-TBO decode graphs to be captured/replayed when decode TBO is disabled.
  • In TboAttnBackend, replay-view child dispatch is skipped for decode/target-verify when decode TBO is disabled.

• Preserved the default behavior:

  • If SGLANG_DISABLE_TBO_FOR_DECODE is unset or false, decode TBO behavior is unchanged.

Accuracy Tests

On other machines (non-NVIDIA GPUs), the DeepSeek model achieves 0.97 accuracy on the MATH-500 benchmark.

Speed Tests and Profiling

Due to limited machine resources, the num_hidden_layers of the DeepSeek-R1-0528 model was changed to 41. The server was deployed on a single node with 8x H20 GPUs as follows:

export SGLANG_DISABLE_TBO_FOR_DECODE=1

nohup python -m sglang.launch_server \
    --model-path /workspace1/models/DeepSeek-R1-0528 \
    --trust-remote-code \
    --host 0.0.0.0 \
    --port 30000 \
    --tp 8 \
    --ep 8 \
    --mem-fraction-static 0.8 \
    --dp-size 8 \
    --enable-dp-attention \
    --moe-a2a-backend deepep \
    --deepep-mode auto \
    --enable-two-batch-overlap \
    --tbo-token-distribution-threshold 0.0 \
    --disable-prefill-cuda-graph \
    --max-prefill-tokens 8280 \
    --disable-radix-cache \
    --cuda-graph-max-bs 32 \
    --max-running-requests 256 \
    --page-size 64 \
    --chunked-prefill-size -1 \
    > "$LOG_FILE" 2>&1 &

Benchmark script:

python3 -m sglang.bench_serving \
    --host 127.0.0.1 \
    --port 30000 \
    --backend sglang \
    --model /workspace1/models/DeepSeek-R1-0528 \
    --dataset-name random-ids \
    --num-prompts 128 \
    --max-concurrency 128 \
    --random-input-len 2048 \
    --random-output-len 256 \
    --random-range-ratio 1.0 \
    >> log_128_2K_256.log 2>&1

Benchmark Results Summary

Configuration	Mean TTFT (ms)	Mean TPOT (ms)
TBO disabled	13166.72	134.41
TBO enabled	13278.05	177.64
TBO enabled + SGLANG_DISABLE_TBO_FOR_DECODE=1	13884.69	134.57

On other machines (non-NVIDIA GPUs), with tp size=16, ep size=16, dp size=16: when TBO is enabled, TTFT is 80% of the TBO-disabled baseline. After enabling the SGLANG_DISABLE_TBO_FOR_DECODE environment variable, TPOT matches the TBO-disabled baseline; without it, TPOT increases by 50%–100% compared to the TBO-disabled baseline.

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 #28035201549
Latest PR Test (Extra): ❌ Run #28035199442

[gemini-code-assist]

Code Review

This pull request introduces the ability to disable Two-Batch Overlap (TBO) specifically for decode and target verification phases via a new environment variable SGLANG_DISABLE_TBO_FOR_DECODE. The changes integrate checks across batch preparation, attention backend dispatch, and CUDA graph runner components. The review feedback suggests caching the environment variable lookup in is_decode_tbo_enabled() to avoid performance overhead in the hot path, and importing the helper function from the public sglang.srt.layers.moe interface for consistency.

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.

[Qin-sx]

This is quite helpful for reducing TTFT in non-PD deployments. Could someone help take a look at this PR when you have a moment? Thanks! @merrymercy @Ying1123 @hnyls2002 @Fridge003 @ispobock @Qiaolin-Yu @hebiao064 @HaiShaw @ch-wan @BBuf @Edwardf0t1 @wisclmy0611 @zijiexia @JustinTong0323 @sogalin

[Qin-sx]

This is quite helpful for reducing TTFT in non-PD deployments. Could someone help take a look at this PR when you have a moment? Thanks! @merrymercy @Ying1123 @hnyls2002 @Fridge003 @ispobock @Qiaolin-Yu @hebiao064 @HaiShaw @ch-wan @BBuf @Edwardf0t1 @wisclmy0611 @zijiexia @JustinTong0323 @sogalin