Performance 1/7: shared engine, conductor and component changes by dghosef · Pull Request #858 · mstar-project/mstar

dghosef · 2026-10-24T05:41:19Z

This is the shared slice of the perf stack: the engine, conductor, worker and component changes that the model slices (#859, #860) build on, and a few that pay on their own. Data-parallel requests went to a random replica, warmup reached one replica at batch size 1, and several pieces had a copy per model: the AR stop check, pinned host-to-device staging, head-shard sizing in every resource builder. This PR routes each request to the least-loaded replica, warms every replica, gives AR submodules one stop path (TokenStop), one pinned H2D helper and one head-shard count, and adds the engine pieces the model slices need: a one-row-per-step piecewise lease, ragged attention as a custom op, a softmax top-k MoE router kernel, a capturable Qwen3 text encoder and threaded request video decode. Against main, least-loaded routing gains up to +45% (Wan2.2, 2 replicas, c=4) and threaded decode +28% (V-JEPA2, c=4); on Qwen3.5 the stop-path changes are neutral.

Routing and warmup. The conductor sends a request to the replica of each node group that serves the fewest live requests (ties at random). Load is counted per (node group, replica), so groups that share a worker don't skew each other's picks. WorkerGraph.replica_ranks is the one place a worker graph's replicas come from, read by the conductor (its instance table included), warmup, the parallel config and placement validation. Each warmup entry is sent once per replica, as concurrent copies. When node groups have different replica counts, the entry first runs once alone, so a group with fewer replicas still compiles its batch-1 shape. Ids are warmup-<spec>.<copy> and warmup-<spec>.solo; a failing copy lets its siblings finish, then the first failure is re-raised. Spreading is best effort: copies over the admission cap queue and can land on a replica that is already warm. Stop path. An AR submodule sets self.token_stop = TokenStop(...) in __init__ and gets check_stop and check_stop_batched from one host copy of the step's tokens; one per-row verdict (_token_stops) holds the walk test, and HostRows.first_values is a single row-map path. Whisper, Canary, Canary-Qwen, Qwen3-ASR (main's TokenStepNode) and Qwen3.5's decode go through it, Qwen3.5's own copy is gone, and the engine no longer passes step_walk. TokenStop.prefill_walks covers a token sampled at prefill, before the loop exists: a stop id there, or a one-token budget, returns main's END_PARTITION (#794), so the conductor ends the partition instead of starting decode. last_prefill_flag names the step-metadata flag of the pass whose token is the reply's, for Qwen3.5's prompt prefilled over several passes. Qwen3.5, Qwen3-ASR, Canary, Canary-Qwen, GLM-5.2 and Command A+ declare their prefill walks, so #794's per-model prefill checks are gone again for the one declaration; Whisper declares none, since an align walk can follow its decode. The token budget is the loop cap alone (get_loop_caps, max(1, max_output_tokens - 1) since #794): unlike main's stop checks, TokenStop doesn't also count max_tokens. Head sharding. EngineResourceInfo.head_shards is how many ways a node's ranks split a resource's heads: all tp × sp ranks of its lockstep group (1 without one). Every builder (KV, cross, ragged, linear attention, recurrent state) sizes by it rather than reading the group itself. The one per-resource exception is main's RaggedAttentionConfig(shard_over_sp=False), which splits a ragged attention's heads over TP only, decided in _build_ragged alone; Wan2.2 in #860 uses it to keep every head on each SP rank. Engine pieces for the model slices. one_row_per_step lets a packed lease_before_step region (a vision tower) lease, check and replay a whole step's requests as one captured row; the runner applies it in its own bucket lookup, the step's segments all count against max_segments_per_request, and it is refused at config time where it means nothing (#859's Qwen3.5 vision and Omnilingual ASR use it). With RaggedAttentionConfig(compile_op=True), the resource's run traces as mstar::ragged_attention, so a compiled region keeps one graph across attention; the op body calls the manager's _run, and the op's manager registry holds managers weakly. It is off by default because Z-Image with the op took +2.6 GiB of capture memory and a slower startup, and only Wan2.2 (#860) turns it on. RaggedAttentionCallable is gone: callers bind the label on resource.run. TopKRouter(fused_softmax_topk=True) routes a softmax router through the Triton top-k kernel, passed to the MoE block through router=; ties go to the lower expert id and bf16 logits tie often, so a model opts in after counting divergences (Qwen3-Omni in #859). Tuned tile tables cover Qwen3-Omni's two expert shapes on H200. The Qwen3 text encoder keeps one rotary table per device at the longest length, sliced per prompt, and a caller that knows whether the batch is padded passes padded, skipping a host sync; this makes klein's and Z-Image's encoders capturable (#860). rms_norm / add_rms_norm (function forms, for norms whose weights live elsewhere) and image_output_format (lets a model check output_format at intake) change no behaviour on their own. Host work. to_device_nonblocking asks the device's pin-memory capability instead of checking for CUDA, and takes a sequence plus dtype, built straight into pinned memory in one allocation; main's tensor path is kept, and an already-pinned contiguous source is copied from in place. The sampler's keep-mask, wav2vec2's frame lengths, and the Nemotron duplex, Parakeet, FastConformer and audio-feature callers use it instead of their own copies. PinnedStager.to_device is deleted, and its callers, the DSA page tables and GLM-5.2's seq_lens, use to_device_nonblocking (this edits #387's GLM-5.2 and DSA files). torchcodec decodes through Model.load_video (and Cosmos3 and V-JEPA2) use min(8, usable_cpus()) FFmpeg frame threads, read once per process; main's usable_cpus() reads the CPU affinity and the cgroup quota, and the count is not tied to the API server's torch pool, which main sizes down from the CPU quota. Qwen3.5 (transformers' loader) and Qwen3-Omni don't go through it. Ride-alongs: benchmark/waypoint/benchmark_streaming.py and test/waypoint/serve_rollout.py match the server's internal ids (with their #n suffix) to client ids, and test_submodule_device_cache checks that a compiled forward reading main's cached device does not recompile.

Hardware: Eras p5en H200s, torch 2.12.1+cu132, FlashInfer 0.7.2.post1, MSTAR_COMPILE_LEVEL=1, SHM arena, Rust graph runtime and control mesh unless a row says otherwise. The per-change gains are push-1 A/Bs against main on the model each change was built for, 2–4 rounds on fresh servers. The "at ad5f644d3" rows are against main fadeb2759 (ed31270d1 differs only in a test docstring): one GPU for Qwen3.5-0.8B and Whisper large-v3, two for Qwen3.5-9B TP2 (GPUs 0 and 1 of a shared 4-GPU slice), four for Command A+ TP4 (command_a_plus_tp4.yaml with its KV pool at 2048 pages, the size the file gives for 141 GB GPUs); throughput there is benchmark.runner, closed loop, a fresh Inductor cache per boot, rounds A B B A A B (plus 3 in reverse order at 0.8B c=64), output tok/s per boot, failed=0 everywhere.

Change	Check	Result
Least-loaded replica	Throughput vs main	Wan2.2, 2 replicas: c=2 +20%, c=4 +45%. Qwen3.5-0.8B, 2 replicas: c=32 +22%, c=64 +4.5%. Greedy outputs on 2 replicas all in main's set, 8/8. CPU: `test_conductor_replica_assignment` (alternation, lockstep instances balanced whole, groups sharing a worker balance independently).
Warmup on every replica	Cold replica, startup	Cosmos3, 2 replicas: a cold replica's first request 10.5 → 6.2 s. Qwen3.5-0.8B DP2 logs one warmup per replica, same startup time (111 vs 117 s). CPU: `test_startup_warmup` (every replica reached, the solo run, a failing copy).
One-row-per-step lease	CPU	`test_piecewise_can_lease`: a real runner leases and replays the step as one row; refused for batched or non-leasing regions.
Ragged attention custom op	Parity	FLUX.2-klein's exact config bit-exact against the pre-change tree. `test_ragged_attention` (CUDA: graph kept, callable built inside a compiled region, rebuilt resource); `test_ragged_attention_cpu` (a dropped manager leaves the registry).
`head_shards`	—	No behaviour change: each builder gets the count it computed from the group before. No test of its own.
`TokenStop`, stop path (`8f3de9809` vs `37ad4f58f`)	Qwen3.5-0.8B c=32, 640×256, 3 rounds + `worker_phases` profile	12995 / 12661 / 13278 vs 12609 / 12724 / 13103 tok/s, within spread. CPU-bound (`event_sync` ~0.09 ms); `check_stop` +9 µs/step.
`TokenStop`, ASR (`8f3de9809`)	WER, 200 LibriSpeech test-clean clips, c=32	Whisper large-v3 1.77% vs main's 1.77–1.81%. Qwen3-ASR (with #859 on top) 1.58–1.60%, same as main. Measured on the stack's top (#972 at `0d242a90a`), not on this PR's own head: WER on 200 LibriSpeech clips, stack vs main: Whisper 1.72 vs 1.74%, Canary 5.52 vs 5.52%, Qwen3-ASR 1.79 vs 1.84%; at c=32, 196–198 of 200 transcripts identical. With `ignore_eos`, Whisper and Canary return `max_output_tokens` + 1 tokens on both arms, as before this PR (Qwen3-ASR's count is #859's).
`TokenStop`, token counts at `ad5f644d3`	`max_tokens` 1/2/4/16, both graph runtimes	2/3/5/17 chunks on both arms for 0.8B and Whisper. A 12-request mixed-budget burst gives budget+1 on every row, and 9B TP2 matches main across 7 boots (the +1 was main's; #794 has since made it N).
`TokenStop`, streaming at `ad5f644d3`	Streamed vs non-streamed	Equal chunk for chunk on every Qwen3.5 request (0.8B over 12 boots, 9B over 7); PR equals main on each runtime, and Rust equals Python on each arm.
`TokenStop`, Whisper at `ad5f644d3`	24 short clips, 6 boots, c=1 and c=8	Identical text, words and segment tokens. Word timestamps differing out of 544 (min/median/max): main vs main 13/19/26, PR vs main 15/19/22, PR vs PR 20/23/25. Against HF (527 words): main 346–348 exact, 18.6–18.9 ms mean error; PR 345–353, 18.41–18.82 ms.
`TokenStop`, Whisper word timestamps	Large-v3, `verbose_json`, c=1; 24 clips (544 words) plus a 90 s and a 605 s clip; 8 main boots, 7 PR boots, each with its own compile cache	Main's word timestamps change from boot to boot by up to 30 of 544 words, each by one 20 ms frame at a word boundary. Inductor's coordinate-descent autotuning picking different reduction configs per boot is one cause; compile level 0 with deterministic Inductor still differs on 9 words. The PR falls within that range, and two of its boots are bit-identical to a main boot across all 26 clips. Short-clip transcripts and segments never differ; the long-form clips vary on main too. Against HF (527 words): main 347–352 exact, 18.41–18.63 ms; PR 345–352, 18.37–18.71 ms. The earlier gap came from comparing a single cache state per arm.
`TokenStop` `prefill_walks`, Command A+ TP4 at `ad5f644d3` (before the move onto `END_PARTITION`)	Rust, Python, `MSTAR_TP_ASYNC_SCHED=1`	Byte-identical to main on all three. A prompt whose prefill samples the stop id returns one chunk, the stop's own, streamed or not, with the prefix cache off, and on a resend. Every `max_tokens` budget, with and without `ignore_eos`, returns the expected count. Greedy 64-token replies byte-identical. 96 mixed requests at c=8: 0 failures, 0 wrong counts, KV pool drained to 0 pages.
`TokenStop`, CPU	Linux CPU	`test_canary_model` (batched agrees with per-request), `test_whisper_model` and `test_qwen3_asr_model` (batched path taken), `test_whisper_decoder_node`, `test_glm52_model` (a stop id or a one-token budget at prefill ends the partition, with and without `ignore_eos`), `test_glm52_engine_cycle`, Command A+'s `test_integration`, `test_qwen3_5_check_stop_batched` (only the last prefill pass's token ends the partition), `test_inline_tokens` (a withheld step skips the fused route-and-send).
Throughput at `ad5f644d3` vs main	0.8B text c=32	12000 / 11718 / 12214 vs 11972 / 11914 / 11953, −0.3%
	0.8B text c=64	14528 / 14613 / 14047 / 13747 / 14209 / 14027 vs 14242 / 13724 / 13585 / 14454 / 14054 / 13995, −1.3%
	0.8B image c=32	7956 / 8064 / 7824 vs 8220 / 8258 / 7331, −0.1%
	9B TP2 c=16	2661 / 2671 / 2668 / 2645 vs 2661 / 2715 / 2697, +1.1% (main's first round overlapped another job's FlashInfer JIT, so a fourth was added)
	9B TP2 c=64	7054 / 7002 / 7023 / 7004 vs 6871 / 7092 / 6985, −0.5%. Every delta is within the round-to-round spread; TTFT and ITL p50 overlap. 0.8B is CPU-bound (`event_sync` 0.12 of a 2.46 ms step at c=32, 0.19 of 3.26 ms at c=64), 9B TP2 GPU-bound (3.64 of 6.01 ms at c=16, 5.00 of 8.77 ms at c=64); step time equal across arms within 0.04 ms.
Pinned H2D helper (`088a61d37` vs `53a60d15c`)	Qwen3.5-0.8B, cold boots	Within the cold-boot floor (41–47 of 96 texts identical vs a same-code floor of 45), recompiles equal (427). Whisper's log-mel lengths through the merged helper match main at `ad5f644d3`. CPU: `test_accelerator_graph_backend` (pins where the capability says, one pinned allocation, CPU device), `test_h2d_to_device` (an already-pinned tensor is copied from in place, contiguous or not).
Fused MoE router, H200 tiles	Qwen3-Omni text→text vs main; parity	c=1 +4%, c=16 0%. Router and fused-MoE tests pass on H200 (`test_moe_router`: matches the torch router, compiles, off by default; `test_fused_moe_decode`). The tiles aren't bit-exact (different K blocking); the Omni greedy divergence they cause is within batching noise.
Qwen3 text encoder	Parity	klein's exact config bit-identical. CPU: `test_diffusion_numerics` (one table across prompt lengths).
Threaded video decode	V-JEPA2 ViT-L vs main; parity	c=4 +28%, c=1 latency −26%. A 720p clip's full decode 3.0 → 0.9 s. Frames `torch.equal` to the single-thread decode on two boxes; served V-JEPA2 outputs byte-identical. CPU: `test_media_decode_limits`.
`ruff`; CPU Core vs main `9f84f18cd`	Eras CPU box	`ruff` clean. CPU Core against main `9f84f18cd`: 8751 passed, 1007 skipped, 2 failed, the same 2 as main (the `test_eval_compare_multigpu` dump_stacks pair); none fail only on this branch. GitHub CI green at `a4b5fcacb` (CPU Core, TP lockstep sim, rust-transport).

Not yet run on GPU at the head: everything after ad5f644d3 (the move of prefill_walks onto #794's END_PARTITION; main merges through 165b403b7, among them #810's per-node dtype, #849, #850, #876, #794 and #950; the head-sharding change; skipping the fused route-and-send on a withheld step); GLM-5.2 (not enough hardware) and the merged H2D helper's DSA and GLM-5.2 callers; the loop-cap budget on Canary-Qwen on this PR's own head; LTX-2.5 SP2 and TP2 with main's per-node dtype merged in; and, since push-1, DP2 routing and multi-replica warmup, the sampler keep-mask H2D, ragged compile_op (#860), the MoE router on Omni, and V-JEPA2's seeking loader's speed.

🤖 Generated with Claude Code

… perf push The parts of the overnight perf push (draft #853) outside the per-model folders, so the speech/language and generation model changes stack on top: - Engine: ragged attention `shard_heads` (CFG-parallel nodes keep every head) and opt-in `compile_op` (`run` traces as the `mstar::ragged_attention` custom op); managers carry their resource key. `RaggedAttentionCallable` is deleted: callers bind the label with `functools.partial(resource.run, label=...)`. Piecewise `one_row_per_step` leases a packed region at batch size 1. - Submodule base: `TokenStop` declaration on `ARNodeSubmodule` (gives `check_stop` / `check_stop_batched`), `loopback_input_ids`, cached `get_device`. - Conductor: least-loaded data-parallel replica pick (`_worker_load`). - API server: warmup `concurrency` (int or list) and one copy per DP replica; `api_switch_interval_s` config. - Components: `rms_norm` / `add_rms_norm` helpers, `TopKRouter` `fused_softmax_topk` through the fused-MoE `router_topk` (now with a softmax mode), Qwen3 text encoder caches one rotary table and takes `padded`, `image_output_format`, `to_device_nonblocking`, threaded video decode, `DenoiseLoopSubmodule.ragged_for` via `resource.run`. - fused_moe H200 tile configs for E=128. - LTX-2.5: the one call site of the deleted `RaggedAttentionCallable`. - Docs (serving, adding_models), diffusion-scaffold skill, waypoint benchmark log parsing, and the tests for all of the above. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Conflicts: - mstar/engine/resources/attn/ragged/flashinfer.py: main (#863) moved the manager machinery into ragged/base.py's _RaggedManagerBase for the FA3 dense backend. Took main's layout and moved this branch's resource-keyed mstar::ragged_attention op, the _MANAGERS registry and the compile_op `run` into _RaggedManagerBase, so both backends get compile_op; the FlashInfer managers just pass `key` through. - test/modular/test_conductor_tick_exception.py: imports; kept main's `re` and this branch's `Counter`. Converged on main: - mstar/utils/h2d.py: main (#890) added to_device_nonblocking too; kept main's and deleted this branch's duplicate (same signature; main's pins only for CUDA where this branch asked DeviceCapability). - mstar/model/components/token_step.py: main's TokenStepNode (#890) carried its own stop checks, the same logic as this branch's TokenStop. It now declares `token_stop` and inherits ARNodeSubmodule's checks; its _stops, check_stop and check_stop_batched are deleted. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The replica pick summed live requests per worker across every node group, so a group sharing a worker with another (a DP2 group on ranks 0 and 1 beside a single-replica group on rank 0) saw the other group's requests as its own load and kept sending new requests to one replica. Each request now records its replica per group (RequestData.dp_replicas), and the load is counted on those keys. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The flag was applied at the engine's early lease only, while the runner's run and can_run derived the batch size from the request ids or segment lengths they were passed, so a caller had to repeat real_bs=1 and the first request id by hand or replay against a bucket the lease never promised. The runner now resolves a one_row_per_step region as one row in _resolve (lease, can_run, run) and trims run's request ids to the first, so the lease and the replay agree by construction and the engine passes the step's batch size. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…ate (P3-3) With generated_key set, the stop check read the count through request_state, which creates a state for a request it doesn't have one for; a check that runs after the request's state was dropped left a fresh one behind for a finished request. It reads request_states directly now. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The mixin derived token_stop from stop_token_ids through a cached_property, while the docs and every other AR node set token_stop in __init__. Canary and Canary-Qwen now set it in __init__, the mixin drops stop_token_ids and the property, and the module docstring no longer claims the stop checks. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The module-level key -> manager map the mstar::ragged_attention op reads was never purged, so it kept every ragged manager (and its plan wrappers and workspaces) alive past the engine that built it. A WeakValueDictionary drops an entry with the engine's last reference; there is no resource teardown hook to deregister from, and the op can only name a manager by key. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…-interval rows (R1.2, R1.3, P3-4) adding_models.rst documents one_row_per_step (and that the whole step's segments then count against max_segments_per_request), shard_heads and compile_op, says the latter two are the model's declaration rather than YAML overrides, and names RaggedAttentionConfig where it said "the config". serving.rst says warmup sends each entry once per DP replica and that concurrency copies can, not will, share a step, and no longer claims the Qwen3.5 configs set api_switch_interval_s (they do in #859, not here). The ragged config comments stop naming stacked-PR adopters. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…source (R1.7, P3-6) The helper only reads the sampler's slot masters and the step's routed ids, so it belongs to the resource that owns those masters rather than to every AR submodule. Its comment no longer claims the row-index copy never blocks: off CUDA, to_device_nonblocking is a plain copy. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…ne (P3-5, P3-6) TopKRouter's fused_softmax_topk docstring says bf16 logits tie often, so an opting-in model should count how often its outputs diverge; the router test's docstring describes the tests it holds; mstar/utils/device.py matches main. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…nity (R1.1) VIDEO_DECODE_THREADS was min(8, os.cpu_count()) fixed at import: it counted the host's CPUs rather than the ones a container's cpuset gives the process, and ignored preprocess_torch_threads, the API server's preprocessing budget. media_limits.video_decode_threads() now bounds a decode by torch's intra-op pool (which that knob sets), by len(os.sched_getaffinity(0)) where the platform has it, and by 8. V-JEPA2's own loader asks it too. Qwen3.5 decodes through transformers' load_video, which takes no thread count; that stays a follow-up. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The Whisper decoder is the third EncoderDecoderNode, missed by the previous commit: with the mixin no longer deriving token_stop from stop_token_ids, it had no stop declaration and its stop checks fell back to the base, which never stops (test_whisper_decoder_node caught it). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…_ids (R2.a) The decode branch of preprocess kept its own copy of the routed-row index_copy_, whose row index was a blocking torch.tensor(..., device=). The sampler's helper does the same copy with the index staged through pinned memory. The test's fake sampler binds the real helper, so the routed-row case now exercises it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…ler names (R2.b) Canary-Qwen built its TokenStop from the module constants while Canary and Whisper used self.DECODE_LOOP; all three now take both the loop and the sampler from the class attributes, so a subclass that renames them stops on the right loop. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

… CPU quota (R2.c) video_decode_threads reads sched_getaffinity, so it honours a cpuset but not a CFS quota (a container --cpus limit). The docstring and the adding_models note now say so. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

A one_row_per_step config refuses capture sizes above 1, so the test runner now captures only (1, 256). can_run(3, ...) then passes only because the runner resolves the step as one row; with the old per-call batch size it finds no bucket. Also drops a stray third blank line in test_conductor_replica_assignment.py. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

… docs paragraph Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Main brings #849 (Host overhead 2/N): the API server hands the data worker one InlineTokens item per result_tokens frame, Model.text_output / postprocess_inline map inline token ids straight to bytes (raw token bytes for Qwen3.5 and Whisper), and the Rust frontend relays chunks in batches (APIServer.relay_chunks_to, ChunkSink). One conflict, the typing import in mstar/api_server/entrypoint.py: both sides' names (NamedTuple for DataParallelReplicas and _SinkBatch, Protocol for ChunkSink). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The row still said the knob would be re-measured once #849 lands. It has landed and changed the message-thread and data-worker path the knob tunes, so the numbers are from before that change and the re-measure is still pending. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

It was there to document the removed splits_heads attribute; what is left repeats the RaggedAttentionConfig row. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Head sizing now equals main's, so the six cases pinned nothing this branch changes. The shard_over_sp=False test stays; the helper loses its recurrent branch, which only the deleted test used. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Main already sizes a shard_over_sp=False ragged attention over the TP group, so the test pinned nothing this branch changes. The file's coverage goes back to main's plus the manager registry test. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Brings in #850 (worker speculates as columns, fused route+send, skips unread publishes) and #974 (only expert parallelism syncs MoE routing). #850's fused route+send (route_and_send_inline) builds its route input from the host column as plain i64s, where this branch's inline values are per-row optional so a withheld output's row carries None. The Rust body now builds Some per row, and the worker sends a step whose stop check withheld an output through the general path, which gives that row None: the fused call reads every row's token off the column and would send it inline. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…-shared Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Conflict in test/modular/test_ragged_attention_cpu.py: main's mixed-wrapper assertions (a sink row on each side of the prompt rows, the tensor-core decode half on the FA2 paged prefill kernel) replace this branch's copy of the old ones; this branch's op-registry test is kept beside them. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Conflicts: - conductor.py: imports; Collection is gone with main's last use of it, Counter stays for the replica load. - audio_features.py: imports; both main's stft_as_real and this branch's DeviceCapability are used. - nemotron_duplex/submodules.py, components/conformer.py: main moved the duplex encoder to components/fastconformer_stream.py and deleted conformer.py. Its device_counts helper is gone on both sides; the stream's two count copies now hand to_device_nonblocking the host list directly, as this branch's other call sites do. Main moved the PNG thread count onto the CPU plan, so video_decode_threads follows the plan's unpinned CPUs too rather than usable_cpus(). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Gaurav-Shah05

On H100s I ran main (9f84f18cd) against this PR with Qwen3.5-0.8B, closed loop:

	main	#858
2 replicas, 32 requests	10.8k tok/s	12.8k tok/s	1.18x
2 replicas, 64 requests	17.4k tok/s	19.0k tok/s	1.09x
2 replicas, TTFT p50 at 32 / 64 requests	18.1 / 27.0 ms	18.1 / 27.1 ms
1 GPU, 64 requests	17.5k tok/s	17.6k tok/s	1.01x

The replicated rows use replies of 16 to 512 tokens, drawn per request, and your least-loaded gain reproduces there with TTFT unchanged. The 1-GPU row has the stop-path changes alone, and they're neutral, as you found.

The description says GLM-5.2's and Command A+'s caps become max_output_tokens - 1, but both were already max(1, n - 1) on main, so neither budget changes.

lgtm overall and my comments are only minor.

Gaurav-Shah05 · 2026-10-31T20:51:42Z

+    # first step (``ARNodeSubmodule.ended_at_prefill``), which withholds its
+    # ``buffer`` (``StopCheck.withheld``), so the loop's client edge must
+    # carry ``buffer``
+    prefill_walks: frozenset[str] = frozenset()


With prefill_walks, a stop id sampled at prefill keeps the request alive for one more decode step. This token is then withheld from both the inline frame and the tensor path.

#794 ends the same request at the prefill by returning END_PARTITION, which the conductor already turns into request done, so I think prefill_walks should mean that walk's row returns END_PARTITION.

ended_at_prefill and the withheld inline row in the worker and both graph runtimes can go then.

Done in 03f7758: a prefill_walks row with a stop id (or a one-token budget) now returns END_PARTITION, and ended_at_prefill and the None inline value in the worker and both runtimes are gone. Qwen3.5, Qwen3-ASR, Canary and Canary-Qwen declare their prefill walks in place of #794's per-model checks.

Gaurav-Shah05 · 2026-10-31T20:52:25Z

    torch.set_num_threads(int(threads))


+def _set_switch_interval(config: dict) -> None:


On #859 I measured 1 ms against Python's default 5 ms on Qwen3.5-0.8B at c32 and c64 and saw no difference beyond the spread between sessions. Also on main #939 replaced the message thread's 1 ms sleep, which this knob was compensating for, with a wait on results.

Agreed, and with #939 gone the 1 ms sleep it made up for is gone too. Dropped the knob in bffb067.

Gaurav-Shah05 · 2026-10-31T20:53:55Z

    files: list[WarmupMedia] = field(default_factory=list)
    output_modalities: list[str] = field(default_factory=lambda: ["text"])
    model_kwargs: dict | None = None
+    concurrency: tuple[int, ...] = (1,)


Is there any config in this stack that sets concurrency? If not its parsing, docs and tests cover a dead path. It would be needed when a config needs a batched warmup.

None does, so it's gone in a25694a; ids are now warmup-..

Gaurav-Shah05 · 2026-10-31T20:54:23Z

-                "warmup request %d/%d failed after %.1f s: %s",
-                index + 1, len(specs), time.perf_counter() - t0,
-                getattr(exc, "detail", exc),
+        if replicas.fewest < replicas.most:


Each warmup copy reaches its own replica only because the least-loaded pick spreads concurrent requests. #920 already had to turn the prefix cache off on a replicated deployment's warmup copies to keep them apart. What do you think about the conductor placing warmup copy i on replica i of each group, so every replica warms under any routing policy?

Agreed it's sturdier than relying on the least-loaded spread. It needs a placement hint the conductor and #920's routing both honour, so I'll do it as a follow-up on top of #920.

Gaurav-Shah05 · 2026-10-31T20:54:42Z

    @classmethod
    def build(cls, spec: KVSpec, info: EngineResourceInfo):
+        # before the cache is allocated: it is sized off the head counts
+        spec.config.shard(info.head_shards)


With the shard moved from __init__ to build, the tests that construct KVManager directly with a two-rank group no longer shard, so test_cfg_tp_publish.py's TP2 managers each hold both KV heads.

Those helpers should go through KVManager.build, or shard the config themselves, so the test still moves one rank's shard.

Done in 641e06c; the TP2 helper shards its config the way build does.

Gaurav-Shah05 · 2026-10-31T20:56:02Z

+    # ``mstar::ragged_attention`` op, rather than breaking there per layer. For
+    # a compiled, uncaptured region whose fusions around attention pay; off by
+    # default, since tracing a whole transformer as one graph per capture
+    # bucket costs startup time and capture memory (Z-Image: +2.6 GiB).


Does anything in the stack turns compile_op on for Z-Image?

No, Z-Image leaves it off; the +2.6 GiB there is why it defaults off. Wan2.2 (#860) and Qwen3.5 (#969) turn it on.

Gaurav-Shah05

meant to approve

Conflicts: RouteInput / RouteArg keep this branch's optional inline values beside main's per-row log-probs (Python and Rust alike); the delivery parity test runs this branch's None inline row with main's log-prob script; Qwen3-ASR's imports take main's encoder move plus TokenStop. Main's Qwen3-Omni ragged-backend stand-in takes the manager key this branch passes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

A case with a pinned CPU fails if video_decode_threads goes back to usable_cpus(); the MSTAR_STEP_THREAD_CORE row names FFmpeg's per-decode threads among what the pins shrink. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Its gain was measured before #849 and #939 changed the message-thread and data-worker path, and the re-measure it was waiting on never ran. The serving-config key, its validation, docs rows and test go; workers keep main's MSTAR_PY_SWITCH_INTERVAL_SEC unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

KVManager shards its config in build, so the test's direct constructor shards it the same way and each rank moves its own shard. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

No config sets it, so its parsing, docs and tests covered a dead path. Each entry still goes out once per data-parallel replica, now as warmup-<spec>.<copy>. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

TokenStop.prefill_walks rows now return #794's END_PARTITION (a stop id, or a one-token budget), with last_prefill_flag for a prompt prefilled over several passes; ended_at_prefill and the None inline value in the worker and both graph runtimes are deleted. Qwen3.5, Qwen3-ASR, Canary and Canary-Qwen declare their prefill walks in place of #794's per-model checks. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

#794's cap-versus-stop-check agreement tests for Canary and Canary-Qwen assumed a stop check that counts max_tokens; TokenStop leaves the budget to the loop's cap, so they pin the cap instead. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…LM-5.2 test comment names the loop cap Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

mstar-review Bot added area/api area/conductor area/engine area/model-shared area/resources area/utils model/ltx2_5 size/L labels Oct 24, 2026

dghosef force-pushed the perf/split-1-shared branch from 39ee03e to df1af58 Compare October 24, 2026 05:55

This was referenced Oct 24, 2026

Performance 2/7: speech and language models #859

Merged

Performance 3/7: image, video and world models #860

Merged

Performance improvements across 18 models #853

Closed

stephen-dwq added this pull request to stack #865 October 24, 2026 23:37

Gaurav-Shah05 mentioned this pull request Oct 27, 2026

DiT scaffold: draw seed noise at ingest and queue the schedule without a host sync #904

Merged

dghosef and others added 15 commits October 27, 2026 01:06

dghosef and others added 10 commits October 29, 2026 23:24

Name the TP x SP head-sharding test for the kinds it covers; reflow a…

50401ad

… docs paragraph Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Drop the head-sharding paragraph from the submodule attributes section

82d9fb5

It was there to document the removed splits_heads attribute; what is left repeats the RaggedAttentionConfig row. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Conductor: build its instance table from WorkerGraph.replica_ranks

bbe03b6

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Merge origin/main (bb0e0a8: #962, #941, #944) into perf/split-1-shared

0e495ae

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Merge origin/main (c96e02b: #968, #856, #778, #876) into perf/split-1…

7ee5eb4

…-shared Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

dghosef mentioned this pull request Oct 30, 2026

Performance 7/7: image, video and world models #972

Merged

dghosef and others added 2 commits October 30, 2026 19:29

Gaurav-Shah05 reviewed Oct 31, 2026

View reviewed changes

Gaurav-Shah05 approved these changes Oct 31, 2026

View reviewed changes

dghosef and others added 9 commits October 31, 2026 19:14

Tests: the TP2 publish test's KV managers hold one rank's head

641e06c

KVManager shards its config in build, so the test's direct constructor shards it the same way and each rank moves its own shard. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

API server: drop warmup_requests' concurrency option

a25694a

No config sets it, so its parsing, docs and tests covered a dead path. Each entry still goes out once per data-parallel replica, now as warmup-<spec>.<copy>. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Merge remote-tracking branch 'origin/main' into land/858

34fbf94

Docs: last_prefill_flag for a prompt prefilled over several passes; G…

f991121

…LM-5.2 test comment names the loop cap Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

dghosef mentioned this pull request Nov 1, 2026

Performance 5/7: Qwen3.5 and Qwen3-Omni #970

Merged

Gaurav-Shah05 mentioned this pull request Nov 1, 2026

Speed up MoE decode with three fused kernels #984

Open

NSagan271 mentioned this pull request Nov 1, 2026

Host overhead 6/N: declared stop policies, checked in the graph runtime #857

Merged

dghosef mentioned this pull request Nov 1, 2026

Qwen3-Omni: fit the encoder captures on an 80 GB GPU #1119

Open

vasilevklart mentioned this pull request Nov 1, 2026

Load a Qwen3-VL checkpoint's language model as a diffusion prompt encoder #987

Closed

merceod mentioned this pull request Nov 3, 2026

Stop at max_output_tokens, not one past it #604

Merged

		torch.set_num_threads(int(threads))


		def _set_switch_interval(config: dict) -> None:

Conversation

dghosef commented Oct 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Gaurav-Shah05 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Gaurav-Shah05 left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

dghosef commented Oct 24, 2026 •

edited

Loading