Repository navigation
Conversation
The lane refused any node holding resources, so #978's codec, with its codec_state pool, can't run on it. A side step now admits, plans and commits its node's own recurrent pools inline on the GPU thread, under the lane's stream. One step in flight, so the next one plans after this one has landed and nothing needs double-buffering. - StepRunner knows the resources only side steps declare. Such a step runs between a chain step and its staged successor while the plan thread may still be staging it, so it neither promotes nor drops that stage (_beside_stage). - A refused admit goes back through _handle_admit_failure, as the chain's does, instead of postprocessing empty outputs. - _remove_request waits for the lane's admit to settle first. - A side node's resource teardowns run on the lane's stream (Engine.set_side_stream): a released slot is zeroed there, so the next stream's first write can't land ahead of the zero. - _side_lane_keys opens only recurrent pools of the node's own with no KV-keyed snapshots; KV, attention, positions, samplers, shared pools and resident walks still refuse.
The four colocated configs set side_stream: true on the Codec group again; #978 took it out because the lane couldn't take codec_state. The split configs can't set it (the codec is alone on its worker).
mstar-review
Bot
added
area/engine
area/resources
area/worker
model/qwen3_tts
size/M
labels
dghosef
left a comment
dghosef
left a comment
There was a problem hiding this comment.
Nice to have the codec back beside the chain. One crash to fix before this goes in (inline).
If you get a chance, a quick H200 run and a short soak with clients cancelling mid-stream would be nice. The cancel soak is what would catch the inline one.
Working on a followup that will generalize this.
| if stream is None: | ||
| self._resources[key].remove_request(rid) | ||
| continue | ||
| with stream: |
There was a problem hiding this comment.
torch.Stream's __enter__ keeps its state on the stream object, not per thread, and a second enter raises ("Stream's context should not be initialized"). The GPU thread holds this same object entered for the whole side step (side_steps.py:76), so a teardown that overlaps it (a cancel, a failed request, or the remove after a refused chain admit) raises here, and the main loop then fails every Talker row in flight and leaves the request half torn down. Can we not enter the lane's stream from the main thread here? torch.accelerator.set_stream around the remove (and restoring the previous stream after) is enough for now. Would also be good if the fake Stream in test_side_steps.py refused a second enter, so a test catches this.
Gaurav-Shah05
left a comment
Gaurav-Shah05
left a comment
There was a problem hiding this comment.
On H200, running the codec beside the Talker is only faster while the host is the limit, and the change that removes that host cost turns the lane into a loss.
Setup
One H200, configs/qwen3tts_1p7b.yaml, 800 Seed-TTS English sentences with voice vivian, closed loop, two boots per arm alternated in one job:
| Lane off | Lane on | ||
|---|---|---|---|
main, c64 |
443.7, 438.6 | 493.7, 495.3 | 1.12x |
main, c128 |
580.7, 575.3 | 501.8, 498.0 | 0.86x |
| one codec step for every width, c64 | 535.7, 540.9 | 481.4, 477.3 | 0.89x |
| one codec step for every width, c128 | 670.2, 673.1 | 615.6, 614.9 | 0.92x |
On main at c128, first audio p50 goes from 120 ms to over a second with the lane on. Seeded audio is byte-identical in every arm, and WER doesn't change.
Analysis
mainis host-bound at c64 because the codec takes one step per chunk width, about 1.75 codec steps for every Talker step, and each step costs the worker's main thread a fixed amount. The lane's gain comes from taking those steps off the chain. Batching every width into one codec step removes the same cost.- I ran a local copy similar to Joe's 0447fd7, and with it the GPU is the limit at 92 to 95% busy. From there the lane has no idle GPU time to fill. In phase timing at c64, a Talker decode step takes 7.0 ms of GPU time with the codec beside it, against 4.8 ms on the chain, and the lane runs 0.86 codec steps for every Talker step, against 0.48.
Your B300 runs stop at c32, where main is still host-bound. Could you rerun lane on against lane off, on top of Joe's 0447fd7, on B300 from c8 to c128?
If the lane wins there, it's a choice for the Blackwell configs. If it doesn't, I don't think the engine should take on side lanes that hold resources until some node shows a win from one.
Puts the Qwen3-TTS codec back on the side lane. #1073 ran codec windows beside the Talker's chain, but the lane refused any node with resources, and #978's incremental codec keeps per-stream decoder state in a
codec_staterecurrent pool, so #978 tookside_streamout of the configs and the codec has been on the chain since. This makes a side step admit, plan and commit a recurrent pool of its own node, so the two gains stack.Why it needed more than dropping the refusal
The engine already runs declare, admit, plan, forward and commit for a side step, inline on the GPU thread under the lane's stream. Four things the chain gets around that didn't hold for a step off the chain:
StepRunnerkeeps one staged pre-plan, the chain's next step, which the plan thread stages after N's commit. The side step launches at the top of the pass, so on the GPU thread it runs after N and before N+1, often while the plan thread is still staging N+1. Any step that isn't the staged one drops the stage on its admit (_drop_stale_preplan->clear_preplanon every resource) and resets the record on its plan. Done by the side step, that races the plan thread and can leave some resources staged and others not._handle_admit_failure(push back, hold, limit). The lane's finish would have postprocessed the empty outputs.RecurrentStatePool.reset_request->_zero_slot) on the main thread's stream, the chain's, behind whatever chain steps are queued. The free list is LIFO, so the next new stream's first window can take that slot and run on the lane's stream before the zero lands, which would wipe that stream's state after its first window.The stage race and the zero race are read off the code paths above, not seen in a run; the tests below pin each fix, and no run here logged a dropped stage.
Not needed: a double buffer. The lane has one step in flight and launches the next only after the last one landed, so the next plan's copies can't overwrite buffers still being read. Plan stays inline.
What changed
mstar/engine/resources/runner.py:StepRunnertakesside_resources, the resources only side steps declare._beside_stage(step): a step whose keys are all such resources neither promotes nor drops the stage inadmit/plan, and its refusal doesn't clear it.set_side_stream(keys, stream), andremove_requesttears those keys down under that stream.mstar/engine/engine.py:load_modelhands the runner the side nodes' resources;Engine.set_side_stream(node, stream). Call sites otherwise unchanged.mstar/worker/side_steps.py:launchregisters the lane's stream with the engine when it makes it;finishsends a refused step to_handle_admit_failureand routes nothing.mstar/worker/worker.py:_settle_side_admit(), called in_remove_requestbefore the engine teardown, waits for the lane's step to admit (_await_admit_settled).mstar/worker/engine_manager.py:_side_lane_keyslets through a recurrent state pool that only the side node declares and that keys no snapshots on a KV cache (_lane_holds). Everything else still refuses, with the reason in the error: a KV cache, attention, positions, a sampler, a pool another node shares or one keyed on a KV cache, resident walks, plus everything it refused before (TP/SP, alone on its worker, no streams, chain walks on the node, combined walks, spans or stop checks, a second side node).docs/adding_models.rst,docs/serving.rst,docs/models.rst.Where the threads touch the pool now: the GPU thread admits, plans, commits and gathers/scatters, all under the lane's stream; the main thread ingests, reserves at readiness (
ask/reservethroughAdmission) and tears down (after the lane's admit settled, on its stream); the plan thread never stages it, since lane keys are never speculated or yield-away targets.Config and guards
configs/qwen3tts.yaml,qwen3tts_1p7b.yaml,qwen3tts_base.yaml,qwen3tts_voicedesign.yaml:side_stream: trueon the Codec group again. Drop it to put the codec back on the chain. The split configs can't set it._handle_admit_failurelike the chain's.Gate
2x B300 (sm_103), torch 2.12.1+cu132,
-O2, SHM arena, GPU 0, server on 8 pinned cores and the client on 4 others with idle spinners on both, a fresh server and Inductor cache per session, 200 Seed-TTS EN sentences, closed loop, 5 warmup waves. Main is origin/main 9c6f1ad (codec on the chain, as shipped); lane on is this branch's config; lane off is this branch withside_streamdropped. Sessions alternated main / on / off, then off / on / main. ITL is the gap between audio chunks at the client. Audible first audio (first chunk with a sample above -40 dBFS) matched first-byte TTFA within 2 ms in every cell, so it isn't a separate column.No failed request, nothing over 60 s of audio. *One request in that cell (sentence 183, normally ~6 s) ran to 34.7 s of audio, and a c=32 cell is only ~3 s of wall time, so it took the cell from ~300 to 232. That's lane off, i.e. the codec on the chain with no side stream; I haven't traced it.
Inline plan cost, from
MSTAR_PHASE_TIMINGon one untimed server per arm (GPU 1, c=8 and c=32, 2 passes each): the lane admits and plans each codec window on the GPU thread, 0.02 ms admit + 0.11 ms plan. Lane off pre-admits and pre-plans the same window on the plan thread for 1707 of its 1710 codec steps (0.012 + 0.122 ms) and promotes it in 0.006 ms. At c=8 there were 1798 codec steps over 4000 passes, so the inline plan is ~0.05 ms of a 5.2 ms pass. That's the most a pre-plan for the lane could get back, which is why it isn't built (an estimate, not measured).Soak and bytes
Tests
On 1b6e3f6 (the head minus a docstring rewording), main 9c6f1ad alongside:
ruff check .: clean.test_side_steps.py(48 -> 60): a side step on a realRecurrentStatePoolandStepRunnerthrough the engine's own declare/admit/plan/forward/commit takes a slot, plans it under the lane's stream, commits it (the next window reads what the first wrote), and on teardown frees it with the zero on the lane's stream; it leaves a staged chain pre-plan in place and the chain's step then promotes it; a raise in its admit or its plan fails only its rows and leaves the stage; a refused step goes to_handle_admit_failureand routes nothing; a full pool parks the row at readiness while the chain schedules, and it runs once a slot frees; a teardown waits for the lane's admit; the refusal list (KV-style resource, a pool next to it, a shared pool, a KV-keyed pool, resident walks) and a node with its own pool taking the lane.test_qwen3_tts_model.py: the four colocated configs take the lane, the split one doesn't, and the Codec holds only what the lane admits.test_eval_compare.pyandtest_eval_stress.pyleft out on both arms): 12316 tests, 11269 passed, 13 failed, 1034 skipped; main 12301 / 11254 / 13 / 1034. The same 13 ids fail on both, none only here. The watermark tests pass on both on this box. The serial pair (test_tp_async_lockstep_sim.py,test_piecewise_config_signature.py): 27 passed on both.test_side_steps,test_streams,test_accelerator_graph_capture,test_qwen3_tts_model,test_qwen3_tts_packed_replay,test_qwen3_tts_reference_audio,test_code_predictor_kernels,test_recurrent_gather_scatter,test_recurrent_padding_rows,test_recurrent_preplan: 264 passed.Not in this PR
side_stream_walks.side_streamis the switch if a live run there says otherwise.No new environment variable or dependency.