Skip to content

Qwen3.5 on ROCm: twin GDN state slots, a decode-only conv kernel and the packed prologue - #1139

Draft
merceod wants to merge 55 commits into
users/atindra/pr-46-t12-rope-fp32from
users/atindra/pr-47-t14-small-step
Draft

merceod wants to merge 55 commits into
users/atindra/pr-46-t12-rope-fp32from
users/atindra/pr-47-t14-small-step

Conversation

@merceod

@merceod merceod commented Nov 2, 2026

Copy link
Copy Markdown
Collaborator

Stacked on #1137 (pr-46-t12-rope-fp32). Picks the commits of #1090, #1096, #1102 first; they drop out of this diff as those merge.

Why

After #1121 a Qwen3.5 0.8B decode step at c=64 on MI325X is GPU-bound: untraced, the replay span plus the gap between replays is the ITL (2368 + 21 us). The step ran 270 kernels, so the easy guess was launch overhead at about 5 us a kernel. It isn't. Inside a captured graph a small kernel costs 1.7-2.1 us whatever sits around it, so 100 fewer kernels would buy about 8 %, not 20. The largest item, the GDN rule (the recurrent state update, 777 us of the step traced), pays for state traffic: a kernel with its reads and writes and no math takes 45.2 of its 46.1 us at 64 rows. What it paid on top of the traffic was writing the state back over the slot it had just read. Writing a second slot instead takes 22-48 % less time at 16-64 rows, same kernel, same values. That is most of this PR; three smaller changes come with it.

What changed

