[Diffusion][Perf] Add Qwen-Image QK RoPE Triton path

[dongbo910220]

Purpose

This PR adds a Qwen-Image specific CUDA fast path for Q/K RMSNorm + RoPE.

In QwenImageCrossAttention, query/key currently go through RMSNorm and Qwen-Image interleaved RoPE as separate operations. This PR fuses that path into a Triton-backed custom op while keeping the existing implementation for torch.compile and unsupported inputs.

Correctness

python -m pytest -q tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py tests/diffusion/layers/test_fused_qk_norm_rope.py

Result:

all passed

Performance Methodology

Item	Configuration
GPU	NVIDIA RTX PRO 6000
Baseline	main
PR	this branch

Operator Benchmark

Packed-QKV-derived view:

Seq len	Native ms	PR ms	Speedup
4096	0.5921	0.1584	3.74x

E2E Benchmark

Qwen-Image, 50 steps:

Mode	Main ms	PR ms	Main / PR
--enforce-eager	26242.49	25602.75	1.0250x
torch.compile	23533.71	23497.84	1.0015x

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

Omni ReviewBot: no human activity for 29 days

@dongbo910220 this pull request has had no human commit, comment or review since 2026-08-10. Per repository policy it may be closed if it stays inactive.

To keep it moving, any one of these is enough: push an update, reply to the open blocker, or post the current plan and timeline.

[Dong1017]

Thanks for the PR. I verified the kernel on an L20 (torch 2.11, Triton 3.6, head 9642c99): 42 configs (seq 1→4096 × bf16/fp16/fp32 × head_dim 128/32, packed-QKV views) match a reference implementation of F.rms_norm + RotaryEmbedding(is_neox_style=False) with zero mismatches, and the op-level speedup holds (5.9x @ seq 1024, 7.6x @ seq 4096 vs eager).

A few findings:

1. Constexpr strides → ~220 ms recompile per new seq_len — vllm_omni/diffusion/models/qwen_image/fused_qk_norm_rope.py:114-136

All Q/K/out strides are tl.constexpr, so the specialization key includes q_stride_b = seq_len * heads * head_dim. Measured on L20: first call at a new seq_len costs ~220 ms vs 0.4 ms cached, and the kernel cache grows one entry per distinct seq_len. Qwen-Image requests vary in image resolution / text length, so each new shape hits this on first request. Fix: pass strides as runtime args; keep only head_dim / num_q_heads / head_block constexpr. (The sibling kernel in vllm_omni/diffusion/layers/fused_qk_norm_rope.py already does this.)

2. Third fused RMSNorm+RoPE implementation — vllm_omni/diffusion/models/qwen_image/fused_qk_norm_rope.py:1

Base main already has vllm_omni/diffusion/layers/fused_qk_norm_rope.py (#5990) with near-identical plumbing, and vllm_omni/model_executor/models/omnivoice/fused_qkv_rope.py:33 documents the policy: "Lift it out when a second model wants it." The RoPE conventions differ (interleaved + separate cos/sin vs half-split + packed table), so a full merge is a refactor — but could you either extend the shared module with an interleaved mode or explain in the PR why a third copy is preferable?

3. _apply_interleaved_rope duplicates an existing helper — vllm_omni/diffusion/models/qwen_image/fused_qk_norm_rope.py:44

Bitwise-identical to apply_rotary_emb_torch(..., interleaved=True) (vllm_omni/diffusion/layers/rope.py:20), which this PR's own benchmark already imports. Suggest deleting the local copy.

4. Unrelated header rewrite — qwen_image_transformer.py:2

The "vLLM project" → "vLLM-Omni project" copyright flip is unrelated to the change; please revert.

Minor: the PR description's test command lists tests/diffusion/layers/test_fused_qk_norm_rope.py (the pre-existing #5990 test, not part of this PR) — likely a rebase artifact.

[dongbo910220]

Thanks @Dong1017 for the detailed review and the L20 validation. I addressed the findings in efd3516c and f88a4b54.

1. Stride specialization: all Q/K/output/cos/sin strides and eps are now runtime kernel arguments. Only code-generation inputs remain tl.constexpr. On an NVIDIA RTX PRO 6000, the initial seq_len=128 call compiled in 434.23 ms; after that, the first seq_len=257 call took 0.090 ms and its repeat took 0.068 ms, confirming that a new sequence length no longer triggers a shape-specific compilation pause.

2. Shared implementation: I removed the Qwen-private implementation and moved the interleaved kernel, eager fallback, CustomOp registration, fake implementation, and validation into diffusion/layers/fused_qk_norm_rope.py as fused_qk_norm_rope_interleaved. The existing packed half-split MiniMax-H3/NPU path is unchanged. The shared interleaved API supports 3-D or 4-D Q/K tensors, different Q/K head counts, strided projection views, and activation-dtype or FP32 cosine/sine tables. I also added a shared-layer regression test using Boogu-Image's head_dim=120, 28/7 Q/K heads, FP32 RoPE tables, and a packed-QKV-derived view. This allows [Kernel][Boogu-Image] Fuse Q/K RMSNorm + interleaved RoPE via fused_qk_norm_rope #6982 to rebase onto this implementation, drop its duplicate shared-kernel changes, and retain the Boogu model adapter and integration.

3. Duplicate eager helper: the private helper has been removed; the fallback now reuses apply_rotary_emb_torch(..., interleaved=True).

4. Header: the unrelated copyright rewrite in qwen_image_transformer.py has been reverted to the pre-existing vLLM project line.

5. Test command: after the shared-layer refactor, tests/diffusion/layers/test_fused_qk_norm_rope.py is now part of this PR's diff and contains the new shared-layout and Boogu-geometry coverage, so it remains in the correctness command intentionally.

The focused local suites and the latest-main NVIDIA RTX PRO 6000 suites all passed. At seq_len=4096, the shared operator measured 0.1637 ms versus 0.5864 ms for the eager reference (3.58x), and the profiler shows both vllm_omni::fused_qk_norm_rope_interleaved and _qk_norm_rope_interleaved_kernel. Same-machine Qwen-Image E2E measured 26254.24 ms to 25688.95 ms in eager mode (1.0220x). Compile mode remains on the existing compiled path and measured 23598.02 ms versus 23597.36 ms, with identical output hashes.

[Dong1017]

Follow-up on the verification above: I ran the full test suites on an L20 (torch 2.13+cu130, Triton 3.6, head 9642c99, vllm 0.28.0 matching the branch's rebase target):

• tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py: 32/32 passed (includes fp32 eager fallback, torch.compile fullgraph capture, non-default head_dim 32, seq 1→4096, packed-QKV views)
• tests/diffusion/layers/test_fused_qk_norm_rope.py ([Kernel] Fuse Q/K RMSNorm and RoPE #5990 sibling kernel): 3/3 passed — the two custom-op registrations coexist without conflict
• full tests/diffusion/models/qwen_image/ directory: 63/63 passed

Update after the two fix commits (efd3516c0, f88a4b541): re-ran both test files at the new head — 38/38 passed — and re-measured the recompile behavior: first call at a new seq_len dropped from ~220 ms to ~0.3 ms, confirming the strides are now runtime arguments. The interleaved kernel now lives in the shared vllm_omni/diffusion/layers/fused_qk_norm_rope.py and the model-specific module is gone. All findings above are addressed.

LGTM 🚀

[princepride]

LGTM

[princepride]

@dongbo910220 CI failed, PTAL

[dongbo910220]

Thanks @princepride. I investigated all three failed Buildkite checks and found no failure attributable to this PR:

CUDA: only test_read_ownership_through_terminal_and_deferred_paths[unknown-False] failed on a five-second NIXL synchronization wait. The exact case passed repeatedly locally, and the complete NIXL connector test file also passed.
Intel: the agent ran out of disk space while cloning the repository, so no tests were executed.
NPU: the A5 layer-test job never acquired an agent and produced no test output; the corresponding A2, A3, and 310P jobs passed.

I merged the latest upstream main .