Tensor communication by NSagan271 · Pull Request #2 · mstar-project/mstar

NSagan271 · 2026-02-26T20:37:25Z

Initial work on tensor communication abstractions via mooncake transfer engine and ZMQ

Three independent bugs surfaced by Qwen3-Omni audio-mode generation under Phase 3 speculation. Two distinct visible failure modes were in play simultaneously: long garbage continuation past <|im_end|>, and mid-sequence token doubling like "of of France France is is Paris Paris". Each fix is required; none of them subsumes another. 1. Speculation re-ingests stale loop-back into a freshly-reset queue (worker.py) Spec submits iter K+1 before iter K's slow_post detects EOS. K's stop arrives in _pending_stops while K+1 is already in flight on the GPU thread. K+1's _fast_postprocess applies the stop (register_finished -> queue done -> queue.reset), then routes K+1's outputs via process_node_outputs. The freshly-reset queue's waiting Thinker GraphNode wants text_inputs, ingests K+1's loop-back tensor, schedules another Thinker iter, which samples another <|im_end|>, fires another stop, and the cycle repeats forever -- producing thousands of garbage tokens after the first valid response. _apply_pending_stops_to_batch now returns the stopped loop's loop-back input names. _fast_postprocess merges them into speculation_consumed_loop_back so the routing-filter the spec path already uses (drop loop-back edges already consumed by spec) drops them too. Loop terminates cleanly at EOS. 2. Stop signals lost when popped by the wrong partition's batch (worker.py) _pending_stops was keyed by rid only. Apply-time pops them when the current batch contains the rid -- regardless of which partition that batch is on. On Q3-Omni audio mode, a Thinker stop sitting in _pending_stops[rid] could be popped by a subsequent Talker batch with the same rid; stop_loops walks Talker partition's queues, finds no thinker_decode_loop, and silently no-ops. The stop is gone, Thinker keeps generating past EOS forever, and fix #1 never fires because register_finished never reaches Thinker's queue. Re-key as (rid, partition_name) and only pop entries whose partition matches batch_partition. Stops on other partitions stay pending until a batch on the matching partition runs. Fix #2 is a prerequisite for #1 to work in audio mode. 3. Triton/FlashInfer sampling cross-stream sync race (sampling.py) Sampler._consume_autotune_sync_budget synced for the first 64 sample calls then no-op'd. The intent was to bridge Triton autotune's first kernel launches (which can land on a non-default stream while autotune is timing configs) and FlashInfer's top_p_sampling_from_probs read of the probs buffer. After autotune was assumed to settle, both kernels would land on the default stream and natural FIFO ordering would suffice. That assumption holds for Orpheus (single autotune key: V~128k, APPLY_PENALTY=False, INCLUDE_GREEDY=False), but NOT for Q3-Omni: Thinker (V~151k, no rep-penalty) and Talker (V~8k, rep-penalty=1.05) hit two distinct autotune keys. Once the 64-sync budget runs out (after iter 64 of any kind), subsequent sample calls skip the sync. FlashInfer then reads partially- written probs and samples from the previous iter's distribution -- producing the "X X Y Y" mid-sequence doubling. Drop the budget; sync unconditionally. Cost is ~10us per sample, negligible against 3-5ms iter time. Orpheus iter_total p50 unchanged at 3.2ms after the change. Plus two non-bug-fix tweaks bundled with the above: - ar_engine: env-gated MMINF_WORKSPACE_BUFFER_MB (default 512, was hardcoded 256). Q3-Omni Thinker prefill (bs=4 x num_tokens=2048) needs ~274 MB of FlashInfer workspace; the old 256 MB cap caused 3 of 40 captures to fail with workspace overflow. - Defensive .clone() in _sample_and_remap (cuda_graph_runner) and _sample_decode_outputs (ar_engine). Ensures iter-N's sampled tensor is on fresh storage, immune to any FlashInfer-internal output-buffer reuse across calls. Not load-bearing for the duplication fix (Triton sync above is what fixes it) but removes an aliasing risk for future FlashInfer changes. pyproject: separate, restructure CUDA-backend extras as cu128 / cu130 so default install picks up CUDA 13 wheels for torch 2.11. Verification: 50/50 Q3-Omni audio-mode runs across 5 prompts produce clean text. Orpheus Phase 3 perf preserved (iter_total p50 3.2 ms). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Three independent bugs surfaced by Qwen3-Omni audio-mode generation under Phase 3 speculation. Two distinct visible failure modes were in play simultaneously: long garbage continuation past <|im_end|>, and mid-sequence token doubling like "of of France France is is Paris Paris". Each fix is required; none of them subsumes another. 1. Speculation re-ingests stale loop-back into a freshly-reset queue (worker.py) Spec submits iter K+1 before iter K's slow_post detects EOS. K's stop arrives in _pending_stops while K+1 is already in flight on the GPU thread. K+1's _fast_postprocess applies the stop (register_finished -> queue done -> queue.reset), then routes K+1's outputs via process_node_outputs. The freshly-reset queue's waiting Thinker GraphNode wants text_inputs, ingests K+1's loop-back tensor, schedules another Thinker iter, which samples another <|im_end|>, fires another stop, and the cycle repeats forever -- producing thousands of garbage tokens after the first valid response. _apply_pending_stops_to_batch now returns the stopped loop's loop-back input names. _fast_postprocess merges them into speculation_consumed_loop_back so the routing-filter the spec path already uses (drop loop-back edges already consumed by spec) drops them too. Loop terminates cleanly at EOS. 2. Stop signals lost when popped by the wrong partition's batch (worker.py) _pending_stops was keyed by rid only. Apply-time pops them when the current batch contains the rid -- regardless of which partition that batch is on. On Q3-Omni audio mode, a Thinker stop sitting in _pending_stops[rid] could be popped by a subsequent Talker batch with the same rid; stop_loops walks Talker partition's queues, finds no thinker_decode_loop, and silently no-ops. The stop is gone, Thinker keeps generating past EOS forever, and fix #1 never fires because register_finished never reaches Thinker's queue. Re-key as (rid, partition_name) and only pop entries whose partition matches batch_partition. Stops on other partitions stay pending until a batch on the matching partition runs. Fix #2 is a prerequisite for #1 to work in audio mode. 3. Triton/FlashInfer sampling cross-stream sync race (sampling.py) Sampler._consume_autotune_sync_budget synced for the first 64 sample calls then no-op'd. The intent was to bridge Triton autotune's first kernel launches (which can land on a non-default stream while autotune is timing configs) and FlashInfer's top_p_sampling_from_probs read of the probs buffer. After autotune was assumed to settle, both kernels would land on the default stream and natural FIFO ordering would suffice. That assumption holds for Orpheus (single autotune key: V~128k, APPLY_PENALTY=False, INCLUDE_GREEDY=False), but NOT for Q3-Omni: Thinker (V~151k, no rep-penalty) and Talker (V~8k, rep-penalty=1.05) hit two distinct autotune keys. Once the 64-sync budget runs out (after iter 64 of any kind), subsequent sample calls skip the sync. FlashInfer then reads partially- written probs and samples from the previous iter's distribution -- producing the "X X Y Y" mid-sequence doubling. Drop the budget; sync unconditionally. Cost is ~10us per sample, negligible against 3-5ms iter time. Orpheus iter_total p50 unchanged at 3.2ms after the change. Plus two non-bug-fix tweaks bundled with the above: - ar_engine: env-gated MMINF_WORKSPACE_BUFFER_MB (default 512, was hardcoded 256). Q3-Omni Thinker prefill (bs=4 x num_tokens=2048) needs ~274 MB of FlashInfer workspace; the old 256 MB cap caused 3 of 40 captures to fail with workspace overflow. - Defensive .clone() in _sample_and_remap (cuda_graph_runner) and _sample_decode_outputs (ar_engine). Ensures iter-N's sampled tensor is on fresh storage, immune to any FlashInfer-internal output-buffer reuse across calls. Not load-bearing for the duplication fix (Triton sync above is what fixes it) but removes an aliasing risk for future FlashInfer changes. Verification: 50/50 Q3-Omni audio-mode runs across 5 prompts produce clean text. Orpheus Phase 3 perf preserved (iter_total p50 3.2 ms).

…torage, and batch it (#322) * Key tensor bookkeeping by uuid, split it from storage, and batch it Tensor uuids are globally unique, so the tensor store no longer keys by request. TensorStore keeps the tensors and which request owns each one (only for teardown). The reference state (refcounts, persist, mem_registered and the descriptor) moves behind a TensorBookkeeping interface that sees only integer uuids. That interface has a Python backend and a Rust one (rust/src/tensors.rs). The Rust backend is a drop-in, covered by parity tests run against both, but nothing selects it by default: it copies descriptors in and rebuilds them on the way out, which pays off only for a Rust-side caller holding the same state. Until one exists, it is reachable only by passing it to TensorStore. The manager gains batched forms, each one boundary crossing under the Rust backend: - store_and_return_tensor_info_batch returns flat columns for a whole output batch - increment_ref_batch_uniform, dereference_batch(_uniform) and cleanup_collectable - register_for_send_uuids, which the arena stages in one pass with one D2H sync The worker, data worker and graph nodes move onto these. The worker still routes per request in Python, so a small helper regroups the batch store's columns into each request's output edges. A model output under a name no output edge carries is no longer stored. Nothing can read it, and under the batch store it would take no reference, so it would be held until the request is torn down. cargo test now runs with --no-default-features (pyo3's extension-module feature is optional), and ParallelList joins utils/containers. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * avoid sending one string per tensor across the Python <> Rust boundary with the SHM arena * Restore the cross-worker transport tests _send still called store_and_populate_graph_edges, removed in this commit's own branch when the store moved to uuid keys. All 12 tests in test/distributed/test_tensor_transport.py have been failing since -- the only coverage of the sharded fanout/fanin paths -- so the helper now does the three things that call did: store, populate the edge, take a reference per outgoing edge. Also corrects the _Placements docstring: ParallelList does NOT share its "always truthy" edge, because it defines __len__. Co-Authored-By: Claude <noreply@anthropic.com> * Respond to PR comments: prevent leaks on raise, make SHM arena registration linear instead of quadratic --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

* Combined graph walks: batch several walks of a node as one A model declares {node: {combined_walk: {walks}}} via Model.get_combined_graph_walks(); the worker builds a validated (node, real walk) -> combined walk map for the scheduler. - MicroScheduler keys scans, backlog, round-robin and caps on the combined walk, pops each real walk separately, and records per-rid walks in ScheduledBatch.request_walks. A batch whose rows share one real walk is labelled with that walk, so uniform steps keep their CUDA graphs. A combined key is capped at the smallest cap of its walks. - StepContext.request_walks / NodeInputs.graph_walk carry each row's real walk to prepare_inputs, the prefix cache, declare_step and KV publish. - The worker stores, routes and sends outputs once per real walk, and stops loops per walk. - TP: ScheduleTPNode.request_walks; followers pop per walk, all or none. - Speculation: speculate_node takes columnar (walks, sample_rids) and returns the first walk with a target; only that walk's rids continue. - EngineManager refuses split_batches_by_capture_key on a combined walk. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Orpheus: batch prefill and decode under one combined walk LLM's prefill and decode rows share a "mixed" walk. Mixed steps replay the packed prefill CUDA graphs, where a decode row is a one-token segment. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Combined walks: cap a combined key by the combined walk's own cap The smallest cap of the constituent walks throttled decode-only batches to the packed-prefill capture size (4 on BAGEL / Qwen3-Omni). The model now declares the combined walk's cap through its captures or max_batch_size. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL: batch prefill_text and decode under one combined walk LLM's prefill_text and decode share "mixed_text". Guidance-off mixed steps replay the packed prefill_text captures, now captured up to batch 64 so a step of decode rows plus a prompt still replays. Prompt tokens are tracked per prefill_text row; the guidance pre-fork applies when a guided prefill row is present, as it already did for a mixed-guidance prefill batch. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni: batch Thinker text/audio prefill and decode together The Thinker's prefill_text, prefill_audio and thinker_decode share "thinker_mixed", which runs the packed prefill path and replays its captures (now up to batch 32). prefill_vision stays apart: it does not batch at all yet, and its deepstack features and MRoPE advance are built for a single request. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * FIFO builder: cap a combined step by the walks it actually holds Under a combined walk the scheduler passes each real walk's cap. The FIFO builder grows the step's walk set in order of first appearance and keeps the set that schedules the most rows (ties keep more walks), so a walk with a small cap no longer shrinks a step it would only join. The oldest row's walk is always in, and a row left out goes to the backlog, where it is first in line next step. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Combined walks: route a relabelled step with its own walk's outputs A combined batch merged from several walks' pops kept the first walk's output signals. When the scheduled step turned out to be one walk and was relabelled, it routed with the wrong names: a decode step routed with prefill's lost its loop-back edge, and the next decode step found no input. ScheduledBatch now keeps each walk's output signals and relabelling takes the matching ones; the worker uses them for a mixed step too. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Worker: set a batch's in-flight flag in every worker graph it spans A combined walk's prefill and decode rows live in different worker graphs, but two places flagged every rid under the first rid's worker graph: _set_in_flight_flag (submit / completion) and _settle_speculation (a speculative batch). The first left other graphs' rows flagged after a mixed step; with backlogged rows flagged too, Orpheus requests hung in a c=48 burst. Fixing only that one made it worse: a speculative batch still flagged a fresh prefill row on its decode node (the first row's graph), the per-graph clear never reached it, and the row never decoded. Both now go per worker graph. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Combined walks: cap a step by the walks it holds, not the combined walk A combined key was capped at the combined walk's own cap, which applied before a uniform step was relabelled, so a decode-only Orpheus step stopped at the packed prefill capture size (16) instead of decode's (64). A combined key now takes its largest walk's cap; a one-walk step takes that walk's and a mixed step the smallest of its walks' and the combined walk's. A speculation merge counts its continuing rows' walk, and room_for_continuing caps by the walks the step would actually hold. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * StepContext.force_prefill: plan prefill attention for packed captures FlashInfer picked its wrapper from the step's shape: a bucket with as many tokens as rows (or an eager step of one-token rows) got the decode wrapper, which has no qo_indptr, so a packed forward's select_last_hidden failed. A packed capture of bs == num_tokens (e.g. 64 rows x 64 tokens) could not be captured. The engine now marks a step leased on a packed capture, at capture and at replay, and FlashInfer plans prefill for it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Orpheus: capture packed prefill at batch sizes 1, 4, 16, 64 A packed graph's work depends on its token count; the batch size only pads host-side rows, so coarser batch buckets reaching 64 cost nothing on the GPU and let a mixed step hold all the decode rows (the mixed cap was 16, and a mixed step left the rest of decode for another step). 24 packed graphs instead of 30. Orpheus A/B vs cp/3, 2 restarts per side: c=48 +0.1% req/s (was -25% before this and the walk-set cap fix), c=16 +0.4%. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Combined walks: one pop per step, all spec targets, FIFO-by-row caps Review follow-ups on #362, mostly cutting per-step Python work under a combined walk: - speculate_node returns every walk's targets (each names its walk) and the caller picks; the async check is per walk in the runtime, so a batch may speculate when any of its walks is async (Python and Rust). - pop_walk_rids pops rows of several walks of one node in one runtime call, from columns (RequestWalks: rids, walks, walk_idx); all or none under check_ready. Both scheduler pop sites use it, and the TP head carries its walks as the same columns. A follower's batch now gets each walk's output signals from the pop rather than asking per walk at routing. - ScheduledBatch caches rids_by_walk; every row removal goes through discard_rid, which resets it. - The FIFO builder composes a combined step row by row: a row joins unless its walk (or the combined walk's cap, once the step mixes walks) would cap the step at or below the rows already in it. It takes raw caps plus the caller's own row count and walk, instead of a separate mixed cap. - The scheduler asks the engine for each (node, walk) cap once. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * cleanup and CPU perf * refuse_combined_walks_across_tp_groups + more cleanup * Combined walks: one postprocess round trip per step across walks complete_and_route_batch takes per-row walks (walks + rid_walk_idx), so a mixed step does one store, route, register and send; the completion keeps each row's walk and worker graph. set_in_flight and commit_speculation take per-row worker graph ids, and pending_loop_stop_rids a list of walks, so neither is called once per worker graph or walk. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni: batch Thinker vision prefill into the mixed walk prefill_vision joins thinker_mixed. preprocess packs each deepstack layer over every row of the step, with zeros beside rows that have no vision features; declare_step advances each row's MRoPE position by its own rule (a vision row by its 3D-grid extent, others by their span). A mixed step holding a vision row replays its own captures (cg_key_info "deepstack", batch sizes 4/16/32 at the text token buckets, 15 graphs), since the plain mixed captures take no deepstack; mixed steps without vision are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni: stamp the deepstack capture key in declare_step The engine checks that a step declares the cg_key_info it was leased under; a mixed step holding a vision row leased the deepstack captures but declared none, so it failed at replay. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3-Omni: one deepstack capture for vision and mixed steps Like prefill_text's captures, which also replay thinker_mixed, the vision capture now replays both prefill_vision and thinker_mixed under the deepstack key. It captures batch sizes 1/4/16/32, so prefill_vision's cap matches the mixed walk's and vision rows join mixed steps under the existing cap rule. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Compiled forwards: keep request ids out of what dynamo traces An eager step recompiled whenever its rows or their ids changed: the sampler entry and the per-rid output dicts made dynamo guard on len(request_ids) and on the ids themselves. Combined walks make such steps common (a long prefill riding in a decode step misses every capture). SamplerResource.sample runs outside the compiled frame, and the Orpheus and BAGEL LLMs return row_outputs for the engine to split. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: batch text prefill and decode in one walk The LLM's mixed walk replays the packed prefill_text captures, now to batch size 32 so decode rows fit. The GDN backend takes the chunked kernel under a packed capture whatever the spans, as attention does: a bucket with one token a row is captured as prefill and replays mixed steps. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: only a decode row's token can stop check_stop_batched read every row's sampled token, so an earlier prefill step's (which postprocess drops) could stop the request on EOS or count toward its budget. Both stop paths now check decode rows only: a prefill row is not in decode_loop yet, so its stop could not end the loop anyway. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Scheduling: a backlog that fills the step runs without a ready scan With more ready rows than a step holds, a split set sits in the backlog across steps. The worker's own scheduling pass still scanned the ready queues, checked and popped their rows, and measured them, only for the builder to park them again: no fresh row joins a step its backlog fills. Skip the scan then, as the speculation merge already does; a backlog step that comes out empty (every row blocked) still takes fresh rows. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * prefill starvation fix * bad cleanup of speculatively scheduled loops that have already finished fix * Combined walks: refuse TP-group spans by the walks a worker runs refuse_combined_walks_across_tp_groups keyed on the worker's nodes, so a P/D config at TP2 on both sides -- which the docs allow -- failed to start on every rank: the prefill and decode halves of a combined walk resolve to different sharding groups even though no rank runs both. Key on the (node, walk) pairs the worker actually runs instead; a worker holding one half can never build a step that mixes them. Co-Authored-By: Claude <noreply@anthropic.com> * Combined walks: park other walks without re-checking our own rows _park_other_walks scanned every ready row of the key and dropped its own walk's afterwards, paying a check_ready per row to discard it. Pass (node, own_walk) as the scan's exclusion instead: both runtimes match exclude_target against the real walk, so the same rows park. Measured at 0.26 -> 0.14 ms of speculate a step on Orpheus at c=128. Co-Authored-By: Claude <noreply@anthropic.com> * Fix the loop-counter test's worker fake for the scheduler label _assemble_speculation has called scheduler.label_for since the first combined-walks commit, which updated seven test files but not this one; its SimpleNamespace worker has no `scheduler`, so all three tests have errored since. They pass on main, so this is a regression the stack introduced, not a pre-existing failure. Binds the real MicroScheduler.label_for over an empty combined-walk map rather than stubbing a return, so the fake cannot drift from it again. Co-Authored-By: Claude <noreply@anthropic.com> * Fix the drain test's batch fake for the per-walk rid grouping _postprocess_batch groups a stopped speculative batch's rids by real walk and discards through the batch, so the SimpleNamespace stand-in no longer answers for it: both parametrisations have errored on AttributeError since the first combined-walks commit. They pass on main. Uses the real ScheduledBatch, which is a plain dataclass here, rather than widening the namespace. Co-Authored-By: Claude <noreply@anthropic.com> * Combined walks: the per-walk cap's left-out rows still park The returned-rows change in #354 made this test assert an empty backlog, which holds only where the plain cap is what pushed a row out. Under a combined walk the per-walk composition leaves rows out instead, and those cannot run this step whichever walk holds them -- handing them back would re-pop them every pass, which is what the backlog exists to avoid. They park under the combined key, which is what this test is here to pin. Co-Authored-By: Claude <noreply@anthropic.com> * Bound the backlog scan by the step's cap, not the caller's leftover Integrating #368's bounded readiness scan, I passed `_room_left` as the bound. On this branch that is the caller's own walk's leftover and returns None when the caller holds no rows -- the ordinary case -- so the scan stayed unbounded and checked every backlogged row, which is what #368 set out to stop (96 sweeps at 128 in flight). `_scan_limit` is an upper bound instead: the builder caps a mixed step at the min over the walks actually in it, which nothing knows before the scan, and the max over the key's caps bounds that from above. Also carries walks/rid_walk_idx onto RouteArg, which #368 split into RouteArgRequired plus an optional half; the rebase left the fields on the required side while take_send_plan reads them off RouteArg, so the extension did not compile. Co-Authored-By: Claude <noreply@anthropic.com> * Rust: a (node, walk) nobody runs is an error, not an empty pop pop_walk_rids returned None for three different things: a rid that is not ready yet, a walk another worker owns, and a (node, walk) no worker graph anywhere runs. The scheduler treats None as "come back later", so the third case -- a structural error -- became an endless retry rather than a failure. PythonGraphRuntime has always raised on it, and AUTO prefers Rust, so the silent side was the one that ran. The reason the two runtimes drifted is that Python keys its index on the conductor's deployment-wide all_worker_graph_ids_to_* maps while Rust's node_owner holds only local graphs, so Rust could not tell somebody else's walk from nobody's. It does have the remote metadata, though -- the constructor appends it to all_worker_graphs -- so runs_anywhere draws the same line Python gets for free, and only a pair that is genuinely nowhere raises. The two caller-side invariants that also returned None (rids and walk indices differing in length, a walk index past the walks given) now raise too, matching speculate_node a few lines down. Checking the index range once up front also lets both loops index instead of silently skipping rows. Co-Authored-By: Claude <noreply@anthropic.com> * Prefix chains: the batch form takes a walk per row extend_prefix_chains_batch took one graph_walk for the whole step, so a combined walk could not use it and fell through to a call per row. Nothing in it was ever shared across the batch except the lock -- the walk is only read inside a per-rid _keyed_label -- so passing step_context.walk_of instead keys each row under its own walk and the mixed step takes the batched path like any other. That removes the gate and its slow fallback rather than keeping a second path that only mixed steps exercise. A lookup, not a dict built per step: the per-row path already called walk_of for every row, so this costs nothing it did not already cost. Co-Authored-By: Claude <noreply@anthropic.com> * BAGEL: guidance labels are per row, not per step requires_cfg is read with any() over the batch, and the labels it picks were applied to every row. That was right while a step held a single walk: every row of it was guided or none was. Under the combined walk an unguided decode row can share a mixed_text step with a guided prefill, and it was handed that row's cfg_img segments -- so its tokens ran in a branch it never asked for. With T2T requests decoding beside seeded T2I, merceod saw one of six images and one of six greedy texts change, where on main both matched the alone runs. The labels are now derived per row from that row's own resource_step_info, and `segments` keeps only the (row, label) pairs the row owns. segments is the ordering authority for KVStep, AttentionStep and PositionStep alike, so all three are scoped by this one change; per_label_pos_ids had to move to the same `label in mine` predicate, since those ids are concatenated positionally against segments and a disagreement would have a row reading another's positions. Not fixed here: the main -> cfg_text pre-fork still reserves and copies for every row, because scoping it needs the fork_rids field #365 adds to KVStep. That is waste, not corruption -- `_reserve_fork` runs per rid independently of segments, so the target stays correctly sized and nothing reads a cfg branch the row now has no segments in. Co-Authored-By: Claude <noreply@anthropic.com> * more instrumentation [temporary] + allow inline output for mixed walks * Prefill wrapper: stage an inline plan's host arrays page-locked A packed (prefill or mixed) step plans inline on the gpu thread, behind the step still in flight on the same stream. FlashInfer's prefill plan copies the host index arrays out of pageable memory, and the page indices blocking regardless, so the plan waited for that step to finish: 1.8-4.3 ms of attention plan per mixed step on Qwen3.5-0.8B, growing with the in-flight step's length. Stage them in the wrapper's own page-locked buffers, copy the indices into the captured buffer non-blocking as the decode wrapper does, and pass the host lengths so the kv-lens copy is asynchronous too. An event recorded after the plan guards the staging buffers against the next restage. Attention plan per mixed step: 0.7-1.1 ms. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * GDN: run a mixed step's decode rows on the decode kernel A mixed step ran every row through the chunked GDN path, which gathers and scatters each row's whole recurrent state (~1 MB a row on Qwen3.5-0.8B) and runs the chunked kernel on rows needing one recurrent update: 18 x ~345 us per mixed step at 118 decode rows against 18 x ~93 us for a decode step. GDNMixedWrapper splits the step over static shapes: the decode kernel takes every row (one-token rows that carry state update in place by slot; the rest point at the null slot), and the chunked path takes at most `split_max_prefill_rows` prompt rows, moved to the front by a planned permutation; the conv splits the same way. A one-token prompt with no state stays on the chunked path, whose zeroing the decode kernel lacks. Index uploads go through page-locked staging, as the plan runs inline. The mixed walk gets its own captures (its decode-rows-plus-a-prompt layout, via PackedCudaGraphConfig.token_layout) instead of replaying prefill_text's, and the scheduler caps a mixed step's prompt rows at that bound (max_rows_in_combined -> BatchBuildRequest.walk_row_caps). Per mixed step at c128 (~113 rows): GPU time 22.6 -> 10.5 ms, the postprocess wait on it 13.4 -> 0.5 ms. Conv output and state are identical to the chunked path; the delta-rule output differs by ~2 bf16 ulps. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Scheduling: speculation stays on its own walk; a mixed step takes the combined walk's cap A speculation chain merges only rows of its own real walk; another walk of the same combined walk is another target, which the fairness peek reports and the chain yields to. A yield-away from a walk leaves that walk's rows parked under the combined walk for the chain's next merge. Steps mix walks only when the scheduler builds one fresh. On Qwen3.5-0.8B (node2, trickled arrivals) this brought the branch from -5..-7% to within ~1% of main at c32/c128: prompts accumulate again before a step, and decode rows are no longer pushed out of the chain by a 32-row mixed cap. A mixed step takes the combined walk's own cap where it has one, else the smallest cap of the walks in it, so Qwen3.5's prefill ladder goes back to main's [1, 2, 4, 8, 16] beside the mixed walk's own captures. Removes the measurement-only MSTAR_SPEC_EXACT_WALK switch. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Attention: plan a mixed step's one-token rows on the decode wrapper FlashInferMixedWrapper: the decode wrapper takes every row of a mixed bucket (a one-token row attends its own pages, any other reads one sink-page token and its output is dropped), the prefill wrapper at most split_max_prefill_rows prompt rows moved to the front by a planned permutation, plus a tail row over the step's other tokens so its qo_indptr ends at the forward's token count. Host arrays are staged page-locked, as the plan runs inline. Output matches the packed prefill path within one bf16 ulp. Plan ~13% cheaper (493 vs 568 us at ~110 decode rows); attention kernel per full-attention layer 138 -> 73 us at 118 decode rows. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * GDN: build a mixed step's split with numpy The plan walked every token of the bucket in Python to build the permutation; vectorized, same outputs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Postprocess: a step of one walk skips the per-row walk lookups Prefix-chain extend takes the step's walk instead of a per-row walk_of, check_stop_batched gets the step's walk (Qwen answers nothing for a non-decode walk and reads no per-row walk for decode), and inline emit takes every row's value in one slice when the host rows are in the batch's order. At c128 decode these were +33 us (route_outputs), +28 us (check_stop_tail) and +18 us (check_stop) per step over main; after, +3 / +9 / +2. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Mixed steps: a row whose walk loops back on the device keeps doing so A mixed step asked which signals stay on the device under its combined label, which no walk of the submodule's answers, so every decode row in it routed its token as a tensor, and its next decode step read that tensor in and freed it instead of taking the sampler's slot master. Rows whose own walk keeps every output on the device now carry zero tensors in the route, are recorded for loop-back under that walk, and arrive empty in a speculated next step, as they would in a step of their walk alone. At c128: input tensors freed per decode step 33.3 -> 0.34 (main 0.46), per mixed step 44.5 -> 2.1. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Drop the temporary measurement instrumentation The per-walk phase labels, the forward stage/replay timers, the MSTAR_COMBINED_WALKS control and the gpu/spec/loop counters were for finding the combined walks' costs; the inline emit per real walk stays. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5: sort the cuda_graph_config import Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Scheduling: a serving config can let a speculation chain mix walks speculation_mixes_walks: per node, one bool for every walk of its combined walks, or per walk. On for a walk, a speculation chain on it merges ready rows of the other walks of its combined walk, the fairness peek and a yield-away from it skip the whole combined walk, a merge at its cap parks the other walks' rows for the next one, and its room counts every backlogged row under the mixed step's cap. Off by default: the chain keeps to its own walk, and steps mix walks only when the scheduler builds one fresh. Refuses nodes with no combined walk, walks outside them, and non-bool values. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Lazy publish: snapshot and finish each row of a mixed step under its own walk The eager publish already took each row's walk under a combined walk; the lazy path (on by default) handed every row the combined label, so a publisher keyed on the walk (BAGEL's per-(node, walk) KV labels) published nothing for a mixed step. The snapshot now runs per walk group and its finish per row's walk; a step of one walk takes one group, as before. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * In-flight flags: a step of one walk names its worker graph once cp/1 sent one worker graph id per row (a mixed step's rows can sit in different worker graphs), a 128-entry list built and extracted twice a step where main sent one id. A node in one walk lives in one worker graph, so a step of one walk sends that id once again, and set_in_flight (both runtimes) takes one id as every row's. A mixed step still sends one per row. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Steps of one walk skip the per-row walk bookkeeping Per-step costs cp/1 added on a plain decode step, measured or found in a pass over the new code, each kept for a step that mixes walks: - complete_and_route_batch keeps no per-rid walk map or index for one walk (route -3.3 us at 128 rows; now within noise of main) - KV: commit's prefix indexing reads the step's walk, and the chain extend keys every row with one (node, walk) - the recurrent pool hands its host slot lists over as built, not copied - Qwen3.5 declare_step builds the per-row walks only for a non-decode step - prepare_inputs decides the prefix probe once for a step of one walk - a speculation chain of one walk builds no rid -> walk map - a TP follower pops a head of one walk with the plain pop - the yield check skips a backlog all of the excluded walk in one test Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Mixed steps: check a row's cap only against the walk it moves The composition built the set of walks twice and checked every walk's row cap for every row. Only the row's own walk count moves, unless the row brings a new walk, which still checks them all. Same rows left out as before; composing a 128-row mixed step on the main thread 90 -> 29 us. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Mixed splits: size the prefill half from the scheduler's caps A combined walk that runs its one-token rows on the decode kernel and its prompt rows on a prefill kernel is declared as a `CombinedSplit`: the combined walk and the member walks that bring prompt rows. The prefill half is captured for as many rows as the scheduler lets those walks hold in one step (`max_rows_in_combined`, an uncapped walk counting the bucket's rows), which the engine hands each resource in `EngineResourceInfo`. This replaces `split_walks` and `split_max_prefill_rows`, a second copy of the scheduler's cap that had to match it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Shorten the long comments the combined-walk work added Each keeps the fact a reader needs at that line and drops the history of how it came to be. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Packed captures: a shared layout of decode rows then one prompt `decode_rows_then_prompt` sits beside `distribute_tokens` in the CUDA graph config, for any mixed capture whose kernels split rows by span; Qwen3.5's mixed capture names it instead of its own lambda. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * KV: scope a step's forks to some of its rows KVStep.fork_rids names the rows a step's pre/post forks apply to (None is every row, as before), at admit reservation, plan and commit. A mixed step forks for its prompt rows only, beside decode rows that forked long ago. The preplan key includes it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL: a mixed step forks cfg_text for its prompt rows only A guided decode row sharing a step with a guided prompt was forked too, so its cfg_text was overwritten with main (a think_mode request's 19 tokens became 205). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Docs: combined walks and how a mixed step is capped Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * BAGEL: a partly guided step runs guidance on its guided rows only With per-row guidance labels, a step whose rows are only partly guided (a T2I prompt beside T2T decode rows) planned cfg_img for the guided rows but ran every row's tokens through it, and FlashInfer refused the step. Such a step now takes its own capture key, `PARTIAL_CFG`, which no capture holds, so it runs eagerly, and its cfg_img pass runs on the guided rows' tokens. An image prefill's commit-time fork is scoped to its guided rows too. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Revert "BAGEL: a partly guided step runs guidance on its guided rows only" This reverts commit 8ec2fa4. * BAGEL: guidance labels are per step again, as on main Reverts 0f9c41c. A step with guided and unguided rows runs every row under the guidance labels, as a decode step mixing them already did on main; scoping the labels per row is left for its own change. The pre-fork stays scoped to the step's prompt rows (`fork_rids`), which is what keeps a decode row sharing a mixed step from having its cfg_text overwritten. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * CI: run the combined-walk CPU tests in CPU Core Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * scheduling fixes * CI: CPU Core can't import Qwen3-Omni without flashinfer * Cosmos3 test: a real StepContext, which carries the per-row walks the engine reads Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Stephen Duan <s129652@cs.washington.edu> Co-authored-by: Gaurav-Shah05 <73552538+Gaurav-Shah05@users.noreply.github.com>

