[AuxOutput] Add block-keyed storage for routed-expert outputs

[xhx1022]

AuxOutput Connector: block-keyed storage for routed-expert outputs

Summary

This PR introduces the first AuxOutput Connector implementation for vLLM.
The first supported auxiliary output is routed-expert output (R3): the expert IDs selected for
every token executed by an MoE model.

R3 cannot remain a request-local GPU buffer when prefix caching and async
scheduling are enabled:

• a prefix-cache hit skips model execution, so cached tokens produce no new R3;
• physical GPU/KV slots are reused and are not stable object identities;
• speculative decoding can execute rows that are later rejected;
• one result may contain cached blocks, new blocks, and a partial tail.

The connector stores immutable R3 blocks using KV-compatible block hashes. KV
and R3 for the same prefix can therefore be reused together. This PR provides a
bounded local SHM backend; distributed storage and more auxiliary output types are
follow-ups.

Why an AuxOutput Connector?

The model capture hook only obtains a tensor from the current forward pass. A
separate lifecycle layer must decide:

• which logical tokens are accepted;
• which blocks are reusable across requests;
• when temporary GPU/CPU buffers can be released;
• where immutable blocks are stored; and
• how cached and new blocks form one ordered result.

These responsibilities also apply to future DSA and logprobs auxiliary outputs. The
connector keeps them out of the model runner and gives storage backends one
common object lifecycle. It is not a second KV Connector or a generic Scheduler
sidecar.

Architecture

+-------------------------------------------------------------------+
| API / OutputProcessor                                             |
| Current: return complete R3                                       |
| Follow-up: return ordered auxiliary output keys                           |
+-------------------------------+-----------------------------------+
                                |
                                v
+-------------------------------------------------------------------+
| EngineCore                                                        |
| Scheduler -> AuxOutputSchedulerConnector                           |
|              - logical request progress                           |
|              - KV-compatible hashes                              |
|              - terminal events                                   |
+-------------------------------+-----------------------------------+
                                | per-step metadata
                                v
+-------------------------------------------------------------------+
| GPU Worker                                                        |
| Model Runner -> router capture -> async GPU-to-CPU copy            |
|              -> AuxOutputWorkerConnector                           |
|                 - accepted-row assembly                           |
|                 - block publication/materialization               |
+-------------------------------+-----------------------------------+
                                |
                  +-------------+-------------+
                  |                           |
                  v                           v
     +--------------------------+  +--------------------------+
     | Current                  |  | Follow-up                |
     | Bounded local SHM        |  | Mooncake distributed     |
     | background publication   |  | immutable objects        |
     +--------------------------+  +--------------------------+

The Scheduler connector owns logical progress and hashes, not payloads. The
Worker connector owns accepted R3 rows and publication. BlockObjectStore owns opaque
immutable blocks and retention, but knows nothing about tokens or speculative
decoding.

Auxiliary output types

An execution auxiliary output is an auxiliary model output that is aligned to a stable
logical coordinate and may need to outlive the worker step that produced it.

Auxiliary output	Coordinate	Reuse	Status
Routed experts (R3)	Executed token	Prefix block	This PR
DSA routing/index data	Executed token	Prefix block	Follow-up
Logprobs/prompt logprobs	Predicted token	Prefix block	Follow-up
Multimodal auxiliary data	Request/object	Request only	Exploratory
Sampling candidate IDs	Sampling step	Request only	Exploratory

Prefix-block auxiliary outputs are content-addressed and reusable with a matching
prefix. Request-only auxiliary outputs use the same storage contract but do not
participate in prefix hits. Reward values, tool results, finish reasons, and
Scheduler status remain normal request/control-plane metadata.

Current PR scope

This PR implements:

• Model Runner V2 R3 capture;
• block-keyed SHM publication and prefix materialization;
• async scheduling, chunked prefill, speculative rejection, and preemption;
• background publication, bounded capacity, and KV-aligned retention;
• terminal-only R3 delivery for streaming output;
• fail-closed checks for unsupported configurations.

It does not implement Mooncake, public auxiliary output keys, additional auxiliary output
fields, PD disaggregation, or new parallel topologies.

When disabled, no auxiliary output capture buffer, Store, per-step metadata, or R3 copy
is created on the normal serving path.

The existing --enable-return-routed-experts flag and routed_experts response
field are unchanged. The new optional configuration is exposed as
--aux-output-config (AuxOutputConfig in Python).

