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

vllm_musa/patches/manifest.py

@@ -209,6 +209,15 @@ def _series_entries() -> list:
         required=False,
         intent="route vllm._custom_ops rms_norm/rotary to MUSA-safe dflash fallbacks",
     ),
+    DivSpec(
+        id="vllm__model_executor__models__qwen3",
+        category="6",
+        path="vllm_musa/patches/vllm__model_executor__models__qwen3.patch.py",
+        upstream_path="vllm/model_executor/models/qwen3.py",
+        apply_phase="runtime",
+        required=False,
+        intent="fuse Qwen3 dense-decode q-norm+RoPE into one kernel (opt-in VLLM_MUSA_QWEN3_FUSED_QKNORM)",
+    ),
 ]
 
 # census of the runtime-override / vendored shadow modules (vllm_musa

vllm_musa/patches/vllm__model_executor__models__qwen3.patch.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MUSA cat-6 object patch: fuse Qwen3 dense-decode q-norm + neox RoPE into one
+TileLang kernel (opt-in via ``VLLM_MUSA_QWEN3_FUSED_QKNORM=1``).
+
+Dormant unless the env flag is set. When enabled it registers the opaque fused
+custom op and rebinds ``Qwen3Attention.__init__``/``forward``; the attention
+backend keeps ``reshape_and_cache`` (no backend surgery)."""
+
+import os
+
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+PATCHES: list = []
+
+
+def apply() -> None:
+    if os.environ.get("VLLM_MUSA_QWEN3_FUSED_QKNORM", "0") != "1":
+        return
+    # MUSA: the fused op is a MUSA-only TileLang kernel; ignore the flag on any
+    # non-MUSA runtime so an accidentally-set env var cannot break CPU/CUDA runs.
+    if not (hasattr(torch.version, "musa") and torch.version.musa is not None):
+        logger.debug("Skipping Qwen3 fused qk-norm patch: MUSA unavailable")
+        return
+    try:
+        from vllm.model_executor.models.qwen3 import Qwen3Attention
+
+        from vllm_musa.qwen3_jit.qk_norm_rope_kv import register_and_patch
+    except Exception as e:
+        logger.debug("Skipping Qwen3 fused qk-norm patch: %s", e)
+        return
+
+    logger.info("Enabling Qwen3 fused qk-norm+RoPE decode kernel (opt-in)")
+    register_and_patch(Qwen3Attention)

vllm_musa/qwen3_jit/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Lazy JIT helpers for Qwen3 dense MUSA decode fast paths."""