Repository navigation
perf: reuse speculative output buffers and combine status D2H - #2670
Merged
Merged
Conversation
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
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>
9 tasks done
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>
This was referenced Oct 29, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.
D2HQueueowns FIFO ordering, slot reuse and completion. Capacity comes frommax_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.SpecStatusQueuespecializes 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.ScheduledBatchOutputdetaches 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, inrocm/atom-dev:nightly_202610151507, using only physical GPU 4:64 → 63 → 64 → 1 → 128verifies 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.