Skip to content

[V1] Fix jump-forward decoding correctness and add tests (follow-up to #36142) - #1

Closed
HelloWorldU wants to merge 0 commit into
HelloWorldU:jump-forward-decodingfrom
vllm-project:main
Closed

[V1] Fix jump-forward decoding correctness and add tests (follow-up to #36142)#1
HelloWorldU wants to merge 0 commit into
HelloWorldU:jump-forward-decodingfrom
vllm-project:main

Conversation

@HelloWorldU

@HelloWorldU HelloWorldU commented Mar 15, 2026

Copy link
Copy Markdown
Owner

Purpose

This PR builds upon vllm-project#36142 by @FredericOdermatt, addressing correctness issues
identified in the bot reviews and adding test coverage for jump-forward decoding.

Changes

Bug Fixes

1. Fix pending_ff_tokens cleared globally each step
Previously, all pending ff_tokens were cleared every schedule step, even for
requests not scheduled in that step (preempted or skipped due to token budget).
Those requests would lose their deterministic tokens permanently, corrupting
continuation state. Fixed by only clearing tokens for requests present in
num_scheduled_tokens.

2. Fix missing stop checks after appending ff_tokens
ff_tokens were appended without calling check_stop, allowing requests to bypass
EOS/stop-token/max-token termination. Fixed by checking stop conditions during
ff_token application, consistent with _update_request_with_output behavior.

Unit Tests Added

New tests in tests/v1/core/test_scheduler.py:

  • ff_tokens correctly injected for scheduled requests
  • ff_tokens preserved for unscheduled requests across steps
  • ff_tokens truncated correctly when stop token is encountered

Updated tests/v1/core/utils.py to support enable_jump_decoding config via
proper StructuredOutputsConfig → VllmConfig → Scheduler.__init__ chain.

Test Results

pytest tests/v1/core/test_scheduler.py -v -k "jump"
4 passed in 81.48s

(vllm) root@autodl-container-m414yxjsh7-f3a0a1c2:~/autodl-tmp/vllm# pytest tests/v1/core/test_scheduler.py -v -s -k "jump_forward"
================================================================================= test session starts =================================================================================
platform linux -- Python 3.12.3, pytest-9.0.2, pluggy-1.6.0 -- /root/miniconda3/bin/python
cachedir: .pytest_cache
rootdir: /root/autodl-tmp/vllm
configfile: pyproject.toml
plugins: asyncio-1.3.0, anyio-4.10.0
asyncio: mode=Mode.STRICT, debug=False, asyncio_default_fixture_loop_scope=None, asyncio_default_test_loop_scope=function
collected 98 items / 94 deselected / 4 selected                                                                                                                                       

tests/v1/core/test_scheduler.py::test_jump_forward_tokens_injected INFO 03-15 21:23:42 [model.py:533] Resolved architecture: OPTForCausalLM
INFO 03-15 21:23:42 [model.py:1580] Using max model len 2048
INFO 03-15 21:23:42 [scheduler.py:231] Chunked prefill is enabled with max_num_batched_tokens=8192.
INFO 03-15 21:23:42 [vllm.py:748] Asynchronous scheduling is disabled.
PASSED
tests/v1/core/test_scheduler.py::test_jump_forward_tokens_stop_eos '(ReadTimeoutError("HTTPSConnectionPool(host='huggingface.co', port=443): Read timed out. (read timeout=10)"), '(Request ID: 32dbbb1f-9610-4428-bd8c-eca0606d7e91)')' thrown while requesting HEAD https://huggingface.co/facebook/opt-125m/resolve/main/config.json
[2026-03-15 21:23:58] WARNING _http.py:320: '(ReadTimeoutError("HTTPSConnectionPool(host='huggingface.co', port=443): Read timed out. (read timeout=10)"), '(Request ID: 32dbbb1f-9610-4428-bd8c-eca0606d7e91)')' thrown while requesting HEAD https://huggingface.co/facebook/opt-125m/resolve/main/config.json
Retrying in 1s [Retry 1/5].
[2026-03-15 21:23:58] WARNING _http.py:329: Retrying in 1s [Retry 1/5].
'(ReadTimeoutError("HTTPSConnectionPool(host='huggingface.co', port=443): Read timed out. (read timeout=10)"), '(Request ID: 50ea8229-1f5d-43e8-9d39-526ddb172b90)')' thrown while requesting HEAD https://huggingface.co/facebook/opt-125m/resolve/main/preprocessor_config.json
[2026-03-15 21:24:12] WARNING _http.py:320: '(ReadTimeoutError("HTTPSConnectionPool(host='huggingface.co', port=443): Read timed out. (read timeout=10)"), '(Request ID: 50ea8229-1f5d-43e8-9d39-526ddb172b90)')' thrown while requesting HEAD https://huggingface.co/facebook/opt-125m/resolve/main/preprocessor_config.json
Retrying in 1s [Retry 1/5].
[2026-03-15 21:24:12] WARNING _http.py:329: Retrying in 1s [Retry 1/5].
INFO 03-15 21:24:20 [model.py:533] Resolved architecture: OPTForCausalLM
INFO 03-15 21:24:20 [model.py:1580] Using max model len 2048
INFO 03-15 21:24:20 [scheduler.py:231] Chunked prefill is enabled with max_num_batched_tokens=8192.
PASSED
tests/v1/core/test_scheduler.py::test_jump_forward_tokens_stop_max_tokens '(ReadTimeoutError("HTTPSConnectionPool(host='huggingface.co', port=443): Read timed out. (read timeout=10)"), '(Request ID: 8255237f-d720-4e77-8a08-902dc701a3ea)')' thrown while requesting HEAD https://huggingface.co/facebook/opt-125m/resolve/main/preprocessor_config.json
[2026-03-15 21:24:38] WARNING _http.py:320: '(ReadTimeoutError("HTTPSConnectionPool(host='huggingface.co', port=443): Read timed out. (read timeout=10)"), '(Request ID: 8255237f-d720-4e77-8a08-902dc701a3ea)')' thrown while requesting HEAD https://huggingface.co/facebook/opt-125m/resolve/main/preprocessor_config.json
Retrying in 1s [Retry 1/5].
[2026-03-15 21:24:38] WARNING _http.py:329: Retrying in 1s [Retry 1/5].
INFO 03-15 21:24:44 [model.py:533] Resolved architecture: OPTForCausalLM
INFO 03-15 21:24:44 [model.py:1580] Using max model len 2048
INFO 03-15 21:24:44 [scheduler.py:231] Chunked prefill is enabled with max_num_batched_tokens=8192.
PASSED
tests/v1/core/test_scheduler.py::test_jump_forward_tokens_retained_for_unscheduled_requests INFO 03-15 21:24:52 [model.py:533] Resolved architecture: OPTForCausalLM
INFO 03-15 21:24:52 [model.py:1580] Using max model len 2048
INFO 03-15 21:24:52 [scheduler.py:231] Chunked prefill is enabled with max_num_batched_tokens=8192.
WARNING 03-15 21:24:52 [scheduler.py:273] max_num_batched_tokens (8192) exceeds max_num_seqs * max_model_len (1600). This may lead to unexpected behavior.
PASSED

================================================================================== warnings summary ===================================================================================
<frozen importlib._bootstrap>:488
  <frozen importlib._bootstrap>:488: DeprecationWarning: builtin type SwigPyPacked has no __module__ attribute

<frozen importlib._bootstrap>:488
  <frozen importlib._bootstrap>:488: DeprecationWarning: builtin type SwigPyObject has no __module__ attribute

../../miniconda3/lib/python3.12/site-packages/torch/jit/_script.py:362: 14 warnings
  /root/miniconda3/lib/python3.12/site-packages/torch/jit/_script.py:362: DeprecationWarning: `torch.jit.script_method` is deprecated. Please switch to `torch.compile` or `torch.export`.
    warnings.warn(

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
============================================================== 4 passed, 94 deselected, 16 warnings in 81.48s (0:01:21) ===============================================================
sys:1: DeprecationWarning: builtin type swigvarlink has no __module__ attribute

Pending / Future Work

  • xGrammar support via find_jump_forward_string() (requires safe FF-string to
    token conversion, non-trivial due to token boundary issues)
  • E2E tests with real model inference
  • Performance benchmarks

Acknowledgements

Thanks to @FredericOdermatt for the original implementation, @benchislett for
the thorough review, @mmoskal for the llguidance fast-forward documentation,
and the vLLM team for the excellent v1 architecture.

@github-actions

Copy link
Copy Markdown

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant