[BLANKETusers]

Purpose

Splits the MoE routing half out of #9527. The #7318 symptom there (greedy decode stopping after ~6 tokens under cudagraph_mode: FULL) is fixed by the runtime metadata refresh alone; this PR carries the separate routing defect.

The AR backbone routes tokens between the understanding FFN and gen_mlp (layers 14–27, moe_type = "ffn-14:28") from input_ids >= gen_vocab_start_index. That dispatch reads the mask on the host (.any() / .all() / .item()), which cannot be recorded into a device graph: dummy capture inputs are all zeros, so the branch taken at capture time is replayed for real inputs.

Two distinct problems:

1. The generation expert never runs. The multimodal graph path passes inputs_embeds only, so input_ids is None and gen_token_mask does not exist. Every token takes und_expert, in every mode — including enforce_eager: true, which is the stock mammoth_moda2.yaml. The checkpoint ships 42 gen_mlp.* tensors covering exactly layers 14–27 that are loaded and then never used.
2. No graph-safe dispatch exists under capture, so the recorded branch is whatever the dummy mask produced.

What this changes

• dispatch capturing streams through a torch.where blend that runs both experts and selects per token; eager fast paths untouched
• keep input_ids on the mm graph path via requires_raw_input_tokens on both AR wrappers and the top-level module
• fail closed when capture runs without a mask, instead of silently baking und_expert
• clamp both embedding lookups so the discarded side of the blend cannot index out of range
• warn once when an NPU/XPU backend exposes no capture probe, rather than always reporting "not capturing"

Why this is the correct behaviour

The reference implementation does the same thing, inside the model:

• modeling_mammothmoda2_qwen2_5_vl.py (initial commit df2154e, the arch this checkpoint declares — llm_config.model_type = mammothmoda2_qwen2_5_vl) and modeling_mammothmoda2_qwen3_vl.py both derive gen_token_mask = input_ids >= self.gen_vocab_start_index in model.forward whenever inputs_embeds is not supplied, then route every gen-vocab token through gen_mlp for the whole sequence.
• Neither variant resets the mask in prepare_inputs_for_generation, so decode keeps routing on token ids — the same semantics as this PR.
• The t2i prompt ends in gen-vocab ids (<|image start|>, <|image token|>), so prompt positions already route to gen_mlp in the reference.

main's all-und_expert behaviour is therefore the deviation, not the baseline.

Behaviour change and cost — needs to be weighed

• Default-path numerics change. With requires_raw_input_tokens, the stock eager config stops being identical to main: prompt gen tokens now go through gen_mlp. The first differing hidden row is the first gen-vocab prompt token. A reference-output comparison is still the open question; this PR does not claim the new image is better, only that it matches the reference code path.
• The blend costs about +25% AR decode. H200 (SM90, FA3, cudagraph_mode: FULL over mixed prefill-decode), greedy seed 42, 1024²: 36.0 s vs 28.7 s with the same tree minus the routing change. Both experts run for every graphed token, since the captured graph must have no data-dependent control flow. That dual-expert tax also hits understanding-only traffic under FULL.
• Short-term (this PR): keep und-only on eager (mammoth_moda2_ar.yaml). PIECEWISE does not avoid the dual-expert MoE path when MoE is captured; opt into FULL only via mammoth_moda2_ar_graph_full.yaml (graph settings only — sampling stays at base defaults). See the CUDA graphs section in recipes/MammothModa2/MammothModa2.md.
• Follow-up: mask-keyed all-und / all-gen / blend split graphs in [Perf][MammothModa2] Mask-keyed FULL cudagraph split for und/gen MoE experts #9865 (not dual-thread overlap inside one FULL blend).

Relationship to #7319

#7319 raises on capture with a non-None mask and relies on input_ids=None. That is a different routing design and the two cannot land together — one of them has to be chosen. This PR is the graph-safe-blend option; direction should be settled there.

Test plan

