Repository navigation
Resident decode planning: decode rows on device tables - #1127
Conversation
A steady decode step only moves lengths by one and crosses a page now and then, so keep each request's block table row, length, GDN slot and position offset on the device. The captured step reads them by a row of slots and commits its own lengths; the host sends what changed and stays the truth. Mixed and eager steps keep the host plan and mark their rows. Behind knobs: resident_planning (auto|off), Qwen3.5 decode opts in, fa2 keeps its per-step plan.
Picks up #852. No textual conflicts.
The planner's stage runs before runner.plan sets plan_phase, so a fresh stage's laps landed under res.plan. A resident step doesn't read or write the recurrent addressing memo; the cached KV plan's np.insert still feeds fa2's attention plan on a resident step.
#980 has the runner drop its stage when an admit is refused, since the unwind gave back the pages it planned into. The planner's stage names the same pages in its crossing cells, so with no commit in between a retry over the same rows promoted cells pointing at a page the row no longer held. The KV pool now drops it wherever it gives a stage's pages back.
The knob lists in knobs.py and docs/serving.rst conflicted with #988's compile_deterministic entry; both sides kept.
#859 declares device loop-back (DeviceLoopback) and builds Qwen3.5's decode rotary tables in the graph. The engine enables the sampler's reader the way main does now and then builds the planner; Qwen3.5 keeps only resident_walks, since the loop-back hooks come from the declaration.
|
Merged main again (#908, #859, #860) at the head above. Two conflicts: the engine now enables the sampler's loop-back reader the way #859 does ( |
|
Pushed de9824c, the #988 wiring: a resident step's capture gathers each row's used page columns and length into the fmha wrapper's own table (one Triton kernel), and a bucket that gathers that way takes no host plan at all. The branch as pushed before ran #988's decode without it and read 7-20% lower at 0.8B than the "+#988" column: on this head, 793 / 10112 / 26059 / 33047 tok/s at c=1/16/64/256 without, 993 / 11309 / 27902 / 34310 with (2B c=16/64 8759 / 22320 vs 8901 / 23253). The "#988" section of the description says the same. |
#820's LoRA lease (a base-only capture key when no row names an adapter) goes first, and the planner is matched to whichever lease it returns. Imports merged.
|
Merged main once more for #820: its LoRA lease (the base-only capture key when no row names an adapter) runs first and the planner is matched to the lease it returns; a node with a LoRA pool stays host-planned, since a resource without tables refuses the planner. Imports merged, nothing else conflicted. |
|
@merceod this overlaps with my #933 and #966 (draft PRs, I'm working on getting them ready), so here's where they stand relative to #1127: Like #1127, they keep decode-step state resident on the device so the host plans less. Some differences: - Which state: #1127 keeps host-truth tables (KV block table and lengths, GDN slot, position offset). #933/#966 give every resource on the walk a per-row share, including state that lives only on the device: the sampler's last token, RNG offset and loop-back states.
Results so far (1× H200, Qwen3.5 0.8B, prefix cache off, two sessions per arm): #1127 measures +1–3% over main (c=64: 20.9k / 20.7k vs 19.8k / 20.6k tok/s; c=128: 27.0k / 27.6k vs 26.6k / 27.3k). Its fmha decode path is Blackwell-only, so it doesn't run here (and also I'd assume the model to be more host-bound on Blackwell than Hopper, which would affect the results). Let me know if you think I should drop #933 and #966 in favor of #1127? Or, if the A/B favors them, whether there are parts of #1127 (the fmha wiring, the debug check of device rows against streams) you'd want on top. |
|
Thanks for running it, and for laying the two out side by side. Where the gain is. c=64/128 is the regime where both of us are inside the spread: on the B300, #1127 reads +2% at c=64 and c=256 against How I read the two. #966 is the general form and #1127 the narrow one, and the line between them is device-only state. Everything on #1127's tables has host truth, which is why mark-stale is enough: a mixed step changes nothing on the device, and the next resident step rewrites the row from the host. Once the sampler's last token and the RNG offset sit only on the device, a mixed step has to copy out first, which is your lend. That buys the host less work (admit, declare and the per-row commits, which #1127 leaves in Python: 0.08 ms of plan thread at 0.8B c=16, 0.23 at c=256) and costs the moves and lends. The per-step floor looks alike: a steady resident step spends 0.04-0.06 ms of plan thread on the slot row and the staging H2D (the Suggestion. I wouldn't drop #933/#966. I'd land #1127 first as the narrow step: it sits on Same harness, both boxes. I'll run |
|
Sounds good. I will hold the B300 run for the refined #966 then, and run the four arms when you say it is ready. Thanks for taking the review. |
NSagan271
left a comment
There was a problem hiding this comment.
The overall mechanism is good and very much needed, but I have some concerns about how row allocation interacts with admission and offload (in particular, this implementation seems to raise if we run out of space in the device table, as opposed to applying backpressure in admit, with potential to offload or evict).
| # per-request state off tables on the device; each owner of that state | ||
| # keeps its own table. A resource overrides these only once it has one. | ||
|
|
||
| def resident_refusal(self) -> str | None: |
There was a problem hiding this comment.
Nit: I wonder if it would be cleaner to put these resident functions into their own mixin that Resource uses, just to group the methods together.
There was a problem hiding this comment.
Done in e0cee953b. The hooks moved to ResidentHooks in mstar/engine/resources/resident_hooks.py with their docstrings, and Resource inherits it. The module imports resident.py only under TYPE_CHECKING, so there's no cycle.
| if ready: | ||
| self._cooling = [(slot, seq) for slot, seq in self._cooling if seq > done] | ||
| self._free.extend(ready) | ||
| if not self._free: |
There was a problem hiding this comment.
Running out of slots raises in the middle of plan, after every resource's admit has already succeeded. On the exec path that reaches Worker._handle_main_loop_error, which fails every in-flight request, not just the one that needed a slot. Per invariant 1(4), an over-subscribed finite pool should be backpressure the scheduler can resolve.
Elsewhere the contract is that whatever owns a finite thing refuses at admit with an AdmitFailedReason. The runner then rolls back every resource the step reached, and the worker re-queues, evicts or fails depending on the reason (e.g. RecurrentStatePool._refused → AllocationFailed). Positions and the sampler hold nothing finite per request, so their admit is always ADMIT_OK. The planner isn't a Resource and isn't in the runner's admit chain, so it has no way to refuse. A host-plan fallback doesn't work either, since every captured graph of a resident walk has the gather recorded.
There's also a reason eviction can't help as written: an offload marks the row stale but keeps its slot, which goes back only at REMOVE_REQUEST. So an AllocationFailed from here would have _try_offload_cold_request offload a victim without freeing anything. A reload already rewrites the row whole, so releasing the slot on offload looks cheap. Once that's done, I think held slots are bounded by rows resident on the device, which the owners' own pools already cap. That gives two options:
- Size the slots at startup from those caps (e.g. the recurrent pool's usable slots, plus slack for slots cooling behind in-flight steps), so exhaustion can't happen. On 0.8B that would also take the store from 4096 rows to ~130.
- Take the slot at admit and refuse there.
(1) seems simpler if the bound holds for every resource; otherwise this probably needs some admission / eviction mechanism.
There was a problem hiding this comment.
Went with your first option, sized by the batch cap as in your follow-up. The planner can't run out any more: 2B+1 rows, and a step short of free slots displaces the least recently run requests. Those just take a fresh slot and a whole-row rewrite the next time they run. Nothing in the planner refuses, so there's no admit-time path to add. The offload half is fixed too: Engine.offload_request releases the slot, and so does a TP follower's replay of the leader's offload (apply_resident_delta, which doesn't go through offload_request). So an AllocationFailed that offloads a victim now frees its slot as well as its pages. Marking the row stale on offload and reload is gone, since the row has no slot left to mark. test_a_refused_admit_offloads_a_victim_and_its_reload_is_written_whole runs that through _handle_allocation_failure. The bound's argument is in the reply above and in ResidentPlanner's docstring.
One thing differed from your sketch: the recurrent pool can't be the refusal that triggers the offload. It refuses with AllocationFailed, but it doesn't evict (supports_eviction is False), so _try_offload_cold_request finds nothing evictable for gdn_state and holds the named row instead. test_a_recurrent_pool_refusal_holds_the_request_since_the_pool_cannot_evict pins that, and the offload test goes through a KV refusal.
| from mstar.engine.resources.kv.resident import KVResidentTables | ||
|
|
||
| page_size = self._groups[0].page_size | ||
| self._resident = KVResidentTables(planner, page_size, -(-self.config.max_seq_len // page_size)) |
There was a problem hiding this comment.
The table is ceil(max_seq_len / page_size) columns wide. The KV config's max_seq_len here is the model's max_position_embeddings (262144 on Qwen3.5), not the serving config's max_seq_len: 32768. So each of the 4097 rows has 2048 columns, which is the whole pool (max_num_pages defaults to 2048).
A stream's pages are distinct pool pages, so min(ceil(max_seq_len / page_size), pool pages) is a safe width. If the serving max_seq_len caps request length, that would give 256 columns, ~8 MiB instead of 64. With the current formula, a 1M-context model on 16-token pages comes to ~2 GiB at the default slot count.
There was a problem hiding this comment.
Right, the max_seq_len here is the pool's (max_position_embeddings, 262144). The serving config's top-level key isn't read anywhere but the config check, so it can't cap a request. Width is now min(ceil(max_seq_len / page_size), request_pages): a stream's pages are distinct pool pages, so the pool bound always holds, and a smaller resources: kv_cache: max_seq_len narrows it. With the slot count at 257 the store is 4.0 MiB on Qwen3.5, from 64.1. A 1M-context model on 16-token pages with a 2048-page pool gets 2048 columns too, since the pool caps it, not the context. KVResidentTables.row raises if a stream holds more pages than the width, and so does a crossing past it, which would otherwise write into the next row. Neither is reachable with the pool bound.
| self._last_rids = None | ||
| if self._staged is not None and rid in self._staged.rids: | ||
| self._staged.commits = -1 | ||
| seq = int(self._last_read[slot]) |
There was a problem hiding this comment.
Under TP each rank has its own planner. Whether a released slot is free or still cooling depends on how release (main thread, at REMOVE_REQUEST) interleaves with commit (GPU thread, advancing _committed_seq). Followers replay removes in the same gap between steps, but I don't think that pins this interleaving. Near exhaustion, one rank could raise in _take_slot while the other runs into the step's collective. This goes away if the slot count is bounded at startup. An admit-time refusal would need to be decided identically on every rank.
There was a problem hiding this comment.
With the bound, no rank can come up short whatever the interleaving. Whether a released slot is free or cooling when the next step plans only changes which slot a row gets on that rank, and slot numbers stay inside a rank's planner: no collective, message or capture decision reads one. Victims are ordered by plan seq, ties by the order rows took their slots, so ranks with the same history displace the same requests (test_ranks_with_the_same_history_displace_the_same_requests). Where a release lands on the other side of a commit, one rank may take a free slot where another displaces, which only changes which row gets a host rewrite and one H2D (test_a_rank_whose_release_lands_on_the_other_side_of_a_commit_still_runs). 9B TP2: 40 requests at c=8, 0 failed, and the row check read 0 mismatches over 1200 steps on each rank.
| if ctx.resident is not None: | ||
| ctx.plan_phase = "res.pre_plan" | ||
| t0 = phase_start() | ||
| ctx.resident.plan(step, ctx) |
There was a problem hiding this comment.
Question: if resident.plan raises here (out of slots), pre_admit has already reserved, and this block has only a finally. The worker's pre-plan wrapper logs it and calls _reset_skip_plan_flags. Does anything roll back what pre_admit took before exec re-drives the step?
There was a problem hiding this comment.
Nothing leaks, on main or here. pre_admit puts its pages on the streams and marks them in flight. A raise anywhere after it in that block, from resident.plan or from the runner's own pre_plan, unwinds to Worker._preplan_spec's except, then _reset_skip_plan_flags, reset_pre_plan_for_batch, StepRunner.clear_preplan. That drops _pre_admitted and batch.step, so exec declares and admits afresh, and KV's clear_preplan clears the in-flight marks and removes any labels the stage invented. The reserved pages stay with their streams, where exec's admit finds them: the decode row's pages already cover its token, so it acquires nothing (kv/manager.py, the len(page_indices) * page_size >= seq_len check). If the row isn't in the retried step, the page waits on its stream for the row's next step or its teardown. ctx.resident.plan sits inside the same try as the runner's pre-plan, so it gets the same unwind, and there's a comment at the call saying so. test_a_raise_in_the_preplan_leaves_the_reservation_for_exec runs the real _preplan_spec and _pre_plan_for_batch with a raise from each, then checks the free page count and the stream's pages across the re-driven step.
What could leak was the planner's own marks: a build that raised after taking rows off _stale lost them. _build now hands them back on any raise (test_a_plan_that_raises_leaves_its_rows_to_the_next_step).
|
|
||
| self._lock = threading.Lock() | ||
| # rid -> slot. Handle-keyed: `release` purges it. | ||
| self._slot_of: dict[int, int] = {} |
There was a problem hiding this comment.
_slot_of, _stale and _moved are handle-keyed and purged in release, which looks right. Could you add them to the list of known handle-keyed state in mstar/worker/rid_table.py's docstring?
There was a problem hiding this comment.
Added: ResidentPlanner.{_slot_of, _stale, _moved} (purged by release). Displacement drops all three for the displaced request too.
| - config's ``knobs: resident_slots``, else ``4096`` | ||
| - Request rows of those tables, allocated at startup because a captured | ||
| graph holds their addresses; a step that needs a row past the last | ||
| fails. Each row costs a cell per page of ``max_seq_len`` plus a few, 8 |
There was a problem hiding this comment.
"a cell per page of max_seq_len": readers will compute this from the serving config's max_seq_len, but it's the KV pool's (the model's native context), 8x larger on Qwen3.5. Worth saying which, or changing the width (see the comment on KVManager.bind_resident).
There was a problem hiding this comment.
The MSTAR_RESIDENT_SLOTS row is gone with the knob. The MSTAR_RESIDENT_PLANNING row now says how the tables are sized, and that a row's cells count pages of the KV pool's max_seq_len (the model's context, or resources: <kv>: max_seq_len), not the config's top-level key, capped at the pool's pages.
| return | ||
| with self._lock: | ||
| if len(values) > self._numel: | ||
| self._grow(len(values)) |
There was a problem hiding this comment.
Minor: the first flush over 64Ki cells (e.g. a mixed step marking a few hundred long rows stale at once) grows the stager here. _grow synchronizes every ring event, on the GPU thread right before the replay. It's a one-time host sync, but it lands in the step path. Could the planner size its stager at build time?
There was a problem hiding this comment.
The planner builds its stager at its largest flush, the widest capture's rows at full stride plus a cells row (B × 2050 + 4B cells on Qwen3.5, 16 MiB of pinned memory over the ring's 8 buffers). So _grow doesn't run in the step path. test_a_flush_of_the_largest_step_stages_without_growing checks the size on CPU, and its CUDA twin flushes B full rows through the real stager with _grow patched to fail.
| _assert_row(planner, live, 1) | ||
|
|
||
|
|
||
| def test_running_out_of_slots_says_how_to_fix_it(): |
There was a problem hiding this comment.
This checks that it raises, but nothing checks what the batch or the scheduler sees when it does. If this becomes an admit refusal, a test through _handle_admit_failure (re-queued, not failed) would pin it down.
There was a problem hiding this comment.
With no refusal there's nothing to re-queue, so the test pins the cache instead:
test_more_live_requests_than_slots_run_off_the_cache, four seeds: B=4, so 8 slots and the trash row. 20 live requests take turns over 300 steps, with pre-plans, abandoned pre-plans, host-planned mixed steps and teardowns with recycled handles mixed in. Nothing raises, and every row a step reads matches its stream. Each displaced row is written whole when it next runs, and the victims are always the least recently read. The slot ledger (free, cooling, held) partitions the slots after every step.test_a_displaced_row_is_written_whole_with_its_lengths,test_a_step_keeps_its_own_oldest_row_and_a_victim_leaves_no_marks,test_a_step_does_not_displace_its_own_rows_or_the_step_behind_it,test_plans_that_never_ran_hold_no_slots.- Through the worker:
test_a_refused_admit_offloads_a_victim_and_its_reload_is_written_whole(KV refuses, the real_handle_allocation_failureand_try_offload_cold_requestoffload the victim, its slot is released, the retry runs, and the reload is written whole at its own lengths), plus the recurrent-pool case above. Alsotest_a_follower_replaying_an_offload_gives_the_slot_up_too. - The knob's tests are gone, and one checks the knob is unknown now.
Live test: resident decode planningWorks and helps a little on H200. One thing to fix in the description: no shipped config turns the feature on. H200 · tested at No shipped config enables itResident planning covers the decode step only when For everything below I set Throughput (ABBAAB, 3 restarts per arm), output tok/s
c=8, c=32 and i2t are higher on every restart. c=1 and c=128 overlap with main, and the direction is the same as on B300 but smaller. 0 failed requests and 0 tracebacks. Correctness (2 restarts per arm, plus a PR server with
|
| Check | main | PR | PR, resident off |
|---|---|---|---|
| GSM8K, 100 greedy | 55, 58 | 57, 57 | 52 |
| Needle at 2k / 8k / 24k, depth 0.1 / 0.5 / 0.9 | 9/9, 9/9 | 9/9, 9/9 | 9/9 |
| 32 short prompts concurrent vs solo, identical | 8, 7 | 8, 8 | 10 |
| Abort storm (24 dropped, 16 fresh, ×3) | all OK, fresh ≤ 0.35 s | all OK, ≤ 0.33 s | all OK, ≤ 0.35 s |
/health, tracebacks |
200, 0 | 200, 0 | 200, 0 |
The GSM8K scores fall in one band (0.8B greedy changes across restarts on main too). Agreement across arms is at the same level as within one arm.
Still to come: KV pressure at 96 pages (t2t c=32, GSM8K, storm) on both arms, MSTAR_RESIDENT_SLOTS=16 at c=64 to exercise slot exhaustion, and Qwen3.5-9B TP2.
|
@merceod Update on how #933/#966 relate to this, after more measurement on 2× H200 (Qwen3.5-0.8B, worker_phases).
I think, since the host cost of #966 is larger than its benefit, I'll just keep #933/#966 as drafts for reference and not land them (though #932 is worth landing, cc. @stephen-dwq). Two things #966 did that this PR doesn't, which I'd like to add in a follow-up PR on top of this one:
On my earlier comment about capacity/backpressure:
|
|
Thanks, all five were real, and they're one bug: per-step state the planner held once per node. It's keyed by resource key now: 0af5053, plus a test rename, a docstring and two merges of What's keyed how.
Read-only pools. A pool the walk declares at span 0 (a cross-attention context) keeps rows and lengths on its table, and the capture gathers both. There's no write addressing, no in-graph advance and no crossings: What still refuses. Everything from before except FlashInfer cross attention: bounded and ring KV, dense and MLA attention, Mamba2, the MLA layout, sliding-window pools and FlashInfer sliding windows, block positions, a pool with no sink, a pool shared by nodes or read across nodes. Also cross attention with its context in the query's pool, positions over a pool the step only reads, and a recurrent pool declared at span 0. Nothing reads cross attention's pages off the tables, so a node with cross attention still plans attention on the host every step. On Whisper that means the KV write addressing and positions come off the tables, but no bucket is host-plan-free. Whisper. I didn't port #1143. To check the planner takes what it needs, a run-local hook gives The gate. B300 GPU 0, resplan3's harness and setup, this change (0af5053) against the head you reviewed after resplan3's rework (7d2a6cd), 0.8B, 0 failed in every cell. Medians: c=1 1030 vs 1013 tok/s (+1.6%), c=16 11049 vs 11358 (-2.7%), c=64 27844 vs 26818 (+3.8%). That's three rounds, order swapped each time, and five at c=16, where the head read lower in four. Paired by round the sign flips at c=1 (+3 / -8 / +4%) and c=64 (+7 / -5 / +3%), and at c=16 it's -1 / -3 / -6 / -2 / +3%, with the two arms' ranges overlapping (10607-11442 vs 10638-11635). What the gate is for is the plan thread, and the per-key records don't show there. The Also run on the change: greedy 32/32 against After each merge of |
The engine asks each submodule for its resident walks at warmup, before sizing; #1142's test builds its submodule as a SimpleNamespace.
|
For the record after the merge (8741219): gated the sizes this turns on that had not been measured, and the merged head. Setup: B300 GPU 0, resplan3's harness,
|
The merge kept #1127's copy_runs_ from the old ring stager, which reads _numel, _next, _views and _bufs. None of those exist on #932's segmented PinnedStager, so every resident flush on a GPU raised AttributeError (CPU takes the early torch-copy return, which is why CI didn't see it). It now claims one contiguous run of the arena with the same _claim copy_ uses, then issues one cudaMemcpyAsync per destination (torch copy off CUDA), and notes its stream on the segment like copy_ does. test_runs_past_the_staging_buffer_grow_it is sized off _seg_numel, since a 4-element stager now starts at a 32 KiB segment.
The #1127 tests that broke on the merge, all in their fakes: - page lists go through StreamPages.assign/clear instead of replacing CacheStream.page_indices with a list, which admit_rows can't take - _decoder_step passes the bucket's capture length to _static_plan_state - the stager test reads _seg_numel - the pre-plan raise test's engine/batch/runner fakes get _pre_plan_done, uniform_inputs and pad_inputs(uniform=), and torch.Event is faked on CPU, so the raise comes from the planner or the runner as intended instead of from the fakes before pre_admit. It also checks the event is recorded on the raise. - test_preplan_error_paths' step context gets resident=None
The merge put #1127's resident stage ahead of _order_after_pre_plans on the exec path, so an inline plan's cells and slots copies were issued before the wait. A dropped pre-plan's copy into the same bucket buffers on its plan stream could then land after them. The wait now comes first, as it already did on the pre-plan path.
…mns, one-copy staging, a segmented stager (#932) * Host overhead 9/N: per-step planning in the Rust page arena, one-copy staging, a segmented stager The KV page pool's bookkeeping moves into a Rust arena (rust/src/page_arena.rs, kv/arena.py): the free list, owner counts, seals and every stream's page list, in each of the pool's groups. A plan label is built in one pass over its rows' stream ids, and one call (kv_plan.rs) computes its FlashInfer arrays, host lengths and decode write locations; views are built only when a consumer asks. PythonPageArena keeps the previous implementation for a build without the extension, and is the arena the conductor's routing mirror runs. Also: captured steps stage their write addressing in one copy per slot buffer and none when the buffer already holds zeros; eager writes' addressing is computed on the host (Rust, numpy fallback) and uploaded in one copy; CUDA-core decode wrappers plan directly after their first plan; cross attention plans off the context label's columns; the sampler's scalar rows stage in one copy (ScalarRows); PinnedStager is a segmented pinned arena copied with cudaMemcpyAsync on CUDA; every step is ordered after earlier pre-plans' copies; worker_phases keys step phases by node and walk. Squashed from the review branch onto PR 8 (#908), with main's KV groups (#950), routing mirror (#920), chunked prefill and #852 integrated. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * KV: a step's admit, plan and commit pass over its rows in one arena call On a step that isn't a capture, the page arena reads each row's stream off the manager's streams in one call: admit_rows finds the rows whose pages already cover their span (the rest reserve per row as before), plan_rows builds the first group's plan arrays from each stream's length, generation and page list, and commit_rows moves each row's counts. The calls read and write the CacheStream fields themselves (stored_len, generation, step_in_flight, offloaded), so CacheStream is slotted; both arenas implement them, the Python one row by row. Each pass hands back what it can't do: a row with no stream yet, an offloaded row, a row short of pages. Admit takes the passes only in a pool of one group and with no chunked row reserving ahead; commit only with no prefix index, retention or debug checks; plan only for the first group. Everything else takes the per-row path. An abandoned pre-plan clears its rows' in-flight marks with one call over its segments (set_rows_in_flight) rather than a list of (rid, label). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test: a CUDA-core decode wrapper's later plans go through the planner too Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * LoRA graph test: the bare engine carries its pre-plan map * Resident gather and its test pass the bucket's capture length * copy_runs_ on the segmented stager The merge kept #1127's copy_runs_ from the old ring stager, which reads _numel, _next, _views and _bufs. None of those exist on #932's segmented PinnedStager, so every resident flush on a GPU raised AttributeError (CPU takes the early torch-copy return, which is why CI didn't see it). It now claims one contiguous run of the arena with the same _claim copy_ uses, then issues one cudaMemcpyAsync per destination (torch copy off CUDA), and notes its stream on the segment like copy_ does. test_runs_past_the_staging_buffer_grow_it is sized off _seg_numel, since a 4-element stager now starts at a 32 KiB segment. * Resident tests on StreamPages, capture lengths and the pre-plan event The #1127 tests that broke on the merge, all in their fakes: - page lists go through StreamPages.assign/clear instead of replacing CacheStream.page_indices with a list, which admit_rows can't take - _decoder_step passes the bucket's capture length to _static_plan_state - the stager test reads _seg_numel - the pre-plan raise test's engine/batch/runner fakes get _pre_plan_done, uniform_inputs and pad_inputs(uniform=), and torch.Event is faked on CPU, so the raise comes from the planner or the runner as intended instead of from the fakes before pre_admit. It also checks the event is recorded on the raise. - test_preplan_error_paths' step context gets resident=None * Inline plan: wait on pre-plans before the resident stage The merge put #1127's resident stage ahead of _order_after_pre_plans on the exec path, so an inline plan's cells and slots copies were issued before the wait. A dropped pre-plan's copy into the same bucket buffers on its plan stream could then land after them. The wait now comes first, as it already did on the pre-plan path. * Side-step test: the fake worker carries a phase period --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: merceod <atindra@mstarlabs.ai>
|
The sampler follow-up you asked for is #1164: each request's sampling scalars are planner cells, read in the graph by one Triton kernel, and the slot-index row uploads only when the planner's slot row changes (49-64 uploads per 1,000 steps at c=64/128 against every step before). On a B300 at 0.8B c=128 the sampler's host time per captured step goes from 139-158 µs to 83-91; at c=64/128 that comes back as |
One conflict, in mstar/engine/engine.py: mstar-project#1127's resident planning and this branch's weight-reload methods were both added in the same gap between the build loop and paged_kv_caches. Both are new methods with nothing in common, so both are kept. The branch's net diff against main is unchanged by the merge.
Conflicts: - kv/manager.py: main's plan states (#932, #1127): every eager packed step's addressing is built on the host, and decode takes the resident rows. #388's small-step host path and its thresholds go, with KVPlanOutput.max_span and test_kv_host_plan.py. correct_len now uses tokens_per_page, cache.copy_pages and StreamPages' del/extend, and refuses a cache with KV groups. - glm52/submodules.py: main's TokenStop (#794) is the plain decode's stop. Under MTP, check_stop counts a verify step's tokens, and a prefill stop or one-token budget returns END_PARTITION, which replaces the ended_at_prefill latch. The prefill groups read PREFILL_TOKEN_BUCKETS. - fused_moe/decode.py: main's kernel registry (#986). GLM-5.2's H200 tiles move from _H200_GLM52 to kernels/tables/moe_decode/NVIDIA_H200.json, so they now apply on H200 only. - mla.py, adding_models.rst: both sides' additions. - test_glm52_dsa_engine.py, test_fused_moe_decode.py: #388's tests on the merged code. Also, to fit main: - GLM-5.2's per-step index upload goes through to_device_nonblocking, since PinnedStager no longer has to_device. - prefill_token_buckets is a named Glm52Model argument, which the shipped-ladder test checks. - #388's tests take main's capture_into_graph stream, the arena's owners/is_sealed, and swap the KV manager's arena (now a read-only Rust object) rather than patch it. - test_glm52_capture_gpu's harness ends decode at Glm52Model.get_loop_caps, as the runtime does: since #857 check_stop stops on EOS only, so main's own copy of this test streams 15 tokens for 12 and fails.
Keeps each decode request's block table row, stored length, GDN slot and position offset on the device, so a captured decode step reads its rows by a row of slots and commits its own lengths. The host sends only what changed. On a B300 against
main53c47cb, which already builds decode's rope in the graph (#859) and runs #988's fmha decode, this is +32% at 0.8B c=1, +12% at c=16, +7% at c=64 and +10% at c=256, and flat at 2B (+1% at c=16 and c=64). On Blackwell every decode bucket then needs no host plan at all.Why
0.8B is host-bound: its step is the three step threads' CPU. Two steps of one batch differ by "every length +1" and "a row crossed a page", yet the plan thread rebuilds each row's write location, position and GDN slot every step. #852 made that rebuild cheaper. This stops doing it for a captured decode step, and every other step keeps #852's cheaper rebuild.
What changed
mstar/engine/resources/resident.py:ResidentPlanner, one per node. A slot is one row of a fixed-address int64 store, and slots are a cache, not something a request can run out of:2 * B + 1rows, B the widest capture of a resident walk (128 on every Qwen3.5 config, so 257 rows, 4.0 MiB), the +1 the trash row. A step whose new rows find no free slot takes those of the least recently run requests (by plan sequence number, ties in the order they took their slots) that aren't in the step and that no uncommitted step read. A displaced request takes a fresh slot when it next runs and is written whole, as a new request is: every cell mirrors host state. 2B is always enough, on every rank and whatever the commit timing. A step's rows are at most B, and besides them only the rows of the one earlier plan that may not have committed are out of reach, at most B, because a plan is built only behind the commit of the step before it. So nothing here refuses at admit or raises, and the size is derived, not configured.Engine.offload_requestand in a TP follower's replay of the leader's offload: a reload is written whole anyway, so the old "mark it stale on offload and reload" is gone.kv/resident.pya pool's block table and lengths,recurrent/pool.pythe GDN slot,position/manager.pythe position offset.planon those resources returns right away on a resident step, and the graph fills the same static buffers the model already reads, so the forward doesn't change. The hooks they override sit in a mixinResourceinherits (resources/resident_hooks.py);bind_residenthands each owner its resource key.ResidentStream: the label of its own segments, or the step's when it names none, and whether it appends a token a row or only reads), and every step of that bucket plans with the same records. Every owner reads its own label.ResidentWrites.kv_lenhas a stored length per pool, and positions take their offset against, and gather with, the pool theirPositionConfig.kv_cachenames.CapturedResidentStepholds the store once and each pool's block table, andResidentBuffersa length buffer per pool per (bucket, slot), so an fmha decode wrapper reads its own pool's table and lengths. A capture refuses two labels, or spans other than all 1 or all 0, on one resource.host_plan_free, still one per bucket) only if no attention needed one: cross attention reads both its pools' KV plans, so a node with it plans every step. FlashInfer cross attention over a context pool of its own now opts in, host-planned as before; with its context in the query's own pool it refuses.min(ceil(max_seq_len / page_size), request_pages)block table cells: the pool'smax_seq_len(the model's context unlessresources: <kv>: max_seq_lenoverrides it; the serving config's top-level key doesn't reach the pool), capped by the pool, since a stream's pages are distinct pool pages. On Qwen3.5 that's the pool's 2047 pages, one span.accelerator_graph_runner.py: the gather before the forward and the length commit after it are recorded in the capture.engine.py: builds the planner at warmup, from the runner's widest capture of the node's resident walks and before any capture records the store's address; plans, flushes rows before the replay, commits. Nothing reads back: the host advances its copy by the same rule.SINK_PAGE, stored length 0 (attention length 1),SINK_SLOT.submodule_base.py:resident_walks(), empty by default. Qwen3.5 namesdecodeon every config: since Performance 2/7: speech and language models #859 a captured decode step builds its rope in the graph whateveringraph_decode_ropesays, so itspreprocessreads nothing the tables fill (only the loop-back tokens, off the sampler).Config and guards
knobs: resident_planning: auto | off, defaultauto, envMSTAR_RESIDENT_PLANNING. On only where the submodule names walks, a capture covers them and every resource of the node can; otherwise the worker log says why. There's no knob for the tables' size. Resources refuse throughResource.resident_refusal, and the base refuses, so a new resource stays host-planned until it has a table. Tested refusals: a resource with no table, a pool shared by nodes, a pool read across nodes, MLA layout (GLM-5.2), sliding-window pools, a FlashInfer sliding window, block positions, a pool with no sink, cross attention whose context is in the query's own pool, bounded KV, ring KV, dense and MLA attention, Mamba2. FlashInfer cross attention over a context pool of its own no longer refuses: its context pool is read only on the walk and it plans on the host. A capture raises if a resident walk declares, on any one resource, more than one stream or spans other than all 1 (one token a row) or all 0 (read only); or forks, retention, a sink, explicit positions orcommit=False. It also raises if positions aren't one stream with the pool they name, appending, and if a recurrent pool is only read. A stream with released tokens raises at its row, and so does one with more pages than a row holds, which the pool bound makes unreachable.Every Qwen3.5 config turns it on: the LLM node names
decodeand all six of its resources can. Run here on 0.8B, 2B, 4B, 9B on one GPU and 9B TP2. Each logs256 slots of 2050 cells, 4.0 MiB(per rank on TP2), and on this B300 the capture log has every decode bucket, bs 1 to 128, host-plan-free. 27B TP4 needs four GPUs and this box has two, so it hasn't run. Whisper's decoder (a self-attention pool, a context pool and cross attention) names no resident walk on this branch; that's #1143's one line. With it named run-locally, large-v3 builds a planner (resident planning on decoder: 128 slots of 19 cells, 0.0 MiB, walks ['decode']: 3 narrow cells, 4 self-attention pages, 12 context pages), every decode bucket keeps its attention host plan (cross attention reads both pools' KV plans), and c=1 gives 200/200 transcripts againstmainwith the row check at 0 mismatches over 5,600 resident steps.On top of #852
#852 is in
mainand this branch merges it (merge commits, no rebase). Git merged the ten files both touched without a conflict, and the branch's diff against 9c05c56 is line for line the diff it had against its old base, plus one commit for the places they meet:np.insert. It skips_setup_plan_states, the write addressing the graph now fills from the table. Host-planned steps keep all of Host overhead 5/N: cheaper resource planning, admit and commit #852: a captured prefill or mixed step stages its addressing through_stage_packed_plan_state, decode through_stage_decode_plan_state._build_addressing, so it neither reads nor refreshes the memo, and a steady decode step doesn't bump_version(commit bumps it only whenhas_stateturns on). The next host-planned step over the same rows, a mixed step say, still finds the memo. The pool's table and the memo read the sameSlotState.index, which for a live request is set once, at the alloc before its first decode step, so the two can't disagree.StepRunner.plansetsctx.plan_phase, so the engine sets it first and the stage's laps land underres.pre_plan.resident,res.plan.residentorres.plan_fresh.resident, beside the resources' own.#852's five test modules (
test_gdn_pinned_stage.py,test_kv_packed_plan_state.py,test_flashinfer_direct_plan.py,test_qwen3_5_mixed_ids.py,test_recurrent_preplan.py) andtest_kv_decode_plan_cache.pyrun unchanged: they build resources with no planner, and a step without one is the host path they test.test_resident_planning.pygains three tests for the seams: a resident run leaves the memo for the next host step, a resident step over #852's cached KV plan (crossings included) hands attention the same indptrs a host plan does, and the stage's laps follow the step's plan phase.mainmoved again while this was measured (#857, #977, #980, #916), so the head merges it a second time. One conflict, inStepContext, where #916's chunk fields and the planner fields both landed: kept both. #916's chunked prefill runs on prefill and mixed walks, which stay host-planned, so a chunked prompt's row is written whole at its first decode step, as a prefill's always was; its up-front reservation adds pages past the row's committed columns, which no table reads.test_chunked_prefill.py's engine stub neededresident=None, as other stubs here already have. On top of that, one fix where #980 and the planner meet. #980 has the runner drop its stage when an admit is refused, since the unwind gave back the pages the stage planned into. The planner keeps a stage of its own whose crossing cells name those pages, and it promoted that stage whenever a retry had the same rows, bucket and slot with no commit in between, so the retried step wrote its token into a page the row no longer held. The KV pool now drops it wherever it gives a stage's pages back (rollback_admit, an unpromotedclear_preplan).test_a_refused_admit_drops_the_stage_naming_its_pagesfails without it: the step writes at page 3, its stream at page 4. It runs only on refusals and abandoned pre-plans, so the gate below, measured without it, still describes a steady step.That merge, fa20064, was checked on its own against
main4079d7f at the time. Greedy, prefix cache off and on (both passes): 32/32 at 0.8B and 2B,mainagainstmain32/32; KV pressure withinmain's own spread, 0 assertions. #916's chunked prefill with the 512-token gate prompts at a 200-token budget (min_chunk_tokens: 64), so every prompt fills in three or more host-planned chunks before its first decode step: 32/32 at c=1 at 0.8B and 2B. At c=8 the outputs follow the batches:mainagainstmain21-28/32 over four servers, this branch againstmain16-28/32 over five, and the diverging outputs I read are coherent and about their prompts, the earliest a first-token flip between two of the model's usual openings. To check the tables rather than the text, a debug hook compared each resident step's device rows, lengths and write location with the streams after the replay: 0 mismatches over 4539 steps under chunked prefill, 629 steps under KV pressure with 27 refused admits, and 1200 steps on each rank of 9B TP2 (40 requests, 0 failed). Whisper large-v3 c=1 200/200 transcripts and Qwen3-TTS seeded c=1 200/200 WAV bytes againstmain.This head merges
maintwice more, up to 53c47cb (#859, #821, #1110, #1129) and then 2394b4d (#1074), with no conflicts. #859 builds decode's rope in the graph on every config, which is what lets Qwen3.5 drop itsingraph_decode_ropecondition.What stays on the host
Performance
2x B300 (sm_103), GPU 0, torch 2.12.1+cu132, FlashInfer 0.7.2.post1,
--rust-frontend -O2,MSTAR_SHM_ARENA=1, 512-token prompts, 256 out, prefix cache off. Every arm runsknobs: step_thread_core: autoandsampler_ingraph_scatter: 1, pluskv_plan_cache: 1andgc_gen0_threshold: 50000;ingraph_decode_ropeis no longer set, sincemaindoesn't need it for decode. Server mask 32-47 + 160-175: step threads on 35,163, API on 44. Arms alternated per round, one Inductor cache per round so they replay the same kernel choices, two runs a session.As before, #852 runs its direct planner only on FlashInfer 0.7.2, and every arm widens that set to this host's 0.7.2.post1 at startup, outside the tree. It matters less now: on this GPU the decode buckets are fmha, and only prefill and mixed steps plan through FlashInfer.
Arms:
mainis 53c47cb. "head" is d279b7a. The pushed head 7d2a6cd adds 254e0fd (an abandoned stage restages its row, capture refusescommit=False, a node with no table builds), none of which a steady step runs, and a merge ofmain2394b4d (#1074, the conductor's wait). #988's wiring is in the head, and every decode bucket of 0.8B, 2B and 9B TP2 (bs 1 to 128) is host-plan-free on this GPU, so the old "+#988" column is the head itself.Per-key state, against 7d2a6cd. The change keying the planner's per-step state by resource key (0af5053) against the head before it, same setup, 0.8B, 0 failed: c=1 1030 vs 1013 tok/s (+1.6%), c=16 11049 vs 11358 (-2.7%), c=64 27844 vs 26818 (+3.8%), medians over three rounds (five at c=16), and paired by round the sign flips at c=1 and c=64 and is -1 / -3 / -6 / -2 / +3% at c=16. On the plan thread the planner's stage (
residentbelow) reads 0.042 / 0.047 vs 0.041 / 0.044 ms at c=1, 0.052 / 0.047 vs 0.049 / 0.051 at c=16 and 0.058 / 0.074 vs 0.056 / 0.056 at c=64, two phase-timing sessions each (in the 0.074 one every plan-thread phase ran 20-30% slow), and in one process per tree pinned to one coreResidentPlanner.planover the same 64-row steps is 43.0 vs 43.8 µs. The c=16 reading has no mechanism I can find; my hypothesis is session noise.Median over runs; each cell is tok/s [min-max] · ITL p50/mean ms · TTFT p50 ms.
mainin both: c=1 +40 / +32%, c=16 +23 / +10%, c=64 +9 / +6%, c=256 +20 / +2%. At c=256main's own two sessions split 30.1k / 35.3k as before, so the size there isn't resolved; the direction is.ARRIVAL, so rows run to completion and the rows holding slots stay near B. Displacement costs are below, under Correctness.Plan thread CPU by phase, ms per step, phase timing on, one server per arm, same placement:
main's step: the KV plan, KV staging, attention, GDN and positions are all near zero on a host-plan-free bucket. What's left is the planner's own stage (0.04-0.06 ms), declare, admit and padding, and the sampler.event_syncdoubles at c=16 and c=64 (0.074 -> 0.153, 0.132 -> 0.264 ms): the host now waits on the GPU more.The merged head against
main7cc4b04.mainis 7cc4b04, themainat launch (#1131, #1142, #1154 and #725 past 7ff4aec). The head is eb9500b with thatmainmerged locally (cf466990b). While this ran,mainmoved on to f1c35ef, and the pushed head merges that instead (Correctness, below). Same setup, knobs and placement as above. Two sessions per arm, arms alternated, one Inductor cache per round. A third 4B round at c=1 and c=16 reused round 1's cache. 0 failed in every cell, and no recompile-limit line or traceback in any gate server.main53c47cb read +32 / +12 / +7% at 0.8B c=1/16/64 and +1 / 0% at 2B. At 0.8B c=64 the merged head is +8 / 0% by round: level with or abovemain, by how much isn't resolved.resident planning on LLM: 256 slots of 2050 cells, 4.0 MiB, walks ['decode']. In the capture log, all 62 resident captures, decode bs 1 to 128, are host-plan-free, as at 0.8B and 2B.main: c=16 is -1% in each of three rounds, while c=1 reads -1 / -4 / 0%. Both are GPU-bound here, and the head's decode graph is a little longer (below).Where a 4B and 9B step goes, c=16, phase timing on, one server per arm (ms per step):
The same cells under
nsys(graph-level trace, gate placement, one 3 s window in a steady c=16 run, and c=1 at 4B), µs per decode step:main,event_sync(the host waiting on the GPU) is 1.44 ms of a 2.99 ms step at 4B c=16 and 2.60 of 4.45 at 9B c=16, and the GPU is idle outside the graph 63 and 36 µs a step. llm_kernels1 measured 4B c=16 as host-bound on an oldermain(bb0e0a8:event_sync0.04 ms, 635 µs idle a step). On thismainit isn't.event_sync+0.41 and +0.40 ms), and the step doesn't move. The main thread's CPU rises with the wait (2.00 -> 2.38 ms at 4B). Hypothesis, not checked: the wait spins._gather_column_kernelthree times (1.4 µs each) and_advance_column_kernelonce (2.1 µs), withmain's_split_combine_kerneland_scatter_gather_elementwise_kernel(3.9 µs together) gone. The kernels sum to the same (2409.7 vs 2408.5 µs a step); the seven more launches add about 5 µs of gaps inside the graph. At a GPU-bound size that's the cost: 4B c=16 reads -0.6%, about that graph's 0.5%. At c=1 the head's mean step undernsysis shorter thanmain's (less idle outside the graph), so -0.8% there is withinmain's own spread (392-414 tok/s over six runs).Correctness
Greedy, the runner's 32 text prompts at c=1, natural EOS, 128-token cap, GPU 1, against
mainon one Inductor cache per size. Yardstick:mainagainstmainon the same cache.Prefix off was run on the gated commit, on 254e0fd and (0.8B) on the pushed head 7d2a6cd: 32/32 each. Under KV pressure the outputs follow the batches, so
maindoesn't match itself either; the head against itself reads 20/32 at 0.8B and 23/32 at 2B. Those servers ran withMSTAR_KV_DEBUG_ASSERTS=1(0 assertions) and a debug hook that syncs after every resident step and compares each row's device cells (stored length, used block table columns), the replay's write location and the trash row with the streams: 0 mismatches over 600 steps (about 3,900 rows) on each of six servers, the pushed head's among them. Every offload gave its slot up: 22-28 offloads a server, and as many slot releases in the hook's log.Slot pressure. Decode captures capped at 8 (
MSTAR_QWEN35_DECODE_BUCKETS=1,2,4,8, so 16 slots) at c=64, 512-token prompts, 256 out, the row check on. Under Qwen3.5's stockARRIVALorder rows run to completion and nothing is displaced (29,600 checked steps, 0 mismatches). A run-local hook gives decodeTURNSorder instead, withspeculation_mixes_walks: {LLM: false}(below). That displaces 195,311 rows over 24,456 steps, eight rows nearly every step, and the row check reads 0 mismatches over 30,200 steps (240,590 rows), 0 failed of 1,280 requests, on the gated commit, on 254e0fd and on 7d2a6cd alike. Throughput with the check off, one server each, under the same hook: head c=64 4378 / 4425 tok/s againstmain's 4373 / 4704, and c=8, where rows keep their slots, 6228 / 5641 against 5052 / 5101. So with every row rewritten every step, c=64 is level withmain(-3%, inside its own spread).Forcing
TURNSwith walk mixing at its default fails decode rows with anIndexErrorinprepare_inputs(inputs["text_inputs"][0]), in clumps of four: 34 of 640 at c=64 onmainand 52 of 640 on the head, with device loopback on or off. With mixing off it's 0 on both. This isn't this PR, and it only matters for a model that setsTURNSon a decode walk that has a mixed walk. Hypothesis, not checked: a chain merges a prompt row whose token hasn't been routed yet.9B TP2 on both GPUs: 40 requests at c=8, 0 failed. The row check read 0 mismatches over 1,200 steps (9,477 rows) on each rank, and each rank logs 256 slots of 2050 cells.
4B and 9B on one GPU, on the gated head (eb9500b with
main7cc4b04 merged) againstmain7cc4b04, same greedy setup, one Inductor cache per size:The 9B prefix-off pair: the head matches the first
mainserver 32/32 and the second 28/32, and the twomainservers differ from each other in those same four requests (from token 35 to 126). That'smainagainst itself, not the tables.MSTAR_KV_DEBUG_ASSERTS=1, the row check): 32 of 32 ok, 0 mismatches over 600 resident steps (3,934 rows), 29 offloads and 29 slot releases, 0 assertions. With 10 pages a row is 12 cells (256 slots of 12 cells).16 slots of 2050 cells, 0.3 MiB; c=64, theTURNShook,speculation_mixes_walks: {LLM: false}): 195,189 rows displaced over 24,453 steps, and the row check reads 0 mismatches over 30,100 steps (239,992 rows), 0 failed of 800 requests. With the check off, one server each: c=64 2916 / 2943 tok/s againstmain's 2981 / 2967 (-1.5%), c=8 2912 / 2911 against 2919 / 2916.mainf1c35ef merged, plus the stub below) againstmainf1c35ef at 4B: greedy 32/32 with prefix caching off, and KV pressure with the row check at 0 mismatches over 600 steps (3,953 rows), 34 offloads and 34 slot releases, 0 assertions.The per-key change (0af5053) and each merge of
mainafter it (6496f14, then 7ff4aec) were rerun: greedy 32/32 at 0.8B and 2B with prefix caching off, and on 0af5053 with it on (both passes); KV pressure with the row check at 0 mismatches over 600-632 steps (27-30 offloads, as many slot releases, 0 assertions); Whisper with #1143's line run-locally 200/200 with 0 mismatches over 5,600 steps on both pools; Qwen3-TTS seeded 200/200 WAV bytes on 0af5053.Whisper large-v3 c=1: 200/200 transcripts against
mainrun on the same GPU and Inductor cache, WER 1.35% on both.mainon the other GPU with a fresh cache differs from that samemainin three transcripts (case and punctuation), so that difference comes from the environment. Qwen3-TTS 1.7B seeded c=1: 200/200 WAV bytes, TTFA p50 20 ms on both. Neither runs on tables; they share the code.Tests
ruff check .clean.mainmerged up to f1c35ef: Attention workspaces grow to what a plan asks for #1142'stest_the_engine_seals_between_sizing_and_capturebuilds its submodule as aSimpleNamespace, and the engine asks every submodule forresident_walks()at warmup, so it failed on the merged head. cf77040 gives the stub one; nothing else changed. CPU Core on cf77040 againstmainf1c35ef: 10718 passed, 15 failed vs 10616 and 14, the same 14 plustest_eval_compare::test_stopping_compare_stops_the_job_and_its_servers, which fails onmaintoo when run alone. The gated tree (main7cc4b04 merged, with the stub): 10637 passed, 15 failed vsmain's 10537 and 12, the extra three all eval_compare tests that fail alone onmain7cc4b04 too. TP lockstep sim 11 passed on both. Focused modules 737 passed, 76 skipped (main634 and 70).CUDA_VISIBLE_DEVICES=, xdist 12), againstmain53c47cb's 10421 passed, 46 failed, 906 skipped: the gated d279b7a 10501 passed, 46 failed, 911 skipped, the same 46 by id; 254e0fd 10506 passed and 44 of them failed, the other two being eval_compare timeouts. The pushed head 7d2a6cd, withmain2394b4d merged: 10525 passed, 45 failed, 911 skipped,main's 46 less one of the eval_compare timeouts. All environmental here: a Rust bookkeeping attribute the extension build lacks (32), libibverbs (5), eval_compare subprocess timing (5), prometheus_client, flash_attn, one HTTP status. The first pass, on 00ab09a, failed 13 more, all intest_resident_delta.py: its engine stub bindsEnginemethods by name and lacked the new_release_resident_slot, which d279b7a adds. The per-key change 0af5053 againstmain2394b4d: 10545 passed, 45 failed vs 10443 and 44, the extra an eval_compare hang-timing test that fails onmaintoo when run alone. After each merge ofmain(6496f14, then 7ff4aec), against thatmain: the same failing ids (14, then 13; Make the Rust graph runtime the only graph runtime #1130 cleared the Rust-attribute ones). TP lockstep sim: 11 passed on every head.main: 601 and 58, withouttest_resident_planning.py).test_resident_planning.py: 102 passed, 4 skipped on CPU, 106 with the CUDA ones on a GPU; 21 of them cover the two-pool and read-only cases. Each of 15 parts of the slot cache, and each of 15 parts of the per-key state, undone on its own fails a named test. Four older tests build the planner's records by hand, so their stubs follow the new row and stream records; their assertions didn't change.test_kv_offload.py::test_oom_still_reports_allocation_failed, which fails onmaintoo, since the refusal is anAdmitRuntimeErrorthere now. CI has no GPU to see it.#988
#988 landed first, so its wiring is in this branch (de9824c). Inside a resident step's capture, a small Triton kernel copies each row's used page columns and its length off the tables into the table and lengths the fmha wrapper holds the addresses of (
FmhaDecodeWrapper.gather_resident). A bucket whose attention gathers that way takes no host plan at all (host_plan_free:FlashInferManager.planand the KV plan return at once). On this B300 that's every decode bucket, bs 1 to 128, on 0.8B, 2B and 9B TP2 (from the capture log), so the gate's head column is that path. fmha is Blackwell-only (fmha_decode_applies, capability 10.x): on an H200 the head plans fa2 decode per step, through #852's cached KV plan.Not in here
gather_static), and Hopper's fa2 decode indices built from the block table.main(f1c35ef: Performance 5/7: Qwen3.5 and Qwen3-Omni #970, Performance 6/7: OmniVoice, and a Qwen3-TTS decode fix #971, Performance 7/7: image, video and world models #972, Per-model perf and serving fixes for Qwen3.5, Qwen3-Omni, Fish, OmniVoice and others #1140). The gate below carriesmain7cc4b04. f1c35ef's Qwen3.5 changes are the prefill ladder (nowmodel_kwargs.prefill_token_buckets), the vision buckets and the GC knob in the shipped configs, none on the decode path. That merge is checked for correctness at 4B (below), not gated.DIRECT_PLAN_FLASHINFER_VERSIONSto 0.7.2.post1 in the tree: Host overhead 5/N: cheaper resource planning, admit and commit #852's call, not this PR's.TURNS-with-mixingIndexErrorabove: it's onmain, and separate.resident_walks().The Rust graph runtime, transport and speculation rows are untouched: the branch changes nothing under
rust/(git diff 53c47cb8c HEAD -- rust/is empty). No new dependencies.Gate per session
0.8b c=1
0.8b c=16
0.8b c=64
0.8b c=256
2b c=16
2b c=64
Merged head against
main7cc4b04Rounds marked * ran beside a CPU test run on the client's socket.
4b c=1, against
main7cc4b044b c=16, against
main7cc4b044b c=64, against
main7cc4b049b c=1, against
main7cc4b049b c=16, against
main7cc4b049b c=64, against
main7cc4b040.8b c=1, against
main7cc4b040.8b c=16, against
main7cc4b040.8b c=64, against
main7cc4b042b c=16, against
main7cc4b042b c=64, against
main7cc4b04