Skip to content

Resident decode planning: decode rows on device tables - #1127

Merged
merceod merged 27 commits into
mainfrom
users/atindra/resident-plan
Nov 3, 2026
Merged

merceod merged 27 commits into
mainfrom
users/atindra/resident-plan

Conversation

@merceod

@merceod merceod commented Nov 2, 2026 •

Copy link
Copy Markdown
Collaborator

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 main 53c47cb, 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 + 1 rows, 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.
  • An offload gives the slot up too, in Engine.offload_request and 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.
  • Each state's owner keeps its table: kv/resident.py a pool's block table and lengths, recurrent/pool.py the GDN slot, position/manager.py the position offset. plan on 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 mixin Resource inherits (resources/resident_hooks.py); bind_resident hands each owner its resource key.
  • Per-step state is keyed by resource key, so a node can hold any number of KV pools (Whisper's decoder has two). A capture records each resource's stream once per bucket (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_len has a stored length per pool, and positions take their offset against, and gather with, the pool their PositionConfig.kv_cache names. CapturedResidentStep holds the store once and each pool's block table, and ResidentBuffers a 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.
  • A pool the walk declares at span 0 (a cross-attention context) is read only. Its row holds the stream's pages and length, and the capture gathers the lengths. There's no write addressing, no in-graph advance and no crossings, so its commit and host advance do nothing. A host step can move it, so for the rows a host step advanced the pool compares its own stream, and rewrites the row whole only if that changed. The mode comes from the declared step at capture, not a setting.
  • Attention reports at capture how it read its pages: off the tables (fmha) or through a host plan. A bucket skips the host plan (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.
  • A row is the narrow columns, then one span per KV pool of min(ceil(max_seq_len / page_size), request_pages) block table cells: the pool's max_seq_len (the model's context unless resources: <kv>: max_seq_len overrides 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.
  • Crossings and refreshed lengths go in one pinned H2D into a staging row (room for a length and a column a row per pool), and the graph's first kernel scatters them. A whole row (new, displaced or stale) is a run per span, each one H2D on the node's stream right before the replay; a run that starts where the last one ended joins it, so the narrow cells and the first block table are one copy and a one-pool node like Qwen3.5 still writes a row in one. So every table write is in step order on one stream, and an abandoned pre-plan leaves the tables alone. The stager is sized at build for the largest flush a capture can issue (B full rows plus a cells row), so it never grows on the GPU thread.
  • Padding rows sit on row 0, which stays zero: SINK_PAGE, stored length 0 (attention length 1), SINK_SLOT.
  • A freed slot comes back after the commit of the last step that read it, since that commit still advances the slot's host copy.
  • submodule_base.py: resident_walks(), empty by default. Qwen3.5 names decode on every config: since Performance 2/7: speech and language models #859 a captured decode step builds its rope in the graph whatever ingraph_decode_rope says, so its preprocess reads nothing the tables fill (only the loop-back tokens, off the sampler).

Config and guards

knobs: resident_planning: auto | off, default auto, env MSTAR_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 through Resource.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 or commit=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 decode and all six of its resources can. Run here on 0.8B, 2B, 4B, 9B on one GPU and 9B TP2. Each logs 256 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 against main with the row check at 0 mismatches over 5,600 resident steps.

On top of #852

#852 is in main and 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:

#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) and test_kv_decode_plan_cache.py run unchanged: they build resources with no planner, and a step without one is the host path they test. test_resident_planning.py gains 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.

main moved again while this was measured (#857, #977, #980, #916), so the head merges it a second time. One conflict, in StepContext, 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 needed resident=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 unpromoted clear_preplan). test_a_refused_admit_drops_the_stage_naming_its_pages fails 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 main 4079d7f at the time. Greedy, prefix cache off and on (both passes): 32/32 at 0.8B and 2B, main against main 32/32; KV pressure within main'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: main against main 21-28/32 over four servers, this branch against main 16-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 against main.

This head merges main twice 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 its ingraph_decode_rope condition.

What stays on the host

state where it is why
decode attention plan nothing on Blackwell: every decode bucket is fmha and reads its rows off the tables. On other GPUs FlashInfer's fa2 decode plan, per step, through #852's cached KV plan and direct planner fa2's split-KV schedule comes from host lengths
streams, lengths, pages, prefix index, eviction host the source of truth; the tables mirror it
admit, declare, per-row commit, sampler as they were 0.06 ms of the plan thread at 0.8B c=1, 0.09 at c=16, 0.16 at c=64
prefill, mixed and eager steps host plan, #852's their rows are marked on commit and written by the next resident step that reads them

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 runs knobs: step_thread_core: auto and sampler_ingraph_scatter: 1, plus kv_plan_cache: 1 and gc_gen0_threshold: 50000; ingraph_decode_rope is no longer set, since main doesn'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: main is 53c47cb. "head" is d279b7a. The pushed head 7d2a6cd adds 254e0fd (an abandoned stage restages its row, capture refuses commit=False, a node with no table builds), none of which a steady step runs, and a merge of main 2394b4d (#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 (resident below) 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 core ResidentPlanner.plan over 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.

cell main head vs main
0.8b c=1 788 [688-811] · 1.19/1.21 · 14 1042 [1021-1068] · 0.91/0.90 · 14 +32.1%
0.8b c=16 9995 [8500-10204] · 1.42/1.48 · 30 11196 [11098-11369] · 1.24/1.29 · 31 +12.0%
0.8b c=64 26405 [25782-26756] · 1.92/2.20 · 42 28308 [28156-28516] · 1.73/2.02 · 48 +7.2%
0.8b c=256 32920 [29575-35478] · 3.07/3.70 · 1016 36079 [35582-36318] · 2.70/3.35 · 933 +9.6%
2b c=16 8767 [7909-9051] · 1.61/1.69 · 33 8869 [8757-9032] · 1.60/1.67 · 32 +1.2%
2b c=64 22827 [21169-23347] · 2.27/2.57 · 47 22894 [20981-23793] · 2.24/2.54 · 54 +0.3%
  • 0.8B two sessions per arm, 2B three (main's second 2B session read 8% under its other two at c=16, so a third round). 0 failed requests in every cell, no recompile-limit line and no traceback in any gate server.
  • Paired by round at 0.8B the head is above main in both: c=1 +40 / +32%, c=16 +23 / +10%, c=64 +9 / +6%, c=256 +20 / +2%. At c=256 main's own two sessions split 30.1k / 35.3k as before, so the size there isn't resolved; the direction is.
  • 2B pairs: c=16 +1 / +10 / 0%, c=64 +2 / 0 / 0%. At 2B the step isn't plan-bound: the plan thread's drop doesn't reach throughput.
  • No request was displaced in any gate server, c=256 included: Qwen3.5's decode order 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:

cell arm step main GPU plan KV plan KV stage attention GDN+linear positions resident sampler declare+admit+pad rest
0.8b c=1 main 1.414 0.450 0.515 0.490 0.080 0.106 0.106 0.066 0.028 0.000 0.007 0.057 0.040
0.8b c=1 head 1.063 0.428 0.515 0.147 0.000 0.000 0.002 0.002 0.002 0.039 0.006 0.058 0.038
0.8b c=16 main 1.695 0.619 0.603 0.544 0.086 0.112 0.108 0.075 0.032 0.000 0.011 0.079 0.041
0.8b c=16 head 1.537 0.742 0.694 0.210 0.000 0.000 0.002 0.003 0.002 0.054 0.013 0.091 0.045
0.8b c=64 main 2.436 1.017 0.873 0.630 0.093 0.112 0.106 0.086 0.040 0.000 0.013 0.139 0.042
0.8b c=64 head 2.335 1.237 0.989 0.281 0.000 0.000 0.002 0.003 0.002 0.061 0.014 0.156 0.044
cell plan thread main -> head (ms) head takes of main's step
0.8b c=1 0.490 -> 0.147 0.343 (70%) 24.2%
0.8b c=16 0.544 -> 0.210 0.334 (61%) 19.7%
0.8b c=64 0.630 -> 0.281 0.348 (55%) 14.3%
  • The resident step takes 0.33-0.35 ms off the plan thread at 0.8B, 55-70% of it, 14-24% of 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_sync doubles at c=16 and c=64 (0.074 -> 0.153, 0.132 -> 0.264 ms): the host now waits on the GPU more.
  • At c=1 the step drops by what the plan thread does (1.41 -> 1.06 ms). At c=16 and c=64 the main and GPU threads' CPU rises by two-thirds to all of it (c=16: main 0.62 -> 0.74, GPU 0.60 -> 0.69 ms), so the step drops less (1.70 -> 1.54, 2.44 -> 2.34 ms). Hypothesis, as before and not checked: the three threads share two CPUs, and the freed time goes to waiting hot.

The merged head against main 7cc4b04. main is 7cc4b04, the main at launch (#1131, #1142, #1154 and #725 past 7ff4aec). The head is eb9500b with that main merged locally (cf466990b). While this ran, main moved 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.

cell main 7cc4b04 head vs main paired by round
4b c=1 397 [392-414] · 2.44/2.46 · 18 393 [391-398] · 2.46/2.47 · 19 -0.8% -1 / -4 / +0
4b c=16 5151 [5133-5191] · 2.78/2.91 · 52 5119 [5029-5163] · 2.79/2.93 · 51 -0.6% -1 / -1 / -1
4b c=64 12278 [12092-12450] · 4.15/4.82 · 84 12244 [11993-12392] · 4.17/4.83 · 93 -0.3% -0 / -0
9b c=1 295 [290-299] · 3.30/3.32 · 21 295 [291-299] · 3.30/3.31 · 21 -0.1% +0 / -0
9b c=16 3492 [3480-3501] · 4.09/4.29 · 75 3495 [3491-3500] · 4.09/4.29 · 76 +0.1% +0 / -0
9b c=64 9087 [9045-9113] · 5.36/6.46 · 131 9080 [9060-9109] · 5.36/6.46 · 120 -0.1% +0 / -0
0.8b c=1 778 [750-814] · 1.21/1.23 · 14 1081 [993-1110] · 0.86/0.87 · 14 +38.9% +34 / +40
0.8b c=16 10086 [9939-10183] · 1.40/1.47 · 27 11706 [11672-11749] · 1.19/1.25 · 28 +16.1% +15 / +17
0.8b c=64 26625 [26099-27098] · 1.88/2.17 · 45 27672 [26826-28554] · 1.80/2.09 · 48 +3.9% +8 / +0
2b c=16 8817 [8584-8953] · 1.59/1.69 · 31 8976 [8870-9026] · 1.58/1.66 · 31 +1.8% +3 / +1
2b c=64 23416 [21492-23574] · 2.23/2.52 · 45 23412 [22244-23888] · 2.19/2.48 · 55 -0.0% -3 / +5
  • 0.8B and 2B rows are the rounds that ran alone: 0.8B rounds 3 and 4, 2B rounds 2 and 3. The earlier head against main 53c47cb 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 above main, by how much isn't resolved.
  • I ran a CPU test suite beside four rounds by mistake: 4B and 9B round 1, 0.8B rounds 1 and 2, 2B round 1. It was on cores 108-127, the client's socket. 4B and 9B read the same in those rounds as in the others. At 0.8B, though, each c=64 session split its two runs about 27.5k / 25.4k in both arms, and the c=64 gain read +0.3% (c=1 +24%, c=16 +18%; 2B +3 / 0%). Those rounds are in the per-session blocks, marked.
  • 4B and 9B on one GPU: every head server logs 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.
  • 9B is level at c=1, 16 and 64. 4B is 0.3-0.8% under 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):

cell arm step event_sync main GPU plan KV plan KV stage attention GDN+linear positions resident sampler declare+admit+pad rest
4b c=16 main 2.992 1.444 2.004 0.635 0.559 0.087 0.113 0.109 0.077 0.034 0.000 0.011 0.082 0.045
4b c=16 head 2.988 1.850 2.376 0.630 0.189 0.000 0.000 0.002 0.002 0.001 0.048 0.011 0.082 0.042
9b c=16 main 4.447 2.596 3.235 0.694 0.635 0.098 0.130 0.122 0.088 0.038 0.000 0.012 0.095 0.052
9b c=16 head 4.458 2.999 3.640 0.761 0.233 0.000 0.000 0.002 0.004 0.002 0.061 0.013 0.100 0.051

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:

cell arm graph execution (median) step (median / mean) idle outside the graph (mean)
4b c=1 main 2430 2455 / 2522 90
4b c=1 head 2438 2463 / 2472 31
4b c=16 main 2747 2776 / 2823 63
4b c=16 head 2760 2789 / 2821 51
9b c=16 main 4061 4092 / 4113 36
9b c=16 head 4070 4100 / 4140 53
  • At 4B and 9B the step is the GPU's. On 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 older main (bb0e0a8: event_sync 0.04 ms, 635 µs idle a step). On this main it isn't.
  • The head takes 0.37 ms (66%) off the plan thread at 4B and 0.40 ms (63%) at 9B, about what it takes at 0.8B. The host then waits that much longer on the GPU (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.
  • The head's decode graph is 8 µs longer at 4B c=1, 13 µs at 4B c=16 and 9 µs at 9B c=16 (0.2-0.5%). A node-level trace at 4B c=1 (GPU 1, one server per arm on one cache) shows what was added: _gather_column_kernel three times (1.4 µs each) and _advance_column_kernel once (2.1 µs), with main's _split_combine_kernel and _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 under nsys is shorter than main's (less idle outside the graph), so -0.8% there is within main'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 main on one Inductor cache per size. Yardstick: main against main on the same cache.

setting 0.8B head vs main 0.8B main vs main 2B head vs main 2B main vs main
prefix cache off 32/32 32/32 32/32 32/32
prefix cache on, first pass 32/32 32/32 32/32 32/32
prefix cache on, same prompts again 32/32 32/32 32/32 32/32
KV pressure: 8 at once, 10 pages, offload on, 21-31 evictions a server 19-22/32 20/32 25-26/32 24/32
32 images, prefix cache off 32/32 32/32

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 main doesn't match itself either; the head against itself reads 20/32 at 0.8B and 23/32 at 2B. Those servers ran with MSTAR_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 stock ARRIVAL order rows run to completion and nothing is displaced (29,600 checked steps, 0 mismatches). A run-local hook gives decode TURNS order instead, with speculation_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 against main'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 with main (-3%, inside its own spread).

Forcing TURNS with walk mixing at its default fails decode rows with an IndexError in prepare_inputs (inputs["text_inputs"][0]), in clumps of four: 34 of 640 at c=64 on main and 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 sets TURNS on 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 main 7cc4b04 merged) against main 7cc4b04, same greedy setup, one Inductor cache per size:

setting 4B head vs main 4B main vs main 9B head vs main 9B main vs main
prefix cache off 32/32 32/32 32/32, 28/32 28/32
prefix cache on, first pass 32/32 32/32 32/32 32/32
prefix cache on, same prompts again 32/32 32/32 32/32 32/32

The 9B prefix-off pair: the head matches the first main server 32/32 and the second 28/32, and the two main servers differ from each other in those same four requests (from token 35 to 126). That's main against itself, not the tables.

  • KV pressure, 4B (8 at once, 10 pages, offload on, 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).
  • KV pressure, 9B, same setup, 40 requests: 40 of 40 ok, 0 mismatches over 700 steps (4,892 rows), 44 offloads and 44 slot releases, 0 assertions.
  • Slot pressure, 4B (decode captures capped at 8, so 16 slots of 2050 cells, 0.3 MiB; c=64, the TURNS hook, 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 against main's 2981 / 2967 (-1.5%), c=8 2912 / 2911 against 2919 / 2916.
  • The pushed head (cf77040: main f1c35ef merged, plus the stub below) against main f1c35ef 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.
  • 27B TP4 needs four GPUs; this box has two.

The per-key change (0af5053) and each merge of main after 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 main run on the same GPU and Inductor cache, WER 1.35% on both. main on the other GPU with a fresh cache differs from that same main in 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.
  • With main merged up to f1c35ef: Attention workspaces grow to what a plan asks for #1142's test_the_engine_seals_between_sizing_and_capture builds its submodule as a SimpleNamespace, and the engine asks every submodule for resident_walks() at warmup, so it failed on the merged head. cf77040 gives the stub one; nothing else changed. CPU Core on cf77040 against main f1c35ef: 10718 passed, 15 failed vs 10616 and 14, the same 14 plus test_eval_compare::test_stopping_compare_stops_the_job_and_its_servers, which fails on main too when run alone. The gated tree (main 7cc4b04 merged, with the stub): 10637 passed, 15 failed vs main's 10537 and 12, the extra three all eval_compare tests that fail alone on main 7cc4b04 too. TP lockstep sim 11 passed on both. Focused modules 737 passed, 76 skipped (main 634 and 70).
  • CPU Core selection (CUDA_VISIBLE_DEVICES=, xdist 12), against main 53c47cb'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, with main 2394b4d 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 in test_resident_delta.py: its engine stub binds Engine methods by name and lacked the new _release_resident_slot, which d279b7a adds. The per-key change 0af5053 against main 2394b4d: 10545 passed, 45 failed vs 10443 and 44, the extra an eval_compare hang-timing test that fails on main too when run alone. After each merge of main (6496f14, then 7ff4aec), against that main: 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.
  • The resident, Host overhead 5/N: cheaper resource planning, admit and commit #852's five, chunked-prefill, admit-failure (three modules), micro-scheduler, offload (two), resident-delta, h2d and runner modules: 704 passed, 64 skipped on eb9500b (main: 601 and 58, without test_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.
  • On a B300: the resident, h2d, plan-cache, offload and Host overhead 5/N: cheaper resource planning, admit and commit #852's modules 359 passed, 7 skipped (XPU, ROCm), 1 failed: test_kv_offload.py::test_oom_still_reports_allocation_failed, which fails on main too, since the refusal is an AdmitRuntimeError there 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.plan and 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

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

arm session tok/s r1, r2 ITL p50/mean ms TTFT p50 ms sibling busy CPUs failed
main s1 688, 790 1.25/1.30 14.5 2.5 0
main s2 787, 811 1.17/1.20 13.9 0.9 0
br s1 1042, 1021 0.91/0.91 13.7 1.6 0
br s2 1068, 1041 0.91/0.89 13.4 0.7 0

0.8b c=16

arm session tok/s r1, r2 ITL p50/mean ms TTFT p50 ms sibling busy CPUs failed
main s1 9827, 8500 1.46/1.62 31.7 4.2 0
main s2 10164, 10204 1.38/1.44 30.4 1.0 0
br s1 11369, 11098 1.24/1.30 28.8 2.8 0
br s2 11245, 11147 1.24/1.29 32.0 1.3 0

0.8b c=64

arm session tok/s r1, r2 ITL p50/mean ms TTFT p50 ms sibling busy CPUs failed
main s1 25782, 26139 1.92/2.24 46.3 3.5 0
main s2 26756, 26670 1.89/2.18 42.1 2.2 0
br s1 28370, 28245 1.71/2.01 52.5 1.7 0
br s2 28516, 28156 1.74/2.03 45.5 1.7 0

0.8b c=256

arm session tok/s r1, r2 ITL p50/mean ms TTFT p50 ms sibling busy CPUs failed
main s1 30719, 29575 3.08/4.02 1095.2 5.4 0
main s2 35120, 35478 2.99/3.43 949.0 1.3 0
br s1 36318, 36102 2.98/3.33 927.5 1.8 0
br s2 35582, 36056 2.70/3.38 939.7 2.0 0

2b c=16

arm session tok/s r1, r2 ITL p50/mean ms TTFT p50 ms sibling busy CPUs failed
main s1 8768, 8800 1.61/1.69 33.7 1.9 0
main s2 7909, 8199 1.75/1.85 33.2 1.5 0
main s3 9051, 8767 1.59/1.67 30.0 2.1 0
br s1 8870, 8873 1.61/1.67 31.7 1.5 0
br s2 8868, 8858 1.61/1.67 31.6 2.7 0
br s3 9032, 8757 1.57/1.65 34.9 1.6 0

2b c=64

arm session tok/s r1, r2 ITL p50/mean ms TTFT p50 ms sibling busy CPUs failed
main s1 22929, 23347 2.24/2.53 44.2 2.0 0
main s2 21169, 22322 2.37/2.70 50.7 1.4 0
main s3 22726, 23005 2.27/2.57 46.0 2.1 0
br s1 23793, 23576 2.19/2.46 52.5 1.8 0
br s2 22390, 20981 2.24/2.70 54.1 3.8 0
br s3 22874, 22914 2.28/2.54 55.3 1.4 0

Merged head against main 7cc4b04

Rounds marked * ran beside a CPU test run on the client's socket.

4b c=1, against main 7cc4b04

arm session tok/s r1, r2 ITL p50/mean ms TTFT p50 ms sibling busy CPUs failed
main s1* 396, 397 2.44/2.46 17.7 1.8 0
main s2 414, 414 2.34/2.35 18.2 1.1 0
main s3 392, 394 2.44/2.48 17.8 3.6 0
br s1* 392, 391 2.46/2.49 18.3 2.1 0
br s2 398, 398 2.42/2.45 17.8 3.7 0
br s3 394, 393 2.46/2.47 19.7 1.8 0

4b c=16, against main 7cc4b04

arm session tok/s r1, r2 ITL p50/mean ms TTFT p50 ms sibling busy CPUs failed
main s1* 5142, 5166 2.77/2.91 51.1 2.8 0
main s2 5141, 5133 2.78/2.92 52.9 2.1 0
main s3 5161, 5191 2.78/2.90 52.4 1.8 0
br s1* 5100, 5138 2.79/2.93 51.9 1.7 0
br s2 5106, 5029 2.81/2.96 50.4 4.6 0
br s3 5131, 5163 2.79/2.91 50.1 2.1 0

4b c=64, against main 7cc4b04

arm session tok/s r1, r2 ITL p50/mean ms TTFT p50 ms sibling busy CPUs failed
main s1* 12447, 12450 4.09/4.75 84.5 3.2 0
main s2 12110, 12092 4.21/4.90 84.1 2.3 0
br s1* 12385, 12392 4.10/4.75 92.6 1.7 0
br s2 12104, 11993 4.23/4.90 92.8 2.7 0

9b c=1, against main 7cc4b04

arm session tok/s r1, r2 ITL p50/mean ms TTFT p50 ms sibling busy CPUs failed
main s1* 298, 299 3.27/3.28 20.1 4.1 0
main s2 290, 292 3.34/3.36 21.9 2.8 0
br s1* 298, 299 3.26/3.27 21.4 2.1 0
br s2 291, 291 3.34/3.36 22.2 2.3 0

9b c=16, against main 7cc4b04

arm session tok/s r1, r2 ITL p50/mean ms TTFT p50 ms sibling busy CPUs failed
main s1* 3480, 3491 4.07/4.30 74.7 3.3 0
main s2 3494, 3501 4.09/4.28 79.5 2.0 0
br s1* 3491, 3499 4.09/4.29 77.5 2.6 0
br s2 3492, 3500 4.09/4.30 74.2 4.7 0

9b c=64, against main 7cc4b04

arm session tok/s r1, r2 ITL p50/mean ms TTFT p50 ms sibling busy CPUs failed
main s1* 9045, 9099 5.35/6.42 133.9 1.8 0
main s2 9076, 9113 5.36/6.47 121.5 2.8 0
br s1* 9080, 9081 5.36/6.45 132.0 1.6 0
br s2 9109, 9060 5.37/6.49 107.0 1.9 0

0.8b c=1, against main 7cc4b04

arm session tok/s r1, r2 ITL p50/mean ms TTFT p50 ms sibling busy CPUs failed
main s1* 785, 712 1.22/1.28 14.2 2.4 0
main s2* 788, 782 1.18/1.22 14.5 1.6 0
main s3 801, 814 1.17/1.18 14.4 1.6 0
main s4 756, 750 1.24/1.27 14.0 2.5 0
br s1* 938, 1001 0.96/0.98 13.7 2.7 0
br s2* 910, 998 0.90/1.00 13.4 4.5 0
br s3 1074, 1088 0.86/0.87 14.3 2.9 0
br s4 993, 1110 0.85/0.90 12.7 4.2 0

0.8b c=16, against main 7cc4b04

arm session tok/s r1, r2 ITL p50/mean ms TTFT p50 ms sibling busy CPUs failed
main s1* 8592, 8175 1.48/1.78 32.1 3.4 0
main s2* 10254, 9813 1.41/1.48 27.3 2.1 0
main s3 10183, 10124 1.40/1.46 26.9 2.5 0
main s4 9939, 10047 1.41/1.49 28.6 1.3 0
br s1* 10629, 11016 1.29/1.36 30.1 2.6 0
br s2* 11151, 10575 1.18/1.35 28.8 5.2 0
br s3 11702, 11672 1.20/1.25 28.7 3.3 0
br s4 11709, 11749 1.19/1.25 27.7 4.4 0

0.8b c=64, against main 7cc4b04

arm session tok/s r1, r2 ITL p50/mean ms TTFT p50 ms sibling busy CPUs failed
main s1* 27150, 25379 1.94/2.23 42.8 1.4 0
main s2* 27377, 26491 1.87/2.16 44.1 2.7 0
main s3 26802, 26099 1.84/2.20 41.0 3.8 0
main s4 26449, 27098 1.93/2.17 48.3 1.7 0
br s1* 27772, 25445 1.82/2.18 51.4 4.0 0
br s2* 26039, 27766 1.77/2.14 52.6 4.2 0
br s3 28554, 28467 1.78/2.02 47.5 2.6 0
br s4 26878, 26826 1.82/2.16 48.4 5.5 0

2b c=16, against main 7cc4b04

arm session tok/s r1, r2 ITL p50/mean ms TTFT p50 ms sibling busy CPUs failed
main s1* 8811, 8721 1.59/1.70 32.4 1.3 0
main s2 8584, 8807 1.63/1.71 32.1 1.6 0
main s3 8953, 8827 1.55/1.68 29.6 9.2 0
br s1* 9061, 9032 1.58/1.64 31.5 2.1 0
br s2 8978, 8870 1.59/1.67 31.0 2.9 0
br s3 9026, 8974 1.57/1.65 30.1 2.8 0

2b c=64, against main 7cc4b04

arm session tok/s r1, r2 ITL p50/mean ms TTFT p50 ms sibling busy CPUs failed
main s1* 23523, 23049 2.22/2.52 49.8 1.4 0
main s2 23574, 23328 2.23/2.51 44.6 2.8 0
main s3 23503, 21492 2.21/2.64 44.3 7.0 0
br s1* 23149, 23275 2.21/2.51 57.0 2.0 0
br s2 23290, 22244 2.21/2.55 55.5 3.0 0
br s3 23534, 23888 2.18/2.46 52.4 2.0 0

merceod and others added 7 commits November 1, 2026 18:36
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.
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.
Picks up #857, #977, #980 and #916. One conflict, in StepContext: the
resident fields and #916's chunk fields both added, kept both.
#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.
@merceod

merceod commented Nov 2, 2026

Copy link
Copy Markdown
Collaborator Author

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 (resources[loopback.sampler].enable_device_loopback_reader()) and then builds the planner, and Qwen3.5 keeps only resident_walks, since the loop-back hooks the branch carried come from #859's DeviceLoopback declaration. One thing for review: resident_walks still names decode only when ingraph_decode_rope is set, which was there because preprocess read positions before the tables were staged. #859 builds decode's rotary tables in the graph on every config, so that condition could go and the tables would be on for every Qwen3.5 config by default. I left it as gated (the 0.8B config sets the knob); say if you'd rather have it on everywhere, the 2B/4B/9B-TP2 checks above ran with the knob on.

@merceod

merceod commented Nov 2, 2026

Copy link
Copy Markdown
Collaborator Author

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.
@merceod

merceod commented Nov 2, 2026

Copy link
Copy Markdown
Collaborator Author

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.

@NSagan271

Copy link
Copy Markdown
Collaborator

@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.

@merceod

merceod commented Nov 2, 2026

Copy link
Copy Markdown
Collaborator Author

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 main (and main's own c=256 sessions split 31.3k / 35.3k), so your +1-3% agrees with it. Its gain is at c=1 and c=16, where the plan thread is the step: +16% and +9% on fa2 (888 vs 763, 9873 vs 9074 tok/s), +40% / +26% with #988's fmha, where a resident bucket takes no host plan at all. Blackwell being more host-bound is part of it, but low concurrency is host-bound on any GPU: at 0.8B c=1 the GPU is busy 0.4 ms of a 1.3 ms step here. So before deciding I'd add c=1 and c=16 to the H200 A/B for all three arms. Mixed steps matter in the comparison too: #916 makes more of them at high concurrency, and both designs pay on those (your 0.38 ms check-in, my whole-row rewrite on the next resident step).

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 resident column in the phase table), against the ~0.1 ms you give for a steady rows step. Dense rows as the state remove the gather, but the gather runs inside the capture (one of the five kernels), so I don't think that's where the time goes; the extra saving is in the sampler rows, which is the state #1127 leaves on the host.

Suggestion. I wouldn't drop #933/#966. I'd land #1127 first as the narrow step: it sits on main with no Rust change, measured at three sizes and c=1-256, with the device rows checked against the streams (0 mismatches over 4539 chunked-prefill steps, 629 under KV pressure with 27 refused admits, 1200 per rank of 9B TP2). Then #966 generalizes it: the three tables (kv/resident.py, recurrent/pool.py, position/manager.py) become row buffers with the host mirrors as truth, the planner's slot store becomes a RowSpace, and the sampler rows come with #966, which is the follow-up #1127's description names. Whichever order, the pieces worth keeping are the fmha wiring (gather_resident / host_plan_free: the host-plan-free bucket is most of the Blackwell gain), the debug check of rows against streams, the refusal list (resident_refusal and the tested refusals), and the stage drop on a refused admit where #980 meets a device table (a same-key retry wrote into a page the row no longer held). If you'd rather the rows land first, the reduced #1127 is those four on top of #966's KV rows, and I'd re-gate it here.

Same harness, both boxes. I'll run main, #932, #966 and #1127 on the B300 with the harness from the description (alternated arms, one Inductor cache per round, two sessions per arm) at 0.8B c=1/16/64/128/256, plus a phase-timing pass for the plan-thread split and the mixed-step count, and post the table here within a day.

@NSagan271

Copy link
Copy Markdown
Collaborator

@merceod That sounds like a plan! I'm still working on refining #966, so I'll lyk when it's ready for testing on B300s.

I'll review #1127 now then (and expect it to land before #966).

@merceod

merceod commented Nov 2, 2026

Copy link
Copy Markdown
Collaborator Author

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 NSagan271 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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).

Comment thread mstar/engine/resources/base.py Outdated
# 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:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread mstar/engine/resources/resident.py Outdated
if ready:
self._cooling = [(slot, seq) for slot, seq in self._cooling if seq > done]
self._free.extend(ready)
if not self._free:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

  1. 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.
  2. 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.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread mstar/engine/resources/kv/manager.py Outdated
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))

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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])

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread mstar/engine/engine.py
if ctx.resident is not None:
ctx.plan_phase = "res.pre_plan"
t0 = phase_start()
ctx.resident.plan(step, ctx)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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] = {}

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