Follow-ups

This PR
  R3 lifecycle + local SHM
        |
        v
Unified object/profile contract + ordered keys
        |
        v
Direct Mooncake backend
        |
        v
Joint KV + mandatory-auxiliary output prefix readiness
        |
        v
DSA, logprobs, and additional topology support

The main planned work is:

1. define a backend-independent object/profile contract and ordered key output;
2. add a direct Mooncake backend for external trainer/JAX consumers;
3. admit remote prefix hits only when KV and mandatory auxiliary outputs are ready;
4. add DSA and logprobs through small field adapters, not new lifecycle code;
5. validate CPU KV offload, PD, PP/DCP/PCP, and multi-writer deployments before
   enabling them.

Validation

Naming-update validation: 96 passed, 1 skipped across configuration, store,
capture, Scheduler, output, and Ray tests. An AST comparison verified that all
24 renamed files preserve the original logic, allowing renamed identifiers,
messages, and sorted imports. Additional smoke checks passed for the existing
R3 flag, the new JSON/Python configuration, compilation hashes, serialization,
object keys, and R3 bytes. The latest naming-only update was also exercised on GB200; results and
limitations are below.

VLLM_TARGET_DEVICE=cpu .venv/bin/python -m pytest -q \
  tests/config/test_aux_output_config.py \
  tests/distributed/aux_output_connector/test_store.py \
  tests/model_executor/test_routed_experts_capture.py \
  tests/v1/executor/test_ray_utils.py tests/v1/core/test_scheduler.py \
  tests/v1/core/test_async_scheduler.py tests/v1/engine/test_output_processor.py \
  -k 'aux_output or routed_experts or pending_output or store or worker or detach_zero_copy'
pre-commit run

GB200 rename A/B (2026-09-08)

Compared 85f100032e (before rename) with f99d53b521 (after rename) on the
same two GB200 nodes (8 GPUs), using the locally cached DeepSeek-V4-Flash-Base
FP8 weights. Both arms enable R3, MRV2, MP, TP1/DP8/EP8, async scheduling and
prefix caching. They use identical native libraries/dependencies; Python source
injection was verified by per-file SHA checks.

Workload: 64 concurrent requests, each with 262,144 input tokens and one output
token; synthetic token IDs, with output constrained to token 42. GPU memory
utilization is 0.90, max batched tokens 16,384, block size 256, and batch-invariant
mode is off. This is a serving/R3 regression test, not a model-quality evaluation.

Measurement	Before rename	After rename
Total throughput (input + output tokens/s)	82,741.42	82,568.30
Total duration	202.768 s	203.193 s
Preemptions	0	0
Prefix-hit tokens in the timed workload	0	0

Throughput changed by −0.21% in this single A/B: no obvious regression, not a
statistical significance claim. Startup and warmup are excluded. Both arms still
reported the same first-long-request CombineTopkSwaIndicesKernel.kernel JIT
warning, so this is not a strictly all-kernels-warm measurement.

Correctness checks:

• All 128 long responses passed token-count, R3 row-count, shape and dtype checks.
• Separate cold/hot tests cover eight 1K-prompt/16-decode requests and eight
  64K-prompt/1-decode requests per version. The latter exercises multi-step
  chunked prefill. All requests' reused-prefix R3 SHA hashes match exactly:
  768 rows/request for the short case and 65,280 rows/request for the long case.
• All eight short cold requests have identical full R3 hashes across versions.

Limitation: long-context full-output bitwise equivalence did not pass. None
of the 64 full 256K R3 hashes matched across the two server runs. In the saved
request-0 sample, differences begin at row index 2051 and affect 55.38% of
elements. The same candidate version also differs when comparing the identical
first 64K input prefix in 256K versus 64K cold requests (first difference at
2051; 46.44% of elements). Recomputed decode R3 can differ between cold/hot runs
as well, already before the rename. These observations do not establish the
underlying cause or prove long-output equivalence; they must not be presented
as a full R3 correctness pass.

This run validates the response structure, exact cached-block reuse and short
cross-version outputs. It does not add new forced-preemption, speculative
decoding or free-generation quality coverage. Slurm job 24176 completed with
exit code 0; test containers and node allocations were released.

Development validation covers focused unit tests, async scheduling, forced
preemption, chunked prefill, speculative rejection, cold/hot prefix reuse,
streaming aggregation, SHM capacity/LRU behavior, pre-commit, and GB200/H200
throughput experiments reported in the PR discussion.

