Merge branch 'main' into test-cpu-32-hk-verify

Author: JavaPythonAIForBAT

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",))