Allow Multiple Sampling Configs Per Node

[NSagan271]

What does this PR do?

Closes #199. Replaces the broken sample_with_config workaround for the Qwen3-Omni code predictor with the ability for a model to define one or more "auxiliary" samplers, which are instantiated as full Sampler or CudaGraphableSampler classes, and have their own (derived) seed and config.

Migrated Qwen3-Omni and Qwen3-TTS over to the new sampling logic.

How was it tested?

Added new unit test.

Integration test with Qwen3-Omni, Qwen3-TTS, and Orpheus, with and without TP where applicable: outputs are correct, and spot-checked Qwen3-Omni benchmarks.

Checklist

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

[stephen-dwq]

I think we can remove a lot of self.main is not None checks in handling MultiSamplingConfig if we add guard against that in the construction site in ./mstar/model/base.py, rather than permitting a propagation of None downstream from there, and adding spot-handling everywhere.

Otherwise, lgtm. Generally, I think we will have to abstract a lot of the structure as we bring in more models; some stuff (like the Sampling that was changed here) seems a little bit tight to a single model/architecture.

[NSagan271]

The Qwen3-TTS PR also got merged, so I migrated Qwen3-TTS over to the new sampler logic (and removed the custom fallback sampling logic)

[stephen-dwq]

LGTM.

In the future, especially with Rust migration, we may consider a rewrite to a general NSampler with separate gather and view, rather than combining the two in gather_for_request_ids which seems to be doing both. This would really screw with the API, so it would be a very low priority if we consider this improvement.

[merceod]

I had an observation from serving legs (not this PR's fault): qwen3-tts with an explicit seed kwarg is not reproducible across repeats even within one serverinstance. On this branch AND on main identically (different byte lengths per run). Since "seed" implies determinism to users, probably worth its own issue.

[merceod]

BLOCKING/CRITICAL: this call now always passes sampler_buffers=, but VJepa2ACRolloutPredictorSubmodule.get_piecewise_cuda_graph_configs (mstar/model/vjepa2/submodules.py lines 919) still has the old signature and this call sits outside the per-label try/except below, so nothing swallows it. Serving configs/vjepa2_ac.yaml on this branch kills the worker at engine warmup with "TypeError: got an unexpected keyword argument 'sampler_buffers'". The identical serve on main starts and answers requests. The base NodeSubmodule and qwen3_tts overrides were updated while vjepa2's was missed.

Fix should be pretty easy (accept and ignore the kwarg).

Also another note: no modular test builds vjepa2 piecewise runners, so the full suite is green on a branch that can't serve the model. Inthink a cheap signature-conformance test over get_piecewise_cuda_graph_configs overrides would close that gap.

[NSagan271]

This (and the "This call is for its side effect " comment below) revealed a design flaw in incorporating the sampler into the piecewise cuda graph runner. get_piecewise_cuda_graph_configs is not the natural place to inject the sampler; instead, I added a flag to the piecewise cuda graph config for whether sampling is used, and if so, the cuda graph runner gathers the sampler and passes it into the captured function directly. The Qwen3-TTS submodules.py now doesn't have to handle sampler buffer/gathering logic; the captured function just gets a sampler it can use.

I still added the test for future reference.

[merceod]

The aux SamplerBuffers are allocated once at engine build from the default config, and MultiSampler.set_config /MultiSamplerBuffers.update_request_config silently intersect/skip unknown labels. So a model whose aux label SET depended on model_kwargs would have late labels silently ignored (sample_aux does assert at use, but only then). The declare-at-startup contract is fine but we should be stating it here in the docstring (something like"the label set must not depend on model_kwargs, only the values may") as it would prevent a future issue.

[NSagan271]

Good catch, added the comment.

[merceod]

This call is for its side effect (staging params into thebuffers the captured samplers read) and discards the returned sampler. Itt almost reads like someone forgot to use the result. This is correct currently I think (the per-bucket capture samplers alias the same gather buffers, and the padded tail is stale-but-sliced-off as the comment says), but a stage_-named wrapper or a one-line comment on the buffer aliasing would make the intent clear.

Also subtle: padded_bs=batch_size stages only the real rows while the bucket samples its full captured width which is fine per the slicing argument, just worth the note.

[NSagan271]

Same fix as #201 (comment)

[merceod]

About the doctrine (nitpick): the signature changed to MultiSamplerBuffers -> MultiCudaGraphableSampler, so this is no longer compatible with old callers. Either drop the "shim" framing or delete the function if nothing in-tree still uses it.

[NSagan271]

Fixed, just removed the function because it was only used in one place.

[Gaurav-Shah05]

LGTM pending atindra's comments; the existing reviews cover everything I'd raise.

[merceod]

LGTM!

Two comments:

(1) needs a rebase since there is one trivial conflict in mstar/cli/main.py's DEFAULT_CONFIGS (wan22 vs sampling_test entries)

(2) that same CLI map registers "sampling_test": "sampling_test.yaml" but no such config file exists in the tree (at least I could not find it though I might be wrong lol) so mstar serve sampling_test would fail