fix(omni): GLM-Image noise in dynamo disaggregated path

[ptarasiewiczNV]

Summary

Three commits fix GLM-Image in dynamo's disaggregated omni path:

• refactor(omni) — colocate glm_image.yaml under examples/backends/vllm/launch/stage_configs/ (matching single_stage_llm.yaml, qwen2_5_omni_pd.yaml) instead of reaching into the installed vllm_omni package. Removes a leaky abstraction that breaks the launch script against vllm-omni builds that don't ship that specific yaml.
• fix(omni) — attach mm_processor_kwargs={target_h, target_w} to the stage-0 OmniTextPrompt for image/video generation requests, so OmniInputPreprocessor._process_text routes through the HF multimodal processor path. Without this, AR-based image-gen models (e.g. GLM-Image) never emit their image-generation scaffold and the DiT stage denoises a collapsed token stream into textured noise.
• fix(omni) — place target h/w on the original_prompt dict (both as mm_processor_kwargs for the post-feat: JailedStream #3034 ar2diffusion and as top-level height/width for the 0.19.0rc1 ar2diffusion shipped in the dynamo runtime). Stage processors read this to upsample AR-generated prior tokens; without it they fall back to the 1024x1024 default and decouple from the requested size.

Evidence

zai-org/GLM-Image, prompt "a red apple on a white table", 2×A6000:

Default size (1024×1024) — MD5-identical output across the native and fixed disaggregated paths:

Path	Setup	Output MD5
dynamo disagg (before fix, 3 runs)	same container, same mount stack	36421ed1d1cfb07499fd166141f7998c — red striated noise
vllm-omni serve --omni (native)	same container, same local vllm-omni install	fa91343423d032e053327eb6047459b4 — coherent apple
dynamo disagg (after fix)	byte-identical to the native path	fa91343423d032e053327eb6047459b4 — coherent apple

Non-default size (512×512):

Path	Before first fix	After route image-gen through multimodal processor (commit 2)	After pass target h/w to stage processor (commit 3)
dynamo disagg	noise	RuntimeError at glm_image_transformer.py:883 (tensor a (1024) vs tensor b (4096) — DiT at requested 512 but AR prior upsampled to 1024-scale)	✅ coherent apple, ~60s

Root cause

The dynamo disagg path is a set of separate worker processes glued together by a custom router. It bypasses vllm-omni's OpenAI chat entrypoint entirely (goes through dynamo.vllm.omni.stage_router → dynamo.vllm.omni.stage_worker → AsyncOmni directly), so the upstream vllm-omni#3034 fix doesn't reach it. Two separate pieces of size-metadata have to make it across the stage boundary in the dynamo path and neither did:

1. Stage 0 (AR) — needs mm_processor_kwargs={target_h, target_w} on the engine prompt so OmniInputPreprocessor._process_text takes the multimodal branch and the HF processor emits GLM-Image's scaffold (<|image|>PROMPT<sop>H W<eop><sop>h w<eop><|dit_token_N|>). Without the scaffold, AR produces a handful of repeated VQ codes and DiT denoises them into noise.
2. Stage 1 (DiT via the ar2diffusion custom processor) — needs the target size on original_prompt (as mm_processor_kwargs["target_h"/"target_w"] on post-feat: JailedStream #3034 vllm-omni, or top-level height/width on 0.19.0rc1) so it slices and upsamples the AR prior token grid to the right latent shape. Without it, it defaults to 1024×1024 and produces a 64×64 prior regardless of the requested size, which then mismatches the DiT hidden-state shape for any non-1024 latent.

Scope

Minimal, two touched files (utils.py for the preprocessing plumbing, disagg_omni_glm_image.sh + new yaml for the colocation refactor). Image/video-generation paths only; chat / text / audio branches of parse_omni_request are untouched. Models whose HF processor and stage processor ignore target_h/target_w and height/width are unaffected.

Test plan

• Manual: disagg_omni_glm_image.sh + /v1/images/generations at 1024x1024 — output MD5-identical to vllm-omni serve.
• Manual: same at 512x512 — coherent image (was RuntimeError before).
• Regression: 1024x1024 re-tested after the 512 fix — still MD5-identical to the native baseline.
• Regression (reviewer): qwen omni agg / single-stage paths — confirm the chat/text branch of parse_omni_request and the audio handler are unchanged (not touched in the diff).
• CI: pre-commit hooks pass locally (black / ruff / codespell / yaml / shebangs all clean across all three commits).

🤖 Generated with Claude Code

[ayushag-nv]

Thanks for working on this