[BugFix] Restore vllm_c IR op priority and torch.nn.RMSNorm for Qwen-Image (#4964)

[NumberWan]

Purpose

Fixes #4964

After the v0.24 rebase (#4709), Qwen-Image nightly perf regressed for two independent reasons. This PR lands both fixes in one place:

Part	File	What it fixes
Part 1	vllm_omni/platforms/cuda/platform.py	Platform IR priority: inductor-active → native-only RMSNorm dispatch
Part 2	vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py	Model backend: #3588 accidentally reverted #4074's torch.nn.RMSNorm

Short Summary

](Qwen-Image vs main, local H200): After this PR, serving latency_mean −26.9% (1.7263s → 1.2613s; 512², 10 steps, C=1) and test_qwen_image_matches_diffusers accuracy recovers (SSIM 0.9477 → 0.9786, PSNR 26.20 → 31.09 dB; FAIL → PASS).

---

Part 1 — Platform: restore vllm_c IR op priority

Problem

CudaOmniPlatform.get_default_ir_op_priority() followed upstream inductor-aware CUDA defaults. When inductor compilation is active, rms_norm / fused_add_rms_norm priority becomes ["native"] only.

For Qwen-Image diffusion, this forces repeated RMSNorm calls in the denoise loop onto a slower native decomposition path instead of vLLM CUDA (vllm_c) kernels, causing large nightly perf regressions across Qwen-Image workloads.

This part restores Omni's CUDA default IR op priority to ["vllm_c", "native"] while preserving existing VLLM_USE_OINK_OPS behavior.

Before / After Rebase

Before #4709, Omni CUDA platform always preferred vllm_c kernels, regardless of whether inductor compilation was active:

# pre-#4709 (e.g. c45ac74d)
default = ["vllm_c", "native"]  # force using vllm_c kernels

	Pre-#4709	Post-#4709 (main)	This PR (Part 1)
get_default_ir_op_priority()	["vllm_c", "native"] always	["native"] when inductor active	["vllm_c", "native"] always (= pre-#4709)
RMSNorm dispatch	_C::rms_norm	aten::mean decomposition	_C::rms_norm

#4709 copied upstream vLLM 0.24 logic: inductor on → prefer native for op fusion (reasonable for LLM AR). Qwen-Image already ran torch.compile pre-#4709 but kept vllm_c — diffusion denoise loops are RMSNorm-heavy and vllm_c is faster. Part 1 removes the using_inductor branch #4709 added.

Scope: CUDA platform default for all diffusion models using vLLM RMSNorm under inductor, not Qwen-Image-only. Models that need native-only already have per-model overrides (e.g. Cosmos3).

Root cause (Part 1)

	Pre-#4709	Post-#4709 (current main)
inductor active?	yes (compile on by default)	yes (same)
IR priority when inductor active	["vllm_c", "native"]	["native"] only
RMSNorm dispatch	_C::rms_norm / vllm::rms_norm_kernel	native decomposition (aten::mean x ~2000)
nightly perf	within baseline	large regression (#4964)

Part 1 — Test plan and result

Local validation on H200, Qwen-Image-2512, 512×512.

Kernel dispatch（profiler，4 steps，offline）

Variant	stage_0_gen_ms	aten::mean
a560ed18 (broken)	1288 ms	2042
Part 1 only (patched)	1080 ms (-16%)	114

Profiler shows RMSNorm dispatch switching from native decomposition back to _C::rms_norm.

E2E serving（C=1，20 steps，n=10）

Variant	latency_mean
c45ac74d (v0.23)	2.2847s
a560ed18 (v0.24, broken)	2.4960s (+9.2%)
Part 1 only (patched)	2.4502s (-1.8% vs broken; +7.2% vs v0.23)

Part 1 recovers the kernel path and most of the profiler gap, but E2E serving on v0.24 does not fully return to v0.23 levels — likely additional v0.24 rebase overhead outside RMSNorm.

Qwen-Image-2512 pixelwise vs diffusers (3 variants; skip P1+P2)

test_qwen_image_2512_matches_diffusers_pixelwise — 1664×928, 50 steps, seed=42. Thresholds: mean_abs_diff ≤ 3e-2, p99_abs_diff ≤ 4e-1.

Variant	pytest	mean_abs_diff	p99_abs_diff
v0.23 (c45ac74)	PASS	2.53e-2	0.322
main / post-#4709 (a560ed1)	PASS	1.72e-2	0.204
Part 1 only	PASS	1.80e-2	0.251

All three pass pixelwise thresholds vs diffusers. Part 1 IR priority change does not materially shift 2512 output vs main.

Non-Qwen A/B — FLUX.1-dev (justify global platform default)

Part 1 only — isolates platform IR priority (native-only vs vllm_c + native).
FLUX.1-dev uses vllm.model_executor.layers.layernorm.RMSNorm (same IR dispatch path as #4964).
Local validation: H200, v0.24 venv, single GPU, C=1.

Serving perf (512², 8 steps, n=5):

Variant	IR priority (inductor active)	latency_mean	Δ vs A
A — post-#4709 / current main	["native"] only	0.6166s	—
B — Part 1 (this PR)	["vllm_c", "native"]	0.5402s	−14.1%

A reproduces the post-#4709 platform default (#4709 inductor-aware logic). For diffusion RMSNorm-heavy models this path is slower than vllm_c.

Accuracy (FLUX.1-dev offline cross-variant, 512², 8 steps, seed=42): main vs Part 1 — SSIM 0.992, PSNR 41.27 dB (output unchanged).

Other models on the same vllm.model_executor.layers.layernorm.RMSNorm IR path (Part 1 applies; local A/B not run unless noted):

FLUX.1-dev, FLUX.2 / FLUX.2 Klein, Wan22S2V, HunyuanVideo15, SD3, OmniGen2, LongCat, Ernie, Ovis, LTX-2 / LTX-2.3 (local serving A/B: parity).

tests/platforms/test_cuda_ir_op_priority.py (CPU):

• CompilationMode.NONE / VLLM_COMPILE / STOCK_TORCH_COMPILE → rms_norm == ["vllm_c", "native"]
• VLLM_USE_OINK_OPS=1 → rms_norm == ["oink", "vllm_c", "native"]

---

Part 2 — Model: restore torch.nn.RMSNorm for Qwen-Image

Problem

Separate from Part 1: Qwen-Image transformer norm layers were switched back to vllm.RMSNorm in #3588 while adding AutoRound W4A16 — an accidental revert of #4074.

When	PR / issue	RMSNorm in Qwen-Image
Earlier regression	#3933	switched to omni RMSNorm → perf + accuracy issues
Fix	#4074	torch.nn.RMSNorm — fixed SSIM (#4029) and improved latency
W4A16 quant	#3588	incidentally switched back to vllm.RMSNorm while adding AutoRound support

#3588 did not require vLLM RMSNorm on norm layers themselves (W4A16 quant targets Linear weights). Part 2 restores #4074's backend while keeping #3588 quant paths intact.

torch.nn.RMSNorm bypasses vLLM IR / vllm_c dispatch on Qwen-Image's RMSNorm-heavy denoise loop — independent of Part 1's platform priority.

Change (Part 2)

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py only:

• norm_q / norm_k
• norm_added_q / norm_added_k
• txt_norm

# before (since #3588)
from vllm.model_executor.layers.layernorm import RMSNorm
self.norm_q = RMSNorm(...)

# after (Part 2, same as #4074)
self.norm_q = nn.RMSNorm(...)

Part 2 — Test plan and result

Local validation: H200, v0.24 venv, Qwen-Image, 512×512.

Serving benchmark (#4074-style: 10 steps, n=10, C=1)

Variant	latency_mean	throughput_qps	latency_p50
main (before this PR) — vllm.RMSNorm, native IR (#4709)	1.7263s	0.5792	1.7258s
P1+P2 (this PR) — torch.nn.RMSNorm, vllm_c IR	1.2613s	0.7927	1.2666s

This PR vs main: latency_mean −26.9% (1.7263s → 1.2613s).

Part 2 isolation (Part 1 held fixed; supplementary)

Variant	latency_mean	throughput_qps	latency_p50
vllm.RMSNorm (P1 fixed; Part 2 reverted)	1.4037s	0.7123	1.4142s
torch.nn.RMSNorm (P1+P2)	1.2613s	0.7927	1.2666s

Holding Part 1 constant, Part 2 alone: latency_mean −10.1%.

Accuracy vs diffusers (test_qwen_image_matches_diffusers, 512², 20 steps)

Variant	SSIM	PSNR (dB)	pytest
main (before this PR) — vllm.RMSNorm, native IR (#4709)	0.9477	26.20	FAIL
P1+P2 (this PR) — torch.nn.RMSNorm, vllm_c IR	0.9786	31.09	PASS

Main regresses on both Part 1 (IR path) and Part 2 (vllm.RMSNorm since #3588). This PR fixes both.

Part 2 isolation (Part 1 held fixed; supplementary)

Variant	SSIM	PSNR (dB)	pytest
vllm.RMSNorm (P1 fixed; Part 2 reverted)	0.9734	29.52	FAIL
torch.nn.RMSNorm (P1+P2)	0.9786	31.09	PASS

Holding Part 1 constant, Part 2 alone recovers accuracy vs diffusers (PSNR 29.52 → 31.09).

Unit / regression tests (Part 2)

Test	Result
test_qwen_image_matches_diffusers	PASSED
test_qwen_image_autoround_w4a16_load	PASSED
test_qwen_image_autoround_w4a16_generate	PASSED

Confirms W4A16 checkpoint still loads and generates correctly with torch.nn.RMSNorm.

---

How Part 1 and Part 2 relate

Qwen-Image denoise loop
│
├─ Layers still using vllm.RMSNorm
│   └─ Part 1: platform IR priority → prefer _C::rms_norm over aten::mean
│
└─ norm_q/k, norm_added_q/k, txt_norm  (5 sites)
    └─ Part 2: torch.nn.RMSNorm → bypasses vLLM IR entirely

• Part 1 = primary [Bug]: Nightly CI, Qwen-image-*, performance degradation exceeding 10% compared to the baseline was observed in some scenarios. #4964 mechanism fix (inductor → native-only); profiler evidence above.
• Part 2 = restore [BugFix]Qwen-Image performance regression by using torch RMSNorm(RMSNorm backend) #4074 accuracy/latency path; orthogonal to platform IR priority.
• Combined = management request to land both in #5009.

CI test plan (combined)

• tests/platforms/test_cuda_ir_op_priority.py (Part 1, local CPU)
• test_qwen_image_matches_diffusers (Part 2, local)
• test_qwen_image_autoround_w4a16_load / _generate (Part 2, local)

BEFORE SUBMITTING: read CONTRIBUTING.md and run the precheck-pr skill with the code agent for a self-check against project conventions.
(anything written below this line will be removed by GitHub Actions)

[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.

[NumberWan]

@SamitHuang @Gaohan123 PTAL

[Gaohan123]

Please add a UT to protect it

[Gaohan123]

Please supplement performance comparison with P1+P2 and main

[NumberWan]

Please supplement performance comparison with P1+P2 and main

The supplement performance has been added .
A short summary:
After this PR, Qwen-image latency_mean faster 26.9% and increased accuracy in (test_qwen_image_matches_diffusers).

Details have been added in the PR description

[Gaohan123]

LGTM. Thanks