Skip to content

perf: reuse speculative output buffers and combine status D2H - #2670

Merged
valarLip merged 3 commits into
mainfrom
codex/final-speculative-output-d2h
Oct 22, 2026
Merged

valarLip merged 3 commits into
mainfrom
codex/final-speculative-output-d2h

Conversation

@valarLip

@valarLip valarLip commented Oct 21, 2026 •

Copy link
Copy Markdown
Collaborator

Deferred speculative outputs retain pinned PyTorch storage through NumPy views, and their D2H destinations are recreated each step. Return independently owned NumPy arrays while reusing pinned destinations and completion events. Both prefill and verification now copy the two speculative status counts in one D2H operation.

  • D2HQueue owns FIFO ordering, slot reuse and completion. Capacity comes from max_num_seqs (with the existing two-request block-drafter warmup floor); changing batch size reuses storage and copies only the active region. No per-transfer handle is allocated.
  • SpecStatusQueue specializes only the output snapshot, copying each row once in scheduler order. The runner supplies contiguous [bonus, rejected] counts from the final post-broadcast verdict or a persistent prefill zero buffer, without a GPU packing kernel.
  • ScheduledBatchOutput detaches non-owning NumPy views and retains already-owned arrays. Received outputs survive later transfers and cleanup without retaining pinned Torch storage.

Validation on main 65a199ea8, in rocm/atom-dev:nightly_202610151507, using only physical GPU 4:

  • 433 passed, no failures or skips: transfer/storage, deferred speculative status, prefill postprocessing, DSpark, scheduler and pipeline-parallel tests, including real asynchronous copies.
  • Independent CPU tests with GPU visibility disabled and AITER imports prohibited: 6 passed, two GPU-only tests skipped as intended.
  • Profiling 64 → 63 → 64 → 1 → 128 verifies no new destination allocation after warmup and one status copy for both prefill and verification. Black, Ruff and diff checks pass.

The naming cleanup changes no runtime operations (checked by normalizing identifiers and comparing ASTs). Regression tests are limited to three new modules: transfer invariants with CPU doubles and real GPU copies, runner wiring/capacity, and output-storage ownership.

A same-process transfer microbenchmark against the earlier PR revision (2fe751209) used 31 randomized A/B pairs, 200 send/receive operations per block. Prefill paired median time decreased 34.2% at BS64 (95% bootstrap interval: −37.2% to −24.3%) and 19.4% with changing batch sizes (−24.4% to −15.5%). Draft and verified-status differences were inconclusive; their intervals include zero. These measure synchronous D2H send/receive round trips, not end-to-end inference throughput or latency, and compare against the previous PR implementation rather than main.

@github-actions

Copy link
Copy Markdown
Contributor

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

  • ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

  • ✅ Run after the PR is approved and Pre Checkin passes
  • ✅ Run immediately when an approval review is submitted
  • ✅ Can be requested before approval with labels
Label Tests
ci:full Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom Run native ATOM model accuracy tests
ci:vllm Run ATOM vLLM OOT model accuracy tests
ci:sglang Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2670 --add-label <label>

@valarLip valarLip changed the title perf: reuse speculative output buffers and combine MTP status D2H perf: reuse speculative output buffers and combine status D2H Oct 21, 2026
@zufayu
zufayu requested a review from yitingw1 October 22, 2026 01:13
@valarLip
valarLip merged commit 4438322 into main Oct 22, 2026
@valarLip
valarLip deleted the codex/final-speculative-output-d2h branch October 22, 2026 02:47
demonsan added a commit that referenced this pull request Oct 22, 2026
…copy (#2670)

After the rebase onto origin/main 4438322 the constructor copies non-owning NumPy views (_copy_array_view), which the ast-extracted class in test_gemm_fusion_r12_abort_path needs a numpy stand-in for. Also pins that the two coexist: a view is copied, an owning array kept, an output-less batch allocates nothing, and aborted_req_ids rides the same constructor.
jiayyu added a commit that referenced this pull request Oct 22, 2026
postprocess called `_settle_mixed_verdict` on every verified step and let
it return early for a pure one. main's #2670 tests drive postprocess on a
stub runner with no mixed path, which then failed on the missing method.
Gate the call on `is_mixed` so a pure step never reaches the mixed path;
the method keeps its own guard for direct callers.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
rbrugaro-amd added a commit to rbrugaro-amd/ATOM that referenced this pull request Oct 23, 2026
An all-greedy batch always verifies by exact match, so postprocess no
longer resolves the lossless inputs for it. Besides keeping the default
T=0 path untouched by the feature, this stops postprocess from reading
batch.req_ids and the drafter on that path, which broke
test_output_transfer_runner.py::test_postprocess_submits_prefill_or_final_broadcast_counts
after the rebase onto ROCm#2670.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
ZLkanyo009 added a commit that referenced this pull request Oct 26, 2026
Conflict in atom/model_engine/model_runner.py with #2670 (perf: reuse
speculative output buffers and combine status D2H), which replaced
send_to_cpu_async_draft / send_mtp_status_to_cpu_async with D2HQueue-backed
send_draft_token_ids_to_cpu and spec_status_queue.send(spec_counts, ...).

Resolution: take main's queue-based paths and move the hold into
D2HQueue.send, which now records each CUDA source on its stream after
queuing the copy. That covers both queue users on async_copy_stream (the
spec status, a view of the sampler's fresh verdict, and the draft ids).
send_to_cpu_async keeps _hold_for_copy_stream for the sampled ids and
logprobs. The GPU test is ported to the new API; its positive control now
disables Tensor.record_stream and runs for every path.

Co-authored-by: Cursor <cursoragent@cursor.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants