[ZLkanyo009]

Summary

Keep the GPU sources of the side-stream D2H copies on tokenIDProcessor.async_copy_stream alive until those copies are known to have completed. Each source is held next to the completion event that already exists for its copy, and released only after that event has been waited on. Without this, the caching allocator can hand the block to the next step's allocation on the default stream while the copy is still queued, and the host reads the next step's bytes as this step's MTP status / sampled ids / draft ids / DSpark ell.

Under MTP this is what produces the GDN _causal_conv1d_update_kernel illegal memory access seen with in-process LMCache offload (#2636): a stale num_bonus of 16384 becomes num_accepted_tokens = 16385, which the conv kernel adds to the state offset unchecked.

No record_stream, no new CUDA events: the earlier _hold_for_copy_stream / record_stream version is gone, and every hold reuses an event the code already had.

Motivation

https://github.com/wuhuikx/LLM-Booster/issues/157 — with lmcache_offload + MTP 3 on Qwen3.8-Flash-Next (#2636), AgentX conc 32/64 crashed within 13–54 min with an illegal memory access in the GDN conv update kernel. The faulting offset (slot * 6 + num_accepted_tokens - 1 = 16487, 16487 mod 6 = 5) is impossible for any valid num_accepted_tokens (1..4), so the bad value is num_accepted_tokens, not the slot.

num_accepted_tokens comes from num_bonus + 1, and num_bonus comes from the spec-status D2H queued on async_copy_stream (now SpecStatusQueue, a D2HQueue). Its source is a view of the rejection sampler's verdict, a fresh torch.empty on the default stream every step, and nothing holds it after postprocess: once the last reference is dropped the block is free for the default stream's next allocation. If the side-stream copy runs after that allocation is written, the host gets the new contents.

The copy runs late whenever async_copy_stream, or a HIP hardware queue it shares with other streams (default GPU_MAX_HW_QUEUES=4), has work in front of it. In-process LMCache adds its own pack/copy streams with kernels on them, which is why this only showed up with offload on; without MTP there is no status D2H, which is why spec 0 never crashed. The bug itself is on main, so the fix is proposed against main.

Caught in production with host-side checks added only to a debug build (not in this PR), c32 + 160 GiB CPU pool, #2636 + #2606:

[L1DBG-MTP-STATUS] num_bonus=[16384, 1, 1, ...] prev_batch.req_ids=[757, 798, ...]
[L1DBG-GDN-OOB] bad_nat_rows=[0] row=0 req=757 nat=16385 num_bonus_sched=16384 state_slots[0]=[5]

16384 is --max-num-batched-tokens: the next step's prefill metadata landed in the freed verdict block. The step was a decode step immediately followed by a prefill step while LMCache was storing 32k tokens on its own stream. The request's state load had succeeded, so offload correctness is not involved.

A micro-benchmark (bs 32, the same send/recv code path) reproduced it deterministically (measured with the first, record_stream version of this fix):

queued ahead of the side-stream copy	current code	with the fix
nothing	0/200	0/200
one 256 MiB D2H on async_copy_stream	200/200 stale	0/200
16 streams × 10 4096² matmuls (GPU_MAX_HW_QUEUES=4)	99/100 stale	0/100
same, GPU_MAX_HW_QUEUES=1	0/100	—

Changes

Contract (commented at async_copy_stream): every D2H queued on that stream keeps its GPU source referenced until the copy is known to have completed. The five sources on that stream:

source	where the reference lives	released after
sampled token ids	cpu_tensor_handle entry (cpu_tensor, copy_done, (gpu_tensor, gpu_logprobs))	recv_async_output waits on copy_done, then pops the entry
sampled logprobs	same entry (same copy, same copy_done)	same
spec status (verdict[:2], the crash path)	_D2HSlot.source of SpecStatusQueue	D2HQueue.receive after copy_done.synchronize()
draft ids	_D2HSlot.source of the draft-id D2HQueue	same
DSpark ell (record_ell)	_ell_slot_source[slot], next to _ell_slot_event[slot]; one-off buffers: _ell_oneoff_inflight as (event, ell) with the copy's own event	_ell_stage reuses the slot after its event completed; one-off entries are pruned once their event has completed

• ell is deliberately not held through _ell_pending: that list evicts entries whose copy may still be in flight (del self._ell_pending[:-_MAX_ELL_INFLIGHT]), and ell_by_req = {} clears it outright.
• D2HQueue.drain() waits for every pending transfer, then drops its output and source. tokenIDProcessor.clean() waits on every queued copy_done and drains both queues before it rebuilds them, so dropping the handles cannot free a source under an in-flight copy.
• Removed tokenIDProcessor._hold_for_copy_stream and both record_stream calls.

Tests

• tests/test_side_stream_d2h_source_lifetime.py (GPU only, skipped on CPU): parks the copy stream with a torch.cuda._sleep spin on a stream that does not share the default stream's hardware queue, queues the copy, drops the source, allocates and writes the next step's tensors, and checks the host got this step's values. Paths: spec status, sampled ids, logprobs, draft ids, DSpark ell (ring slot), DSpark ell (one-off buffer), and clean() with all copies still in flight. The positive control monkeypatches the holds away (for clean(): skips its wait, holds intact) and requires every path to race, so the test cannot pass vacuously.
• CPU unit tests: test_output_transfer.py (a slot's source is alive at its copy_done wait and released after receive / drain), test_dspark_scheduler.py (ell outlives its _ell_pending entry and the setter, and is released only when its slot is reused after the event; one-off ell held until its own event lands), test_output_transfer_runner.py (handle owns ids and logprobs until recv_async_output has waited; clean() waits on every queued copy before dropping it; needs AITER to import the runner, like the rest of that file).

Accuracy

GSM8K (shared/scripts/gsm8k.sh: 5-shot, chat, max_tokens=12288, num_concurrent=64, temperature 0), MTP 3 with real verification (no --spec-decode-acceptance-length):

tree	offload	strict-match	MTP acceptance
main 40db24c2 + this PR (hold-refs version)	off	0.978	71.6%
main 5fea029d + this PR (record_stream version)	off	0.9795	72.5%
main 5fea029d (control)	off	0.9788	72.3%
#2636 + #2606 + this PR (record_stream version)	lmcache_offload (160 GiB CPU)	0.9788	72.2%
reference: #156, MTP 3	off	0.9757

Performance

Measured with the first, record_stream version of this fix; the hold-refs version keeps the same blocks out of the allocator for about the same time (until the host waits on the copy, which it does the next step anyway), so no difference is expected.

AgentX (inferencex-agentx-mvp, seed 42, 3600 s profiling, 1 GPU per point, MTP 3 with --spec-decode-acceptance-length 2.32 as in the original crash runs), effective_total_throughput per GPU:

conc	CPU pool (state)	tree	tok/s/GPU	requests	errors	TTFT mean	ITL p90	crashed
32	160 GiB (32)	#2636 + #2606 + this PR	81,141	2,786	0	15.4 s	32.1 ms	no
48	224 GiB (64)	#2636 + #2606 + this PR	69,959	2,621	0	46.1 s	38.0 ms	no
64	288 GiB (64)	#2636 + #2606 + this PR	57,369	2,277	0	92.5 s	47.0 ms	no
64	288 GiB (64)	same, repeat	61,901	2,458	0	85.8 s	43.4 ms	no
32	96 GiB (32)	same	37,324	1,463	0	52.9 s	75.7 ms	no
32	160 GiB (32)	#2636 + #2562 + #2606 + this PR	82,938	2,853	0	14.5 s	30.1 ms	no
32 / 48 / 64	no offload	#2636 + #2606 + this PR	31,373 / 31,662 / 30,401	1,288 / 1,286 / 1,246	0			no

"#2636" here includes its 183cfec2 (one LMCache lookup socket per server); before that commit, servers started together in one container could query each other's LMCache and those numbers are not comparable.

Without this PR, 4 of 5 one-hour offload runs on the first machine crashed 13–54 min in (c32 + 160 GiB at about 70.8k tok/s/GPU before the crash, c64 + 288 GiB at about 55.7k; c32 + 96 GiB crashed twice), and on this machine the debug check caught num_accepted_tokens = 16385 12 min (c32 + 160 GiB) and 17 min (c64 + 288 GiB) into profiling. With it, every point above ran the full 3600 s. The two c64 runs differ by 8%. The first overlapped with up to 6 other servers on the host and the repeat ran its last 30 min with only one, so host-side LMCache copy bandwidth is the likely cause (not measured); treat c64 as ~57–62k.

How to use

No flags. Nothing changes for users other than the crash going away.

Test plan

Hold-refs version (88a00dd9, main 40db24c2 merged), MI355X, rocm/atom-dev:latest + CI aiter wheel amd_aiter-0.1.1.dev1+g508736740:

• black 26.10.1 / ruff 0.17.2 clean on changed files.
• tests/test_side_stream_d2h_source_lifetime.py: 14 passed, 5/5 runs on one GPU plus one run on a second GPU. Control counts, 50 trials per path: with the holds 0/50 stale on every path; without them 50/50 stale on every path (spec status, sampled ids, logprobs, draft ids, ell slot, ell one-off, clean()).
• Same test on unmodified main 40db24c2: all 7 paths fail, 20/20 trials each read the next step's tensor.
• GPU run of test_output_transfer.py, test_output_transfer_runner.py, test_spec_deferred_status_queue.py, test_dspark_scheduler.py, test_gdn_num_accepted_tokens.py: 111 passed, 2 skipped (opt-in) vs 103 passed, 2 skipped on main; the 8 extra are the new CPU tests.
• CPU unit suite (.github/scripts/run_unit_tests.sh, CPU torch, no aiter): 3 failed / 9814 passed / 1007 skipped vs 3 / 9808 / 1006 on main 40db24c2. Same 3 failures (test_graph_marker x2, test_sparse_indexer_side_effect_op); the per-test diff is exactly the 6 new CPU tests and the GPU-only file skipping.
• Qwen3.8 MXFP4 server on this branch (no offload, MTP 3, real verification): GSM8K above.

Earlier (record_stream version, same paths held):

• [Qwen 3.8] LMCache offload for QSA + GDN + PLE state #2636 + fix(qwen4_exp): load unscaled Quark MTP weights as BF16 #2606 + this PR: AgentX c32 + 160 GiB, c48 + 224 GiB, c64 + 288 GiB (twice) and c32 + 96 GiB, each for the full 3600 s: no crash, 0 errors (previously c32 crashed in ~19 min, c64 in ~20 min).
• GSM8K on the same tree with lmcache_offload + MTP 3 real verification: 0.978, acceptance 72.1%.
• [Qwen 3.8] LMCache offload for QSA + GDN + PLE state #2636 + [ATOM] accelerate lm cache: load the CPU prefix when the Qwen3.5 HBM hit is off the chunk grid #2562 + fix(qwen4_exp): load unscaled Quark MTP weights as BF16 #2606 + this PR: AgentX c32 + 160 GiB, full 3600 s, no crash, 0 errors. [ATOM] accelerate lm cache: load the CPU prefix when the Qwen3.5 HBM hit is off the chunk grid #2562 does not touch this path and is not required.

Limitations

• This fixes the stale host copy. The GDN conv kernel still trusts num_accepted_tokens; a bounds assert there would have turned a silent illegal access into a clear error, but it costs a sync per step and is left out.
• Unrelated to the crash, but found while investigating it on [Qwen 3.8] LMCache offload for QSA + GDN + PLE state #2636: about 1.6–41% of state-tier loads fail (c32 + 160 GiB 1.6%, c48 8.6%, c64 15–18%, c32 + 96 GiB 41%) because state images and KV share one LMCache LocalCPU pool and LMCache's LRU evicts state images under KV pressure behind the state codec's back (the KV leg reports Retrieved, the pair still fails and recomputes). That is a capacity issue in [Qwen 3.8] LMCache offload for QSA + GDN + PLE state #2636 and is tracked there, not here.

[github-actions]

🏷️ 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 2654 --add-label <label>

[yitingw1]

Looks correct — I traced the three record_stream call sites and they're the only async_copy_stream D2H users in tokenIDProcessor. One sibling instance isn't covered though: atom/spec_decode/dspark_verify.py::record_ell queues a D2H of ell (a fresh per-step GPU tensor from compute_ell) on this same async_copy_stream, with no record_stream — same race, not touched here. Worth a one-line follow-up (self.runner.tokenID_processor._hold_for_copy_stream(ell) after queuing the copy), either in this PR or a quick follow-up.

[ZLkanyo009]

Merged origin/main at 5fea029 (#2606) into this branch (merge fb6bc96), and added one commit on top (07c728d). No rebase or force-push.

Conflict: atom/model_engine/model_runner.py. #2670 replaced send_to_cpu_async_draft / send_mtp_status_to_cpu_async with D2HQueue (output_transfer.py); the spec status is now a single copy of verdict[:2]. I took main's queue-based code and moved the hold into D2HQueue.send: after queuing the copy it calls record_stream(self.stream) on CUDA sources. That covers both queue users, the spec status (a view of the fresh verdict) and the draft ids, plus any future ones. send_to_cpu_async still holds the sampled ids and logprobs through _hold_for_copy_stream.

@yitingw1 thanks for spotting record_ell. 07c728d adds self.runner.tokenID_processor._hold_for_copy_stream(ell) after the copy is queued, and the GPU test now covers that path too.

GPU test: test_side_stream_d2h_source_lifetime.py is ported to the new API and covers four paths: spec status, sampled ids, draft ids and DSpark ell. Its positive control replaces Tensor.record_stream with a no-op and requires every path to race. It was order-dependent before: HIP maps streams onto a few hardware queues, and when the copy stream shared the default stream's queue, the next step's write waited behind the parked copy. The test now picks a copy stream that has its own queue and parks it with a spin kernel rather than a large D2H. It passed 5/5 on two GPUs.

Verified (MI355X, rocm/atom-dev:latest + CI aiter wheel):

• black/ruff (CI versions) clean.
• CPU unit tests (.github/scripts/run_unit_tests.sh): 3 failed / 9721 passed. These are the same 3 failures as unmodified main; the new GPU test skips on CPU.
• GPU: the test above plus test_output_transfer*.py, test_spec_deferred_status_queue.py, test_dspark_scheduler.py and test_gdn_num_accepted_tokens.py: 111 passed, 2 skipped (opt-in). Main gives 103 passed, 2 skipped on the same files without the new test.
• Qwen3.8 GSM8K, full 1319, MTP 3 with real verification: 0.9795 (main + fix(qwen4_exp): load unscaled Quark MTP weights as BF16 #2606 control: 0.9788). Acceptance rate 0.725 (main 0.723).

[ZLkanyo009]

@yitingw1 Changed it as you suggested, in 88a00dd9 (main 40db24c2 merged first, no rebase / force-push):

• send_to_cpu_async now puts the GPU ids and logprobs into the cpu_tensor_handle entry with cpu_tensor and copy_done. recv_async_output waits on copy_done, then pops the entry, and the pop drops the GPU references.
• Same pattern in D2HQueue (the spec status, the path that crashed, and the draft ids): each slot holds its source, and receive clears it after copy_done.synchronize().
• DSpark ell: held next to its ring slot's _ell_slot_event and released when _ell_stage reuses the slot after that event. A one-off buffer keeps ell with that copy's own event. _ell_pending can evict an entry while its copy is still running, so it no longer holds the source.
• clean() waits on every queued copy_done and drains both queues before it drops them.
• _hold_for_copy_stream and both record_stream calls are removed. No new CUDA events: every hold reuses an event that was already there.

Verification (MI355X, CI aiter wheel):

• GPU lifetime test, 7 paths (spec status, sampled ids, logprobs, draft ids, ell slot, ell one-off, clean() with copies in flight): 14 passed, 5/5 runs, plus one run on a second GPU. Positive control, 50 trials per path: 0/50 stale with the holds, 50/50 stale with them removed, on every path. On unmodified main all 7 paths fail, 20/20 each.
• Related GPU unit files: 111 passed, 2 skipped (main: 103 / 2; the 8 extra are new CPU tests). CPU suite: 3 failed / 9814 passed, the same 3 failures as main; per test, the only difference is the 6 new tests.
• Qwen3.8 MXFP4, MTP 3 real verification, GSM8K full 1319: 0.978, acceptance 71.6% (reference 0.9757–0.978).

The PR description is updated. Since the mechanism changed after your approval, could you take another look?

[ZLkanyo009]

@valarLip could you take a look at this one?

Short version: side-stream D2H copies on async_copy_stream read sources that the default stream frees at the end of the step, so the caching allocator can hand that block to the next step before the copy runs. The host then reads the next step's data. Under MTP, a stale bonus count becomes num_accepted_tokens, and the GDN conv kernel goes out of bounds; we hit this crash under MTP + LMCache load.

This revision follows the earlier review suggestion: no record_stream and no new events. Each source is held next to the copy's existing copy_done event and released only after that event is waited on, which covers all 5 sources on that stream: spec status (verdict[:2]), draft ids, sampled ids, logprobs and the DSpark ell.

Validation is in the comment above. On unpatched main, the deterministic lifetime test fails on all paths (20/20 trials read the next step's tensor); on this branch it passes in 5 of 5 runs. GSM8K with real MTP-3 verification is 0.978, and the CPU suite matches main.