[Bugfix] Wrong Answer in PD disaggregation [P w/o SPS | D w/ SPS]

[Ch3ngY1]

Motivation

After SGLang v0.5.2, when deploying a PD-disaggregated MHA model (e.g., GLM4.5 Air FP8), if speculative decoding parameters are not configured on the Prefill node but are configured on the Decode node, the service may return incorrect responses.

Modifications

Investigation showed that KV transfer caused incorrect Value cache on the Decode (D) node, so we updated the corresponding logic accordingly.

Accuracy Tests

Deploy Scripts:

# prefill
python3 -m sglang.launch_server \
    --model-path /dev/shm/GLM-4.5-Air-FP8 \
    --disaggregation-mode prefill \
    --tp-size 4 \
    --mem-fraction-static 0.8 \
    --kv-cache-dtype fp8_e4m3 \
    --host 0.0.0.0 \
    --attention-backend fa3 \
    --port 30001 \
    --chunked-prefill-size $((16384 * 4)) \
    # --speculative-algo EAGLE --speculative-num-steps 2 --speculative-eagle-topk 1 --speculative-num-draft-tokens 3
# decode
python3 -m sglang.launch_server \
    --model-path /dev/shm/GLM-4.5-Air-FP8 \
    --disaggregation-mode decode \
    --tp-size 4 \
    --mem-fraction-static 0.8 \
    --kv-cache-dtype fp8_e4m3 \
    --host 0.0.0.0 \
    --attention-backend fa3 \
    --port 30011 \
    --speculative-algo EAGLE --speculative-num-steps 2 --speculative-eagle-topk 1 --speculative-num-draft-tokens 3

Running evaluation with

python benchmark/gsm8k/bench_sglang.py     --host http://127.0.0.1     --port 30000     --num-questions 1319     --parallel 256

Results

# P with Speculative Decoding
100%|████████████████████████████████████████████████████████████| 1319/1319 [00:56<00:00, 23.24it/s]
Accuracy: 0.920
Invalid: 0.000
Latency: 57.106 s
Output throughput: 2379.669 token/s

# P without Speculative Decoding without this PR
100%|████████████████████████████████████████████████████████████| 1319/1319 [04:54<00:00,  4.48it/s]
Accuracy: 0.011
Invalid: 0.313
Latency: 294.957 s
Output throughput: 2248.717 token/s

# P without Speculative Decoding with this PR
100%|████████████████████████████████████████████████████████████| 1319/1319 [01:02<00:00, 21.20it/s]
Accuracy: 0.916
Invalid: 0.000
Latency: 62.603 s
Output throughput: 2246.093 token/s

Benchmarking and Profiling

The bench serving results for the multi-turn chat are as follows. We first use the first 10,000 samples for KV-cache warmup, and then benchmark the remaining 10,000 samples. The results are shown below.

• Device = H200 * 8
• P = 2xTP2
• D = 2xTP2

# Prefill with speculative decoding: 
# Round1 --> Sample ID [0, 10000)

============ Serving Benchmark Result ============
Backend:                                 vllm      
Traffic request rate:                    40.0      
Max request concurrency:                 not set   
Successful requests:                     10000     
Benchmark duration (s):                  282.80    
Total input tokens:                      25873815  
Total generated tokens:                  1876513   
Total generated tokens (retokenized):    1869016   
Request throughput (req/s):              35.36     
Input token throughput (tok/s):          91490.07  
Output token throughput (tok/s):         6635.37   
Total token throughput (tok/s):          98125.43  
Concurrency:                             784.53    
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   22186.77  
Median E2E Latency (ms):                 23569.19  
---------------Time to First Token----------------
Mean TTFT (ms):                          19182.07  
Median TTFT (ms):                        20462.93  
P99 TTFT (ms):                           26864.65  
---------------Inter-Token Latency----------------
Mean ITL (ms):                           34.28     
Median ITL (ms):                         33.80     
P95 ITL (ms):                            43.10     
P99 ITL (ms):                            62.13     
Max ITL (ms):                            992.52    
Mean TPOT (ms):                          16.09     
Median TPOT (ms):                        16.28     
P95 TPOT (ms):                           18.65     
P99 TPOT (ms):                           21.74     
Max TPOT (ms):                           34.69     
==================================================

# Round2 --> Sample ID [10000, 20000)

============ Serving Benchmark Result ============
Backend:                                 vllm      
Traffic request rate:                    40.0      
Max request concurrency:                 not set   
Successful requests:                     10000     
Benchmark duration (s):                  275.27    
Total input tokens:                      25997995  
Total generated tokens:                  1876513   
Total generated tokens (retokenized):    1868842   
Request throughput (req/s):              36.33     
Input token throughput (tok/s):          94446.92  
Output token throughput (tok/s):         6817.10   
Total token throughput (tok/s):          101264.02 
Concurrency:                             195.73    
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   5387.78   
Median E2E Latency (ms):                 5333.12   
---------------Time to First Token----------------
Mean TTFT (ms):                          2439.15   
Median TTFT (ms):                        2320.33   
P99 TTFT (ms):                           4927.51   
---------------Inter-Token Latency----------------
Mean ITL (ms):                           33.58     
Median ITL (ms):                         33.55     
P95 ITL (ms):                            41.72     
P99 ITL (ms):                            53.83     
Max ITL (ms):                            481.44    
Mean TPOT (ms):                          15.79     
Median TPOT (ms):                        16.04     
P95 TPOT (ms):                           17.97     
P99 TPOT (ms):                           18.68     
Max TPOT (ms):                           20.51     
==================================================

# with this PR, Prefill without speculative decoding: 
# Round1 --> Sample ID [0, 10000)

============ Serving Benchmark Result ============
Backend:                                 vllm      
Traffic request rate:                    40.0      
Max request concurrency:                 not set   
Successful requests:                     10000     
Benchmark duration (s):                  272.63    
Total input tokens:                      25873815  
Total generated tokens:                  1876533   
Total generated tokens (retokenized):    1845595   
Request throughput (req/s):              36.68     
Input token throughput (tok/s):          94905.18  
Output token throughput (tok/s):         6883.12   
Total token throughput (tok/s):          101788.30 
Concurrency:                             235.95    
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   6432.56   
Median E2E Latency (ms):                 4249.37   
---------------Time to First Token----------------
Mean TTFT (ms):                          3562.08   
Median TTFT (ms):                        1218.30   
P99 TTFT (ms):                           15838.37  
---------------Inter-Token Latency----------------
Mean ITL (ms):                           35.46     
Median ITL (ms):                         34.48     
P95 ITL (ms):                            45.19     
P99 ITL (ms):                            71.50     
Max ITL (ms):                            3195.92   
Mean TPOT (ms):                          15.47     
Median TPOT (ms):                        15.12     
P95 TPOT (ms):                           20.75     
P99 TPOT (ms):                           24.14     
Max TPOT (ms):                           37.20     
==================================================

# Round2 --> Sample ID [10000, 20000)

============ Serving Benchmark Result ============
Backend:                                 vllm      
Traffic request rate:                    40.0      
Max request concurrency:                 not set   
Successful requests:                     10000     
Benchmark duration (s):                  271.78    
Total input tokens:                      25997995  
Total generated tokens:                  1876323   
Total generated tokens (retokenized):    1664673   
Request throughput (req/s):              36.79     
Input token throughput (tok/s):          95658.17  
Output token throughput (tok/s):         6903.83   
Total token throughput (tok/s):          102561.99 
Concurrency:                             131.76    
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   3580.96   
Median E2E Latency (ms):                 3495.13   
---------------Time to First Token----------------
Mean TTFT (ms):                          448.86    
Median TTFT (ms):                        423.63    
P99 TTFT (ms):                           940.30    
---------------Inter-Token Latency----------------
Mean ITL (ms):                           36.30     
Median ITL (ms):                         34.45     
P95 ITL (ms):                            49.14     
P99 ITL (ms):                            77.42     
Max ITL (ms):                            3771.21   
Mean TPOT (ms):                          16.94     
Median TPOT (ms):                        16.28     
P95 TPOT (ms):                           23.44     
P99 TPOT (ms):                           28.65     
Max TPOT (ms):                           41.04     
==================================================

deploy	TTFT-R1	TPOT-R1	TTFT-R2	TPOT-R2
P with SPS	19182.07	16.09	2439.15	15.79
P without SPS	3562.08	15.47	448.86	16.94

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 Process

1. Ping Merge Oncalls to start the PR flow. 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.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.

[gemini-code-assist]

Summary of Changes

Hello @Ch3ngY1, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request resolves a critical bug affecting PD-disaggregated SGLang deployments where speculative decoding is active solely on the decode stage. By correcting the KV cache transfer logic, particularly for Value cache pointers, the change ensures the model produces accurate responses and maintains performance, as demonstrated by the provided accuracy and benchmark results.

Highlights

• Bug Fix for PD-Disaggregation with Speculative Decoding: Addresses an issue in SGLang v0.5.2 where PD-disaggregated MHA models produced incorrect responses when speculative decoding was enabled only on the Decode node. The problem stemmed from incorrect Value cache handling during KV transfer.
• KV Cache Transfer Logic Update: Modified the get_mha_kv_ptrs_with_pp function in conn.py to correctly index dst_v_ptrs when the Decode node has more KV layers than the Prefill node (e.g., due to draft model KV). This prevents an empty Value Cache and ensures accurate data transfer.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request addresses a critical bug in PD disaggregation where incorrect responses were returned when the Decode node used speculative decoding while the Prefill node did not. The root cause was identified as an incorrect Value cache transfer. The fix introduces a specific logic branch to handle this scenario by correctly calculating the offsets for the destination KV pointers, ensuring the Value cache is populated correctly. The provided accuracy tests and benchmarks demonstrate that the fix is effective, significantly improving accuracy from 0.011 to 0.916 in the problematic configuration. The change is well-targeted and appears correct. I have no further comments.

[ShangmingCai]

/rerun-stage unit-test-backend-8-gpu-h20

[github-actions]

✅ Triggered unit-test-backend-8-gpu-h20 to run independently (skipping dependencies).

[github-actions]

🔗 View workflow run

[ShangmingCai]

/tag-and-rerun-ci

[ShangmingCai]

PDxPP-related tests have passed. Let me run some basic tests, if CI is flaky, we can bypass since this is a clear PR with simple change.

[ShangmingCai]

All related tests has passed. GB200 failure is irrelevant.