[aoshen02]

Purpose

Split from #62044 (by @faresobeid): only the scheduler-side hash fix. The store eviction fix is #62392.

The scheduler sent a request's block hashes to the worker as soon as they existed. A prompt's hashes all exist before its first chunk runs, so with chunked prefill the worker pinned one store object for every prompt block, including blocks whose KV had not been computed yet. When those objects were still stored from an earlier turn, the extra pins could push the KV-sized store over capacity and make put() raise BlockObjectStoreError in the writer thread.

AuxOutputSchedulerConnector now sends only the hashes of blocks covered by num_computed_tokens + num_scheduled_tokens (or num_computed_tokens for hash-only updates and finished requests). It takes the hash block size from the scheduler to do that.

As a side effect, when a request finishes and its last sampled token completes a block, that block's hash is no longer sent unless a step that computes that token is already in flight: otherwise the token's KV is never computed, so the worker held a reference to a key with no stored object.

Test Plan

pytest tests/v1/core/test_scheduler.py tests/distributed/aux_output_connector

• test_routed_experts_hashes_follow_scheduled_tokens runs a 160-token prompt with a 64-token budget and 20 decode steps through the real scheduler, and checks the number of hashes sent after each step against fixed values: 4, 8 and 10 during prefill, then 11 once the decode tokens complete the next block.
• test_scheduler_connector_sends_only_computed_hashes_off_schedule (hash-only update, finish) schedules a request's first 4-token block, then lets a settled token complete a second hash block while only 7 tokens are computed. Neither a hash-only update nor the terminal event may send the second hash; the hash-only update sends it once the eighth token is computed.
• test_worker_does_not_pin_unscheduled_prompt_blocks drives the scheduler connector and the worker with a 4-object store. A first request stores 4 prompt blocks and finishes. A second request with the same prompt runs its first 4-token chunk next to a request that stores a new block.

Test Result

• With this change: the new tests pass; test_scheduler.py and tests/distributed/aux_output_connector pass (288 tests).
• On main, with the new constructor argument removed from the tests, all new tests fail:
  • the scheduler sends all 10 prompt hashes at the first step ([10, 10, ...] != [4, 8, ...]);
  • after the first block is sent, the hash-only update and the terminal event both send the second hash although only 7 of its 8 tokens are computed;
  • the worker pins all 4 stored prompt blocks, so storing the new block raises BlockObjectStoreError: auxiliary output store cannot retain the requested batch: limit=4 objects, the writer-thread failure seen in production.
• Reverting only one of the three bounds to "send every hash" fails only the tests for that path: the scheduled-step bound fails the scheduler and worker tests, and the hash-only and terminal-event bounds each fail their case of test_scheduler_connector_sends_only_computed_hashes_off_schedule. The scheduler test cannot see the terminal-event bound, because its request finishes with 179 computed tokens and 11 complete hash blocks either way.

🤖 Generated with Claude Code

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #95918 for commit 6c2a48586d61.

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #95932 for commit 11ecb60224b7.

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #95936 for commit edbb3085f9bc.

ℹ️ This PR is 2 commits behind upstream main. Merged upstream main into this branch as edbb3085f9bc before starting CI. Pull the branch before pushing more changes.

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #95941 for commit 6d7df130dcc5.

ℹ️ This PR is 2 commits behind upstream main. Merged upstream main into this branch as 6d7df130dcc5 before starting CI. Pull the branch before pushing more changes.