[arun-elr]

Fixes #445

What does this PR do?

Found while auditing the opt-in KV host optimizations from #396 (both off by default). Either one, turned on, breaks plans with more than one label:

• Index ring (MSTAR_KV_PINNED_INDPTRS=1). A 3-label plan overwrote its first label's arrays and KVManager.plan() raised IndexError. The ring now keeps one buffer set per plan.
• Plan cache (MSTAR_KV_PLAN_CACHE=1|2). A cached decode step counted query rows across labels (main + cfg_img gave [0,1,2] per label). It now counts each label's own rows.

Single-label plans are byte-identical, with no timing change. A small padding guard and a docstring fix are separable (see "What changed"). #438 keeps BAGEL image_gen on one combined plan label, so it reaches neither bug.

Schematic of PinnedIndexRing and the plan cache's per-label qo_indptr; both paths are opt-in and off by default.

How was it tested?

On ffc8a13 (head), CI's install, GPU hidden: cpu-core.txt 1308 passed, 55 skipped; the two KV test files 26 passed; ruff check . clean. 15 of the 23 tests in the new test/modular/test_kv_multi_label_plans.py fail on main. With the GPU visible they pass too. On an H200, BAGEL think-mode decode logs the plan-cache mismatch on main and not here.

Not tested

CUDA async-copy lifetime of the ring buffers (every consumer's copy issued before its set is reused, across preplan, promotion, discard and inline plans) was checked by source tracing and with mocked events, not by observing real async copies.

What changed

Index ring (MSTAR_KV_PINNED_INDPTRS=1). Each plan label takes 4 of the ring's 8 host buffers, so a 3-label plan overwrote its first label's arrays and KVManager.plan() raised IndexError. The ring now keeps one buffer set per plan, so every array lives until the plan after next, as single-label plans already did.

Plan cache (MSTAR_KV_PLAN_CACHE=1|2). A cached decode step counted query rows across labels (main + cfg_img gave [0,1,2] per label). It now counts each label's own rows. On an H200, BAGEL think-mode decode logs the mismatch on main and not here.

• kv/plan.py: PinnedIndexRing holds depth=2 buffer sets, one per plan. begin_plan() records the finished plan's event on its stream and waits on the event of the set it reuses; take() fills the next buffer of the current set, growing the set past four takes if needed. One event record per plan instead of per take, so a single-label plan does no more work than before.
• kv/manager.py: plan() calls begin_plan() after the preplan promotion return (a promoted plan reuses its staged arrays). _fast_decode_plan uses len(rows.rows) + n_pad.
• Separable padding guard in _derive_rows: a hit now requires each label's padding rows to be the step's padding rows, row for row (request, label, span), not just as many. In a combined label, a later source label's real rows can follow an earlier one's padding, so a count alone let real rows pass as padding. A hit rebuilds each label's padding from all of the step's padding rows, so padding under only some labels gave the other labels extra views (verification mode flagged it; no shipped model lays padding out that way). One condition; easy to drop if you prefer the strict scope.
• Plan cache hit check (_fast_decode_plan): a step is also planned from the streams when its combined labels differ from the cached step's, or when its padding rows change label. The cached views place each label's rows for one grouping and one padding order, so either change gave a combined label the wrong rows in mode 1 (mode 2 caught it).
• mstar/utils/knobs.py and docs/environment_variables.rst: the two knobs' docstrings said "default 1" (the code says 0), and the ring's cost now reads one CUDA event per plan instead of event bookkeeping per take. Docs only.

Tests and evidence

• New test/modular/test_kv_multi_label_plans.py (in cpu-core.txt, about 0.3 s): the two audit regressions; cached against full plans for 1, 2 and 3 labels and a mixed set, each with the ring on and off; padded multi-label steps under verification mode; padding under some labels only; one plan's arrays intact while the next is built. 15 of its 23 tests fail on main, including combined labels with padding, padding that changes label between steps, and a step that starts combining labels.
• test_kv_pinned_index_ring.py updated to per-plan rotation.
• On ffc8a13 (the head), CI's install, GPU hidden: cpu-core.txt 1308 passed, 55 skipped; the two KV test files 26 passed; ruff check . clean.
• With the GPU visible (pinned buffers, real CUDA events), the new tests and the ring and plan-cache tests pass.
• Single-label plans hash the same on main and here (cache modes 0, 1 and 2, ring on and off, with and without padding); plan() timing at 32 rows is unchanged within noise.
• BAGEL on one H200 (bagel_single_gpu.yaml, MSTAR_KV_PLAN_CACHE=2 MSTAR_KV_PINNED_INDPTRS=1), three think-mode text-to-image requests (one alone, two concurrent): main logs "the plan cache disagrees with the full plan for 2 segments"; this branch logs none and has no tracebacks, and the images hash the same on both (verification mode serves the full plan).

Checklist

• ruff check . passes
• Added or updated tests / docs where relevant

AI assistance: written with AI coding agent

[arun-elr]

/review

[github-actions]

Automated review

This PR changes PinnedIndexRing to keep one buffer set per plan instead of one buffer per take. KVManager.plan() now calls begin_plan() after the promotion return, so a plan with three labels no longer overwrites its own arrays. The decode plan cache also changes. Each label now builds qo_indptr from its own rows. A cache hit now requires the same combined_labels, padding rows whose labels don't change, and each label's padding to match the step's padding row for row. The code matches the description.

Findings: none.

I traced the ring's lifetime through plan(): promotion, inline plans after a discarded preplan, and mode 2, which takes from the ring for both the cached and the full plan. A set is reused two plans later, as single-label plans already were on main. The event is recorded at the same point as before, so the async-copy window is no worse than on main. In mode 2 it is better, because main's per-take ring wrapped after one plan there. The new zip(..., strict=True) in _fast_decode_plan can't raise: the length check earlier in the same or chain short-circuits first, and cache.step_segments always has n_real + n_pad entries. Steps with several separate labels and padding still never hit the cache. That is unchanged from main, because _split_padding already rejects label-major padded layouts. The new padding check only makes the cache more conservative for the synthetic layouts in the tests.

On the description: the claim that plan() timing at 32 rows is "unchanged within noise" doesn't say which harness, hardware or command produced it. AGENTS.md invariant 10 asks for those on a change to attention planning, even an opt-in one. One sentence would cover it.

---

_{Advisory, not a merge gate. Rules: AGENTS.md. Reply to discuss.}

[NSagan271]

🤖 Automated live test, AI-generated (Claude Code). Not checked line by line by a human. Numbers come from one allocation, so treat small deltas as provisional.

Live test: [KV cache] Fix multi-label plans under the host index ring and plan cache

Works as far as served traffic reaches it: with MSTAR_KV_PLAN_CACHE=2 MSTAR_KV_PINNED_INDPTRS=1, BAGEL think-mode decode logs the plan-cache disagreement on main and not on the PR, with identical images. The 3-label index-ring crash isn't reachable from a shipped config, so it rests on the PR's CPU tests (fail on main, pass here). Default-path Qwen3.5 throughput unchanged.

1x H200, bagel_single_gpu.yaml, Qwen3.5 0.8B · PR head 3c03bedb vs main merge-base 641162ab

Claims

Claim	Result
Plan cache counted query rows across labels; BAGEL think-mode logs the mismatch on main, not here	✅ MSTAR_KV_PLAN_CACHE=2 (verify mode) + ring on, 6 BAGEL text-to-image requests (4 think-mode, one alone then concurrent groups, plus 2 non-think). Main: the plan cache disagrees with the full plan for 2 segments once; PR: 0. No tracebacks either arm
Outputs unchanged	✅ all 6 images hash-identical across main/PR and across cache modes 1 and 2 (mode 1 serves the cached plan; it did not change an image here)
3-label plan overwrote the ring's buffers → IndexError	➖ not exercised served (no shipped config builds a 3-label plan: #438 keeps BAGEL image_gen on one combined label). CPU tests cover it (below)
New and updated tests fail on main, pass here	✅ test_kv_multi_label_plans.py + updated test_kv_pinned_index_ring.py: main 18 failed / 8 passed, PR 26 passed
Single-label plans unchanged	✅ Qwen3.5 A/B with default knobs (table)

Performance

Qwen3.5 0.8B, benchmark.runner text_to_text, closed loop, 5×c requests, ABBA, one server per row, default knobs (both opt-in paths off); text tok/s.

Arm	c=8 / 32 / 128
main r1	2877 / 7310 / 13504
main r4	2932 / 7419 / 14358
PR r2	2948 / 7351 / 13921
PR r3	2983 / 7420 / 13943

No failed requests; inside main's restart spread. Host load 22–35.

Commands

MSTAR_KV_PLAN_CACHE={2,1} MSTAR_KV_PINNED_INDPTRS=1 python -m mstar.cli.main serve bagel --config configs/bagel_single_gpu.yaml; python bagel_client.py (/generate, output_modalities=image, model_kwargs {think_mode, seed 100+i}). Qwen3.5: python -m benchmark.runner --model qwen3.5 --model-id Qwen/Qwen3.5-0.8B --request-type text_to_text --ignore-eos --output-len-min 128 --output-len-max 384 --output-len-seed 0 --profiling-type closed_loop --inference-system ours --num-requests $((5*c)) --num-warmup 2 --max-concurrency c

Other

• Not tested: Qwen3.5 perf with the knobs on (the A/B uses the default path); BAGEL CFG-parallel (3 GPUs); real async-copy lifetime of the ring buffers beyond what the served runs exercise.
• Media: media/bagel.cache{1,2}.{base,pr}/ (PNG per request).

[merceod]

Ran 3c03bed (and 446m = main 84fd4a6 + this PR, since the head sits 23 commits behind main) against main on a klone H200 (clock-limited node, 1560 MHz under load).

• BAGEL think mode, MSTAR_KV_PLAN_CACHE=2 with the ring on, seeds 100 to 102, one request alone then two at once. Main logs "plan cache disagrees with the full plan for 2 segments" once, the head and 446m log none, 0 IndexError and 0 tracebacks in all 6 launches, and the images hash the same on every arm and in modes 0, 1 and 2 (mode 1 serves the cached plan). Thinking text identical between main and 446m. Same as the live test.
• Qwen3.5-0.8B at c=32 in a quiet block, 3 launches per cell. Knobs on, 446m 3019 out tok/s [2955, 3395] vs main 3096 [2913, 3119], knobs off 2814 vs 3083, ring alone and cache alone on main 3270 and 3233, all inside the launch spread, so the knobs cost nothing measurable on either tree and the PR changes nothing there. Mode 2 on Qwen3.5, 0 mismatch lines on both. Greedy outputs agree within the launch-to-launch floor of this stack (54 to 72 of 112 across every pair, the same tree and knobs pair agrees on 60).
• An earlier, busier block had main with the knobs on at 2009 vs 2534 tok/s with two "plan_executor timed out waiting for batch N commit" warnings. That turned out to be a one-time first big batched prefill that passed 10 s in those two launches (19.4 and 11.3 s), not the knobs. 0 warnings in the other 43 Qwen3.5 launches, knobs on or off.
• 26/26 new tests, 18 fail on main (15 multi-label, 3 ring). cpu-core failing set equals main's, ruff clean.
• Reading agrees. Every consumer of a plan's ring buffers runs inside the same plan call on the stream begin_plan recorded on, preplans run under the plan stream, and promotion issues no copies, so a set is reused two plans later with its event waited on. n_rows counts each label's own rows plus padding, the strict zip cannot raise, and no shipped model declares three separate labels in one KV step, so the IndexError stays a CPU test, as you say.

One nit, inline. The kv_pinned_indptrs docstring still says the cost is event bookkeeping per take, it is one event per plan now.

Approving.

[merceod]

The text below still says the cost is "event bookkeeping per take". With this PR the ring records one event per plan, so the per-take wording is stale.

[arun-elr]

Fixed: the docstring and its docs/environment_variables.rst row now say the cost is one CUDA event per plan, which a buffer set waits on before it is rewritten.

[merceod]

Thanks, reads right. Checked ffc8a13 against 3c03bed with each diffed against the same base 641162a, only knobs.py and the docs row changed, the kv/plan.py, kv/manager.py and test hunks are byte-identical, so the approval stands.

[arun-elr]

Thanks for the runs.

The kv_pinned_indptrs docstring still says the cost is event bookkeeping per take, it is one event per plan now.

Fixed in mstar/utils/knobs.py and the matching docs/environment_variables.rst row. That docs-only commit is the only change since the approval; CPU core on the new head: 1308 passed, 55 skipped.