_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?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added: ResidentPlanner.{_slot_of, _stale, _moved} (purged by release). Displacement drops all three for the displaced request too.

Comment thread docs/environment_variables.rst Outdated
- 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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

"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).

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread mstar/utils/h2d.py
return
with self._lock:
if len(values) > self._numel:
self._grow(len(values))

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread test/modular/test_resident_planning.py Outdated
_assert_row(planner, live, 1)


def test_running_out_of_slots_says_how_to_fix_it():

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_failure and _try_offload_cold_request offload 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. Also test_a_follower_replaying_an_offload_gives_the_slot_up_too.
  • The knob's tests are gone, and one checks the knob is unknown now.

@NSagan271

Copy link
Copy Markdown
Collaborator

🤖 Automated live test, AI-generated (Claude Code). Not checked line by line by a human. Partial report: the KV-pressure / slot-exhaustion job and 9B TP2 are still queued, so I'll follow up with those.

Live test: resident decode planning

Works and helps a little on H200. One thing to fix in the description: no shipped config turns the feature on.

H200 · tested at bfcf640 vs merge base a5de953. The head has since moved to 863cf87, which only merges main. · Qwen3.5-0.8B, servers not CPU-pinned

No shipped config enables it

Resident planning covers the decode step only when ingraph_decode_rope is on (mstar/model/qwen3_5/submodules.py:384). #859 removed that from the 0.8B config, so with the shipped configs the planner never engages, and the server logs no resident planning on LLM line. The description's B300 numbers are from a config that still had it.

