[oneraghavan]

Purpose

Fixes #42655.

This PR resolves a scheduler/worker state mismatch in the v1 engine that
crashes the engine with KeyError -> EngineDeadError under a specific
combination of features:

1. Async scheduling (scheduler runs ahead of the GPU worker)
2. Speculative decoding (placeholder spec tokens on requests)
3. Double streaming_update (same live request legally resumed twice)

Root cause

When a streaming session is rebuilt via _update_request_as_session()
(triggered by streaming_update on a live request), stale async
scheduling and speculative decode state (spec_token_ids,
num_output_placeholders) was not cleared.

This caused the scheduler to continue emitting
scheduled_spec_decode_tokens for the rebuilt request. However, the
rebuilt request was removed from and re-added to the worker's persistent
batch, so the worker's prev_req_id_to_index mapping no longer
contained it. When the worker's deferred spec-decode correction tried to
look up the previous batch row for that request, it crashed with a
KeyError (on v0.17.1) or silently corrupted state (on current main
where .get() is used).

Crash chain

1. Request is RUNNING with active speculative decode
2. First streaming_update arrives → queued on streaming_queue
3. Async scheduling continues emitting spec tokens for the request
4. Second streaming_update arrives → also queued
5. Request stops → _handle_stopped_request() rebuilds session via
   _update_request_as_session() without clearing spec/async state
6. Rebuilt request is re-scheduled with stale spec_token_ids
7. Worker's prev_req_id_to_index doesn't contain the rebuilt request
8. KeyError → EngineDeadError

Fix

Clear spec_token_ids and num_output_placeholders in:

• _update_request_as_session() — the core session rebuild path
• _handle_stopped_request() — the WAITING_FOR_STREAMING_REQ path
  (defensive, prevents stale state from leaking into the next chunk)

This matches the pattern already used by _preempt_request() (clears
spec_token_ids) and reset_prefix_cache() (clears
num_output_placeholders).

Test Plan

Added two regression tests in
tests/v1/streaming_input/test_scheduler_streaming.py:

1. test_update_request_as_session_clears_async_spec_state —
   Verifies that spec_token_ids and num_output_placeholders are
   cleared after session rebuild with stale async/spec state.

2. test_handle_stopped_request_clears_spec_state — Verifies that
   stale spec/async state is cleared when a resumable request stops and
   enters WAITING_FOR_STREAMING_REQ.

python -m pytest tests/v1/streaming_input/test_scheduler_streaming.py -v -k "clears"

Test Result

• Both new regression tests pass
• All existing streaming scheduler tests continue to pass (9/9)
• The pre-existing test_streaming_e2e_lifecycle failure is unrelated
  (fails on clean upstream main as well)

[gemini-code-assist]

Code Review

This pull request fixes a KeyError crash in the v1 scheduler by clearing stale speculative decoding and async scheduling state (such as spec_token_ids and num_output_placeholders) when rebuilding a streaming session or handling a stopped request. The reviewer points out a critical issue with this fix: simply resetting num_output_placeholders to 0 will cause state corruption because in-flight async steps from the prior chunk are still active. To resolve this, the cleared placeholders must be accumulated into async_tokens_to_discard so that the engine correctly discards their outputs when they return. Corresponding assertions should also be added to the regression tests.

[gemini-code-assist]

When resetting session.num_output_placeholders to 0 during a session rebuild, any in-flight async scheduled steps from the prior chunk are still active in the pipeline. If we do not accumulate these placeholders into session.async_tokens_to_discard, the engine will not discard their outputs when they eventually return. This will cause stale tokens from the previous chunk to be appended to the rebuilt session, leading to state corruption.

We should accumulate session.num_output_placeholders into session.async_tokens_to_discard before clearing it.

Suggested change

	session.spec_token_ids = []
	session.num_output_placeholders = 0
	session.spec_token_ids = []
	session.async_tokens_to_discard += session.num_output_placeholders
	session.num_output_placeholders = 0

[gemini-code-assist]

Similarly, when a resumable request stops and transitions to WAITING_FOR_STREAMING_REQ, any in-flight async scheduled steps must be marked for discarding. Otherwise, when those steps return, their outputs will be appended to the waiting request, corrupting its state before the next streaming update is applied.

We should accumulate request.num_output_placeholders into request.async_tokens_to_discard before clearing it.

Suggested change

	if request.spec_token_ids:
	request.spec_token_ids = []
	request.num_output_placeholders = 0
	if request.spec_token_ids:
	request.spec_token_ids = []
	request.async_tokens_to_discard += request.num_output_placeholders
	request.num_output_placeholders = 0

[gemini-code-assist]

We should also assert that session.async_tokens_to_discard is correctly updated to match the stale placeholders that were cleared.

Suggested change

	assert session.spec_token_ids == []
	assert session.num_output_placeholders == 0
	assert session.spec_token_ids == []
	assert session.async_tokens_to_discard == 4
	assert session.num_output_placeholders == 0

[gemini-code-assist]

We should also assert that request.async_tokens_to_discard is correctly updated to match the stale placeholders that were cleared.

Suggested change

	assert request.spec_token_ids == []
	assert request.num_output_placeholders == 0
	assert request.spec_token_ids == []
	assert request.async_tokens_to_discard == 3
	assert request.num_output_placeholders == 0

[oneraghavan]

Hi @njhill @WoosukKwon — could you take a look when you get a chance? This fixes a scheduler/worker state mismatch (KeyError → EngineDeadError) that occurs when async scheduling + speculative decoding + streaming session rebuild interact. Fixes #42655. CI needs the ready label to run. Thanks!

[oneraghavan]

Friendly ping — this PR is ready for review. Let me know if there are any changes needed or blockers. Happy to address feedback. Thanks!

[oneraghavan]

@njhill — Friendly ping for review. This fixes #42655 where an async double streaming_update can break the request-ID-to-batch-row mapping and crash the engine. Would appreciate a look when you get a chance. Thanks!

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @oneraghavan.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[sawsa307]

Hi @oneraghavan thank you for the PR! Closing as superseded by #58259: it clears spec_token_ids and num_output_placeholders at the single choke point (_handle_stopped_request) ahead of both session-rebuild paths, and additionally fixes the num_computed_tokens frontier and stale in-flight outputs — a strict superset of this PR. The diff also no longer applies (the async_tokens_to_discard field it edits was removed).