python -m pytest -q \
  tests/model_executor/models/mammoth_moda2/ \
  tests/worker/test_omni_gpu_model_runner.py

144 passed.

CPU regressions added: blend ≡ eager dispatch, per-token expert selection, capture without a mask fails closed, an all-false (dummy-capture) mask still runs both experts, requires_raw_input_tokens present on both wrappers, out-of-range gen ids clamp, and the missing-probe warning fires exactly once.

Evidence

H200, SM90, FlashAttention 3, mode: 0 + cudagraph_mode: FULL, 29 graphs over mixed prefill-decode:

tree	AR tokens	finish_reason	AR gen time
61a6ac259 (no refresh, no routing)	6	stop	0.12 s
routing change removed	4161	length	28.7 s
this PR	4161	length	36.0 s

The middle row is what #9527 now carries; it reaches the full grid without any of these changes. The table is the reason for the split and the source of the +25%.

Not covered here: NPU/XPU graph execution, model_runner: v2, and any comparison against reference-implementation output.

Cost note (understanding-only traffic)

Under FULL capture, layers 14–27 always run both experts via _moe_blend, including pure understanding/chat batches that never contain gen-vocab ids. The +25% AR decode figure measured on H200 t2i therefore understates the cost for und-only traffic on the same path; measuring one understanding-only workload before merge is still open.

Open before merge

• Reference-output comparison against the upstream implementation (even a single prompt / hidden-state or logits diff): stock eager now routes prompt gen tokens through gen_mlp, so default-path numerics change for existing deployments.
• Understanding-only FULL TPOT on H200 (see cost note above).
• NPU/XPU graph capture with a real is_current_stream_capturing probe (missing-probe path now warns once).

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	⚠️ Failed 2026-11-03T01:50:22.006021Z We couldn’t complete this review. Comment @codex review to try again.	cbbc503	Draft marked ready

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/ar_runtime.md, docs/design/module/model_integration.md.

Module owners: @tzhouam @fake0fan @Gaohan123

Routing: @tzhouam via module of the changed files, CODEOWNERS; @fake0fan via module of the changed files; @Gaohan123 via module of the changed files

@BLANKETusers, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

[amy-why-3459]

Thanks for splitting this out of #9527 and for the clear write-up of the cost. The capture-time dispatch reads well: _moe_blend never touches the mask on the host, the eager fast paths are untouched when not capturing, and the fail-closed guard turns the silent all-und_expert bake into a loud error. A few code-level notes:

1. Runner-level coverage for requires_raw_input_tokens. test_ar_wrappers_keep_raw_input_tokens only checks the class attribute. The whole fix depends on the runner actually passing input_ids alongside inputs_embeds on the mm graph path (and the top-level module attribute set in __init__ being the one the runner reads). A small test through tests/worker/test_omni_gpu_model_runner.py that asserts input_ids is non-None in _model_forward for this model, both during dummy capture and on a normal step, would protect against a future runner refactor silently dropping it again, which is exactly the original defect.

2. The fail-closed guard can now abort engine startup. Any capture path that still reaches moe_forward without a mask (for example model_runner: v2, or a PP rank whose intermediate tensors don't carry gen_token_mask) will now raise during graph capture instead of running degraded. That's arguably the right outcome, but it would help to (a) confirm the PP path: make_empty_intermediate_tensors seeds gen_token_mask, so please check that the runner's static intermediate buffers include that key and get it copied on replay, otherwise rank > 0 replays the zeros mask from capture; and (b) make the error message point users to the workaround (enforce_eager: true or a non-FULL cudagraph_mode).

3. Qwen3 variant. MammothModa2Qwen3ARForConditionalGeneration gets requires_raw_input_tokens = True, but the tests and the H100 evidence only exercise the Qwen2.5-VL backbone. Could you confirm the Qwen3 language model derives gen_token_mask from input_ids in the same way, so the new attribute isn't a no-op there and capture doesn't hit the fail-closed branch?

4. Blend cost also applies to understanding-only traffic. Under FULL capture, every graphed batch now runs both experts on layers 14–27, including pure understanding/chat requests that never contain gen-vocab ids. It would be worth stating that in the PR description next to the +25% t2i number, and ideally measuring one understanding-only workload, so the cost is visible to whoever weighs this.

5. Nit: clamp comment. In get_input_embeddings, the comment says the out-of-range side "maps to index 0", but the base lookup clamps gen ids to base_vocab_size - 1, not 0. Only the gen side maps to 0. Harmless since those rows are discarded by where(), but the comment should match.

6. Default-path numerics. As the description notes, the stock eager config now routes prompt gen tokens through gen_mlp, so output changes for every existing deployment. A reference-output comparison (even a single prompt, hidden-state or logits diff against the reference implementation) before this lands would make that change much easier to sign off on.

[BLANKETusers]

@amy-why-3459 thanks for the split review — addressed in cbbc503b2 / PR body update:

1. Runner-level requires_raw_input_tokens coverage. Added test_prepare_mm_inputs_keeps_input_ids_when_requires_raw_input_tokens in tests/worker/test_omni_gpu_model_runner.py. It goes through GPUModelRunner._prepare_mm_inputs and asserts input_ids stays non-None when the flag is set (and is dropped when unset), which is the exact mm-graph contract this fix depends on.

2. Fail-closed abort + PP. Error message now points at enforce_eager: true / non-FULL cudagraph_mode. Confirmed make_empty_intermediate_tensors seeds gen_token_mask (comment + test_empty_intermediate_tensors_seed_gen_token_mask); first rank overwrites it from input_ids when requires_raw_input_tokens holds, so rank > 0 does not replay a missing key.

3. Qwen3 variant. MammothModa2Qwen3ForCausalLM subclasses the Qwen2 MoE host and reuses the same forward / gen_token_mask = input_ids >= gen_vocab_start_index path; docstring + issubclass assert record that. Hardware evidence is still Qwen2.5-VL only; I am not claiming a Qwen3 online A/B here.

4. Understanding-only blend cost. Accepted — called out in the PR description next to the t2i +25% number. Measuring one und-only FULL workload on H200 remains open before merge.

5. Clamp comment. Fixed: base side clamps to base_vocab_size - 1, gen side to [0, num_embeddings - 1].

6. Default-path numerics / reference comparison. Kept as an explicit open item in the PR body. Happy to add a single-prompt hidden-state / logits diff against the reference implementation before this leaves draft; not claiming sign-off on image quality yet.

Relationship to #9527: that PR is the #7318 refresh-only fix (stale generated_len under FULL replay). This PR is the separate routing defect (input_ids dropped → host MoE control flow / wrong expert under capture) with its own behaviour and cost change.

[amy-why-3459]

Re-reviewed only cbbc503b on top of 6367c66e. The error message, clamp comment, Qwen3 docstring and the _prepare_mm_inputs runner test address notes 1, 2b, 3 and 5. Two small follow-ups on the new tests:

1. test_empty_intermediate_tensors_seed_gen_token_mask doesn't exercise the model. It builds an IntermediateTensors dict by hand with a gen_token_mask key and then asserts that the key is there, so it would still pass if _make_empty_intermediate_tensors stopped seeding the mask. Building the tensors through the model's own make_empty_intermediate_tensors factory and asserting on its output would make the test protect the PP path it is named after.

2. The issubclass assertion only covers inheritance. It confirms MammothModa2Qwen3ForCausalLM subclasses the Qwen2 host, but not that the Qwen3 subclass doesn't override forward / get_input_embeddings. An assertion such as MammothModa2Qwen3ForCausalLM.forward is MammothModa2Qwen2ForCausalLM.forward (and the same for get_input_embeddings) would catch a future override that drops the mask derivation.

The remaining open items listed in the PR body (an understanding-only FULL workload measurement and a reference logits or hidden-state comparison for the default eager path) are still the main things worth seeing before this leaves draft.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[BLANKETusers]

Review follow-up (InferMatrixCopilot Direct @ 534ffe82f)

CI is green, but the PR is still CONFLICTING vs main after #9527 landed. The MoE graph-safe blend path itself looks sound; the remaining work is merge hygiene plus the two weak tests already flagged.

Must keep both flags when resolving the conflict in MammothModa2ForConditionalGeneration.__init__: requires_raw_input_tokens (this PR) and supports_omni_decode_step_metadata (#9527). The runner reads both on the top-level module — picking only one side re-breaks either MoE routing or the FULL-replay metadata refresh. Same for tests/worker/test_omni_gpu_model_runner.py: keep both suites.

Inline notes below cover the still-open test pins. Separately still open from the PR body (sign-off, not code blockers): reference logits/hidden vs upstream, understanding-only FULL TPOT, NPU/XPU real capture probe.

[BLANKETusers]

Addressed the review follow-ups in 1581fc699:

1. Rebased onto main ([Bugfix][MammothModa2] Refresh t2i AR-grid constraints under FULL graph replay #9527); top-level module now sets both requires_raw_input_tokens and supports_omni_decode_step_metadata. Runner tests for both paths kept.
2. PP empty-tensor test now goes through make_empty_intermediate_tensors / shared make_empty_intermediate_tensors_with_gen_mask.
3. Qwen3 pins forward / get_input_embeddings identity to the Qwen2 MoE host.

PR is MERGEABLE again (merge state may still show blocked on reviews/CI).

[amy-why-3459]

Re-reviewed the follow-up commits after the rebase onto main (0394f199 and 1ab722aa), on top of my earlier cbbc503b pass. Code paths only; I did not run this.

The Qwen3 identity asserts on forward / get_input_embeddings cover the earlier concern about a silent override, and pulling the seeding into make_empty_intermediate_tensors_with_gen_mask is a nice cleanup.

Two small follow-ups:

1. The PP test still doesn't go through the model's factory. test_empty_intermediate_tensors_seed_gen_token_mask now assigns its own lambda to lm.make_empty_intermediate_tensors that calls make_empty_intermediate_tensors_with_gen_mask directly, so it tests the helper, not what __init__ installs. If __init__ stopped wiring _make_empty_intermediate_tensors to the helper (or went back to a dict without the mask), this test would still pass. Building the LM through __init__ with a tiny config, or asserting on the closure the constructor actually installs, would close that gap.

2. Workaround wording in the capture error. The new message suggests FULL_DECODE_ONLY alongside PIECEWISE as a fallback. Since FULL_DECODE_ONLY still captures the decode step as a full graph, I'd expect it to reach the same moe_forward capture check whenever the mask is missing there (for example the model_runner v2 or PP cases the message mentions). If that's right, recommending enforce_eager or PIECEWISE only would avoid pointing users at a mode that fails the same way.

I didn't re-read the rebased 4ad60f45 / a832b480 beyond these two commits.

[vllm-omni-review-bot]

Omni ReviewBot: superseded

The CI failure noted on 1ab722aacf7d refers to an earlier head; the pull request now points at f7edadd050f6.

[amy-why-3459]

Looked at the new commit f7edadd only (error-message wording and the PP empty-tensor test now going through __init__).

Dropping FULL_DECODE_ONLY from the fail-closed message is right: decode is still captured as a full graph there, so it wasn't a real workaround, and the test now pins that the message offers only enforce_eager and PIECEWISE. Building MammothModa2Qwen2ForCausalLM through its constructor is also a real improvement over the hand-installed lambda, since it now exercises the closure __init__ actually assigns.

Two small notes on the test:

1. factory.__code__.co_names check is brittle. It passes today because the closure references the module-level helper by name, but a harmless refactor (e.g. functools.partial(make_empty_intermediate_tensors_with_gen_mask, hidden_size=...), or a local alias) would break it without any behaviour change. The tensor assertions below already prove the mask is seeded; if you want to keep a wiring check, monkeypatching m.make_empty_intermediate_tensors_with_gen_mask with a spy and asserting it was called with hidden_size == 8 would test the same thing by behaviour.

2. Stub coverage of __init__. The test stubs embeddings, LM head, logits processor, norm, make_layers and get_pp_group, with extra_gen_vocab=False. That's fine for this assertion, but it means a future dependency added to __init__ will fail here with an unrelated-looking error. Consider a one-line comment listing what has to stay stubbed, or covering the is_last_rank=False branch too, since that's the PP case the factory exists for.

The rest of the diff is line-length reflow and looks fine.

[BLANKETusers]

Pushed as 4a0c0e921 (17 graph-safety tests passed).

[amy-why-3459]

Incremental look at 4a0c0e921 (changes since a8940f22, test file only).

Both test notes from my last pass are addressed:

1. The co_names check is gone. The spy on make_empty_intermediate_tensors_with_gen_mask is installed before __init__ runs, and since the closure looks the helper up as a module global at call time, spy_calls == [(batch_size, hidden)] now checks the wiring by behaviour.
2. The docstring lists what has to stay stubbed, and the new (is_first_rank=True, is_last_rank=False) case exercises the non-last PP rank with real PPMissingLayer head/norm.

No further comments on this change. I only reviewed a8940f22..4a0c0e92 and did not run the tests.

[BLANKETusers]

Short-term und-only guidance + follow-up issue:

• Docs/presets in ac1764d82: default AR stays eager; added mammoth_moda2_ar_piecewise.yaml and documented FULL opt-in mammoth_moda2_ar_graph_full.yaml (dual-expert cost). Recipe section explains when to use which.
• Medium-term mask-keyed split graphs: [Core]Add Diffusion executor #865

[amy-why-3459]

Incremental look at ac1764d8 and 5ff71e49 only (docs, deploy presets, and removal of the evidence JPGs). I read the diff; I did not run anything. The routing code is unchanged since my last pass.

Dropping the JPGs from the tree makes sense, and nothing in the recipe still links to them. A few questions on the new guidance:

1. Is PIECEWISE really free of the dual-expert cost? The new recipe section, mammoth_moda2_ar.yaml, and mammoth_moda2_ar_piecewise.yaml all present PIECEWISE as the way to get graphs without FULL's dual-expert tax ("Eager MoE still runs a single expert; PIECEWISE does not require recording both..."). As far as I can tell from moe_forward, _is_graph_capturing() is true for any capture on the current stream, and under PIECEWISE the MoE MLPs sit inside the captured pieces (only attention is split out by default). If so, PIECEWISE capture would also go through _moe_blend and pay the same cost. The other possibility is that torch.compile traces moe_forward before capture and specializes the host-side .any() / .all() branch on the dummy inputs, which would be the original wrong-expert bug again. Could you add a PIECEWISE data point (und-only correctness plus AR decode throughput, compared with eager), or point me to where PIECEWISE avoids the capture path? If PIECEWISE also runs both experts, the docs, the two YAML comments, and the fail-closed RuntimeError text that recommends PIECEWISE should probably say so.

2. Greedy sampling in a user-facing preset. mammoth_moda2_ar_graph_full.yaml sets default_sampling_params: {temperature: 0.0, top_k: 1} "to isolate routing / graph replay". The recipe table recommends this file for "Text-to-image with FULL AR graphs", so anyone following the table would quietly get greedy AR decoding for t2i, which changes outputs compared with the default preset. I'd suggest keeping the preset limited to the graph settings and moving the greedy override to the repro instructions (or a separate repro-only file).

3. Minor: the FULL preset pins compilation_config.mode: 0, while the PIECEWISE preset leaves mode at the default. Since PIECEWISE needs the compile path, a short comment saying it relies on the default compilation mode would stop someone from copying mode: 0 across and ending up with an unsupported or silently downgraded combination.

Nothing blocking on the JPG removal itself.

[amy-why-3459]

Incremental look at 02b53a2f only (drops mammoth_moda2_ar_piecewise.yaml and turns PIECEWISE into a manual setting in the docs and YAML comments). I read the diff; I did not run anything.

Folding the extra preset into a comment is fine, and I don't see any remaining links to the deleted file in the recipe or the other presets. Two small things:

1. The manual instruction is incomplete as written. The recipe and the comments in mammoth_moda2.yaml / mammoth_moda2_ar.yaml now say "set cudagraph_mode: PIECEWISE on the AR stage yourself", but both files keep stage 0 at enforce_eager: true. With eager enforced, setting only cudagraph_mode won't capture any graphs, so someone following the text literally would see no change. The deleted preset had the full recipe (enforce_eager: false plus compilation_config.cudagraph_mode: PIECEWISE); it would help to put that snippet in the recipe section so the manual path is copy-pasteable.

2. My earlier question about whether PIECEWISE actually avoids the dual-expert path still applies, since PIECEWISE is still recommended in the recipe, both YAML comments, and the fail-closed RuntimeError text (and asserted in test_mammoth_moda2_moe_graph_safety.py). Removing the dedicated preset makes this lower-stakes, but if PIECEWISE capture also goes through _moe_blend, the wording should say so. The greedy default_sampling_params note on mammoth_moda2_ar_graph_full.yaml is unchanged in this commit as well.

Nothing blocking in this commit itself.

[BLANKETusers]

@amy-why-3459 thanks — all three land in d0dd1915b:

1. PIECEWISE dual-expert. Agreed. _is_graph_capturing() only probes the current stream, so any capture that runs through MoE (FULL or PIECEWISE pieces that include the MLPs) takes _moe_blend and pays both experts. Docs / YAML comments / the fail-closed RuntimeError were overclaiming; they now point und-only at enforce_eager only and no longer recommend PIECEWISE as a dual-expert workaround. I do not have a PIECEWISE A/B number yet — happy to add one if useful; the correctness path under capture is the same blend either way. Mask-keyed split graphs for und-only under capture stay in [Perf][MammothModa2] Mask-keyed FULL cudagraph split for und/gen MoE experts #9865.

2. Greedy in the FULL preset. Agreed it should not ship in a user-facing t2i preset. Removed default_sampling_params from mammoth_moda2_ar_graph_full.yaml; [Bug] MammothModa2 AR: und/gen MoE routing drops gen_mlp (eager) and bakes wrong expert under cudagraph FULL #9851 / [Bugfix][MammothModa2] Route AR MoE experts graph-safely under cudagraph capture #9764 repros can override temperature/top_k at the client.

3. mode: 0 on FULL. Left mode: 0 on the FULL preset (graphs without torch.compile) and commented why. The separate PIECEWISE deploy file is already gone after the earlier cleanup, so there is nothing left to copy mode: 0 onto by accident.

[amy-why-3459]

Incremental look at d0dd1915 only (PIECEWISE wording, greedy sampling removed from the FULL preset, mode: 0 comment). I read the diff; I did not run anything.

This addresses the three points from my last pass:

1. The recipe, both YAML headers, and the fail-closed RuntimeError now say plainly that any capture going through MoE takes _moe_blend, and they point und-only traffic at eager only. That matches how _is_graph_capturing() works, and it also removes the incomplete "set cudagraph_mode: PIECEWISE yourself" instruction I flagged before.
2. Dropping default_sampling_params from mammoth_moda2_ar_graph_full.yaml looks right; the preset now inherits the base sampling defaults, and the repro note says to override at the client.
3. The mode: 0 comment makes the choice clear.

One small test nit: assert "cudagraph_mode: PIECEWISE" not in message only guards against that exact old phrasing. Since the new message deliberately mentions PIECEWISE ("PIECEWISE is not a workaround"), a positive check such as assert "PIECEWISE is not a workaround" in message would pin the intended wording more directly. Not blocking.

Nothing blocking in this commit.