Skip to content

[Fix] Enable chunked input-logprob processing by default to cap peak memory - #31498

Merged
hnyls2002 merged 9 commits into
mainfrom
lsyin/logprob-chunk-default
Jul 17, 2026
Merged

hnyls2002 merged 9 commits into
mainfrom
lsyin/logprob-chunk-default

Conversation

@hnyls2002

@hnyls2002 hnyls2002 commented Jul 16, 2026 •

Copy link
Copy Markdown
Collaborator

Large return_logprob prefill batches materialize two near-identical fp32 [rows, vocab] logits tensors (~4.4 GiB each at ~7.8k rows on a 152k vocab), which can OOM an almost-full GPU. Behavior change: flip SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK default to True so batches above 2048 input-logprob rows take the existing chunked path (peak memory becomes chunk-sized and batch-independent); batches at or below the threshold are unaffected, and the env var remains as an opt-out. Also keep per-chunk _get_logits calls off the shared graph logits buffer to rule out chunk aliasing, and add forced multi-chunk stitching tests.


CI States

Latest PR Test (Base): 🚫 Run #29608756604
Latest PR Test (Extra): ❌ Run #29608756846

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Warning

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

@hnyls2002

hnyls2002 commented Jul 17, 2026 •

Copy link
Copy Markdown
Collaborator Author

/rerun-test registered/pp/test_pp_single_node.py
(5 tries — previous dispatch cancelled)

@github-actions

github-actions Bot commented Jul 17, 2026 •

Copy link
Copy Markdown
Contributor

Results for /rerun-test registered/core/test_srt_endpoint.py registered/lora/test_lora_moe_tp_logprob_diff.py registered/spec/test_spec_ngram.py:

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

cd test/ && python3 registered/core/test_srt_endpoint.py

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

cd test/ && python3 registered/lora/test_lora_moe_tp_logprob_diff.py

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

cd test/ && python3 registered/spec/test_spec_ngram.py

@github-actions

github-actions Bot commented Jul 17, 2026 •

Copy link
Copy Markdown
Contributor

Results for /rerun-test registered/lora/test_lora_moe_tp_logprob_diff.py:

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

cd test/ && python3 registered/lora/test_lora_moe_tp_logprob_diff.py

@hnyls2002

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci

@github-actions github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Jul 17, 2026
@github-actions

github-actions Bot commented Jul 17, 2026 •

Copy link
Copy Markdown
Contributor

Results for /rerun-test registered/lora/test_lora_qwen3_5_35b_a3b_logprob_diff.py:

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

cd test/ && python3 registered/lora/test_lora_qwen3_5_35b_a3b_logprob_diff.py

@github-actions

github-actions Bot commented Jul 17, 2026 •

Copy link
Copy Markdown
Contributor

Results for /rerun-test registered/lora/test_lora_qwen3_vl_30b_a3b_instruct_logprob_diff.py registered/pp/test_pp_single_node.py:

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

cd test/ && python3 registered/lora/test_lora_qwen3_vl_30b_a3b_instruct_logprob_diff.py

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

cd test/ && python3 registered/pp/test_pp_single_node.py

@hnyls2002
hnyls2002 merged commit 2c856ab into main Jul 17, 2026
95 of 127 checks passed
@hnyls2002
hnyls2002 deleted the lsyin/logprob-chunk-default branch July 17, 2026 20:19
@hnyls2002

Copy link
Copy Markdown
Collaborator Author

/rerun-test registered/pp/test_pp_single_node.py

@hnyls2002

hnyls2002 commented Jul 17, 2026 •

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci
(2 tries — handler checkout flake)

@alisonshao

Copy link
Copy Markdown
Collaborator

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

lora run-ci CI: run the baseline test suite on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants