Repository navigation
feat(spec): capture the Qwen3.8 MTP draft steps in CUDA graphs (MRoPE) - #2663
Conversation
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
|
Hi, it seems this feature enable MTP graph on ATOM standalone server, not sglang-plugin, I am not sure if it will conflict with #2645 |
|
Read it end to end. Clean PR — small, focused, and the accuracy table is the right kind of evidence: acceptance identical bucket-by-bucket (72.36%, 3.171 toks/fwd, 0/1/2/3 distribution all within 0.1pp) proves the draft proposes the same tokens, which a single aggregate score wouldn't. I checked the load-bearing claim and it holds: Nothing blocking. Four small things. 1. 2. 3. Nothing pins the thing being fixed. (a) address stability of first = head._argmax_row_lens(torch.zeros(4, 16))
for rows in range(1, 20): # well past 8
head._argmax_row_lens(torch.zeros(rows, 16))
assert head._argmax_row_lens(torch.zeros(4, 16)).data_ptr() == first.data_ptr()(b) the new gating in _declare_draft_graphs. MRoPE + flag declares, MRoPE without it doesn't — that's the behavior switch of this PR, and it rides on getattr(self.model, "draft_graph_supports_mrope", False). Rename the attribute and the feature stops firing silently with every test still green. Same shape for reuse_draft_graph_step_buffers and draft_graph_hidden_state_shape; the last one is the one I'd worry about, since the generic fallback reads hc_mult and Qwen4Exp's is hc_count, so a typo there yields a plausible-looking shape rather than an error. 4. bs=8/16 are reported as "0" but look like a small consistent cost. Checked and clean, so nobody redoes it: passing end= does not change aiter's backend selection here (_available short-circuits on not fp32, and logits are bf16, so argmax is the only candidate and ragged never enters the decision; _dispatch's argmax arm ignores it too) — semantics match _full_rows exactly. _argmax_row_lens allocates eagerly rather than inside a capture, since DraftGraph.warmup runs _forward_and_epilogue before torch.cuda.graph(...). The other five topk_select call sites (sampler.py ×2, rejection_sampler.py ×3) are on the eager verify/output path, so there's no sibling of this bug to fix. And [3, bs] survives the downstream: _enter_decode_metadata index_selects on the last axis, _advance_decode_metadata turns fuse_mtp off at ndim≠1 and increments all three rows, and qwen4_exp_attn.prepare_mtp_decode reads positions.shape[-1]. |
…m(out=); one amax floor (#6833) * fix(topk): read no length tensor when end is None topk_select built its default `end` from an 8-entry LRU keyed on the row count. Capturing a ninth count evicted, and freed, the tensor an earlier graph had baked in; replaying that graph read reused memory (wrong draft tokens, then a memory fault, ROCm/ATOM#2663). Every backend now takes a full-rows build that reads each row's length from the runtime width, and is handed an empty sentinel (`row_lens.py`) in place of the tensor. Widths that are a whole number of vectors also compile the per-element row test out. Explicit lengths are clamped to the width, so an `end` past the row no longer selects the zeros an out-of-bounds load returns. The streaming selector's launch logic moves out of topk_select into aiter/ops/flydsl/topk/topk_per_row_radix_stream.py, beside the other backends' host wrappers. * feat(gemm): tgemm.mm(out=) writes into the caller's buffer `gemm_a16w16_out` (a mutating custom op) and `hipb_mm(out=)` let a caller whose consumer reads one fixed address take the GEMM result there. Every backend (asm, flydsl, opus, triton, skinny, hipBLASLt) writes into `out` directly; the untuned torch paths and a skinny call with a wider `out` take the result by copy. `-t out` in test_gemm_a16w16 checks each backend lands in `out` and matches its own `out=None` answer. * refactor(quant): one amax floor, aiter::kAmaxFloor, for dynamic quantizers The kernels seeded a group's amax with 1e-6, 1e-8, 1e-10 or 1e-12 depending on the file, so a fused quantizer and the reference one disagreed on all-zero groups (graph warmup, padded rows). They now share `kAmaxFloor = 1e-10` from csrc/include/quant_amax_floor.h, which is also ATOM's AMAX_FLOOR. The per-token KV cache quant floors the amax rather than the inverse scale, so the stored scale is the one used and never 0. The Triton dual RMSNorm + MXFP8 quant takes the same floor as an opt-in `amax_floor` (default 0, unchanged behaviour).
|
@valarLip Thanks for the root fix in ROCm/aiter#6833. It supersedes the The main part of this PR is independent of that bug, though. It lets an MRoPE draft (the Qwen3.8-Flash-Next MTP head) capture draft steps 1..k-1 in CUDA graphs: the From the tables above: a bs=1 decode iteration goes from 11.21 ms to 8.08 ms, AgentX C1 p90 ITL interactivity improves by 31%, and C2 throughput by 35%. Acceptance is identical bucket by bucket. In our MI355X Qwen3.8 runs this is the main reason the low-concurrency points (c1–c8) reach the GB300 curve. Was the close meant for the whole PR? Two options on my side:
|
|
you can keep this PR with some clean, the close should be done by my agent... |
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
… is graph-safe ROCm/aiter#6833 makes `topk_select(end=None)` read no row-length tensor, so the LM head no longer needs to hold one per captured row count. Back to main's `compute_argmax_token`. Also fold the batch-last capacity helper into `buffer()` and cover the Qwen3.8 MTP draft-graph declaration under MRoPE. Co-authored-by: Cursor <cursoragent@cursor.com>
|
@valarLip Reopened and cleaned up as you suggested. What changed
Validation on main
Full numbers are in the updated description. Given the smaller win, I'm fine either way: merge it as a small bs=1 improvement that keeps MRoPE drafts on the same graph path as the others, or close it if you'd rather not carry the batch-last staging. PTAL, thanks. |
|
Reopen + the On where the win went — I think you're missing a name. #2670 ("reuse speculative output buffers and combine status D2H") merged 10-22, the day after you opened this. Combining those D2H status reads removes the same stall draft graphs remove: the host blocks on the readback, the GPU drains waiting for the next launch batch. Two changes, one pool of headroom. That's also why it isn't a subtraction. The win is One thing you didn't account for: the branch itself went 8.08 → 8.75 ms. Most likely the methodology — running main and the branch at the same time on two GPUs of one node shares the host CPU, which is the exact resource under measurement. The one number I'd ask for: bs 8 and 16, from the same microbenchmark you already ran for bs 1/2/4. The draft side replays unconditionally once a size is captured ( Because I think you're underselling this. On main, Separately, and not this PR's problem: two greedy c1 runs of main matching on 17 of 300 questions deserves its own issue. At TP1/c1 the batching is fixed, so that's kernel-level nondeterminism — and every accuracy A/B in this repo is read against it. |
|
@yitingw1 thanks, re-measured with one server at a time. cc @valarLip Agreed on #2670: combining the status D2H removes the same host stall, and with the You were right about the method, too. Last time main and the branch ran at the same moment on two GPUs of one node, so they shared the host CPU, which is exactly what this PR saves. Method.
Decode iteration wall time (ms; median and range over 16 sessions per side; paired Δ = PR / main − 1 within a round):
So bs 8 and 16 are not net negative beyond noise. At 95% the cost is at most about +1.3% (bs 8) and +1.8% (bs 16).
Profile (torch profiler, 3 s of steady decode per bs, µs per iteration; the profiler inflates host time, so compare only within a row pair):
No code change. Restricting replay to some sizes would trade a ≤0.4% structural cost for a new rule, and the data doesn't justify that. I've updated the PR description with these numbers and the single-GPU method. On the c1 nondeterminism (17 of 300 identical between two greedy runs of main): agreed, we'll track it separately. |
|
@valar-cop review this pr |
valar-cop
left a comment
valar-cop
left a comment
There was a problem hiding this comment.
LGTM, but the accuracy jobs are red, please rerun..
any perf number? gsm8k? and no qwen3.8/qwen4 mtp coverage in ci yet
ci Accuracy (DeepSeek-V4-Flash TP4) failed — ci machine/env issue, not the pr (medium): All six jobs of the vLLM accuracy matrix failed in the same run while the same head's ATOM Test and Pre Checkin workflows passed, and each failing job's log is empty (0 bytes). DeepSeek-V4-Flash TP4 runs no speculative decoding, so the diff (ATOM's standalone draft-graph staging plus Qwen3.8 MTP only) cannot reach it; this is a run-level/build-environment failure.
- FAILED job
Accuracy (DeepSeek-V4-Flash TP4)(ATOM vLLM Test), runner linux-atom-pitt-mi350x-8-6nfvf-runner-tf2h4, log .valar/ci/job_123953641936.log, url https://github.com/ROCm/ATOM/actions/runs/41022598167/job/123953641936, regex infra hints: none, same job on main: [], on other recent runs: {}
next: Rerun the ATOM vLLM Test workflow; if it repeats, inspect the 'Build vLLM cell' / container start step of a single job (the log artifacts are empty) and compare against a main-based run. https://github.com/ROCm/ATOM/actions/runs/41022598167/job/123953641936
ci Accuracy (GLM-5.2-MXFP4-MTP TP4) failed — ci machine/env issue, not the pr (medium): Same run-level failure as the rest of the accuracy matrix; the vLLM plugin uses vLLM's own MTP proposer, not atom/spec_decode/draft_graph.py, so this diff is not on the path being tested. Log is empty.
- FAILED job
Accuracy (GLM-5.2-MXFP4-MTP TP4)(ATOM vLLM Test), runner linux-atom-pitt-mi350x-8-6nfvf-runner-wm6p2, log .valar/ci/job_123953641948.log, url https://github.com/ROCm/ATOM/actions/runs/41022598167/job/123953641948, regex infra hints: none, same job on main: [], on other recent runs: {}
next: Rerun with the rest of the matrix; no code action. https://github.com/ROCm/ATOM/actions/runs/41022598167/job/123953641948
ci Accuracy (MiniMax-M3-MXFP4 EAGLE3 TP4) failed — ci machine/env issue, not the pr (medium): Same run-level failure. EAGLE3 here goes through vLLM's spec-decode patch, and the EAGLE changes in this diff are inert for a non-MRoPE target (mrope is False, so the gate and the StagedInput shapes are byte-identical to main). Log is empty.
- FAILED job
Accuracy (MiniMax-M3-MXFP4 EAGLE3 TP4)(ATOM vLLM Test), runner linux-atom-pitt-mi350x-8-6nfvf-runner-fnmwj, log .valar/ci/job_123953641964.log, url https://github.com/ROCm/ATOM/actions/runs/41022598167/job/123953641964, regex infra hints: none, same job on main: [], on other recent runs: {}
next: Rerun with the rest of the matrix; no code action. https://github.com/ROCm/ATOM/actions/runs/41022598167/job/123953641964
ci Accuracy (Kimi-K3 DSpark DCP8 TP8) failed — ci machine/env issue, not the pr (medium): Same run-level failure. DSpark uses DSparkProposer/dspark_proposer.py, whose StagedInputs keep the default batch_last=False; this diff does not touch that file or its layout. Log is empty.
- FAILED job
Accuracy (Kimi-K3 DSpark DCP8 TP8)(ATOM vLLM Test), runner linux-atom-pitt-mi350x-8-6nfvf-runner-zwhdr, log .valar/ci/job_123953641980.log, url https://github.com/ROCm/ATOM/actions/runs/41022598167/job/123953641980, regex infra hints: none, same job on main: [], on other recent runs: {}
next: Rerun with the rest of the matrix; no code action. https://github.com/ROCm/ATOM/actions/runs/41022598167/job/123953641980
ci Accuracy (DeepSeek-V4.1-Flash DSpark TP4) failed — ci machine/env issue, not the pr (medium): Same run-level failure; DeepSeek-V4.1 DSpark goes through the vLLM DSpark speculator and dspark_proposer.py, neither of which this diff changes. Log is empty.
- FAILED job
Accuracy (DeepSeek-V4.1-Flash DSpark TP4)(ATOM vLLM Test), runner linux-atom-pitt-mi350x-8-6nfvf-runner-w7lmb, log .valar/ci/job_123953642038.log, url https://github.com/ROCm/ATOM/actions/runs/41022598167/job/123953642038, regex infra hints: none, same job on main: [], on other recent runs: {}
next: Rerun with the rest of the matrix; no code action. https://github.com/ROCm/ATOM/actions/runs/41022598167/job/123953642038
ci Accuracy (Qwen3.5-397B-A17B-MXFP4 TP4) failed — ci machine/env issue, not the pr (medium): Same run-level failure. This matrix entry passes no --speculative-config, so no drafter and no draft graph is built at all; the diff cannot affect it. Log is empty.
- FAILED job
Accuracy (Qwen3.5-397B-A17B-MXFP4 TP4)(ATOM vLLM Test), runner linux-atom-pitt-mi350x-8-6nfvf-runner-m26wz, log .valar/ci/job_123953642095.log, url https://github.com/ROCm/ATOM/actions/runs/41022598167/job/123953642095, regex infra hints: none, same job on main: [], on other recent runs: {}
next: Rerun with the rest of the matrix; no code action. https://github.com/ROCm/ATOM/actions/runs/41022598167/job/123953642095
| if running_bs > scheduled_bs: | ||
| out[scheduled_bs:].copy_(src[-1]) # broadcasts; no expand needed | ||
| if batch_last: | ||
| out[..., scheduled_bs:].copy_(src[..., -1:]) |
There was a problem hiding this comment.
docstring still says row, now it's the last real sequence for batch_last.. please reword
There was a problem hiding this comment.
@valar-cop I followed your suggestion on atom/spec_decode/draft_graph.py:237. cc @valarLip
The _stage_one docstring now says it tail-repeats the last real sequence, which is the last row, or the last column of a batch_last input (the [3, N] MRoPE positions). The rest of the paragraph says "sequence" instead of "row".
(commit 4b28fdd56, docstring only; the head is now 80bd51de8 after a merge of main. black and ruff are clean. On CPU, tests/test_draft_graph.py passes 28/28 and tests/test_glm5_next_mtp_routing.py passes 2, with 5 skipped for aiter. test_qwen4_exp_mtp, test_dspark and test_region_graph skip on CPU because they need triton, aiter or a GPU. Pre Checkin is green on 80bd51de8.)
There was a problem hiding this comment.
yes, that reads right now, pad repeats the last real sequence in both layouts and tests/test_draft_graph.py:195 pins the batch-last one.
lgtm from me. note the vllm accuracy matrix is still skipped on this head, it needs the approval or a ci:full label to run.
|
@valar-cop thanks for the review. cc @valarLip The docstring is fixed in Accuracy rerun. I reran the failed vLLM accuracy jobs on the old head
On Perf. These are the single-GPU numbers from my earlier reply, which has the method and the full trace breakdown: #2663 (comment)
GSM8K. Full GSM8K (1319 questions), on one MI355X (node 102):
CI coverage. You're right. Qwen3.8 appears only in the SGLang plugin's nightly accuracy catalog, without speculative decoding. The native ATOM and vLLM accuracy matrices have no Qwen3.8 / Qwen4 MTP model, so CI doesn't exercise this PR's draft path. Adding that job is outside this PR, so I've left it as a follow-up. For now the PR's coverage is @valar-cop review this pr |
valar-cop
left a comment
valar-cop
left a comment
There was a problem hiding this comment.
LGTM, but please make the batch-last producer-fill invariant explicit..
any test for draft step at running_bs > scheduled_bs? gsm8k?
ci green.
| ) | ||
| # A producer may have written straight into this exact fixed-storage | ||
| # view. Do not turn that into a self-copy launch. | ||
| if src.data_ptr() != dst.data_ptr() or src.stride() != dst.stride(): |
There was a problem hiding this comment.
batch_last + producer fill gives an overlapping self-copy with wrong strides.. assert it or document it?
There was a problem hiding this comment.
@valar-cop I followed your suggestion on atom/spec_decode/draft_graph.py:234. It's both asserted and documented now. cc @valarLip
You're right that it's real. On CPU I filled buffer("pos", 2) and staged at running_bs=4. That took the copy branch (same data_ptr, strides (2, 1) vs (4, 1)) and scrambled rows 1 and 2 of the positions without any error. 3→8 and 5→6 do the same.
_stage_onenow asserts when the source is a view inside the role's own fixed storage but isn't the exact view it writes. The exact view (same pointer and stride) still skips the self-copy.buffer's docstring states the invariant. A leading-batch producer fillsbuffer(role, scheduled_bs). Abatch_lastproducer fillsbuffer(role, running_bs)[..., :scheduled_bs], which has the staged stride.- I asserted rather than cloned because no caller reaches this. The only
batch_lastinput is the MRoPE positions, which_enter_decode_metadatabuilds with a freshindex_select, and nothing fills them throughbuffer.
(commit 528b871be; new tests test_a_batch_last_producer_fills_the_running_view and test_a_batch_last_fill_of_a_narrower_view_is_refused in tests/test_draft_graph.py, and the second one fails on the old guard. CPU: the draft-related suites pass, 180 passed / 46 skipped; black and ruff are clean.)
There was a problem hiding this comment.
lgtm. the guard now refuses before dst.copy_, and the 0 <= ptr - buf range check catches the narrow-view overlap that used to scramble rows.
checked the callers: only hidden_states (eagle_proposer.py:778) and the dspark/v4 mono anchors fill a buffer, all leading-batch at running_bs, and positions comes from a fresh index_select (eagle_proposer.py:551), so nothing trips it today. assert over clone is fine.
valar-cop
left a comment
valar-cop
left a comment
There was a problem hiding this comment.
LGTM, but the batch_last self-copy guard needs a fix before it bites..
any perf number? gsm8k? and the padded replay / DP path is only shown TP1
ci green.
| ) | ||
| # A producer may have written straight into this exact fixed-storage | ||
| # view. Do not turn that into a self-copy launch. | ||
| if src.data_ptr() != dst.data_ptr() or src.stride() != dst.stride(): |
There was a problem hiding this comment.
stride differs when running_bs != scheduled_bs, so this guard takes the copy branch over overlapping storage.. latent today but hidden_states would corrupt silently. assert no alias or copy via clone?
There was a problem hiding this comment.
@valar-cop I followed your suggestion on atom/spec_decode/draft_graph.py:234: it now asserts there is no alias, in the same commit as the other thread. cc @valarLip
One correction on scope: hidden_states can't hit this. It's a leading-batch input, so buffer("hidden_states", bs) is buf[:bs], and its strides don't depend on bs. At running_bs != scheduled_bs the producer's view is exactly out[:scheduled_bs], so the self-copy is skipped. I checked this on CPU at 2→4, 3→8 and 5→6: same pointer, same strides, and both the real rows and the pad tail come out correct. Only a batch_last input has a batch-dependent stride.
Every source that lives inside a staged buffer today is the exact view:
- EagleProposer's step-0
hidden_states, whichindex_selectwrites intobuffer(..., scheduled_bs). - The next step's
graphed_ids[:scheduled_bs]andret_hidden_states[:scheduled_bs]. - DSpark's anchors, which are filled at
buffer(..., running_bs)and staged atrunning_bs.
So the assert doesn't fire on any current path, and every copy that ran before still runs. A future alias fails loudly instead of corrupting.
(commit 528b871be; test_a_draft_step_wider_than_the_schedule_pads_every_input stages ids, MRoPE positions and producer-filled hidden states at scheduled_bs=3 / running_bs=8, both eager and replayed. It checks the real rows, the pad tail and the fixed addresses.)
There was a problem hiding this comment.
Verified against the code: hidden_states is a leading-batch input (StagedInput(shape=hidden_shape), eagle_proposer.py:159), so buffer(role, scheduled_bs) is always buf[:scheduled_bs], and dst is that same slice at any running_bs — only the batch-last view has a stride that moves with bs. The assert plus the narrow-view test is the right shape.
lgtm, my mistake on hidden_states.
A batch_last view's stride is its batch, so buffer(role, scheduled_bs) overlaps the running_bs view at a different layout. Staging it took the copy branch over that overlap and scrambled the positions without an error. No caller fills a batch_last input today (MRoPE positions are a fresh index_select), so assert instead of cloning, and say in `buffer` which view a producer fills for each layout. Tests: the batch_last fill at running_bs, the refused narrower view, and a draft step at running_bs > scheduled_bs, eager and replayed.
|
@valar-cop thanks. cc @valarLip Both reviews on Batch-last producer fill / self-copy guard. The bug was real but latent. On CPU, a Test at
Perf / GSM8K. I didn't re-measure, because the change can't move either number:
The round-1 numbers still apply (details in #2663 (comment)):
Only TP1 is shown. That's true. Everything I measured for this PR ran at TP1 on one MI355X: wall time, the profile, GSM8K, the serving run and the GPU unit tests. I haven't run Qwen3.8 MTP with padded replay at TP>1 or with DP attention, so I'll leave that as a follow-up, along with the CI coverage from round 1. The replay decision is keyed only on the DP-agreed vLLM accuracy matrix. It's still skipped by the Heavy CI gate on this head. It needs a reviewer to approve again, or a @valar-cop review this pr |
valar-cop
left a comment
valar-cop
left a comment
There was a problem hiding this comment.
LGTM, one minor thing on the stage() diagnostic
any perf number for tp>1? and the atom test on this head is still pending
ci green.
| count that cannot disagree with the tensor it describes. | ||
| """ | ||
| counts = {t.shape[0] for t in srcs.values()} | ||
| counts = {self._batch_of(role, t) for role, t in srcs.items()} |
There was a problem hiding this comment.
undeclared role now raises bare KeyError instead of the old assert message? let's keep the diagnostic..
There was a problem hiding this comment.
Fixed in 2daa454. stage() and buffer() now check the roles before anything else, through one _assert_declared. It names each undeclared role and lists the declared ones (... has no staged input 'missing'; declared roles are ['row']) instead of raising a bare KeyError. test_staging_an_undeclared_role_names_it covers it; on 528b871be it fails with KeyError: 'missing'.
|
@yitingw1 @valarLip your approval (on
valar-cop is LGTM on the current head Could you take another look and re-approve? |
|
ci triage for 19cc0e7:
|
…ffers Builds on #2663, which stages MRoPE positions batch-last and lets a draft model opt into mid-step graphs with draft_graph_supports_mrope. Qwen3_5MTP now sets it, so draft steps 1..k-1 on Qwen3.5 are warmed and captured instead of running eager (about 1.7-1.8 ms of host time each at bs=1). Qwen3_5MTP also opts into reuse_draft_graph_step_buffers, as DeepSeekMTP does: the captured epilogue writes each mid-step's hidden states and argmax ids into the pass's staged inputs, so a replay needs no clones and the next stage no copies. At k=5 the draft mid-steps go from 20 copies to 9 per decode step. Safe for the hybrid model: the draft is one full_attention decoder layer that reads input_ids and hidden_states only at the top of its forward, and it touches no GDN conv/SSM state or checkpoint slot. ATOM_QWEN35_MTP_REUSE_STEP_BUFFERS=0 restores the clone-and-restage path.
…buffer() stage() counted batches through self.inputs[role] before anything checked the role, so a typo surfaced as a bare KeyError. stage() and buffer() now share one assert that names the undeclared and the declared roles. The overlapping self-copy guard in _stage_one is replaced at its source: only buffer() could hand out a batch_last view whose stride differs from the one stage writes, and nothing fills a batch_last input in place. So buffer() refuses batch_last roles, and _stage_one's copy test, buffer()'s capacity and its docstring go back to main's. Also folds _view into its one caller, trims docstrings, and drops tests of the removed API; the batch_last staging test now uses a producer-filled leading input beside the positions.
|
Whole-PR cleanup in Diff against main, before → after:
Reverted.
Deleted or merged.
Verification at
@yitingw1 |
|
@ZLkanyo009 yes, lgtm on the rest. atom test is still pending on this head, wait for it before merge.. |
Keep the change to main's code to what batch_last needs: - stage() names undeclared roles with its own assert; buffer() keeps main's role check, with only the batch-last refusal added after it. - The staging ndim message keeps main's "leading axis is not the batch" wording, so main's test for it keeps its match; its docstring changes only the two sentences batch_last made untrue. - The batch-last disagreement case is its own test again rather than a second case inside main's test. - The _declare_draft_graphs docstring keeps main's paragraph and adds one sentence for draft_graph_supports_mrope.
|
Pushed Diff against main, before → after:
Reverted, by file
No shared interface changed: Not changed
Verification
Waiting on: CI on |
|
@ZLkanyo009 lgtm, the diagnostic survived the revert — still wait for atom test on |
install_mono_decode wraps the loaded Qwen4ExpMTP as well (MonoDraftModel, as V4.1 wraps its DSpark model). A draft step of 1-4 requests runs the drafter's BF16 layer as one mono program (DraftRunner): the first step attends every verify row and runs the MLP on each request's last row only, later steps run one row. Any other step of at most 16 rows runs the fallback draft (fused QSA attention, ATOM's BF16 MLP, one-row CK experts built at startup), anything larger ATOM's layer. The draft ids come from a Group64 W4 screen with a global error certificate refined on the BF16 rows: the BF16 head's argmax for 1, 2 and 4 rows, ATOM's argmax head otherwise. EagleProposer: a draft model that sets draft_returns_head_input returns (state, head input) from forward_with_head (with selected_rows on step 0 of a decode) and maps the head input to ids (compute_draft_ids_from_head), in a declared draft pass and eagerly alike. Qwen3.8 is an MRoPE model, so on main its draft steps 1+ have no declared pass and run eagerly: the program is one launch a layer, so they cost the same as in a captured graph (measured with #2663 applied: request TPOT within 1%). The recipe, the environment variable table and CLAUDE.md name the deployment. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
install_mono_decode wraps the loaded Qwen4ExpMTP as well (MonoDraftModel, as V4.1 wraps its DSpark model). A draft step of 1-4 requests runs the drafter's BF16 layer as one mono program (DraftRunner): the first step attends every verify row and runs the MLP on each request's last row only, later steps run one row. Any other step of at most 16 rows runs the fallback draft (fused QSA attention, ATOM's BF16 MLP, one-row CK experts built at startup), anything larger ATOM's layer. The draft ids come from a Group64 W4 screen with a global error certificate refined on the BF16 rows: the BF16 head's argmax for 1, 2 and 4 rows, ATOM's argmax head otherwise. EagleProposer: a draft model that sets draft_returns_head_input returns (state, head input) from forward_with_head (with selected_rows on step 0 of a decode) and maps the head input to ids (compute_draft_ids_from_head), in a declared draft pass and eagerly alike. Qwen3.8 is an MRoPE model, so on main its draft steps 1+ have no declared pass and run eagerly: the program is one launch a layer, so they cost the same as in a captured graph (measured with #2663 applied: request TPOT within 1%). The recipe, the environment variable table and CLAUDE.md name the deployment. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* feat(qwen38): mono layer programs for Qwen3.8-Flash-Next Whole decoder layers as one resident FlyDSL program each (gfx950): stages (blocks/: linear, hyper-connection, GDN with the classic and the ReplaySSM verify, QSA with its indexer, MoE, select-row, prefetch) composed by recipe.compose_layer with the frozen per-scenario schedules, placed on the 256 resident CTAs (program, layer), handing off through tagged mailbox regions (device/), emitted as one kernel (emit, kernel), and the weight layouts the stages read, checked at bind time (weights/). The geometry is read from the checkpoint's config.json and checked against the one the kernels are written for. The package follows the other mono models (atom/models/<model>/mono/); the model itself stays the single-file qwen4_exp.py. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * feat(qwen38): serve the target's MTP verify steps on mono install_mono_decode (models.registry.mono_decode_installers) wraps the loaded Qwen4ExpForConditionalGeneration: a target MTP3 verify step of 1-4 requests runs its 48 decoder layers as 48 launches (36 GDN, 12 QSA programs; the PLE layer runs ATOM's PLE prelude first), reading and writing ATOM's GDN pools, ReplaySSM records and QSA caches in place, so a request moves between mono and Native steps. Every other step runs ATOM's model. Mono keeps its own A8W4 copy of the routed experts, so the PLE n-gram table (~95 GB BF16) goes to page-locked host memory, read over UVA (ple.py, chosen at the table's construction): in HBM it would leave no room for the per-request state at --max-num-seqs 160. #2771 does the same for every deployment; this copy goes once it merges. config_refusal checks the deployment once, at install: the full 48-layer Quark MXFP4 checkpoint, MI355X with all 256 CUs, TP1, MTP3, FP8 KV with BF16 index caches on 64-token blocks, an fp32 GDN state, BF16 draft caches, either GDN verify (ReplaySSM, ATOM's default, up to 64 records, or classic), no KV connector but LMCache CPU offload (resident_grid_refusals takes the connectors a model admits). Every served shape is built at startup; one that does not build stops the server (ATOM_MONO_ENABLE=0 serves it Native). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * feat(qwen38): run the MTP drafter on mono install_mono_decode wraps the loaded Qwen4ExpMTP as well (MonoDraftModel, as V4.1 wraps its DSpark model). A draft step of 1-4 requests runs the drafter's BF16 layer as one mono program (DraftRunner): the first step attends every verify row and runs the MLP on each request's last row only, later steps run one row. Any other step of at most 16 rows runs the fallback draft (fused QSA attention, ATOM's BF16 MLP, one-row CK experts built at startup), anything larger ATOM's layer. The draft ids come from a Group64 W4 screen with a global error certificate refined on the BF16 rows: the BF16 head's argmax for 1, 2 and 4 rows, ATOM's argmax head otherwise. EagleProposer: a draft model that sets draft_returns_head_input returns (state, head input) from forward_with_head (with selected_rows on step 0 of a decode) and maps the head input to ids (compute_draft_ids_from_head), in a declared draft pass and eagerly alike. Qwen3.8 is an MRoPE model, so on main its draft steps 1+ have no declared pass and run eagerly: the program is one launch a layer, so they cost the same as in a captured graph (measured with #2663 applied: request TPOT within 1%). The recipe, the environment variable table and CLAUDE.md name the deployment. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(qwen38): one REPLAY_MAX, token-major GDN conv pools only, the hub-id install test - REPLAY_MAX is defined once, in kernels/geometry.py (CPU-safe); blocks/gdn.py and config.py import it instead of keeping two copies. - GDNPools takes ATOM's conv pool as it is, token-major [slots, W, conv_dim] (gdn_attn's mamba_k_cache): the branch that transposed a [slots, conv_dim, W] view never ran in serving, and a channel-major pool is refused with the layout mono reads. - test_install_reads_a_hub_id_checkpoint_from_its_local_directory is back, for the install-time drafter wrapping: the resolved directory reaches load_geometry, DraftRunner and TargetRunner, and the drafter is wrapped. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
StagedInput says which axis is the batch (batch_dim, default 0) instead of whether it is last. bind() allocates max_num_seqs at that axis, and stage(), buffer() and the tail repeat narrow along it: one path for every layout, with no branch on it. MRoPE positions are now a [3, max_num_seqs] buffer, viewed per batch as [3, bs] with stride (max_num_seqs, 1). Every reader of the draft's positions indexes them by stride (rows by stride(0), tokens with stride(-1) == 1) or reads only shape[-1]. A batch-1 view is the front of a wider one at the same stride, so buffer() no longer refuses batch_dim=1 roles; a test stages a producer fill of one.
|
Pushed Why. Review feedback on the previous head: the New design. @yitingw1, this is the layout you argued against on 10-22 ("a flat buffer over a
Review-pr self-check (an independent review of the local head before this push): no must-change. Fixed: a docstring sentence in Lines vs merge-base (source code / source comments / test code / test comments): Verification at
Pending: CI running; re-approval; the 4 lgtm'd valar-cop threads still need resolving before merge. |
Summary
Capture the Qwen3.8 MTP draft's steps 1..k-1 in CUDA graphs. On main,
EagleProposer._declare_draft_graphsdeclares nothing for an MRoPE model, so every Qwen3.8 draft step runs eagerly.This removes MRoPE's carve-out from the draft-graph mechanism:
StagedInputassumed a leading batch axis, and MRoPE positions are[3, N].On main
d176cb98a, measured one server at a time on one GPU over 16 alternating rounds:On current main this gain is gone at bs=1. Since #2610 Qwen3.8 decode at bs ≤ 4 runs on mono, whose own draft program already removes the eager draft's host time: at
ab4ada40a, bs=1 is 4.41 ms on main and 4.44 ms with this PR (three alternating sessions each; +0.7%, smaller than the session spread). bs ≥ 5 has not been re-measured on current main. The numbers below are from the older main.Main's eager draft got cheaper since this PR was first opened (11.21 → 8.08 ms then; among others #2688, and #2670 combining the status D2H), so above bs=2 the host time the graphs remove is hidden behind the GPU anyway. Numbers and a profile are below.
Motivation
This is part of the Qwen3.8-Flash-Next (
amd/Qwen3.8-Flash-Next-Quark-MXFP4) work on MI355X, with the goal of matching the GB300 AgentX curve at low concurrency: https://github.com/wuhuikx/LLM-Booster/issues/263At bs=1 each of the three eager draft steps launches about 60 kernels. A torch profile on main still shows the GPU idle for about 3.1 ms per iteration, almost all of it inside the draft and between the draft and the next decode; with this PR that idle is about 0.1 ms. The profiler inflates host time, so without it the exposed part on main is about 0.9 ms per iteration (the bs=1 wall-time difference).
Changes
The recording and replay are main's
DraftGraphonReqGraph(#2688), unchanged. This PR only lets the staging layer take an input whose batch is not its leading axis, and opts the Qwen3.8 MTP head in.atom/spec_decode/draft_graph.pyStagedInputgainsbatch_dim(default 0, so every other drafter's inputs are unchanged): the axis the batch sits on.bind()allocatesmax_num_seqsat that axis, andstage(),buffer()and the tail padding allnarrowalong it. There is one code path for every layout; nothing inDraftGraphbranches on it.[3, n],batch_dim=1) live in a[3, max_num_seqs]buffer; each batch isnarrow(1, 0, bs), a[3, bs]view with stride(max_num_seqs, 1). Padding repeats the last column.buffer()on abatch_dim=1role returns the front of the staged view at the same address and stride, so a producer can fill it in place like any other role.stage()names an undeclared role and the declared ones, asbuffer()already did, instead of raising a bareKeyError.atom/spec_decode/eagle_proposer.py: an MRoPE model gets the mid-step draft graph when its draft model declaresdraft_graph_supports_mrope = True, withpositionsstaged atbatch_dim=1. Other MRoPE models are unchanged and still run the draft eagerly. Qwen3.5 MTP (qwen3_5_mtp.py) is not opted in: its draft graph has not been validated on GPU.atom/models/qwen4_exp_mtp.py:Qwen4ExpMTPsets main's existing draft-graph hooks:draft_graph_supports_mrope = True,reuse_draft_graph_step_buffers = True,draft_graph_hidden_state_shape = (hc_count, hidden_size). The draft carries the whole hyper-connection bundle between steps.tests/test_draft_graph.py:batch_dim=1positions staged as a[3, bs]view of the fixed buffer (stride(max_num_seqs, 1)) with the last column repeated, beside an in-place-filled leading input;buffer()on abatch_dim=1role is the front of the staged view and staging a fill of it pads after it; an undeclared role named in the error;batch_dim=1and leading sources disagreeing on the batch. Main's test for a source whose leading axis is not the batch now matches the message's "batch at axis 0".tests/test_qwen4_exp_mtp.py: the realQwen4ExpMTPclass throughEagleProposer._declare_draft_graphsunder MRoPE (positions declaredStagedInput(shape=(3,), dtype=int64, batch_dim=1),(hc_count, hidden)hidden states, buffer reuse on), and an MRoPE draft that has not opted in still declares nothing.Every reader of the draft's positions takes the strided view: row
stride(0)and tokenstride(-1)go to the kernel, or onlyshape[-1]is read. The audited readers are listed in #2663 (comment). Main's target decode already handsQwen4ExpAttentiona[3, tokens]view with striderunning_tokens(backends.py, pinned bytest_decode_mrope_storage_matches_padded_graph_stride), and an uncaptured batch runs eagerly on the same staged view, so a replay and an eager step see the same layout.Dependency: ROCm/aiter#6833
The draft head's
topk_select(k=1)is recorded in these graphs. Before ROCm/aiter#6833,topk_select(end=None)took its row lengths from an 8-entry LRU, so capturing more than 8 batch sizes freed a tensor an earlier graph had baked in (wrong draft ids, then a memory fault). An earlier revision of this PR worked around that inembed_head.py. #6833 fixes it at the root (no length tensor is read whenendisNone), so the workaround is gone andembed_head.pyis untouched.#6833 is in aiter main since 2026-10-22 and in the wheel ATOM CI installs (
f40355cd5). A capture-then-replay probe oftopk_select(k=1)at 12 batch sizes againsttorch.argmaxgives 0 mismatches on the wheel used below.Setup
MI355X, TP1, image
rocm/atom-dev:nightly_202610221638; FP8 KV, block 64, PTPC, MTP 3, no offload.d176cb98aand this PR merged with it run one after another on the same GPU, 16 rounds in ABBA order.d176cb98ais the newest main the tuned aiter below can run: feat(attention): FlyDSL paged decode for MiniMax-M3 on gfx1250 #2710 imports an aiter symbol that tree doesn't have. The later main commits don't touchatom/spec_decode.f324de48and this PR merged with it, run at the same time on two GPUs of one node.Two aiter builds:
amd_aiter 0.1.1.dev1+g508736740(contains #6833).905967361plus [Config] [Perf] Tune gfx950 BF16 GEMMs for Qwen3.8-Flash-Next (TP1, MTP 3) aiter#6711, #6713 and #6762, the build our Qwen3.8 AgentX runs use (also contains #6833).Accuracy (earlier head, older main
f324de48; for this head see Test plan)GSM8K, real MTP-3 verification (no forced acceptance), 1319 questions, 5-shot chat, greedy:
The outputs are not token-identical between the two, but that is true of main against itself. Even at concurrency 1 with greedy decoding, two runs of main on the first 300 questions give identical text for only 17 of 300; this PR against main gives 19 of 300. The divergence comes from near-ties in the target's verify logits.
To check whether the graphs change the drafts, split acceptance by draft position. Draft token 0 is produced eagerly in both trees; only steps 1..k-1 are graphed.
The ungraphed token 0 differs by as much as the graphed tokens 1 and 2, and two runs of main differ by the same amount. A wrong replay would show up as a drop in tokens 1 and 2 only.
Performance (earlier head, older main
d176cb98a; for this head see Test plan)Forced acceptance length 2.32 (AgentX rules).
Decode iteration wall time, tuned aiter, one server at a time (bs concurrent streams, 64k-token random prompts, 4000 output tokens, median chunk interval while all are decoding). Each value is the median and range over 16 sessions per side. Paired Δ = this PR / main − 1 within a round:
Once a size is captured, the draft replays it unconditionally, so the staging cost is paid at every captured size. That cost is small:
input_ids, about 2 µs on the GPU. The hidden states already live in the staged buffer.Profile, tuned aiter (torch profiler, 3 s of steady decode per bs, averaged per iteration). The profiler inflates host time, so compare within a pair only:
At bs 8 the branch's draft kernels run longer under the profiler: they are the same kernels, back to back, while main's are spread out by profiler-inflated host gaps. Without the profiler the bs 8 wall time is the same.
AgentX C1, 1 hour each (
inferencex-agentx-mvp, seed 42; mainf324de48, both servers running at the same time):Not changed
qwen3_5_mtp.py) and other MRoPE drafts onattention_mhaare not opted in. Before one is,try_mrope_qk_fused's fallbackrotary_embandtriton_fused_qkv_norm_rope_cache(which assumesstride(-1) == 1) need checking for a strided[3, bs].draft_runner.pyno longer runs at the captured bs 1/2/4 (it is a diagnostic, not a correctness check).tests/test_qwen4_exp_mtp.pystill skips on CPU CI (module-level aiter / triton imports); moving the declaration test means reworking those imports.How to use
There is no new switch. Qwen3.8 with MTP captures the draft graph whenever draft CUDA graphs are enabled (the default).
ATOM_DRAFT_CUDAGRAPH=0still disables them.Test plan
Validation of the current head (
19cc0e7ea, after thebatch_dimrefactor). The Accuracy and Performance sections above are from an earlier head on an older main.At
19cc0e7ea(this head), MI355X, imagerocm/atom-dev:nightly_202610221638, aiter wheel4d1980603from the main manifest:black26.10.1 andruff0.17.2 are clean..github/scripts/run_unit_tests.sh): 4 failed / 11195 passed / 1163 skipped. Three failures (test_graph_marker×2,test_sparse_indexer_side_effect_op) are main's on the same host. The fourth,test_multiprocess_metrics, is a spawned metrics process exiting 1 under the suite's load; it passes alone (3 passed) and does not touch draft graphs.test_qwen4_exp_mtp.py,test_draft_graph.py,test_qwen38_mono_prep.py,test_qwen38_mono_ple.py,test_region_graph.py,test_graph_holders.py: 121 passed / 0 failed / 0 skipped. The 18 draft graph / MTP / DSpark / region graph / mono test files (the above plustest_glm5_next_mtp_routing.py,test_dspark_integration.py,test_lm_head_argmax_out.py,test_mtp_draft_warmup_workplan.py,test_mtp_draft_kv_allocation.py,test_gdn_draft_kv_slots.py,test_h2d_draft_publication.py,test_prepare_mtp_decode_contract.py,test_cudagraph_capture_bounds.py,test_dspark.py,test_model_runner_staging_fence.py,test_rollout_memory_manager_sleep.py): 337 passed / 19 skipped / 0 failed (mainab4ada40a: 331 / 19 / 0).DraftGraph.run; GSM8K, concurrency 16, mono on and off, every batch size 1–16 that occurred). Each replayed step was snapshotted and re-run: replayed again, run eagerly on the same staged[3, bs]view, and run eagerly withpositions.contiguous(). Re-running a step is not bit-reproducible on this stack, in eager as much as in replay: two eager runs of one step on the same inputs never give bit-identical hidden states, and give different draft ids in 1–6% of steps. Replay against eager, and the strided view against a contiguous copy, differ exactly as much as these controls (e.g. this head, mono on, bs 5–16: draft ids equal in 672/708 replays against eager, 674/708 strided against contiguous, 672/708 replay against replay, 685/708 eager against eager). The previous head's flat layout (this branch at9cd4ea52e, before the refactor) gives the same numbers.9cd4ea52e) 0.9795, 72.18%, 3.165. Within run-to-run spread (72.1–72.7% across earlier runs).ab4ada40aagainst this head, three alternating sessions each: 4.41 ms (4.34–4.42) against 4.44 ms (4.44–4.46); see the Summary.aea839bcc(9.44 ms) and528b871be(9.46 ms) sit on the older main, before perf(qwen): integrate Mono decode and optimized TP1 MTP3 serving #2610 put Qwen3.8 decode on mono, so they are not comparable with the current tree.f324de48, bs 1/2/4/8/16 decode wall time and profiles on maind176cb98a, as in the sections above.