The latest update also moves R3 GPU-to-CPU transfer onto the existing async
output copy stream, avoiding a synchronous copy on the model execution critical
path.

AI assistance was used for implementation, testing, performance experiments,
and review. The submitting human reviewed the accepted changes and validation.

[mergify]

Documentation preview: https://vllm--45635.org.readthedocs.build/en/45635/

[mergify]

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

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

[mergify]

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

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

[njhill]

/ci retry

[github-actions]

✅ Triggered Buildkite CI #87737 for commit eaeea04329f8, running 4 failed step(s) from Buildkite CI #87665.

[njhill]

/ci run

[github-actions]

✅ Triggered Buildkite CI #87777 for commit cb1872d8cbcf.

[njhill]

@aoshen02 as currently written there's a deadlock when using the uniproc executor (which is the reason for the CI failures)

Suggested fix from K3:

Make the worker connector tolerate step N's output being processed after step N+1's metadata, which is the actual ordering under async scheduling:

• Remove the finished.wait() from begin_step (worker.py:283).
• Replace the single _pending_output slot with a small collection (bounded by max_concurrent_batches = 2 here; each pending holds one GPU snapshot of a few MB), and drop assert self._pending_output is None in prepare_output (worker.py:129). complete() removes its own entry.
• Defer the destructive metadata effects (popping self._requests[request_id], store-key release, and the final-hash delivery for finished requests) until no outstanding pending output covers that request. This matters: a request preempted during schedule(N+1) (scheduler.py:1488 calls request_finished from _preempt) can appear in metadata(N+1)'s finished_requests while its step-N rows are still unprocessed — today the blocking wait papers over that race in mp mode; without the wait it becomes a KeyError in process_output (worker.py:163) unless teardown is deferred.

[aoshen02]

@aoshen02 as currently written there's a deadlock when using the uniproc executor (which is the reason for the CI failures)

Suggested fix from K3:

Make the worker connector tolerate step N's output being processed after step N+1's metadata, which is the actual ordering under async scheduling:

• Remove the finished.wait() from begin_step (worker.py:283).
• Replace the single _pending_output slot with a small collection (bounded by max_concurrent_batches = 2 here; each pending holds one GPU snapshot of a few MB), and drop assert self._pending_output is None in prepare_output (worker.py:129). complete() removes its own entry.
• Defer the destructive metadata effects (popping self._requests[request_id], store-key release, and the final-hash delivery for finished requests) until no outstanding pending output covers that request. This matters: a request preempted during schedule(N+1) (scheduler.py:1488 calls request_finished from _preempt) can appear in metadata(N+1)'s finished_requests while its step-N rows are still unprocessed — today the blocking wait papers over that race in mp mode; without the wait it becomes a KeyError in process_output (worker.py:163) unless teardown is deferred.

Thanks, working on it.

[mergify]

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

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

[mergify]

Hi @xhx1022, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #88715 for commit 0a0b73800b56.

[aoshen02]

/ci run

[mergify]

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

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

[github-actions]

✅ Triggered Buildkite CI #88861 for commit fbe375befe22.

[njhill]

@aoshen02 I don't think the latest changes are the right approach either. I can have a go at refactoring.

[aoshen02]

@aoshen02 I don't think the latest changes are the right approach either. I can have a go at refactoring.

Sure.

[mergify]

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

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

[aoshen02]

@njhill Hi Nick, could you take a look at xhx1022#14? I think it's a simple way to solve the race problem.
The idea is simple: finish processing step N’s R3 before the worker starts step N+1, instead of doing it later in get_output().
For example, request A has 3 buffered R3 rows. Step N produces the 4th row, completing a block. Meanwhile, the scheduler decides to preempt A in step N+1.

EngineCore process                 GPU worker process
------------------                 ------------------
schedule N ----------------------> execute N / sample_tokens N
                                     produce the 4th R3 row
schedule N+1                         start async GPU-to-CPU copy
  decide to preempt A                run draft proposal
  send N+1 ------------------+       wait for the copy
                             |       combine the 3 old rows + 1 new row
                             |       submit the block for storage*
                             |       return
                             |
                             +----> execute N+1
                                     apply A's cleanup
                                     safe: N already used the old rows

The tradeoff is that the copy wait and CPU assembly now sit on the execution critical path.