[NSagan271]

A mixed step, one that holds decode rows beside prompt rows, runs a split walk on two kernels per layer: decode for the one-token rows and prefill for the rest (CombinedSplit). On Qwen3.5 such a step cost more than a decode step followed by a prefill step, so running decode rows in prompt steps (speculation_mixes_walks) lost throughput. This PR removes three costs in that path. #916 (chunked prefill) will stack on it.

What changed

The GDN decode half is sized to its bucket (GDNManager._get_wrapper, mstar/engine/resources/linear_attn/gdn.py). A captured split wrapper took the largest capture's row count, the floor the other wrappers use for their buffers. So every mixed bucket ran the GDN decode kernel on 128 rows, most of them on the null slot. The split's decode half now takes the bucket's rows. The dropped rows only ever addressed the null slot, so every real row's output and state is bit-identical (test_the_decode_halfs_width_changes_no_real_rows_result).

The engine orders a split walk's rows decode first, and the wrappers read them in place.

• GDNMixedWrapper and FlashInferMixedWrapper used to gather every token into prompt-first order and scatter the outputs back, per layer: index_select / index_copy_ over the whole step.
• Now a resource declares its splits through NodeResourceSpec.combined_splits(). The engine records each split's prefill walks at load, and refuses two resources that disagree.
• Engine.order_split_rows puts those walks' rows last. It is called once, in Worker._make_executing_batch, which every path to a step goes through, including the speculative one and TP followers.
• The decode half then reads the leading run of one-token rows as the step's first tokens. A bs-row torch.where writes its outputs over the prefill half's.
• Under capture, the prefill half takes every token, with the run as one null-slot or sink row ahead of the prompt rows. An eager GDN step's prefill half takes only the tokens after the run, so its prefix-cache checkpoints cover only the prompt rows, which is what the eager checkpoint buffer is sized for.
• Order only bounds capacity: a decode row behind a prompt row runs on the prefill half. That is still correct while the step holds no more prompt rows than the bucket was captured for.

Padding rows no longer use up FlashInfer's decode split budget (FlashInferMixedWrapper, mstar/engine/resources/attn/wrappers.py).

