Repository navigation
Conversation
…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.
…replay follows the registers
…mpt's, as without
…cond decode capture
… a top-k/top-p cut needs
… each row's last accepted slot only
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.
…r copied per layer
… request through the fused prologue
…attention keeps replaced row starts
…d head, refuse a context past the rope table
…p an unused availability check
… and writes it to the other
…e cache's whole context
…er pins the platform proxy
This was referenced Nov 3, 2026
This branch has not been deployed
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.
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
GDNDecodeWrappernext to #1090's verify plan, the prologue commit builds onattn/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_slotson 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, andmax_slotsstill 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: trueunderplatform_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 tocausal_conv1d_update.GDNDecodeWrappertakes it onis_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.PackedAttentionCallableasks 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.git diff d82f16aef 732dc085b -- mstaris 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):
Where it came from, 0.8B, cumulative, two launches an arm (the pin alone three); share is points of c=64 ITL:
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
Tests
test/modular/test_recurrent_twin_slots.py, 14 tests: two slots a request withmax_slotscounting 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.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.What is not covered
platform_overrides: rocmon the 0.8B, 2B and 9B configs only, the conv routes onis_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.