For everything below I set MSTAR_INGRAPH_DECODE_ROPE=1 on both arms. The PR logs resident planning on LLM: 4096 slots of 2051 cells. If this is meant to ship as a speedup, a config needs to turn ingraph_decode_rope back on, or the planner needs to cover the default path.

Throughput (ABBAAB, 3 restarts per arm), output tok/s

Cell main PR median Δ
t2t c=1 611, 576, 627 667, 652, 618 +6.8%
t2t c=8 3485, 3634, 3497 3589, 3945, 3701 +5.8%
t2t c=32 7988, 8338, 8402 8641, 8691, 8658 +3.8%
t2t c=128 14683, 15151, 14552 14932, 15287, 14652 +1.7%
i2t c=16 5980, 5769, 6252 6607, 6317, 6404 +7.1%

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 MSTAR_RESIDENT_PLANNING=off)

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.

@NSagan271

Copy link
Copy Markdown
Collaborator

@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:

  • Since every column is a host mirror, I think slots could be sized by the batch cap and the least recently used request evicted when they run out, instead of raising.
  • The evicted request just gets a fresh slot and a full rewrite when it next runs, which is the new-request path. The constraints are the existing cooling rule, plus the same victim on every TP rank (LRU by plan sequence number, not wall clock).

@merceod