• FlashInfer's tensor-core decode splits long-KV rows across CTAs within a budget of 2 CTAs per SM divided by the KV heads (66 rows on an H200 with 4 KV heads). Every planned row costs at least one CTA, padding rows included.
• A 64-row mixed bucket holding about 46 decode rows therefore had nothing left to split its long-KV rows with.
• The decode half now runs the same FA2 paged kernel, planned through the prefill wrapper, where a row with no query contributes no work items (FlashInfer's scheduler.cuh). The run's rows take one token each and the last row carries the rest.
• Micro-benchmark, per layer, 46 decode rows with 15% at 4k–16k KV: 328 µs → 159 µs in a 64-row bucket. The same rows in a 48-row decode-only bucket take 135 µs.
• The existing decode_tensor_cores: false option keeps the CUDA-core decode wrapper.

No new config options or environment variables. speculation_mixes_walks stays off by default: with these changes, mixing ties on Qwen3.5-4B and still loses on 0.8B (below).

Performance

H200, one GPU per server, MSTAR_SHM_ARENA=1, each tree's default Qwen3.5 config (prefix cache on). Closed loop with --ignore-eos. Two rounds per arm, interleaved, one server per arm and round. "mixed" is 400 prompts of mixed length with 128–256 output tokens; "long" is 200 long prompts with 64–128 output tokens. Cells are tok/s, mean of the two rounds with the range in brackets.

Workload	main, mixing on	this PR, mixing on	main, mixing off	this PR, mixing off
4B mixed c48	3384 (3371–3398)	3558 (3545–3570)	3558 (3555–3561)	3536 (3469–3603)
4B long c32	817 (811–823)	861 (861–862)	859 (851–867)	866 (859–873)
0.8B mixed c48	6721 (6714–6728)	6983 (6775–7192)	7377 (7189–7566)	7275 (7051–7500)

• Mixing on: +5.1% on 4B mixed, +5.4% on 4B long and +3.9% on 0.8B mixed over main. On 4B mixed, ITL p95 drops from 18.5–18.6 to 15.7 ms; with mixing off it is 18.6–21.8 ms on both trees.
• Mixing off: unchanged, within run-to-run noise. The row ordering runs only for combined steps, and the other two changes are inside the split wrappers.
• Per step, a mixed step is now cheaper than decode plus prefill. On 4B, a trace priced each mixed step against a decode step and a prefill step taken from a mixing-off run: mixed steps came out 1.2–2.1 ms cheaper. Under nsys, a 64-decode, 256-prompt-token step is 12.0 ms of kernel time against 13.5 ms. These were measured before [Prefix Cache] Qwen3.5 Prefix Caching #877, with the prefix cache off.
• Why that barely shows in throughput (hypothesis, not checked): with mixing on, the remaining decode steps carry fewer rows (40.5 against 43.0 on average on 4B), which costs about what the mixed steps save. That is a scheduling effect, not a kernel one.

Tried and dropped: a denser mixed capture ladder (32/48/96 rows) measured +0.2% for 31 s more startup; CUDA-core decode for the split's decode half was slower than the FA2 path; FA3 for the decode half took about 100 µs per layer against 34 µs on FA2 without long rows.

Correctness

• Tolerance tests, run on CUDA:

  • the mixed wrappers against the paths they replace: packed prefill for attention, the chunked GDN path for GDN;
  • eager and bucket mode;
  • decode-first, interleaved and no-decode layouts, one-token prompt rows and padding rows;
  • a CUDA graph captured on one layout and replayed on another;
  • tensor-core and CUDA-core decode, window_left -1 and 7.

• Prefix-cache snapshots in a mixed step: test_a_mixed_steps_prompt_rows_save_the_snapshots_the_chunked_path_saves (eager and bucket) checks that a prompt row's snapshot is the conv window and state the chunked path saves, with more decode rows than a checkpoint interval.

• Ordering: test/modular/test_split_row_order.py covers prompt rows going last in order, other walks and nodes keeping their order, disagreeing resources being refused, and the worker making a mixed step decode first.

• Live greedy check: Qwen3.5-4B, 32 prompts, temperature 0, prefix cache on, staggered so prompt rows join running decode steps.

  Comparison	Identical outputs
  Same server, sent twice (yardstick)	22–25 / 32
  This PR vs main, mixing on	26 / 32
  This PR vs main, mixing off	24 / 32
  Mixing on vs off, this PR	15 / 32
  Mixing on vs off, main	17 / 32

  The same prompt diverges at its first token between mixing on and off on both trees. No errors in any server log.

Tests

• ruff check .: clean.
• CPU Core (python -m pytest test/modular test/distributed test/chatterbox test/kokoro test/omnivoice test/sampling_test mstar/model/*/tests benchmark/worker_phases, OMP_NUM_THREADS=2): 49 failed, 8511 passed, 946 skipped. On main: 50 failed, 8504 passed, 927 skipped. The branch's failures are the same set as main's, and the one failure only on main is test_eval_compare_multigpu.py::test_a_head_that_hangs_is_cut_at_the_cap_with_stacks_and_fails.
• GPU (H200, torch 2.12.1+cu132, FlashInfer 0.7.2): test_gdn_mixed_split, test_attn_mixed_split, test_attn_decode_tensor_cores, test_combined_split, test_split_row_order, test_qwen3_5_mixed, test_gdn_prefix_snapshots, test_recurrent_prefix_snapshots, test_gdn_cuda_graph, test_micro_scheduler, test_admit_failure_handling, test_accelerator_graph_capture, test_resource_runner, test_ragged_attention_cpu, test_ragged_attention: 365 passed.

Follow-ups

• Why decode steps shrink with mixing on. This is the remaining lever for making mixing beat decode then prefill end to end, and it is on the scheduling side.
• 0.8B. The remaining gap there is probably the host cost of planning a mixed step inline on the GPU thread (hypothesis).

Developed with AI assistance (Claude).

🤖 Generated with Claude Code

[mstar-review]

Automated review

🟡 No findings on 19ee63e, light review

Description

• The end-to-end throughput table doesn't say which harness produced it (benchmark.runner, run_benchmark.sh, or a custom script) or give the commands and flags, so a maintainer can't rerun the 'mixed' and 'long' workloads (AGENTS.md invariant 10).
• The per-layer FlashInfer micro-benchmark (328 µs → 159 µs) and the nsys per-step comparison don't name the script or harness that produced them.

What the review checked

• Worker._make_executing_batch is the only place an ExecutingBatch is built (main, speculative and TP-follower paths). order_split_rows is a stable partition, so it gives the same result when run again and returns the same order on every rank for the leader's rids and walks (invariant 11).
• Later request_ids reassignments (speculation drops, drop_rids, batch_N filtering) only filter the list, so decode-first order survives. read_check_rids and the RequestWalks broadcast are taken from the reordered list.
• Speculation splice (_thread_outputs_to_speculative), loopback tracking (_record_loopback) and the inline-emit and KV row lookups are keyed by rid, not by row position, so reordering rows doesn't misroute outputs.
• Qwen3.5's mixed walk is {prefill_text, decode} and its split's prefill_walks is (prefill_text,). After ordering, only prefill_text rows follow the decode run, so the max_prefill_rows checks in both wrappers stay within the max_rows_in_combined cap of 4.
• GDNMixedWrapper eager edge cases are handled: no decode run (decode half skipped in run/run_conv), all rows decoding (chunked half skipped, _merge returns decode[:run]) and the mixed case (concatenation in token order).
• Captured GDN checkpoint writes for the lead null-slot sequence stay within checkpoint_rows(cg_max_tokens, cg_max_bs, every), because the total token count is unchanged. The eager checkpoint buffer, sized from per-row spans, now matches a chunked half that skips the run.
• Sizing the GDN split's decode half to bucket.bs is consistent: the wrapper is keyed per bucket and slot, its width is max(bucket bs, rows) and _buf raises on any growth under capture.
• FlashInferMixedWrapper tensor-core decode plan: d_qo gives each run row one token and puts the remainder on the last row, so qo_indptr ends at rows = num_rows, matching q[:n] in run(). The prefill half's qo_indptr (lead run row, prompt rows, padding, tail) still sums to the bucket's token count.
• GDNMixedWrapper.takes_strided_inputs stays False, so callers pass contiguous tensors and the q[f:] / x[front:] slices along dim 0 stay contiguous.
• CPU overhead: order_split_rows costs one dict lookup per batch, plus two list comprehensions only when a step mixes walks (invariant 6).
• The docs/adding_models.rst update describes the new CombinedSplit ordering contract (invariant 9).

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

[dghosef]

LGTM - nice speedup. Tested on H200s: +7.6-8.3% on 0.8B/4B and +3.9-5.8% on 9B TP2 with mixing on, mixing off unchanged, TP2 outputs match main. The moved planner costs ~0.055ms per mixed step but the GDN plan gets cheaper by about the same. Pushed some comment/doc fixes (docstrings that still described the pre-reorder version, combined_splits() in adding_models.rst) and a sink-off GPU test. Non-blocking: worth validating CombinedSplit against get_combined_graph_walks() at load since a missed multi-token walk now fails the step, and the "mixing off: unchanged" reason in the description isn't quite right - steps still mix with it off

sorry claude randomly approved this

[dghosef]

ok yea this is good, just merge and then gtg