[Feature] Add rmsnorm dynamic mx quant fusion pass (vllm-project#10730)

### What this PR does / why we need it?
This PR introduces two new patterns to support the fusion of AddRMSNorm
and DynamicMxQuant operators, and two new patterns to support the fusion
of RMSNorm and DynamicMxQuant operators. After replacing the fusion
operators, the avg execution time drops from 7.541 µs to 3.525 µs for
the first combination, and from 8.098 µs to 4.625 µs for the second, and
model accuracy remains unaffected.

Please note that the fused operators introduced in this PR are only
supported on A5. The validation was performed with CANN 9.1.T560.B030
and FrameworkPTAdapter 26.1.0.B060.

### Does this PR introduce _any_ user-facing change?
N/A

### How was this patch tested?


- vLLM version: v0.23.0
- vLLM main:
vllm-project/vllm@967c5c3

---------

Signed-off-by: Sunwish <isunwish@foxmail.com>

Author: Sunwish

vllm_ascend/compilation/passes/norm_quant_fusion_pass.py

@@ -23,6 +23,10 @@
 from vllm.logger import logger
 
 from vllm_ascend.compilation.passes.base_pattern import BasePattern
+from vllm_ascend.device.mxfp_compat import (
+    is_add_rms_norm_dynamic_mx_quant_fusion_available,
+    is_rms_norm_dynamic_mx_quant_fusion_available,
+)
 from vllm_ascend.utils import enable_custom_op
 
 
@@ -474,6 +478,183 @@ def replacement(
         return replacement
 
 
+class AddRMSNormDynamicMXQuantPattern(BasePattern):
+    def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6):
+        super().__init__(vllm_config, eps)
+
+    def get_inputs(self):
+        """
+        Generate example inputs for the AddRMSNormDynamicMXQuant fusion pattern.
+        """
+        rms_norm_input = torch.randn(2, 64, device="npu", dtype=self.dtype)
+        residual = torch.randn(2, 64, device="npu", dtype=self.dtype)
+        rms_norm_weight = torch.randn(64, device="npu", dtype=self.dtype)
+        return [rms_norm_input, residual, rms_norm_weight]
+
+    def get_pattern(self):
+        def pattern(rms_norm_input: torch.Tensor, residual: torch.Tensor, rms_norm_weight: torch.Tensor):
+            """
+            Pattern for AddRMSNormDynamicMXQuant fusion.
+            """
+            output = torch.ops.npu.npu_add_rms_norm(rms_norm_input, residual, rms_norm_weight, self.eps)
+            out0 = output[0]
+            out1 = output[2]
+            quantized_output = torch.ops.npu.npu_dynamic_mx_quant(out0, dst_type=torch.float8_e4m3fn)
+            return quantized_output[0], quantized_output[1], out1
+
+        return pattern
+
+    def get_replacement(self):
+        def replacement(rms_norm_input: torch.Tensor, residual: torch.Tensor, rms_norm_weight: torch.Tensor):
+            """
+            Replacement for the AddRMSNormDynamicMXQuant fusion.
+            """
+            output = torch.ops.npu.npu_add_rms_norm_dynamic_mx_quant(
+                rms_norm_input,
+                residual,
+                rms_norm_weight,
+                epsilon=self.eps,
+                dst_type=torch.float8_e4m3fn,
+            )
+            return (
+                output[0],
+                output[2],
+                output[1],
+            )
+
+        return replacement
+
+
+class AddRMSNormDynamicMXQuantSPPattern(BasePattern):
+    def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6):
+        super().__init__(vllm_config, eps)
+
+    def get_inputs(self):
+        """
+        Generate example inputs for the AddRMSNormDynamicMXQuant fusion pattern.
+        """
+        rms_norm_input = torch.randn(2, 64, device="npu", dtype=self.dtype)
+        residual = torch.randn(2, 64, device="npu", dtype=self.dtype)
+        rms_norm_weight = torch.randn(64, device="npu", dtype=self.dtype)
+        return [rms_norm_input, residual, rms_norm_weight]
+
+    def get_pattern(self):
+        def pattern(rms_norm_input: torch.Tensor, residual: torch.Tensor, rms_norm_weight: torch.Tensor):
+            """
+            Pattern for AddRMSNormDynamicMXQuant fusion.
+            """
+            output = torch.ops.npu.npu_add_rms_norm(rms_norm_input, residual, rms_norm_weight, self.eps)
+            out0 = output[0]
+            out1 = output[2]
+            out0 = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(out0, True)
+            quantized_output = torch.ops.npu.npu_dynamic_mx_quant(out0, dst_type=torch.float8_e4m3fn)
+            return quantized_output[0], quantized_output[1], out1
+
+        return pattern
+
+    def get_replacement(self):
+        def replacement(rms_norm_input: torch.Tensor, residual: torch.Tensor, rms_norm_weight: torch.Tensor):
+            """
+            Replacement for the AddRMSNormDynamicMXQuant fusion.
+            """
+            output = torch.ops.npu.npu_add_rms_norm_dynamic_mx_quant(
+                rms_norm_input,
+                residual,
+                rms_norm_weight,
+                epsilon=self.eps,
+                dst_type=torch.float8_e4m3fn,
+            )
+            mxscale = output[2]
+            quantized_output = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(output[0], True)
+            mxscale = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(mxscale, True)
+            return quantized_output, mxscale, output[1]
+
+        return replacement
+
+
+class RMSNormDynamicMXQuantPattern(BasePattern):
+    def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6):
+        super().__init__(vllm_config, eps)
+
+    def get_inputs(self):
+        """
+        Generate example inputs for the RMSNormDynamicMXQuant fusion pattern.
+        """
+        rms_norm_input = torch.randn(2, 64, device="npu", dtype=self.dtype)
+        rms_norm_weight = torch.randn(64, device="npu", dtype=self.dtype)
+        return [rms_norm_input, rms_norm_weight]
+
+    def get_pattern(self):
+        def pattern(rms_norm_input: torch.Tensor, rms_norm_weight: torch.Tensor):
+            """
+            Pattern for RMSNormDynamicMXQuant fusion.
+            """
+            output = torch.ops.npu.npu_rms_norm(rms_norm_input, rms_norm_weight, self.eps)
+            out0 = output[0]
+            quantized_output = torch.ops.npu.npu_dynamic_mx_quant(out0, dst_type=torch.float8_e4m3fn)
+            return quantized_output[0], quantized_output[1]
+
+        return pattern
+
+    def get_replacement(self):
+        def replacement(rms_norm_input: torch.Tensor, rms_norm_weight: torch.Tensor):
+            """
+            Replacement for the RMSNormDynamicMXQuant fusion.
+            """
+            output = torch.ops.npu.npu_rms_norm_dynamic_mx_quant(
+                rms_norm_input,
+                rms_norm_weight,
+                epsilon=self.eps,
+                dst_type=torch.float8_e4m3fn,
+            )
+            return output[0], output[1]
+
+        return replacement
+
+
+class RMSNormDynamicMXQuantSPPattern(BasePattern):
+    def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6):
+        super().__init__(vllm_config, eps)
+
+    def get_inputs(self):
+        """
+        Generate example inputs for the RMSNormDynamicMXQuant fusion pattern.
+        """
+        rms_norm_input = torch.randn(2, 64, device="npu", dtype=self.dtype)
+        rms_norm_weight = torch.randn(64, device="npu", dtype=self.dtype)
+        return [rms_norm_input, rms_norm_weight]
+
+    def get_pattern(self):
+        def pattern(rms_norm_input: torch.Tensor, rms_norm_weight: torch.Tensor):
+            """
+            Pattern for RMSNormDynamicMXQuant fusion.
+            """
+            output = torch.ops.npu.npu_rms_norm(rms_norm_input, rms_norm_weight, self.eps)
+            out0 = output[0]
+            out0 = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(out0, True)
+            quantized_output = torch.ops.npu.npu_dynamic_mx_quant(out0, dst_type=torch.float8_e4m3fn)
+            return quantized_output[0], quantized_output[1]
+
+        return pattern
+
+    def get_replacement(self):
+        def replacement(rms_norm_input: torch.Tensor, rms_norm_weight: torch.Tensor):
+            """
+            Replacement for the RMSNormDynamicMXQuant fusion.
+            """
+            output = torch.ops.npu.npu_rms_norm_dynamic_mx_quant(
+                rms_norm_input,
+                rms_norm_weight,
+                epsilon=self.eps,
+                dst_type=torch.float8_e4m3fn,
+            )
+            quantized_output = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(output[0], True)
+            mxscale = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(output[1], True)
+            return quantized_output, mxscale
+
+        return replacement
+
+
 class AddRMSNormQuantFusionPass(VllmInductorPass):
     """
     A pass for fusing AddRMSNorm and W8A8 quantization operations on Ascend.
@@ -488,10 +669,29 @@ def __init__(self, vllm_config: VllmConfig):
             logger.debug("Quant fusion not enabled: unsupported dtype %s", dtype)
             return
 
+        dynamic_mx_quant_fusion_available = is_add_rms_norm_dynamic_mx_quant_fusion_available()
+        if not dynamic_mx_quant_fusion_available:
+            logger.debug(
+                "AddRMSNormDynamicMXQuant fusion not enabled: required MX symbols unavailable, or device isn't A5"
+            )
+
+        rms_norm_dynamic_mx_quant_fusion_available = is_rms_norm_dynamic_mx_quant_fusion_available()
+        if not rms_norm_dynamic_mx_quant_fusion_available:
+            logger.debug(
+                "RMSNormDynamicMXQuant fusion not enabled: required MX symbols unavailable, or device isn't A5"
+            )
+
         common_epsilons = [1e-5, 1e-6]
+
         for eps in common_epsilons:
             AddRMSNormDynamicQuantPattern(vllm_config, eps=eps).register(self.pattern_match_passes)
             AddRMSNormDynamicQuantSPPattern(vllm_config, eps=eps).register(self.pattern_match_passes)
+            if dynamic_mx_quant_fusion_available:
+                AddRMSNormDynamicMXQuantPattern(vllm_config, eps=eps).register(self.pattern_match_passes)
+                AddRMSNormDynamicMXQuantSPPattern(vllm_config, eps=eps).register(self.pattern_match_passes)
+            if rms_norm_dynamic_mx_quant_fusion_available:
+                RMSNormDynamicMXQuantPattern(vllm_config, eps=eps).register(self.pattern_match_passes)
+                RMSNormDynamicMXQuantSPPattern(vllm_config, eps=eps).register(self.pattern_match_passes)
             if enable_custom_op():
                 AddRMSNormQuantPattern(vllm_config, eps=eps).register(self.pattern_match_passes)
                 AddRMSNormQuantSPPattern(vllm_config, eps=eps).register(self.pattern_match_passes)

vllm_ascend/device/mxfp_compat.py

@@ -1,6 +1,8 @@
 import torch
 import torch_npu
 
+from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type
+
 # TODO(linfeng): Temporary compatibility shim for MXFP4/MXFP8 because current torch_npu
 # releases do not expose the required dtype attributes yet. Simplify or remove this
 # file after the torch_npu release in March 2026 includes those dtype symbols.
@@ -20,6 +22,10 @@ def _get_missing_symbols(symbols: tuple[str, ...]) -> list[str]:
     return [symbol for symbol in symbols if not hasattr(torch_npu, symbol)]
 
 
+def _is_dynamic_mx_quant_fusion_soc_supported() -> bool:
+    return get_ascend_device_type() == AscendDeviceType.A5
+
+
 def _ensure_symbols_available(feature: str, symbols: tuple[str, ...]) -> None:
     missing_symbols = _get_missing_symbols(symbols)
     if not missing_symbols:
@@ -31,6 +37,22 @@ def _ensure_symbols_available(feature: str, symbols: tuple[str, ...]) -> None:
     )
 
 
+def is_add_rms_norm_dynamic_mx_quant_fusion_available() -> bool:
+    return (
+        _is_dynamic_mx_quant_fusion_soc_supported()
+        and hasattr(torch, "float8_e4m3fn")
+        and not _get_missing_symbols(("npu_dynamic_mx_quant", "npu_add_rms_norm_dynamic_mx_quant"))
+    )
+
+
+def is_rms_norm_dynamic_mx_quant_fusion_available() -> bool:
+    return (
+        _is_dynamic_mx_quant_fusion_soc_supported()
+        and hasattr(torch, "float8_e4m3fn")
+        and not _get_missing_symbols(("npu_dynamic_mx_quant", "npu_rms_norm_dynamic_mx_quant"))
+    )
+
+
 def ensure_mxfp8_scale_dtype_available(feature: str) -> None:
     _ensure_symbols_available(feature, ("float8_e8m0fnu",))