Fuse Qwen3 dense-decode q-norm + RoPE into one kernel (opt-in)

[yeahdongcn]

Summary

Qwen3 dense attention issues three separate GPU launches per decode step —
q_norm, k_norm, and neox rotary_embedding. All three are registered with
direct_register_custom_op (opaque to Dynamo, as rope.py documents), so
Inductor cannot fuse across them: they stay three distinct kernel launches
even under torch.compile + FULL_DECODE_ONLY. Hand-fusing them is therefore
not redundant with the compile path.

This adds one TileLang kernel that does q-norm + k-norm + neox RoPE in a single
launch, writing q and k in place. This is the 3→1 variant that ships wired
here: the attention backend keeps its own reshape_and_cache (no backend
surgery). Wiring is a category-6 runtime object patch, off by default,
enabled with VLLM_MUSA_QWEN3_FUSED_QKNORM=1.

A 4→1 variant that also writes the paged KV cache in the same kernel
(qk_norm_rope_kv_insert, write_kv=True, correctness-verified) ships in the
same module but is not wired — it needs a small backend patch so the
attention backend skips its own reshape_and_cache. It is left as a follow-up
(numbers for both variants below).

Why

On S5000 each tiny decode kernel has an irreducible ~3.5–6.4 µs device floor — a
1-token op moving a few KB cannot amortize grid-launch / scheduling
(8 KB / 1.6 TB/s ≈ 0.005 µs ≪ the ~5 µs floor). Fusing pays that launch floor
once instead of three times. Because the win comes from removing whole
kernel-node floors (not a faster inner loop), it survives CUDAGraph capture —
unlike an absorbed compute win the graph would hide.

Microbench (cold cache, 8 GB L2 flush, mate.bench_kineto, Qwen3-8B 32/8/128, bs1)

Per-launch device time (µs):

launch	q_norm	k_norm	rope	reshape_and_cache	sum
baseline (separate)	3.73	3.92	4.00	5.89	17.55

• Shipped 3→1 replaces q_norm + k_norm + rope = 11.65 µs with one fused
  launch (~9 µs); reshape_and_cache (5.89 µs) stays separate. Net decode-attn
  preamble ~17.55 µs → ~15.0 µs — the fusion pays 3 launch floors once instead of
  three times.
• 4→1 (not wired) fuses all four into one 9.53 µs launch → 17.55 → 9.53,
  1.8× on the op (up to ~2.2× at the op-level wall including launch overhead).

At bs≥16 the fused kernel becomes occupancy-bound and the edge narrows (32B bs16
→ ~1.0×); single-stream decode is the target. The op is launch/occupancy-bound at
M=1 (0.2–2.5 % of the 1600 GB/s ceiling) — the lever is op-count/fusion, not the
inner loop.

Correctness (fused vs torch fp32 reference, bf16 envelope atol 1e-2 rtol 2e-2)

shape	bs	q rel-err	k rel-err	v	busted elems	verdict
Qwen3-8B	1/4/16	0.0015	0.0014–0.0016	bit-exact	0	PASS
Qwen3-32B	1/4/16	0.0013–0.0015	0.0013–0.0015	bit-exact	0	PASS

v is a pure copy (bit-exact). No NaN/Inf, no poison (all-zero/constant) output.

Capture-safety + serve correctness (Qwen3-8B-FP8, TP1, compile ON + FULL_DECODE_ONLY)

• Capture-safe: confirmed. 3/3 FULL_DECODE_ONLY graphs captured, zero
  operation not permitted when stream is capturing, no SKIPPED INLINING. The
  1-op wrapper does no host→device scalar copy and writes q in place; JIT prewarm
  runs during the eager dummy pass before capture.
• Reaches the spawn worker: confirmed. As a cat-6 object patch, apply() runs
  via the vllm.general_plugins entry point in every worker
  (Applied object-patch vllm.model_executor.models.qwen3 in the EngineCore log),
  so it works under normal V1 multiprocessing.
• Serve correct: France→"Paris", Water→"oxygen"; baseline-vs-fused A/B gives
  identical first tokens (one prompt byte-identical). One prompt's long-tail greedy
  divergence is bf16-ULP sensitivity, bounded by the rel-err 0.0015 above.
• Env-gate confirmed dormant when the flag is unset.

E2E — honest read

These are profile-grounded projections, not measured serve A/B numbers —
because the delta is below the serve-benchmark noise floor (3–5 % same-session,
~15 % shared-node drift), so a serve A/B would be noise-dominated. The credible
evidence is the isolated cold-cache microbench. Grounded in the Qwen3-8B-FP8
decode profile (10.62 ms/step):

variant	per-layer saving	over 36 layers	≈ TPOT @ bs1
3→1 (shipped here)	~2–3 µs	~0.08–0.11 ms/step	~0.7–1.1 %
4→1 (follow-up, needs backend patch)	~8 µs	~0.29 ms/step	~2.6 %

So this PR delivers ~1 % TPOT at bs1 — small, positive, and capture-durable.
The larger ~2.6 % needs the 4→1 backend integration. Consistent with the finding
that vLLM-MUSA dense decode is already at/near the S5000 roofline: there is no
single big compute lever here, only a set of ~1–3 % launch-floor fusions, of
which this is one.

Test plan

• py_compile all four files (local + in-container)
• Cold-cache microbench, Qwen3-8B + Qwen3-32B shapes, bs {1,4,16}
• Correctness grid vs fp32 reference (q/k rel-err, v bit-exact, poison check)
• Serve Qwen3-8B-FP8 TP1, compile ON + FULL_DECODE_ONLY — capture-safe + coherent
• Baseline-vs-fused A/B (env-gate on/off) — identical first tokens
• Comment-style clean; manifest loads (cat-6 entry present)

Risk

Low. Off by default. When enabled, one attention class (Qwen3Attention) is
patched; q and k are mutated in place; the attention backend keeps
reshape_and_cache, so there is no backend surgery. Pure-Python + JIT kernel —
no native reinstall. Dormant for every non-Qwen3 model and whenever the env flag
is unset.

[augmentcode]

🤖 Augment PR Summary

Summary: This PR adds an optional fused decode fast path for Qwen3 dense attention on MUSA by combining q-norm, k-norm, and Neox RoPE into a single TileLang kernel launch.

Changes:

• Registers a new category-6 runtime object patch for vllm.model_executor.models.qwen3 in the patch manifest (opt-in).
• Adds a runtime patch module that enables the feature only when VLLM_MUSA_QWEN3_FUSED_QKNORM=1 is set.
• Introduces vllm_musa/qwen3_jit with a TileLang JIT kernel implementing fused RMSNorm (fp32 reduce) + Neox RoPE for Q/K, optionally with KV-cache insertion.
• Registers an opaque custom op (torch.ops.vllm.musa_fused_qk_norm_rope) so Dynamo/Inductor sees a single call and avoids attempting to inline JIT machinery.
• Monkeypatches Qwen3Attention.__init__ to pre-cast the rotary cache to bf16 and Qwen3Attention.forward to call the fused op and then reuse the existing attention backend (reshape_and_cache unchanged).

Technical Notes: The fused path is designed to be capture-safe (single GPU op per call) and remains off by default; it primarily targets launch overhead reduction in 1-token decode.

_{🤖 Was this summary useful? React with 👍 or 👎}

[augmentcode]

Review completed. 3 suggestions posted.

Comment augment review to trigger a new review at any time.