merceod commented Nov 3, 2026

Copy link
Copy Markdown
Collaborator Author

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 main (6496f14, then 7ff4aec), at eb9500bf6.

What's keyed how.

  • Row layout. A row is the narrow columns (each pool's stored length, the position offset, the GDN slot), then one span per block table, and _seal takes any number of wide tables. rewrite builds a row owner by owner. The narrow cells are one run, and each KV owner adds its used pages as a run of its own (ResidentWrites.run), so writes.rows holds (slot, first column, cells) and flush copies one run each. A run that starts where the last one ended joins it, so Qwen3.5's one pool still writes a row in one H2D, as before.
  • Labels. capture_step records a ResidentStream per resource key (label, and whether the step appends a token a row or only reads), from step.steps[key].segments, or the step's default when a resource names none. It needs one label per key, with spans all 1 or all 0, and refuses two labels on one key. The records are built once per bucket and every step of that bucket plans with them, so nothing per key is built per step. Every owner reads writes.streams[its key].label. bind_resident now hands each owner its key.
  • Lengths. ResidentWrites.kv_len is by pool key, set by each KV owner as it rewrites a row. Positions read kv_len[self._kv_cache_name], the pool their PositionConfig.kv_cache names (the same key is their depends_on, so that pool binds and rewrites first), and gather with that pool's buffer. At capture they also check that their stream and the pool's are one label and both append, since the offset only holds then.
  • The captured step. CapturedResidentStep holds store once and block_tables by pool key, and ResidentBuffers.kv_len is a length buffer per KV pool per (bucket, slot). FlashInferManager.resident_gather gives the fmha wrapper its own pool's table and lengths. host_plan_free stays per bucket: each attention says at capture whether it read its pages off the tables or through a host plan, and a bucket skips the host plan only if none needed one. That matters for Whisper, whose cross attention reads both pools' KV plans. The single-cell scatter moved from the KV owner to the planner, so it runs once per replay whatever the number of pools, and only on a node with a pool.

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: resident_commit and resident_advance do nothing for it. Only a host step can move its stream. So for the rows a host step advanced, stage_step compares the pool's own stream and rewrites the row whole only if that stream changed; otherwise it does nothing. The KV pool takes the mode from the declared step (span 0 on its label for the walk) at capture, so there's no setting. FlashInfer cross attention over a context pool of its own now opts in and stays host-planned. With its context in the query's own pool, it refuses (that would be two labels on one key).

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 WhisperDecoderSubmodule #1143's resident_walks() -> {"decode"}. Large-v3 then logs resident planning on decoder: 128 slots of 19 cells, 0.0 MiB, walks ['decode']: decode captures 1 to 64, 3 narrow cells + 4 self-attention pages + 12 context pages. Every decode capture is host-planned (host_plan_free=False), as it should be. At c=1, all 200 transcripts match main (WER 1.35% on both), with no failed request. A debug hook that syncs after each resident step and compares both pools' device rows, lengths and the self pool's write location with the streams read 0 mismatches over 5,600 steps, with cross_kv_cache read only and kv_cache appending.

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 resident column, two phase-timing sessions per arm, reads 0.042 / 0.047 ms a step against 0.041 / 0.044 at c=1, 0.052 / 0.047 against 0.049 / 0.051 at c=16, and 0.058 / 0.074 against 0.056 / 0.056 at c=64. In the 0.074 session every plan-thread phase ran 20-30% slow, and resident was the same 22% share of the plan thread. Main, GPU and plan thread CPU at c=16 are level between the arms too. In one process per tree, pinned to one core, ResidentPlanner.plan over the same steady 64-row steps takes 43.0 µs median against 43.8 (three alternations: 42.9-45.0 vs 42.5-45.8). The new per-step objects are two empty dicts and one more closure; the records are built once per bucket at capture. So I don't have a mechanism for the c=16 reading. My hypothesis is session noise at this sample size, and I haven't shown it.

Also run on the change: greedy 32/32 against main at 0.8B and 2B with prefix caching off and on (both passes). KV pressure: 0 mismatches over 632 resident steps from the row check, 30 offloads with 30 slot releases, 0 assertions. Qwen3-TTS seeded: 200/200 WAV bytes. ruff check . is clean. CPU Core fails main's ids; the one extra in-suite failure, an eval_compare hang-timing test, fails on main too when run alone. The TP sim passes 11/11. test_resident_planning.py has 21 new tests for the two-pool cases, and each of 15 parts of the fix, undone on its own, fails a named one. Four existing tests build the planner's internal records by hand (2-tuple rows passed to flush, the old row/stage_step signatures, the kernel test's namespace), so their stubs follow the new records; their assertions didn't change.

After each merge of main (6496f14, #1130; 7ff4aec, #969 and #697), both clean, I reran against that main: greedy 32/32 at 0.8B and 2B (prefix off), KV pressure with the row check at 0 mismatches (27-29 offloads, as many slot releases, 0 assertions), Whisper 200/200 with 0 mismatches over 5,600 steps on both pools, CPU Core with the same failing ids as main (14, then 13), the TP sim 11/11 and the focused modules. The gate wasn't rerun on the merges.

merceod and others added 3 commits November 3, 2026 02:11
The engine asks each submodule for its resident walks at warmup, before
sizing; #1142's test builds its submodule as a SimpleNamespace.
@merceod
merceod merged commit 8741219 into main Nov 3, 2026
6 checks passed
@merceod
merceod deleted the users/atindra/resident-plan branch November 3, 2026 11:00
merceod added a commit that referenced this pull request Nov 3, 2026
#1071 landed as d74de4f, so this sits on main now. One conflict, in
Engine.warmup: the node's capture stream (_capture_stream_for) and #1127's
runner binding for _build_resident_planner both stay.
@merceod

merceod commented Nov 3, 2026

Copy link
Copy Markdown
Collaborator Author

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, main 7cc4b04 against eb9500b with that main merged, arms alternated, 0 failed in every cell. The description has the tables, and the per-session blocks are under "Gate per session".

  • 4B and 9B on one GPU. The planner turns on (256 slots of 2050 cells, 4.0 MiB), and every decode bucket, bs 1 to 128, is host-plan-free. 9B is level at c=1/16/64. 4B is -0.8 / -0.6 / -0.3%, with c=16 at -1% in each of three rounds. Both sizes are GPU-bound on today's main: at c=16, event_sync is 1.44 of a 2.99 ms step at 4B and 2.60 of 4.45 at 9B, and the GPU idles 63 / 36 µs a step outside the graph, not llm_kernels1's 635. So the 0.37-0.40 ms the head takes off the plan thread becomes GPU wait. Under nsys the head's decode graph is also 8-13 µs longer. A node trace shows two small kernels added (_gather_column_kernel x3, _advance_column_kernel) and two of main's gone; the kernels sum the same, and the extra launches cost about 5 µs.
  • 0.8B and 2B on the merged head, from the rounds that ran alone: 0.8B +39 / +16 / +4% at c=1/16/64 (c=64 +8 / 0% by round, not resolved), 2B +2 / 0%. I ran a CPU suite on the client's socket beside the first rounds by mistake; it made 0.8B c=64 runs bimodal in both arms. Those rounds are marked in the blocks.
  • Correctness at 4B and 9B.
    • Greedy 32/32 against main with prefix caching off and on (both passes). At 9B with prefix off, main matches itself 28/32; the head matched one of those two servers 32/32.
    • KV pressure with the row check: 0 mismatches at both sizes (600 and 700 steps), offloads equal slot releases (29, 44), 0 assertions.
    • Slot pressure at 4B: 195k rows displaced, 0 mismatches over 30k steps, c=64 -1.5% against main with every row rewritten.
  • Found and fixed. With main merged, Attention workspaces grow to what a plan asks for #1142's test_the_engine_seals_between_sizing_and_capture fails: its stub submodule has no resident_walks(), which warmup now calls. That includes the merge of f1c35ef (8922682). cf77040 adds it on top of that merge. On cf77040 against f1c35ef:
    • CPU Core shows main's failing ids plus one eval_compare test that also fails alone on main.
    • The TP sim passes 11/11, and the focused modules 737 passed.
    • At 4B, greedy is 32/32 and the KV row check reads 0 mismatches.
  • Not done. 27B TP4 needs four GPUs. The f1c35ef merge is checked, not gated. Its Qwen3.5 changes are the prefill ladder, vision buckets and the GC knob, nothing on decode.

merceod pushed a commit that referenced this pull request Nov 3, 2026
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.
merceod pushed a commit that referenced this pull request Nov 3, 2026
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
merceod pushed a commit that referenced this pull request Nov 3, 2026
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.
merceod added a commit that referenced this pull request Nov 3, 2026
Brings in #1071 and #1127. The Talker takes #1071's Qwen3Attention layers
and declared rope_theta. The depth loop stays on cp_attend and cp_linear,
so #1071's two-token first pass and its Blackwell gate are dropped, and the
codec stays incremental, so conv_stack.py and its channels-last stack go.
merceod pushed a commit that referenced this pull request Nov 3, 2026
…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>
@merceod

merceod commented Nov 3, 2026

Copy link
Copy Markdown
Collaborator Author

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 event_sync rather than tok/s on this box, and c=1/16 read +3/+5%. The RNG offset, last token, mask and history stay keyed by sampler row, with the reasons in the PR. Greedy 32/32 and seeded bytes 32/32 against main, the row checks at 0 under KV and slot pressure, and TP2 ranks identical step by step. Hopper's in-graph fa2 page indices are still open.

arun-elr added a commit to arun-elr/mstar that referenced this pull request Nov 3, 2026
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.
vasilevklart added a commit that referenced this pull request Nov 3, 2026
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.
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.

2 participants