Skip to content

[Review only][CPU] Assemble auxiliary outputs on the execution path - #14

Open
aoshen02 wants to merge 4 commits into
r3_offloadfrom
codex/aux-cpu-assembly-20260915
Open

aoshen02 wants to merge 4 commits into
r3_offloadfrom
codex/aux-cpu-assembly-20260915

Conversation

@aoshen02

@aoshen02 aoshen02 commented Sep 15, 2026

Copy link
Copy Markdown
Collaborator

Review-only alternative: CPU execution-side assembly

Alternative to #13, not an additional upstream fix. Both target r3_offload;
do not merge the alternatives together. Duplicate-work checks found upstream
vllm-project#45635 and the existing fork #13; this alternative deliberately offers
a different lifecycle tradeoff for review, not a competing upstream PR.

Finish R3 processing inside sample_tokens() before returning. The next worker
RPC cannot clean request state until this completes. Keep the existing async
GPU-to-CPU copy, wait for it after draft proposal, then assemble on CPU. Remove
the pending-output event/wait handshake; output consumption no longer mutates
the connector. No new output type or no-forward wrapper.

Tradeoff: CPU assembly and its copy wait enter the execution critical path.
Background storage remains unchanged.

The runner calls connector-local worker.finish_aux_output(); completion
implementation is no longer added to the public AsyncOutput class. Replace the
retained step metadata object with a presence flag and reuse the computed end
position.

Validation

  • /home/aoshen/vllm/.venv/bin/python -m pytest tests/distributed/aux_output_connector -q: 58 passed (16.83s). Includes completion-order coverage.
  • Full pre-commit hooks on changed files passed during commit.
  • Independent Claude Code review agrees with the execution-order proof.
  • GB200 UniProc with VLLM_GPU_SYNC_CHECK=error: 73 store/config tests and both HTTP R3 tests pass (Slurm 29066).
  • Raw-snapshot-to-HTTP oracle on 0ed49a1: 64 requests, 16,778,176 prompt/decode R3 rows match exactly (Slurm 29059). This does not independently validate model routing computation or D2H.
  • Real MTP: 32 requests complete; 2,715 draft tokens, 93 accepted, 2,622 rejected. R3 shape/range checks pass (Slurm 29061), not an independent MTP value oracle.
  • Forced memory-pressure preemption: 64 requests complete with 167 preemptions and R3 shape/range checks passing (Slurm 29080 CPU stage).
  • Same-node GB200 throughput on 0ed49a1, 256K prompt / concurrency64 / one output token: 82,375 and 82,312 tokens/s, mean82,343, 2.73% below native image R3 (84,652). Actual weights: DeepSeek-V4-Flash-Base; TP1/DP8/EP. This is primarily prefill throughput. Both chunked/short prefix-reuse checks pass; newly computed cold/hot decode rows are not claimed numerically identical.
  • CPU/GPU finish-path profiles completed (Slurm 29093). They identify GPU per-block publication copies as a batching target, but do not isolate CPU's lost async overlap from storage/assembly overhead. Summed worker/event-wait times are not end-to-end regression percentages.

AI assistance (Codex and Claude Code) was used. For human line-by-line
review and hardware validation; not ready to merge.

aoshen02 and others added 4 commits September 15, 2026 23:42
Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>
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