On #1137, with #1090, #1096 and #1102 picked first. These commits were written on the integration tip, which carries the speculative decoding work: the conv route edits GDNDecodeWrapper next to #1090's verify plan, the prologue commit builds on attn/prologue.py, which #1102 adds, and #1102's decode-bucket commit needs #1096's compile_dynamic capture configs, which #1137 doesn't carry. #1103 and #1104 aren't needed.

  • mstar/engine/resources/recurrent/: twin_slots on the GDN state pool. A request holds two slots; a decode step (one token a row) reads the state from one and writes the next state to the other, and commit swaps them. Same kernel, same values, and max_slots still counts requests. Prefill, split walks and the mixed step's decode half still update in place (keep_in_place). The pool refuses twins together with speculative decoding, whose slots already hold a state per token, and only AITER's FlyDSL decode writes the other slot; any other GDN backend refuses the key when it binds.
  • configs/qwen3_5_{0.8b,2b,9b}.yaml: twin_slots: true under platform_overrides: rocm, so only these three on AMD GPUs. The mtp2, mtp3, TP2 and 4B configs serve without twins. docs/models.rst and docs/adding_models.rst say what the key costs.
  • mstar/utils/causal_conv1d/kernels.py: causal_conv1d_decode, the conv state update for one token a row, bit-equal to causal_conv1d_update. GDNDecodeWrapper takes it on is_rocm().
  • qwen3_5: a plain decode replay on one rank takes the packed attention prologue, q/k norm, rope and the KV write in one kernel (the one Speculative decoding on Qwen3.5: a fused draft step and an fp8 draft head #1102's draft passes use), 12 kernels fewer a step on 0.8B. PackedAttentionCallable asks the KV manager to stage flat write slots (page * page_size + offset) for decode replays; where the kernel doesn't load, the layer keeps its own path. The kernel has an AMD implementation only.
  • mstar/engine/resources/attn/aiter.py: a captured decode step picks AITER's paged attention variant for 4096 keys (DECODE_SEGMENT_PIN) instead of the cache's whole context (32768), which launched 128 segments a row for keys that rarely fill two. A longer row still reads every key, the segments just get wider.
  • The last two commits update hand-built test stubs for the pool's write slots and bring comments in line with twin slots; git diff d82f16aef 732dc085b -- mstar is docstrings and comments.

Numbers

Every number below is from MI325X VF, one GPU an arm in a device cgroup on a shared node (other jobs beside it), the two arms alternating in one job, text-chat-256, three timed runs a launch and the median of launches. They are directions, not a table of record. The A/B on the branch (d82f16a, which differs from the head by tests and comments, against fe6c94a):

row, c tok/s delta ITL p50 ms delta
0.8B c=1 882 -> 932 +5.7 % 1.089 -> 1.025 -5.9 %
0.8B c=4 2676 -> 2884 +7.8 % 1.436 -> 1.327 -7.6 %
0.8B c=16 9041 -> 9901 +9.5 % 1.669 -> 1.518 -9.0 %
0.8B c=64 23990 -> 25992 +8.3 % 2.395 -> 2.190 -8.6 %
2B c=1 / c=64 624 -> 646 / 18844 -> 19773 +3.5 / +4.9 % 1.557 -> 1.499 / 3.105 -> 2.913 -3.7 / -6.2 %
9B c=1 / c=64 226 -> 230 / 7338 -> 7611 +1.8 / +3.7 % 4.354 -> 4.282 / 8.221 -> 7.868 -1.7 / -4.3 %
9B TP2 c=1 / c=64 327 -> 333 / 8829 -> 8874 +1.6 / +0.5 % 2.992 -> 2.936 / 6.666 -> 6.629 -1.9 / -0.6 %
9B mtp2 c=1 / c=64 409 -> 409 / 8684 -> 8734 0 / +0.6 % 5.832 -> 5.796 / 14.976 -> 14.931 -0.6 / -0.3 %
  • 0.8B c=64 pooled over 7 launches an arm in three jobs: +8.17 % tok/s and -8.48 % ITL; the single jobs read +7.4, +8.3 and +8.8 %, and ITL is the steadier number (-8.0 to -8.6 % in all three). TP2 and mtp2 serve without twins and stay level. 4B, also without twins, reads +2.9 / +0.7 % at c=1 / c=64, long-1k +0.5 to +4.5 %, Qwen3-Omni RTF level.
  • The A/B on the integrated tip (732dc08 against fe6c94a), same conditions: 0.8B c=64 24031 -> 25798 tok/s (+7.35 %), ITL p50 2.399 -> 2.205 ms (-8.1 %), c=1 ITL 1.120 -> 1.040 ms, three launches an arm that don't overlap (23972-24076 against 25748-25888). One launch an arm for the rest: 2B c=64 +7.5 % (ITL -7.0 %), 9B c=64 +4.0 % (ITL -4.3 %), 9B TP2 8851 -> 8833 (ITL -0.3 %), mtp2 413 / 8841 tok/s at c=1 / c=64, read as level.

Where it came from, 0.8B, cumulative, two launches an arm (the pin alone three); share is points of c=64 ITL:

commit kernels c=64 ITL ms c=1 ITL ms share
fe6c94a 270 2.388 1.079
ff3de91, decode-only conv same 2.354 1.068 -1.4
40483b0, packed prologue -12 2.335 1.049 -0.8
49f4d94 and 878a84d, twin slots same 2.213 1.062 -5.1
d82f16a, the 4096-key pin same 2.196 1.035 -0.7
  • The pin pays most below c=64: over twins alone it takes c=1 ITL -2.9 %, c=4 -2.35 %, c=16 -1.6 %. Its one measured loss is 4 rows at 16k keys, 177 -> 189 us per attention.
  • The per-kernel floor inside a captured graph is about 1.8 us (1.70-2.10, with 1 or 128 distinct kernels, with or without data chains), so cutting launches is worth less here than it looks.
  • The GDN rule is state traffic, not launches or math: the same reads and writes with no math take 45.2 of its 46.1 us at 64 rows.
  • Writing the state in place over what was just read costs more than writing a second slot: 45.5 against 35.4 us at 64 rows, 32.3 against 18.8 at 32, 18.6 against 9.7 at 16, level at one row (one layer's traffic alone). Read-only (16.8-18.1 us) plus write-only (15.5-15.6) add up to about the twin number. Why in place costs the extra 10 us is not known; holding stores back doesn't help.
  • The GEMMs sit at hipBLASLt's floor here. 28 Triton tilings per GEMM size (XCD-local row tiles included), three HIP MFMA GEMM variants, add+norm folded into a Triton GEMM, silu folded into gate/up and a HIP conv on packed loads were measured and rejected; none is in this PR. Of the 2193 us step that is left, the rule is about 600 (bandwidth), GEMMs about 800 plus lm_head 155, add+norm 245, attention 160 and activation 120.
  • Cost: twins double the GDN pool, 0.8B and 2B 2.4 -> 4.8 GiB, 9B 6.3 -> 12.6 GiB (server logs). Ready time is level, 60-69 s on 0.8B and 99-105 s on 9B for both arms.

Window P (exclusive node, job 1058): the integrated tip 732dc08 against the week-3 tip d559763, which has neither #1121 nor #1137 nor these commits, one GPU an arm in device cgroups with nothing else on the node, three timed runs, text-chat-256. 0.8B 894 / 2751 / 10027 / 25818 tok/s at c=1 / 4 / 16 / 64 against 895 / 2669 / 9036 / 23657 (+9.1 % at c=64, ITL p50 2.20 against 2.38 ms); 2B 650 / 2132 / 7501 / 20127 against 626 / 2046 / 6954 / 18648 (+7.9 % at c=64, ITL 2.90 against 3.09); 9B 229 / 788 / 2837 / 7578 against 225 / 768 / 2638 / 7276 (+4.2 % at c=64, ITL 7.94 against 8.26). 9B TP2 8763 against 8760 at c=64 (level, twins and the prologue are off under TP); mtp2 408 / 8693 at c=1 / c=64 with acceptance 0.715 / 0.700; mtp3 440 / 7839. 0.8B c=1 reads level there (894 against 895) where the shared-node A/Bs read +6 %: c=1 moves with the launch, and P ran one launch an arm. 9B c=64 TTFT p50 69.1 against 71.4 ms. All 12 greedy passes PASS with 0 non-tie; long-decode PASS on all three sizes on the tip (9B engine-only 37 of 56, 2B 1 of 6, 0.8B 0 of 3) and FAIL on 9B and 2B for the week-3 tip, which still rounds the rope frequencies #1137 fixes.

Correctness

  • Greedy near-tie: 34 of 34 launches PASS with 0 non-tie divergences in the branch A/B, and 16 of 16 passes on the integrated tip.
  • Long-decode (two servers, 8 prompts to 2048 tokens past EOS) passes on all three sizes: on the branch 9B engine-only 37 against a limit of 50, 2B 1, 0.8B 0; on the integrated tip, with the harness's recalibrated limits, 9B 32 against 56, 2B 1 against 6, 0.8B 0 against 3.
  • mtp2 acceptance is level: c=1 0.7231 -> 0.7159 (the tree's launches 0.7094 and 0.7224), c=64 0.6987 -> 0.7001.
  • The 20-item image set passes with 0 non-tie on 12 of 12 cells (8 on the branch, 4 on the integrated tip), but the answers do move inside the near-tie band. On 9B the exact-match count reads lower, 7.2 -> 5.3 of 20 pooled (5 cells on the base, 3 on this tree), from four items that were steady on the base (img_07, img_09, img_13, img_16). Their gaps are at most 0.25 and the largest gap anywhere doesn't grow, and two servers on one tree agree on a median 11 of 20 items (10 across the two trees). 2B reads 5.0 -> 5.3. So no answer is claimed better or worse by the bar. Twins and the conv kernel change no values, so the move comes from the prologue or the pin, which change the rounding order; it isn't split per commit yet.

Tests

  • test/modular/test_recurrent_twin_slots.py, 14 tests: two slots a request with max_slots counting requests, the read/write split and the swap at commit, only decode steps writing the twin, walks held in place, padding on the sink, a pre-plan after the commit reading what the step wrote, rollback and release zeroing both slots, decode snapshots and forks reading the right slot, the refusals (speculation, in-place kernels) and the YAML override.
  • test/modular/test_conv_decode_route.py, 4 tests: AMD GPUs take the decode kernel where it fits, every other platform the general one, and a plan with write slots takes the decode kernel.
  • test/modular/test_attn_prologue.py, 4 new tests: flat slots are staged only for a layer that asked, and which stacks (rank count, platform) take the packed prologue.
  • GPU tests in test/rocm/: the decode conv bit-equal to the general one (null slots, strided rows, no bias or activation, a captured replay against eager steps, a write slot leaving the read slot alone), FlyDSL decode on twins against in place and a captured twin replay touching no other slot, and the 4096-key pin reading every key. test/rocm: 2032 passed, 10 skipped on one MI325X VF.
  • The hand-built stubs in test_qwen3_5_device_loopback.py, test_qwen3_5_mixed.py and test_recurrent_prefix_snapshots.py carry the pool's write slots.
  • CPU: Core 10960 tests, 34 failing, the same set as the tree under it (10938 / 34; the +22 are the three new files); ruff clean. These counts come from the integrated tip, not this cut branch.

What is not covered

  • Nothing run on CUDA, 27B or TP4. Twins are a YAML key under platform_overrides: rocm on the 0.8B, 2B and 9B configs only, the conv routes on is_rocm(), the prologue has an AMD implementation only and the pin sits in the AITER manager, so a CUDA resolution runs today's kernels. Shared code did change: the pool, the KV manager, the GDN wrappers and the Qwen3.5 attention layer gained branches that stay off there. Read, not run.
  • Twins aren't on 4B, 9B TP2, the mtp configs or the mixed step's decode half, which still updates in place.
  • Twins under forced eviction and client aborts mid-decode are tested on CPU only, not served; so is the prefix cache's decode snapshot order, since AITER keeps snapshots off.
  • Open: why the in-place write costs what it does, and why attention + reduce drops 223 -> 160 us traced with the twin commit, which touches no attention.
  • c=1 depends on the launch (spans of 1060 and 1030 us on one tree), consistent with per-process tuning and not separated from it. 9B c=64 TTFT read 64.1 -> 70.1 ms on one launch an arm on the integrated tip; TTFT moves with admission at c=64.
  • The pin reaches every model on AITER paged attention with max_seq_len above 4096; Qwen3.5 and Qwen3-Omni are the ones checked.

merceod and others added 30 commits November 2, 2026 09:20
…erver

Loads the base model through transformers and the checkpoint's mtp.* head
from the cached snapshot, runs speculative decoding with sequential
verification for k drafts per step in greedy and sampling mode, and reports
per-position acceptance and tokens per step over the sweep's prompt sets.
draft_cost.py times one draft token (head + lm_head) eager and as a HIP
graph. The verify logic has a CPU test on bigram table models.
A bundle is a directory per (GPU product, torch, HIP or CUDA, Triton) key holding one uncompressed tar per cache (Inductor, Triton, MIOpen) and a manifest. Packing appends only new or changed files under a lock; unpacking extracts the bundle for this stack into the cache dirs a server here resolves and leaves a marker that each worker reads at startup to log the key and whether the bundle matches. mstar serve --compile-bundle unpacks before starting.
The runner hints dynamo that each static input's bucket-varying dim is dynamic (maybe_mark_dynamic, so a forward that needs a fixed size still compiles per bucket) and compiles under the platform's overrides for any count. Size 1 keeps a graph of its own. Each bucket is still captured separately. Meant for encoders over fixed-size windows, whose per-bucket Inductor compiles are most of a cold start.
MIOpen keeps each database in one file (find and perf DBs as key=value text, kernels in SQLite), so packing caches warmed separately replaced the bundle's records with the last host's. A changed database is now merged with the bundle's last copy, this host's record winning on the same key, and appended only when the merge adds something. MIOpen files are compared by content: two warm-ups can leave databases of the same size and mtime.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant