[mickqian]

Motivation

On B300 the fused Q/K RMSNorm + complex RoPE + K/V packing kernel of Qwen-Image-2.1 took 89 us per launch inside a 1024x1024 request (1280 launches, ~114 ms of GPU time per request): one warp per 128-wide row, 8-byte accesses and 64-bit work indices kept it near 1.5 TB/s.

Modifications

• A half-warp holds a row: each lane owns 8 elements (16-byte loads/stores), i.e. two of the reference's 4-element groups. The aten mean order is kept exactly: four squares left to right per group, then the same tree (reference lane offsets 16, 8, 4, 2 become lane offsets 8, 4, 2, 1, offset 1 becomes the in-thread add of the two groups).
• 32-bit work indices (host check that the row count fits), 128-thread blocks.
• Tests: odd head counts (6 heads, a partial last warp) added to the existing bit-exactness tests.

Accuracy

Bit-identical. torch.equal against main's kernel on the production pack shape ([1, 4096, 32, 128], 27-token prefix), the kept-buffer edit shape, copy-V, odd heads and the single-tensor entry point; the existing eager-reference tests pass (test_complex_rope.py + test_qwen_image21_cuda.py: 70 passed on B300). End to end, the decoded PNG pixels of every request match main (md5 of the RGBA array, T2I and edit, 6+ requests each).

Benchmark (B300 SXM6, 1100 W, 1 GPU, b200-1gpu-resident-fa profile)

Kernel alone, [1, 4096, 32, 128] + 27-token prefix (CUDA events, 200 launches):

variant	us/launch
main (warp per row, 8 B, int64 indices)	65.5
4 rows per warp	68.8
32-bit indices only	54.0
half-warp per row, 16 B, 256 threads	38.6
this PR (half-warp per row, 16 B, 128 threads)	37.1

In a serving request (nsys): 89.4 -> 52.1 us per launch.

diffusion-bench-framework harness (devbox_run_cases.sh, fresh server per arm, two rounds in reversed order, vLLM-Omni 0da1d831 on the same box):

cell	main	this PR	vLLM-Omni
Qwen-Image-2.1 T2I 1024, single_e2e (s)	2.147 / 2.144	2.119 / 2.121	2.345 / 2.355
T2I throughput, 4 req @ conc 2 (req/s)	0.4671 / 0.4656	0.4744 / 0.4747	0.4362 / 0.4351
Edit 1024, single_e2e (s)	2.608 / 2.603	2.582 / 2.582	2.826 / 2.840

Combined with #47465, #47466 and #47467 (B300 harness, two rounds in reversed order)

final = main 15ef675 + the three PRs; vLLM-Omni 0da1d831 (blackwell-cudnn-static) on the same box; sglang b200-1gpu-resident-fa.

cell	main	final	final + --async-output-save	vLLM-Omni	lead
Qwen-Image-2.1 T2I 1024, single_e2e (s)	2.145 / 2.144	2.093 / 2.094	2.093 / 2.094	2.350 / 2.350	12.3%
T2I throughput, 4 req @ conc 2 (req/s)	0.4685 / 0.4685	0.4789 / 0.4789	0.4805 / 0.4810	0.4348 / 0.4348	10.1% (10.6% with async save)
Edit 1024, single_e2e (s)	2.610 / 2.603	2.525 / 2.527	-	2.827 / 2.835	12.1%

Decoded PNG pixels of every request match main (md5 of the RGBA array).

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ✅ Run #41068062426
Latest PR Test (Extra): ❌ Run #41068062357
Latest PR Test (AMD ROCm 10): ⏳ Run #41068062323

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[mickqian]

/tag-and-rerun-ci