[Triton][Config] Add muls_add triton kernel and refactor AscendCompilationConfig (#5518)

### What this PR does / why we need it?
Add muls_add triton kernel with related fusion pass. What's more, this
PR refactors `AscendCompilationConfig` and delete `NpugraphExConfig`.

### Does this PR introduce _any_ user-facing change?
None

### How was this patch tested?
CI passed with new added test.


- vLLM version: v0.13.0
- vLLM main:
45c1ca1ca1

---------

Signed-off-by: whx-sjtu <2952154980@qq.com>

vllm_ascend/compilation/compiler_interface.py

@@ -30,7 +30,7 @@ from vllm.compilation.compiler_interface import CompilerInterface
 from vllm.config import VllmConfig
 from vllm.config.utils import Range
 
-from vllm_ascend.ascend_config import NpugraphExConfig, get_ascend_config
+from vllm_ascend.ascend_config import AscendCompilationConfig, get_ascend_config
 from vllm_ascend.utils import COMPILATION_PASS_KEY
 
 
@@ -71,7 +71,7 @@ def npugraph_ex_compile(
     example_inputs: list[Any],
     compiler_config: dict[str, Any],
     vllm_config: VllmConfig,
-    npugraph_ex_config: NpugraphExConfig,
+    ascend_compilation_config: AscendCompilationConfig,
     compile_range: Range,
     key: str | None = None,
 ) -> tuple[Callable | None, Any | None]:
@@ -83,7 +83,7 @@ def npugraph_ex_compile(
     config.mode = "reduce-overhead"
     # execute FX graph in eager mode before graph mode to optimize FX graph.
     config.debug.run_eagerly = True
-    if npugraph_ex_config.enable_static_kernel:
+    if ascend_compilation_config.enable_static_kernel:
         config.experimental_config.aclgraph._aclnn_static_shape_kernel = True
         # According to the cudagraph_capture_size configuration, set the shapes
         # that can trigger the compilation of static kernel. If this configuration is
@@ -117,8 +117,8 @@ class AscendCompiler(CompilerInterface):
     name = "AscendCompiler"
 
     def compute_hash(self, vllm_config: VllmConfig) -> str:
-        npugraph_ex_config = get_ascend_config().npugraph_ex_config
-        if npugraph_ex_config.enable:
+        npugraph_ex_enabled = get_ascend_config().ascend_compilation_config.enable_npugraph_ex
+        if npugraph_ex_enabled:
             self.vllm_config = vllm_config
         return vllm_config.compute_hash()
 
@@ -134,11 +134,11 @@ class AscendCompiler(CompilerInterface):
         # see https://github.com/pytorch/pytorch/issues/138980
         graph = copy.deepcopy(graph)
 
-        npugraph_ex_config = get_ascend_config().npugraph_ex_config
-        if npugraph_ex_config.enable:
+        ascend_compilation_config = get_ascend_config().ascend_compilation_config
+        if ascend_compilation_config.enable_npugraph_ex:
             assert hasattr(self, "vllm_config")
             return npugraph_ex_compile(
-                graph, example_inputs, compiler_config, self.vllm_config, npugraph_ex_config, compile_range, key
+                graph, example_inputs, compiler_config, self.vllm_config, ascend_compilation_config, compile_range, key
             )
         else:
             return fusion_pass_compile(graph, example_inputs, compiler_config, compile_range, key)

vllm_ascend/compilation/graph_fusion_pass_manager.py

@@ -64,6 +64,11 @@ class GraphFusionPassManager:
 
             self.passes.append(MatmulAllReduceAddRMSNormPass(config))
 
+        if self.ascend_compilation_config.get("fuse_muls_add", True):
+            from .passes.muls_add_pass import MulsAddFusionPass
+
+            self.passes.append(MulsAddFusionPass(config))
+
         if config.compilation_config.pass_config.enable_sp:
             from .passes.sequence_parallelism import AscendSequenceParallelismPass
 

vllm_ascend/compilation/passes/muls_add_pass.py

@@ -0,0 +1,117 @@
+#
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+# This file is a part of the vllm-ascend project.
+#
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from __future__ import annotations
+
+import torch
+from torch._inductor.pattern_matcher import PatternMatcherPass
+from vllm.config import VllmConfig
+from vllm.config.compilation import Range
+from vllm.logger import logger
+
+from vllm_ascend.compilation.passes.base_pattern import BasePattern
+from vllm_ascend.utils import vllm_version_is
+
+if vllm_version_is("0.15.0"):
+    from vllm.compilation.vllm_inductor_pass import VllmInductorPass  # type: ignore
+else:
+    from vllm.compilation.passes.vllm_inductor_pass import VllmInductorPass
+
+
+class MulsAddPattern(BasePattern):
+    """
+    Pattern that matches an element-wise mul + add sequence:
+        tmp = x * scale
+        out = tmp + y
+    and replaces it with a call to the muls_add_triton kernel.
+    """
+
+    def __init__(self, vllm_config: VllmConfig, scale: float = 1.0):
+        super().__init__(vllm_config)
+        self.scale = scale
+
+    def get_inputs(self) -> list[torch.Tensor]:
+        """
+        Generate example inputs for the MulsAddPattern.
+
+        The exact shapes are not important for pattern matching; they only
+        provide meta information for the pattern matcher.
+        """
+        x = torch.randn(2, 2048, device="npu", dtype=self.dtype)
+        y = torch.randn(2, 2048, device="npu", dtype=self.dtype)
+        # Only tensor inputs are needed here. The scalar scale is stored on the
+        # pattern instance (self.scale) instead of being passed as an input.
+        return [x, y]
+
+    def get_pattern(self):
+        def pattern(x: torch.Tensor, y: torch.Tensor):
+            """
+            Pattern for element-wise x * scale + y.
+            """
+            tmp = x * self.scale
+            out = tmp + y
+            return out
+
+        return pattern
+
+    def get_replacement(self):
+        def replacement(x: torch.Tensor, y: torch.Tensor):
+            """
+            Replacement that calls the muls_add_triton kernel using the
+            class-level scalar self.scale.
+            """
+            return torch.ops.vllm.muls_add(x, y, self.scale)
+
+        return replacement
+
+
+class MulsAddFusionPass(VllmInductorPass):
+    """
+    A fusion pass that replaces simple element-wise x * scale + y patterns
+    with the Triton-based muls_add_triton kernel on Ascend.
+    """
+
+    def __init__(self, vllm_config: VllmConfig):
+        super().__init__(vllm_config)
+        self.pattern_match_passes: PatternMatcherPass = PatternMatcherPass(pass_name="muls_add_fusion_pass")
+
+        # For now we enable this pass for all floating-point dtypes that the
+        # model is configured to use.
+        dtype = vllm_config.model_config.dtype
+        if dtype not in (torch.float16, torch.bfloat16, torch.float32):
+            logger.debug("MulsAdd fusion not enabled: unsupported dtype %s", dtype)
+            return
+
+        # Currently we only register a single pattern instance with a fixed
+        # scalar scale value. If needed, multiple instances with different
+        # scales can be added here in the future.
+        MulsAddPattern(vllm_config, scale=1.0).register(self.pattern_match_passes)
+
+    def __call__(self, graph: torch.fx.Graph) -> None:  # type: ignore[override]
+        self.begin()
+        self.matched_count = self.pattern_match_passes.apply(graph)
+        logger.debug("Fused %s muls_add patterns", self.matched_count)
+        self.end_and_log()
+
+    def is_applicable_for_range(self, compile_range: Range) -> bool:
+        """
+        Check if the pass is applicable for the current configuration.
+
+        For now, muls_add fusion is always allowed for the selected ranges.
+        This hook exists so that we can add more fine-grained range control
+        in the future if needed.
+        """
+        return True