Two plan shapes passed ingest and then stalled or ended early: - a later step, outside a partition a stream moves, whose walk takes no input the conductor sends: the conductor moves a worker onto a step only by sending it inputs, so the worker kept running the old walk; - in a partition a stream moves (Connection.consumer_walk), a last step that runs once per chunk: that partition advances on every report, so its first chunk ended it. Both are model bugs, refused at ingest with the same 500 as the other plan checks. No shipped plan has either shape. Review findings: r1-code-review #1, #2. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

take_shard_edge now asks NodeResourceSpec.splits_steps, as the engine and sequence_parallel_refusals do, so every caller finds the shard one way (abstractions-3 #1). strategy_line's isinstance goes with it in #2. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Drops Resource.sp_strategy, strategy_line, SequenceShard.describe and the ragged _describe_strategy hook, four pieces that only fed one load-log line; the paged and ragged attentions log their strategy in bind_dependencies, where they decide it (abstractions-3 #2). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The ring was the one exchange that ignored the `run` closure: it reached back into the KV resource through a new attention-to-KV reference (`AttentionManager._kv`) and wrote and attended in `_ring`, calling the plan state with a signature of its own. It is now `sequence_shard.ring(run, q, k, v, r)` beside `ulysses` / `all_gather` / `head_slice`: it calls `run(q, k, v, tokens=r.index)`, `AttentionCallable._attend` forwards `tokens` to `write_kv`, and `attn.run` -> `_attend_placed` drives `RingAttentionWrapper.attend(queries, kv_cache_layer, new_keys)`, the same signature as `ContextParallelWrapper.attend`. Drops `_kv` and `_ring`. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…ons-4 #2) `RingAttentionWrapper.attend` now takes `queries` as a tuple (the shared plan-state signature), and the loop rebound that name to each round's partner queries; it is `theirs` now. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

NSagan271 added 6 commits February 25, 2026 17:31

in progress work on tensor communication

a92023a

continue refactor and refine tensor mannager

b01dac7

more refactor

5a71bcc

fix graph bugs

53cf47e

fix dummy model

58160aa

fix bigs in processing done fwd

5d684d3

merceod self-requested a review February 27, 2026 19:16

merceod merged commit 0902565 into main Feb 27, 2026

kamahori deleted the tensor-communication branch March 14, 2026 00:38

NSagan271 mentioned this pull request Oct 9, 2026

LTX-2.5: tensor- and sequence-parallel DiT and the two-stage distilled recipe #376

Merged

dghosef mentioned this pull request Oct 18, 2026

Client streams: the client as a streaming edge into a running request (full duplex 3/6) #593

Merged

NSagan271 mentioned this pull request Oct 18, 2026

[Chunked Prefill #2] Combined Walks #362

Merged

2 tasks done

Navigation Menu

Repository navigation

Tensor communication - #2

Tensor communication#2
merceod merged 6 commits into
mainfrom
tensor-communication

NSagan271 commented Feb 26, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

NSagan271 commented Feb 26, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants