init v0.23.0

Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>

vllm_ascend/compilation/acl_graph.py

@@ -2,9 +2,11 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import dataclasses
+import weakref
+from collections.abc import Callable
 from contextlib import ExitStack
 from dataclasses import dataclass
-from typing import Any, Callable, Optional
+from typing import Any
 from unittest.mock import patch
 
 import torch
@@ -18,18 +20,46 @@ from vllm.forward_context import BatchDescriptor, get_forward_context
 from vllm.logger import logger
 from vllm.platforms import current_platform
 
+from vllm_ascend.ascend_forward_context import _EXTRA_CTX
+
 from ..utils import weak_ref_tensors
 
+_acl_graph_wrappers: weakref.WeakSet[Any] = weakref.WeakSet()
+_STREAM_RESOURCE_ERROR_CODE = "207008"
+_STREAM_RESOURCE_ERROR_MARKERS = (
+    "insufficient_stream_resources",
+    "stream resources are insufficient",
+)
+_STREAM_RESOURCE_GUIDANCE = (
+    "ACL graph capture failed with a known stream-resource exhaustion "
+    "signature. Consider upgrading to a newer HDK/CANN stack, reducing "
+    "cudagraph_capture_sizes, lowering max_cudagraph_capture_size, preferring "
+    "FULL or FULL_DECODE_ONLY for mostly uniform decode workloads, or "
+    "temporarily disabling graph mode to confirm the failure is capture-related."
+)
+
+
+def _is_stream_resource_capture_error(exc: RuntimeError) -> bool:
+    message = str(exc)
+    lowered_message = message.lower()
+    has_error_code = _STREAM_RESOURCE_ERROR_CODE in message
+    has_stream_resource_marker = any(marker in lowered_message for marker in _STREAM_RESOURCE_ERROR_MARKERS)
+    return has_stream_resource_marker or (has_error_code and "stream resource" in lowered_message)
+
+
+def _raise_stream_resource_capture_error(exc: RuntimeError) -> None:
+    raise RuntimeError(f"{_STREAM_RESOURCE_GUIDANCE}\nOriginal error:\n{exc}") from exc
+
 
 @dataclasses.dataclass
 class ACLGraphEntry:
     batch_descriptor: BatchDescriptor
-    aclgraph: Optional[torch.npu.NPUGraph] = None
-    output: Optional[Any] = None
+    aclgraph: torch.npu.NPUGraph | None = None
+    output: Any | None = None
 
     # for aclgraph debugging, track the input addresses
     # during capture, and check if they are the same during replay
-    input_addresses: Optional[list[int]] = None
+    input_addresses: list[int] | None = None
 
 
 class ACLGraphWrapper:
@@ -57,41 +87,49 @@ class ACLGraphWrapper:
     guaranteed when VLLM_LOGGING_LEVEL == "DEBUG".
     """
 
-    def __init__(self,
-                 runnable: Callable,
-                 vllm_config: VllmConfig,
-                 runtime_mode: CUDAGraphMode,
-                 graph_pool: Any = None,
-                 cudagraph_options: Optional[CUDAGraphOptions] = None):
+    def __init__(
+        self,
+        runnable: Callable,
+        vllm_config: VllmConfig,
+        runtime_mode: CUDAGraphMode,
+        cudagraph_options: CUDAGraphOptions | None = None,
+        *,
+        use_eagle: bool = False,
+        enable_enpu: bool = False,
+    ):
         self.runnable = runnable
         self.vllm_config = vllm_config
-        self.graph_pool = graph_pool
         self.runtime_mode = runtime_mode
         self.compilation_config = vllm_config.compilation_config
 
         self.first_run_finished = False
         self.is_debugging_mode = envs.VLLM_LOGGING_LEVEL == "DEBUG"
+        self._runnable_str = str(runnable) if self.is_debugging_mode else None
 
         # assert runtime_mode is not NONE(no aclgraph), otherwise, we don't
         # need to initialize a ACLGraphWrapper.
         assert self.runtime_mode != CUDAGraphMode.NONE
-        if self.graph_pool is None:
-            self.graph_pool = current_platform.get_global_graph_pool()
+        self.graph_pool = current_platform.get_global_graph_pool()
 
         if cudagraph_options is None:
             cudagraph_options = CUDAGraphOptions()
         self.aclgraph_options = cudagraph_options
         # the entries for different batch descriptors that we need to capture
         # aclgraphs for.
-        self.concrete_aclgraph_entries: dict[BatchDescriptor, ACLGraphEntry]\
-                                                                        = {}
+        self.concrete_aclgraph_entries: dict[BatchDescriptor, ACLGraphEntry] = {}
+        self.enable_enpu = enable_enpu
+        self.use_eagle = use_eagle
+        _acl_graph_wrappers.add(self)
 
     def __getattr__(self, key: str):
         # allow accessing the attributes of the runnable.
         if hasattr(self.runnable, key):
             return getattr(self.runnable, key)
-        raise AttributeError(f"Attribute {key} not exists in the runnable of "
-                             f"aclgraph wrapper: {self.runnable}")
+        if self.is_debugging_mode:
+            raise AttributeError(
+                f"Attribute {key} not exists in the runnable of aclgraph wrapper: {self._runnable_str}"
+            )
+        raise AttributeError(f"Attribute {key} not found. Set VLLM_LOGGING_LEVEL=DEBUG for more details.")
 
     def unwrap(self) -> Callable:
         # in case we need to access the original runnable.
@@ -102,8 +140,7 @@ class ACLGraphWrapper:
         batch_descriptor = forward_context.batch_descriptor
         aclgraph_runtime_mode = forward_context.cudagraph_runtime_mode
 
-        if aclgraph_runtime_mode == CUDAGraphMode.NONE or \
-                            aclgraph_runtime_mode != self.runtime_mode:
+        if aclgraph_runtime_mode == CUDAGraphMode.NONE or aclgraph_runtime_mode != self.runtime_mode:
             # CUDAGraphMode.NONE could mean the profile run, a warmup run, or
             # running without aclgraphs.
             # We do not trigger capture/replay if the runtime mode is not
@@ -114,8 +151,7 @@ class ACLGraphWrapper:
 
         if batch_descriptor not in self.concrete_aclgraph_entries:
             # create a new entry for this batch descriptor
-            self.concrete_aclgraph_entries[batch_descriptor] = \
-                ACLGraphEntry(batch_descriptor=batch_descriptor)
+            self.concrete_aclgraph_entries[batch_descriptor] = ACLGraphEntry(batch_descriptor=batch_descriptor)
 
         entry = self.concrete_aclgraph_entries[batch_descriptor]
 
@@ -125,14 +161,11 @@ class ACLGraphWrapper:
                 # capturing is fast, we don't need to log it for every
                 # shape. E.g. we only log it for the first subgraph in
                 # piecewise mode.
-                logger.debug("Capturing a aclgraph on (%s,%s)",
-                             self.runtime_mode.name, entry.batch_descriptor)
+                logger.debug("Capturing a aclgraph on (%s,%s)", self.runtime_mode.name, entry.batch_descriptor)
             # validate that aclgraph capturing is legal at this point.
             validate_cudagraph_capturing_enabled()
 
-            input_addresses = [
-                x.data_ptr() for x in args if isinstance(x, torch.Tensor)
-            ]
+            input_addresses = [x.data_ptr() for x in args if isinstance(x, torch.Tensor)]
             entry.input_addresses = input_addresses
             aclgraph = torch.npu.NPUGraph()
 
@@ -145,22 +178,46 @@ class ACLGraphWrapper:
                     # therefore, we only run gc for the first graph,
                     # and disable gc for the rest of the graphs.
                     stack.enter_context(patch("gc.collect", lambda: None))
-                    stack.enter_context(
-                        patch("torch.npu.empty_cache", lambda: None))
+                    stack.enter_context(patch("torch.npu.empty_cache", lambda: None))
 
                 # mind-exploding: carefully manage the reference and memory.
+
+                # Sync offloader's copy stream before capture.
+                # Ensure any pre-capture prefetches from offloader are complete.
+                from vllm.model_executor.offloader.base import get_offloader
+
+                get_offloader().sync_prev_onload()
                 forward_context.capturing = True
-                with torch.npu.graph(aclgraph, pool=self.graph_pool):
-                    # `output` is managed by pytorch's aclgraph pool
-                    output = self.runnable(*args, **kwargs)
-                    if self.aclgraph_options.weak_ref_output:
-                        # by converting it to weak ref,
-                        # the original `output` will immediately be released
-                        # to save memory. It is only safe to do this for
-                        # the last graph in piecewise aclgraph mode, because
-                        # the output of the last graph will not be used by
-                        # any other acl graph.
-                        output = weak_ref_tensors(output)
+                try:
+                    with torch.npu.graph(aclgraph, pool=self.graph_pool):
+                        # `output` is managed by pytorch's aclgraph pool
+                        output = self.runnable(*args, **kwargs)
+                        # Join offloader's copy stream after forward to avoid
+                        # unjoined stream error. The last layer's start_prefetch
+                        # forks copy_stream, but wait_prefetch only happens in
+                        # the next forward pass.
+                        get_offloader().join_after_forward()
+                        if self.aclgraph_options.weak_ref_output:
+                            # by converting it to weak ref,
+                            # the original `output` will immediately be released
+                            # to save memory. It is only safe to do this for
+                            # the last graph in piecewise aclgraph mode, because
+                            # the output of the last graph will not be used by
+                            # any other acl graph.
+                            output = weak_ref_tensors(output)
+                except RuntimeError as exc:
+                    if _is_stream_resource_capture_error(exc):
+                        _raise_stream_resource_capture_error(exc)
+                    raise
+
+            # here we always use weak ref for the workspaces
+            # to save memory
+            global _graph_params
+            global _draft_graph_params
+            global _draft_graph_prefill_params
+            weak_ref_workspaces(_graph_params)
+            weak_ref_workspaces(_draft_graph_params)
+            weak_ref_workspaces(_draft_graph_prefill_params)
 
             # here we always use weak ref for the output
             # to save memory
@@ -176,57 +233,60 @@ class ACLGraphWrapper:
 
         if self.is_debugging_mode:
             # check if the input addresses are the same
-            new_input_addresses = [
-                x.data_ptr() for x in args if isinstance(x, torch.Tensor)
-            ]
+            new_input_addresses = [x.data_ptr() for x in args if isinstance(x, torch.Tensor)]
             assert new_input_addresses == entry.input_addresses, (
                 f"Input addresses for aclgraphs are different "
                 f"during replay. Expected {entry.input_addresses}, "
-                f"got {new_input_addresses}")
+                f"got {new_input_addresses}"
+            )
 
         logger.info_once("Replaying aclgraph")
+        # In async scheduling or multi-threaded (MT) scenarios, it is possible that
+        # the CPU's record event (from update_attn_params) for the iteration i completes
+        # before the grph replay of iteration i-1.
+        # To ensure proper ordering, we must call synchronize here before replaying,
+        # so that update_attn_params only executes after the previous graph replay has fully completed.
+        # If we do not in main model and in full-graph mode when using merge-eagle-graph,
+        # we do not need to synchronize.
+        # When enable_enpu is on, model_runner orders update vs replay; skip here.
+        # When FULL + EAGLE draft (merge path), replay does not need this barrier.
+        is_draft_eagle = _EXTRA_CTX.is_draft_model and self.use_eagle
+        need_sync = self.runtime_mode == CUDAGraphMode.FULL and not is_draft_eagle
+        if not self.enable_enpu and need_sync:
+            torch.npu.current_stream().synchronize()
         entry.aclgraph.replay()
         return entry.output
 
 
-def update_attn_params(update_stream, forward_context, runtime_shape):
-    graph_params = get_graph_params()
-    # FIXME: Behold! We are using a temporary hack here to update the args
-    # for each layer's attention op in the graph.
-    for key, param, handle, event in zip(
-            forward_context.attn_metadata,
-            graph_params.attn_params[runtime_shape],
-            graph_params.handles[runtime_shape],
-            graph_params.events[runtime_shape],
-    ):
-        (
-            query,
-            key_cache,
-            value_cache,
-            num_kv_heads,
-            num_heads,
-            scale,
-            block_table,
-            seq_lens,
-            output,
-        ) = param
-        # block_table = forward_context.attn_metadata[key].block_tables
-        seq_lens = forward_context.attn_metadata[key].seq_lens
+def weak_ref_workspaces(params):
+    if params is None:
+        return
+    for num_tokens in params.workspaces:
+        if params.workspaces[num_tokens] is None:
+            continue
+        params.workspaces[num_tokens] = weak_ref_tensors(params.workspaces[num_tokens])
 
-        with torch.npu.stream(update_stream):
-            torch.npu.graph_task_update_begin(update_stream, handle)
-            torch_npu._npu_paged_attention(query=query,
-                                           key_cache=key_cache,
-                                           value_cache=value_cache,
-                                           num_kv_heads=num_kv_heads,
-                                           num_heads=num_heads,
-                                           scale_value=scale,
-                                           block_table=block_table,
-                                           context_lens=seq_lens,
-                                           out=output)
-            torch.npu.graph_task_update_end(update_stream)
 
-            event.record(update_stream)
+def update_full_graph_params(
+    attn_backend,
+    update_stream,
+    forward_context,
+    num_tokens,
+    vllm_config,
+    speculative_config=None,
+    num_dcp_pcp_tokens=None,
+    draft_attn_metadatas=None,
+):
+    impl_cls = attn_backend.get_impl_cls()
+    impl_cls.update_graph_params(
+        update_stream,
+        forward_context,
+        num_tokens,
+        vllm_config,
+        speculative_config,
+        num_dcp_pcp_tokens,
+        draft_attn_metadatas,
+    )
 
 
 @dataclass
@@ -237,24 +297,76 @@ class GraphParams:
     attn_params: dict[int, list[tuple]]
 
 
-_graph_params: Optional[GraphParams] = None
+_graph_params: GraphParams | None = None
 
 
-def set_graph_params(aclgraph_capture_sizes: set[int]):
+def set_graph_params(aclgraph_capture_sizes: list[int]):
     global _graph_params
     if _graph_params is not None:
         raise ValueError("Graph parameters have already been set!")
     _graph_params = GraphParams(
-        {size: []
-         for size in aclgraph_capture_sizes},
-        {size: None
-         for size in aclgraph_capture_sizes},
-        {size: []
-         for size in aclgraph_capture_sizes},
-        {size: []
-         for size in aclgraph_capture_sizes},
+        {size: [] for size in aclgraph_capture_sizes},
+        {size: None for size in aclgraph_capture_sizes},
+        {size: [] for size in aclgraph_capture_sizes},
+        {size: [] for size in aclgraph_capture_sizes},
     )
 
 
+def update_graph_params_workspaces(num_tokens: int, workspace: torch.Tensor):
+    global _graph_params
+    if _graph_params is not None:
+        _graph_params.workspaces[num_tokens] = workspace
+
+
 def get_graph_params():
     return _graph_params
+
+
+_draft_graph_params: GraphParams | None = None
+
+
+def set_draft_graph_params(aclgraph_capture_sizes: list[int]):
+    global _draft_graph_params
+    if _draft_graph_params is not None:
+        raise ValueError("DraftGraph parameters have already been set!")
+    _draft_graph_params = GraphParams(
+        {size: [] for size in aclgraph_capture_sizes},
+        {size: None for size in aclgraph_capture_sizes},
+        {size: [] for size in aclgraph_capture_sizes},
+        {size: [] for size in aclgraph_capture_sizes},
+    )
+
+
+def update_draft_graph_params_workspaces(num_tokens: int, workspace: Any):
+    global _draft_graph_params
+    if _draft_graph_params is not None:
+        _draft_graph_params.workspaces[num_tokens] = workspace
+
+
+def get_draft_graph_params():
+    return _draft_graph_params
+
+
+_draft_graph_prefill_params: GraphParams | None = None
+
+
+def set_draft_graph_prefill_params(aclgraph_capture_sizes: list[int]):
+    global _draft_graph_prefill_params
+    if _draft_graph_prefill_params is not None:
+        raise ValueError("DraftGraph preill parameters have already been set!")
+    _draft_graph_prefill_params = GraphParams(
+        {size: [] for size in aclgraph_capture_sizes},
+        {size: None for size in aclgraph_capture_sizes},
+        {size: [] for size in aclgraph_capture_sizes},
+        {size: [] for size in aclgraph_capture_sizes},
+    )
+
+
+def update_draft_graph_prefill_params_workspaces(num_tokens: int, workspace: Any):
+    global _draft_graph_prefill_params
+    if _draft_graph_prefill_params is not None:
+        _draft_graph_prefill_params.workspaces[num_tokens] = workspace
+
+
+def get_draft_graph_prefill_params():
+    return _draft_graph_prefill_params

vllm_ascend/compilation/compiler_interface.py

@@ -0,0 +1,368 @@
+#
+# 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.
+#
+import copy
+import functools
+import os
+from collections.abc import Callable
+from typing import Any, cast
+
+import torch
+import torch.fx as fx
+from torch._dynamo.backends.common import aot_autograd
+from torch._inductor.compile_fx import graph_returns_tuple, make_graph_return_tuple
+from torch._inductor.decomposition import select_decomp_table
+from torch.fx import GraphModule
+from vllm.compilation.compiler_interface import CompilerInterface
+from vllm.config import VllmConfig
+from vllm.config.utils import Range
+from vllm.logger import logger
+
+from vllm_ascend.ascend_config import AscendCompilationConfig, get_ascend_config
+from vllm_ascend.utils import COMPILATION_PASS_KEY
+
+
+def compile_fx(graph: GraphModule, example_inputs: list, inner_compile: Callable, decompositions: dict) -> Callable:
+    recursive_compile_fx = functools.partial(compile_fx, inner_compile=inner_compile, decompositions=decompositions)
+
+    if not graph_returns_tuple(graph):
+        return make_graph_return_tuple(graph, example_inputs, recursive_compile_fx)
+    return aot_autograd(fw_compiler=inner_compile)(graph, example_inputs)
+
+
+def fusion_pass_compile(
+    graph: fx.GraphModule,
+    example_inputs: list[Any],
+    compiler_config: dict[str, Any],
+    compile_range: Range,
+    key: str | None = None,
+) -> tuple[Callable | None, Any | None]:
+    def compile_inner(graph, example_inputs):
+        current_pass_manager = compiler_config[COMPILATION_PASS_KEY]
+        graph = current_pass_manager(graph)
+        return graph
+
+    decompositions = select_decomp_table()
+
+    compiled_fn = compile_fx(
+        graph=graph,
+        example_inputs=example_inputs,
+        inner_compile=compile_inner,
+        decompositions=decompositions,
+    )
+
+    return compiled_fn, None
+
+
+def _compute_decode_cudagraph_batch_sizes(vllm_config: VllmConfig) -> list[int]:
+    num_spec_tokens = vllm_config.speculative_config.num_speculative_tokens if vllm_config.speculative_config else 0
+    uniform_decode_query_len = num_spec_tokens + 1
+    max_num_tokens = vllm_config.scheduler_config.max_num_seqs * uniform_decode_query_len
+    return [
+        x
+        for x in vllm_config.compilation_config.cudagraph_capture_sizes
+        if max_num_tokens >= x >= uniform_decode_query_len
+    ]
+
+
+def _configure_backend(
+    config: Any,
+    ascend_compilation_config: AscendCompilationConfig,
+    vllm_config: VllmConfig,
+    process_kwargs_options: Callable | None = None,
+) -> None:
+    if ascend_compilation_config.enable_static_kernel:
+        # npugraph_ex's static_kernel requires LOCAL_WORLD_SIZE to determine the
+        # physical node topology for creating per-node Gloo groups, which
+        # coordinate static kernel compilation and .run package installation.
+        # vLLM does not set this env var by default (unlike torchrun), so we
+        # compute it from parallel config:
+        #   local_world_size: processes per node for one DP replica
+        #   data_parallel_size_local: number of DP replicas on this node
+        #   actual_local_world_size: total processes on this physical machine
+        if "LOCAL_WORLD_SIZE" not in os.environ:
+            actual_local_world_size = (
+                vllm_config.parallel_config.local_world_size * vllm_config.parallel_config.data_parallel_size_local
+            )
+            os.environ["LOCAL_WORLD_SIZE"] = str(actual_local_world_size)
+            logger.info_once(
+                "Setting LOCAL_WORLD_SIZE=%d for static kernel (local_world_size=%d * data_parallel_size_local=%d).",
+                actual_local_world_size,
+                vllm_config.parallel_config.local_world_size,
+                vllm_config.parallel_config.data_parallel_size_local,
+                scope="global",
+            )
+
+    if process_kwargs_options is not None:
+        # npugraph_ex (both old and new): build options dict and use _process_kwargs_options.
+        # It maps flat option names to nested config paths for old versions,
+        # and directly setattr for new versions with flat CompilerConfig.
+        # force_eager=True: execute FX graph in eager mode before graph capture.
+        # inplace_pass=False: disable reinplace pass to avoid gelu fallback to CPU.
+        options: dict[str, Any] = {
+            "force_eager": True,
+            "inplace_pass": False,
+            "clone_input": False,
+            "clone_output": False,
+        }
+        if ascend_compilation_config.enable_static_kernel:
+            logger.info_once(
+                "enable_static_kernel is enabled, static shape kernel will be used to accelerate aclgraph execution.",
+                scope="global",
+            )
+            options["static_kernel_compile"] = True
+            # Set sym_range to limit static kernel compilation to specified batch sizes.
+            options["_vllm_aclnn_static_kernel_sym_range"] = _compute_decode_cudagraph_batch_sizes(vllm_config)
+        process_kwargs_options(config, {"options": options})
+    else:
+        # torchair (reduce-overhead): use nested config structure directly.
+        # mode="reduce-overhead": use aclgraph mode, avoid fx graph to Ascend IR transformation.
+        config.mode = "reduce-overhead"
+        config.debug.run_eagerly = True
+        # Disable reinplace pass to avoid gelu fallback to CPU causing host-device copy error.
+        config.debug.aclgraph.disable_reinplace_inplaceable_ops_pass = True
+        if ascend_compilation_config.enable_static_kernel:
+            logger.info_once(
+                "enable_static_kernel is enabled, static shape kernel will be used to accelerate aclgraph execution.",
+                scope="global",
+            )
+            config.experimental_config.aclgraph._aclnn_static_shape_kernel = True
+            config.experimental_config.aclgraph._aclnn_static_shape_kernel_sym_value_range = (
+                _compute_decode_cudagraph_batch_sizes(vllm_config)
+            )
+
+
+def npugraph_ex_compile(
+    graph: fx.GraphModule,
+    example_inputs: list[Any],
+    compiler_config: dict[str, Any],
+    vllm_config: VllmConfig,
+    ascend_compilation_config: AscendCompilationConfig,
+    compile_range: Range,
+    key: str | None = None,
+    cache_dir: str | None = None,
+) -> tuple[Callable | None, Any | None]:
+    # Try npugraph_ex first, fall back to torchair for backward compatibility.
+    try:
+        import npugraph_ex as nge
+
+        cache_path = os.path.join(cache_dir, key) if (cache_dir and key) else None
+        torch.npu.set_compile_mode(jit_compile=False)
+        config = nge.CompilerConfig()
+        # _process_kwargs_options exists in both old and new npugraph_ex,
+        # but in different modules: new -> compiler_config, old -> npugraphex_config.
+        try:
+            from npugraph_ex.configs.compiler_config import _process_kwargs_options
+        except ImportError:
+            from npugraph_ex.configs.npugraphex_config import _process_kwargs_options
+        _configure_backend(
+            config, ascend_compilation_config, vllm_config, process_kwargs_options=_process_kwargs_options
+        )
+        import npugraph_ex.npu_fx_compiler as nfx
+
+        _original_get_compiled_gm = nfx._NpuFxCompiler._get_compiled_gm
+
+        def patched_get_compiled_gm(self, graph, example_inputs):
+            compiled_gm = _original_get_compiled_gm(self, graph, example_inputs)
+            if cache_path:
+                py_code = compiled_gm.get_code()
+                if py_code:
+                    # Triton kernel indices (kernel_side_table) are registered in-process
+                    # at compile time and are not serializable across process boundaries.
+                    # Graphs containing triton_kernel_wrapper calls must not be cached,
+                    # because loading the py_code in a new process will hit an
+                    # AssertionError in kernel_side_table.get_kernel().
+                    if "triton_kernel_wrapper" in py_code:
+                        logger.info(
+                            "Skipping npugraph_ex cache for graph containing Triton kernels "
+                            "(kernel_side_table indices are process-local): %s",
+                            cache_path,
+                        )
+                    else:
+                        os.makedirs(os.path.dirname(cache_path), exist_ok=True)
+                        with open(cache_path, "w") as f:
+                            f.write(py_code)
+                        logger.info("Saved compiled graph to cache: %s", cache_path)
+            return compiled_gm
+
+        nfx._NpuFxCompiler._get_compiled_gm = patched_get_compiled_gm
+        backend = nge.get_npu_backend(compiler_config=config)
+        # torch.compile requires the output of the fx graph to be a tuple
+        if not graph_returns_tuple(graph):
+            compiled_fn = make_graph_return_tuple(graph, example_inputs, backend)
+        else:
+            compiled_fn = backend(graph, example_inputs)
+        nfx._NpuFxCompiler._get_compiled_gm = _original_get_compiled_gm
+        return compiled_fn, (key, cache_path)
+    except ImportError:
+        import torchair
+
+        torch.npu.set_compile_mode(jit_compile=False)
+        config = torchair.CompilerConfig()
+        _configure_backend(config, ascend_compilation_config, vllm_config)
+        backend = torchair.get_npu_backend(compiler_config=config)
+        # torch.compile requires the output of the fx graph to be a tuple
+        if not graph_returns_tuple(graph):
+            compiled_fn = make_graph_return_tuple(graph, example_inputs, backend)
+        else:
+            compiled_fn = backend(graph, example_inputs)
+        return compiled_fn, None
+
+
+class AscendCompiler(CompilerInterface):
+    """
+    AscendCompiler is a custom compiler interface for the Ascend platform.
+    This class provides a method to compile a PyTorch FX graph module with
+    specific configurations for graph fusion and decomposition.
+    """
+
+    name = "AscendCompiler"
+
+    # TODO(wxs): add passes related to compilation in compute_hash
+    def compute_hash(self, vllm_config: VllmConfig) -> str:
+        self.vllm_config = vllm_config
+        ascend_compilation_config = get_ascend_config().ascend_compilation_config
+        from hashlib import sha256
+
+        import torch_npu
+
+        factors = {
+            "torch_npu_version": torch_npu.__version__,
+            "enable_npugraph_ex": ascend_compilation_config.enable_npugraph_ex,
+            "enable_static_kernel": ascend_compilation_config.enable_static_kernel,
+        }
+        logger.info("AscendCompiler hash factors: %s", factors)
+        return sha256(str(factors).encode(), usedforsecurity=False).hexdigest()[:10]
+
+    def initialize_cache(self, cache_dir, disable_cache=False, prefix=""):
+        self.cache_dir = cache_dir
+        self.disable_cache = disable_cache
+
+    def compile(
+        self,
+        graph: fx.GraphModule,
+        example_inputs: list[Any],
+        compiler_config: dict[str, Any],
+        compile_range: Range,
+        key: str | None = None,
+    ) -> tuple[Callable | None, Any | None]:
+        # inductor can inplace modify the graph, so we need to copy it
+        # see https://github.com/pytorch/pytorch/issues/138980
+        graph = copy.deepcopy(graph)
+
+        from torch._guards import detect_fake_mode
+
+        current_fake_mode = detect_fake_mode()
+        if current_fake_mode is not None:
+            example_inputs = [
+                current_fake_mode.from_tensor(inp)
+                if (
+                    isinstance(inp, torch.Tensor)
+                    and hasattr(inp, "fake_mode")
+                    and inp.fake_mode is not current_fake_mode
+                )
+                else inp
+                for inp in example_inputs
+            ]
+
+        ascend_compilation_config = get_ascend_config().ascend_compilation_config
+        if ascend_compilation_config.enable_npugraph_ex:
+            cache_dir = None if getattr(self, "disable_cache", False) else getattr(self, "cache_dir", None)
+            logger.info_once(
+                "enable_npugraph_ex is enabled, which will bring graph compilation optimization.",
+                scope="global",
+            )
+            assert hasattr(self, "vllm_config")
+            return npugraph_ex_compile(
+                graph,
+                example_inputs,
+                compiler_config,
+                self.vllm_config,
+                ascend_compilation_config,
+                compile_range,
+                key,
+                cache_dir,
+            )
+        else:
+            return fusion_pass_compile(graph, example_inputs, compiler_config, compile_range, key)
+
+    def load(self, handle, graph, example_inputs, graph_index, compile_range):
+        key, path = handle
+        # Cache file may be absent when the graph was skipped at save time (e.g. it
+        # contained Triton kernels whose kernel_side_table indices are process-local
+        # and cannot be serialized).  Fall back to a fresh compilation so the Triton
+        # kernels are properly registered in the current process.
+        if not path or not os.path.exists(path):
+            logger.info(
+                "npugraph_ex cache miss for key %s (file absent or not saved), recompiling",
+                key,
+            )
+            # Mirror the same pre-processing done in compile(): deepcopy the graph
+            # to prevent make_graph_return_tuple from mutating the caller's copy,
+            # and re-wrap FakeTensor inputs under the current fake mode to avoid
+            # "fake mode mismatch" in aot_module_simplified.
+            graph = copy.deepcopy(graph)
+            from torch._guards import detect_fake_mode
+
+            current_fake_mode = detect_fake_mode()
+            if current_fake_mode is not None:
+                example_inputs = [
+                    current_fake_mode.from_tensor(inp)
+                    if (
+                        isinstance(inp, torch.Tensor)
+                        and hasattr(inp, "fake_mode")
+                        and inp.fake_mode is not current_fake_mode
+                    )
+                    else inp
+                    for inp in example_inputs
+                ]
+            ascend_compilation_config = get_ascend_config().ascend_compilation_config
+            assert hasattr(self, "vllm_config")
+            compiled_fn, _ = npugraph_ex_compile(
+                graph,
+                example_inputs,
+                {},
+                self.vllm_config,
+                ascend_compilation_config,
+                compile_range,
+                key,
+                getattr(self, "cache_dir", None),
+            )
+            return compiled_fn
+
+        from npugraph_ex.npu_fx_compiler import _CompiledFxArtifacts, _CompiledFxGraph
+
+        with open(path) as f:
+            py_code = f.read()
+        artifacts = _CompiledFxArtifacts()
+        artifacts.py_code = py_code
+        logger.info("Loaded npugraph_ex compilation cache from %s", path)
+        compiled_fn = cast(Callable[..., Any], _CompiledFxGraph.load_artifacts(artifacts))
+
+        # The saved code was compiled from the graph after make_graph_return_tuple mutated it
+        # to return a flat tuple. If the original graph didn't return a tuple, we need to
+        # recreate the unflatten wrapper so callers receive the original output structure.
+        if not graph_returns_tuple(graph):
+            _inner_fn = compiled_fn
+
+            def compiled_fn(*args, **kwargs):
+                result = _inner_fn(*args, **kwargs)
+                if isinstance(result, (tuple, list)) and len(result) == 1:
+                    return result[0]
+                return result
+
+        return compiled_fn

vllm_ascend/compilation/graph_fusion_pass_manager.py

@@ -0,0 +1,79 @@
+#
+# 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 torch import fx as fx
+from vllm.compilation.passes.inductor_pass import get_pass_context
+from vllm.compilation.passes.vllm_inductor_pass import VllmInductorPass
+from vllm.config import VllmConfig
+
+
+class GraphFusionPassManager:
+    """
+    A pass manager for graph fusion passes.
+    It handles the configuration and execution of passes.
+    The counterpart in vllm is PostGradPassManager. Since torch_npu
+    does not support triton for now, we define our own pass manager.
+    """
+
+    def __init__(self):
+        self.passes: list[VllmInductorPass] = []
+
+    def __call__(self, graph: fx.Graph) -> fx.Graph:
+        compile_range = get_pass_context().compile_range
+
+        for pass_ in self.passes:
+            if pass_.is_applicable_for_range(compile_range):
+                pass_(graph)
+        graph.recompile()
+        return graph
+
+    def add(self, pass_: VllmInductorPass):
+        assert isinstance(pass_, VllmInductorPass)
+        self.passes.append(pass_)
+
+    def configure(self, config: VllmConfig):
+        from vllm_ascend.utils import is_310p
+
+        # By default, we enable the graph fusion and quantization fusion pass.
+        self.ascend_compilation_config: dict = config.additional_config.get("ascend_compilation_config", {})
+        if self.ascend_compilation_config.get("fuse_norm_quant", True) and not is_310p():
+            from .passes.norm_quant_fusion_pass import AddRMSNormQuantFusionPass
+
+            self.passes.append(AddRMSNormQuantFusionPass(config))
+
+        if self.ascend_compilation_config.get("fuse_qknorm_rope", True):
+            from .passes.qknorm_rope_fusion_pass import QKNormRopeFusionPass
+
+            self.passes.append(QKNormRopeFusionPass(config))
+
+        if self.ascend_compilation_config.get("fuse_allreduce_rms", True):
+            from .passes.allreduce_rmsnorm_fusion_pass import MatmulAllReduceAddRMSNormPass
+
+            self.passes.append(MatmulAllReduceAddRMSNormPass(config))
+
+        if self.ascend_compilation_config.get("fuse_muls_add", True) and not is_310p():
+            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 SequenceParallelismPass
+            from .passes.sequence_parallelism_moe import SequenceParallelismMoePass
+
+            self.passes.append(SequenceParallelismPass(config))
+            self.passes.append(SequenceParallelismMoePass(config))

vllm_ascend/compilation/passes/allgather_chunk_noop_pass.py

@@ -0,0 +1,40 @@
+import torch
+import torch._inductor.pattern_matcher as pm
+from torch._inductor.pattern_matcher import PatternMatcherPass
+from vllm.compilation.passes.vllm_inductor_pass import VllmInductorPass
+from vllm.config import VllmConfig
+from vllm.distributed import get_tensor_model_parallel_world_size, get_tp_group
+from vllm.logger import logger
+
+
+class AllGatherChunkNoOpCleanupPass(VllmInductorPass):
+    """Fold all_gather + sequence_parallel_chunk_impl into identity."""
+
+    def __init__(self, config: VllmConfig):
+        super().__init__(config)
+        self.tp_group = get_tp_group()
+        self.tp_size = get_tensor_model_parallel_world_size()
+        self.patterns: PatternMatcherPass = PatternMatcherPass(pass_name="npu_allgather_chunk_noop_cleanup_pass")
+        self._register_patterns()
+
+    def _all_gather(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.ops.vllm.all_gather(x, dim=0, world_size=self.tp_size, group_name=self.tp_group.unique_name)
+
+    def _empty(self, *args, **kwargs):
+        return torch.empty(*args, dtype=self.model_dtype, device=self.device, **kwargs)
+
+    def _register_patterns(self) -> None:
+        def pattern(input: torch.Tensor) -> torch.Tensor:
+            gathered = self._all_gather(input)
+            return torch.ops.vllm.sequence_parallel_chunk_impl(gathered)
+
+        def replacement(input: torch.Tensor) -> torch.Tensor:
+            return input
+
+        pm.register_replacement(pattern, replacement, [self._empty(8, 16)], pm.fwd_only, self.patterns)
+
+    def __call__(self, graph: torch.fx.Graph) -> None:
+        self.begin()
+        matched_count = self.patterns.apply(graph)
+        logger.debug("AllGatherChunkNoOpCleanupPass replaced %s patterns", matched_count)
+        self.end_and_log()

vllm_ascend/compilation/passes/allreduce_rmsnorm_fusion_pass.py

@@ -0,0 +1,159 @@
+# 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.
+#
+import torch
+from torch._inductor.pattern_matcher import PatternMatcherPass, PatternPrettyPrinter
+from vllm.compilation.passes.vllm_inductor_pass import VllmInductorPass
+from vllm.config import VllmConfig
+from vllm.config.compilation import Range
+from vllm.distributed import get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce
+from vllm.distributed.parallel_state import get_tp_group
+from vllm.logger import logger
+
+from vllm_ascend.compilation.passes.base_pattern import BasePattern
+
+# computation-communication tiling block is 512
+ALLREDUCE_NORM_FUSE_THRESHOLD = 512
+
+
+class MiddleLayerMatmulAllReduceAddRMSNormPattern(BasePattern):
+    """
+    recognizing the Matmul+AllReduce+AddRMSNorm computation pattern
+    AllReduce is optimized in the fusion operator to a two-stage communication of ReduceScatter+AllGather
+    """
+
+    def __init__(self, vllm_config, eps=1e-6):
+        self.vllm_config = vllm_config
+        self.eps = eps
+        device_group = get_tp_group().device_group
+        backend = device_group._get_backend(torch.device("npu"))
+        self.local_rank = torch.distributed.get_rank(group=device_group)
+        self.tp_group_name = backend.get_hccl_comm_name(self.local_rank)
+        self.tp_size = get_tensor_model_parallel_world_size()
+
+    def get_inputs(self):
+        batch_size, seq_len = 2, 4
+        hidden_size = 4096
+        x = torch.randn(batch_size, seq_len, hidden_size, device="npu")
+        weight = torch.randn(hidden_size, hidden_size, device="npu")
+        residual = torch.randn(batch_size, seq_len, hidden_size, device="npu")
+        rms_norm_weight = torch.randn(hidden_size, device="npu")
+        return [x, weight, residual, rms_norm_weight]
+
+    def get_pattern(self):
+        def pattern(x, weight, residual, rms_norm_weight):
+            mm = torch.ops.vllm.unquantized_gemm(x, weight, None)
+            all_reduce_ = tensor_model_parallel_all_reduce(mm)
+            chunked_residual = torch.ops.vllm.maybe_chunk_residual(all_reduce_, residual)
+            output = torch.ops._C_ascend.npu_add_rms_norm_bias(all_reduce_, chunked_residual, rms_norm_weight, None)
+            out0 = output[0]
+            out1 = output[2]
+            return out0, out1
+
+        return pattern
+
+    def get_replacement(self):
+        def replacement(x, weight, residual, rms_norm_weight):
+            out0, out1 = torch.ops._C_ascend.matmul_allreduce_add_rmsnorm(
+                x,
+                weight,
+                residual,
+                rms_norm_weight,
+                self.tp_group_name,
+                self.tp_size,
+                self.local_rank,
+                self.eps,
+                True,
+                False,
+            )
+            return out0, out1
+
+        return replacement
+
+
+class LastLayerMatmulAllReduceAddRMSNormPattern(BasePattern):
+    def __init__(self, vllm_config, eps=1e-6):
+        super().__init__(vllm_config, eps)
+        device_group = get_tp_group().device_group
+        backend = device_group._get_backend(torch.device("npu"))
+        self.local_rank = torch.distributed.get_rank(group=device_group)
+        self.tp_group_name = backend.get_hccl_comm_name(self.local_rank)
+        self.tp_size = get_tensor_model_parallel_world_size()
+
+    def get_inputs(self):
+        batch_size, seq_len = 2, 4
+        hidden_size = 4096
+        x = torch.randn(batch_size, seq_len, hidden_size, device="npu")
+        weight = torch.randn(hidden_size, hidden_size, device="npu")
+        residual = torch.randn(batch_size, seq_len, hidden_size, device="npu")
+        rms_norm_weight = torch.randn(hidden_size, device="npu")
+        return [x, weight, residual, rms_norm_weight]
+
+    def get_pattern(self):
+        def pattern(x, weight, residual, rms_norm_weight):
+            mm = torch.ops.vllm.unquantized_gemm(x, weight, None)
+            all_reduce_ = tensor_model_parallel_all_reduce(mm)
+            chunked_residual = torch.ops.vllm.maybe_chunk_residual(all_reduce_, residual)
+            output = torch.ops._C_ascend.npu_add_rms_norm_bias(all_reduce_, chunked_residual, rms_norm_weight, None)
+            return output[0]
+
+        return pattern
+
+    def get_replacement(self):
+        def replacement(x, weight, residual, rms_norm_weight):
+            out0, _ = torch.ops._C_ascend.matmul_allreduce_add_rmsnorm(
+                x,
+                weight,
+                residual,
+                rms_norm_weight,
+                self.tp_group_name,
+                self.tp_size,
+                self.local_rank,
+                self.eps,
+                True,
+                False,
+            )
+            return out0
+
+        return replacement
+
+
+class MatmulAllReduceAddRMSNormPass(VllmInductorPass):
+    def __init__(self, vllm_config: VllmConfig):
+        super().__init__(vllm_config)
+        self.pattern_match_passes: PatternMatcherPass = PatternMatcherPass(pass_name="allreduce_rmsnorm_fusion_pass")
+
+        MiddleLayerMatmulAllReduceAddRMSNormPattern(vllm_config).register(self.pattern_match_passes)
+        LastLayerMatmulAllReduceAddRMSNormPattern(vllm_config).register(self.pattern_match_passes)
+
+    def __call__(self, graph: torch.fx.Graph):
+        self.begin()
+        self.matched_count = self.pattern_match_passes.apply(graph)
+        pattern_idx = 0
+        for pattern_entry in self.pattern_match_passes.patterns.values():
+            for p in pattern_entry:
+                p_str = PatternPrettyPrinter.run(p.pattern)
+                logger.debug("Pattern %d: %s", pattern_idx, p_str)
+                pattern_idx += 1
+        logger.debug("Replaced %s 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.
+        """
+        applicable = compile_range.start > ALLREDUCE_NORM_FUSE_THRESHOLD
+        return applicable

vllm_ascend/compilation/passes/base_pattern.py

@@ -0,0 +1,63 @@
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+
+import torch
+import torch._inductor.pattern_matcher as pm
+from torch._inductor.pattern_matcher import PatternMatcherPass
+from vllm.config import VllmConfig
+
+try:
+    import npugraph_ex as nge
+except ImportError:
+    import torchair as nge
+
+from vllm_ascend.compilation.passes.utils.npugraph_ex_utils_check import extra_stream_scope_check
+
+# Global set to track registered patterns and prevent duplicates
+_registered_patterns: set[str] = set()
+
+
+class BasePattern(ABC):
+    def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6):
+        self.vllm_config = vllm_config
+        self.dtype = vllm_config.model_config.dtype
+        self.eps = eps
+
+    @abstractmethod
+    def get_inputs(self) -> list[torch.Tensor]:
+        pass
+
+    @abstractmethod
+    def get_pattern(self) -> Callable:
+        pass
+
+    @abstractmethod
+    def get_replacement(self) -> Callable:
+        pass
+
+    def get_extra_stream_scope_check(self):
+        return extra_stream_scope_check
+
+    def register(self, pm_pass: PatternMatcherPass) -> None:
+        # Create a unique identifier for this pattern based on class name and eps
+        pattern_id = f"{self.__class__.__name__}_{self.eps}"
+
+        # Skip registration if this pattern has already been registered globally
+        if pattern_id in _registered_patterns:
+            return
+
+        pattern_fn = self.get_pattern()
+        replacement_fn = self.get_replacement()
+        example_inputs = self.get_inputs()
+
+        pm.register_replacement(pattern_fn, replacement_fn, example_inputs, pm.fwd_only, pm_pass)
+
+        nge.register_replacement(
+            search_fn=pattern_fn,
+            replace_fn=replacement_fn,
+            example_inputs=example_inputs,
+            extra_check=self.get_extra_stream_scope_check(),
+        )
+
+        # Mark this pattern as registered
+        _registered_patterns.add(pattern_id)

vllm_ascend/compilation/passes/muls_add_pass.py

@@ -0,0 +1,110 @@
+#
+# 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.compilation.passes.vllm_inductor_pass import VllmInductorPass
+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
+
+
+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
+
+        routed_scaling_factor = getattr(vllm_config.model_config.hf_text_config, "routed_scaling_factor", 1.0)
+        MulsAddPattern(vllm_config, scale=routed_scaling_factor).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

vllm_ascend/compilation/passes/noop_elimination.py

@@ -0,0 +1,62 @@
+from collections.abc import Iterable
+
+import torch
+import torch.fx
+from torch import SymInt
+from torch.fx.experimental.symbolic_shapes import statically_known_true
+from vllm.compilation.passes.vllm_inductor_pass import VllmInductorPass
+from vllm.logger import logger
+
+
+class NoOpEliminationPass(VllmInductorPass):
+    """Remove no-op view/reshape nodes after pattern rewrites."""
+
+    def __call__(self, graph: torch.fx.Graph) -> None:
+        fx_graph = graph.graph if hasattr(graph, "graph") else graph
+        removed = 0
+        for node in list(fx_graph.nodes):
+            if not self._is_view_like(node):
+                continue
+
+            input_node = node.args[0]
+            if not isinstance(input_node, torch.fx.Node):
+                continue
+
+            input_meta = input_node.meta.get("val")
+            output_meta = node.meta.get("val")
+            if input_meta is None or output_meta is None:
+                continue
+
+            input_shape = getattr(input_meta, "shape", None)
+            output_shape = getattr(output_meta, "shape", None)
+            if input_shape is None or output_shape is None:
+                continue
+
+            if self._all_dims_equivalent(input_shape, output_shape):
+                node.replace_all_uses_with(input_node)
+                fx_graph.erase_node(node)
+                removed += 1
+
+        logger.debug("NoOpEliminationPass removed %s no-op views", removed)
+
+    @staticmethod
+    def _is_view_like(node: torch.fx.Node) -> bool:
+        return (node.op == "call_method" and node.target in {"view", "reshape"}) or (
+            node.op == "call_function"
+            and node.target
+            in {
+                torch.ops.aten.view.default,
+                torch.ops.aten.reshape.default,
+            }
+        )
+
+    @staticmethod
+    def _dims_equivalent(dim: int | SymInt, i_dim: int | SymInt) -> bool:
+        return statically_known_true(dim == i_dim)  # type: ignore[no-any-return]
+
+    def _all_dims_equivalent(self, dims: Iterable[int | SymInt], i_dims: Iterable[int | SymInt]) -> bool:
+        dims_ = list(dims)
+        i_dims_ = list(i_dims)
+        if len(dims_) != len(i_dims_):
+            return False
+        return all(self._dims_equivalent(s, i_s) for s, i_s in zip(dims_, i_dims_))

vllm_ascend/compilation/passes/norm_quant_fusion_pass.py

@@ -0,0 +1,742 @@
+#
+# 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.
+#
+import torch
+from torch._inductor.pattern_matcher import PatternMatcherPass
+from vllm.compilation.passes.vllm_inductor_pass import VllmInductorPass
+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.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
+
+
+class AddRMSNormQuantPattern(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 AddRMSNormQuant fusion pattern.
+        """
+        rms_norm_input = torch.randn(2, 4, device="npu", dtype=self.dtype)
+        residual = torch.randn(2, 4, device="npu", dtype=self.dtype)
+        rms_norm_weight = torch.randn(4, device="npu", dtype=self.dtype)
+        scale = torch.ones(4, device="npu", dtype=self.dtype)
+        scale_reciprocal = torch.ones(4, device="npu", dtype=self.dtype)
+        offset = torch.zeros(4, device="npu", dtype=self.dtype)
+        return [rms_norm_input, residual, rms_norm_weight, scale, scale_reciprocal, offset]
+
+    def get_pattern(self):
+        def pattern(
+            rms_norm_input: torch.Tensor,
+            residual: torch.Tensor,
+            rms_norm_weight: torch.Tensor,
+            scale: torch.Tensor,
+            scale_reciprocal: torch.Tensor,
+            offset: torch.Tensor,
+        ):
+            """
+            Pattern for AddRMSNormQuant fusion.
+            """
+            output = torch.ops._C_ascend.npu_add_rms_norm_bias(
+                rms_norm_input, residual, rms_norm_weight, None, self.eps
+            )
+            out0 = output[0]
+            out1 = output[2]
+            quantized_output = torch.ops.vllm.quantize(out0, scale, scale_reciprocal, offset)
+            return quantized_output, out1
+
+        return pattern
+
+    def get_replacement(self):
+        def replacement(
+            rms_norm_input: torch.Tensor,
+            residual: torch.Tensor,
+            rms_norm_weight: torch.Tensor,
+            scale: torch.Tensor,
+            scale_reciprocal: torch.Tensor,
+            offset: torch.Tensor,
+        ):
+            """
+            Replacement for the AddRMSNormQuant fusion.
+            """
+            output = torch.ops.npu.npu_add_rms_norm_quant(
+                rms_norm_input, residual, rms_norm_weight, scale, offset, epsilon=self.eps
+            )
+            quantized_output = output[0]
+            out1 = output[2]
+            return quantized_output, out1
+
+        return replacement
+
+
+class AddRMSNormQuantPatternWithBias(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 AddRMSNormQuant fusion pattern.
+        """
+        rms_norm_input = torch.randn(2, 4, device="npu", dtype=self.dtype)
+        residual = torch.randn(2, 4, device="npu", dtype=self.dtype)
+        rms_norm_weight = torch.randn(4, device="npu", dtype=self.dtype)
+        rmsnorm_bias = torch.randn(4, device="npu", dtype=self.dtype)
+        scale = torch.ones(4, device="npu", dtype=self.dtype)
+        scale_reciprocal = torch.ones(4, device="npu", dtype=self.dtype)
+        offset = torch.zeros(4, device="npu", dtype=self.dtype)
+        return [rms_norm_input, residual, rms_norm_weight, scale, scale_reciprocal, offset, rmsnorm_bias]
+
+    def get_pattern(self):
+        def pattern(
+            rms_norm_input: torch.Tensor,
+            residual: torch.Tensor,
+            rms_norm_weight: torch.Tensor,
+            scale: torch.Tensor,
+            scale_reciprocal: torch.Tensor,
+            offset: torch.Tensor,
+            bias: torch.Tensor,
+        ):
+            """
+            Pattern for AddRMSNormQuant fusion.
+            """
+            output = torch.ops._C_ascend.npu_add_rms_norm_bias(
+                rms_norm_input, residual, rms_norm_weight, bias, self.eps
+            )
+            out0 = output[0]
+            out1 = output[2]
+            quantized_output = torch.ops.vllm.quantize(out0, scale, scale_reciprocal, offset)
+            return quantized_output, out1
+
+        return pattern
+
+    def get_replacement(self):
+        def replacement(
+            rms_norm_input: torch.Tensor,
+            residual: torch.Tensor,
+            rms_norm_weight: torch.Tensor,
+            scale: torch.Tensor,
+            scale_reciprocal: torch.Tensor,
+            offset: torch.Tensor,
+            bias: torch.Tensor,
+        ):
+            """
+            Replacement for the AddRMSNormQuant fusion.
+            """
+            output = torch.ops.npu.npu_add_rms_norm_quant(
+                rms_norm_input, residual, rms_norm_weight, scale, offset, epsilon=self.eps, beta=bias
+            )
+            quantized_output = output[0]
+            out1 = output[2]
+            return quantized_output, out1
+
+        return replacement
+
+
+class AddRMSNormQuantSPPattern(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 AddRMSNormQuant fusion pattern.
+        """
+        rms_norm_input = torch.randn(2, 4, device="npu", dtype=self.dtype)
+        residual = torch.randn(2, 4, device="npu", dtype=self.dtype)
+        rms_norm_weight = torch.randn(4, device="npu", dtype=self.dtype)
+        scale = torch.ones(4, device="npu", dtype=self.dtype)
+        scale_reciprocal = torch.ones(4, device="npu", dtype=self.dtype)
+        offset = torch.zeros(4, device="npu", dtype=self.dtype)
+        return [rms_norm_input, residual, rms_norm_weight, scale, scale_reciprocal, offset]
+
+    def get_pattern(self):
+        def pattern(
+            rms_norm_input: torch.Tensor,
+            residual: torch.Tensor,
+            rms_norm_weight: torch.Tensor,
+            scale: torch.Tensor,
+            scale_reciprocal: torch.Tensor,
+            offset: torch.Tensor,
+        ):
+            """
+            Pattern for AddRMSNormQuant fusion.
+            """
+            output = torch.ops._C_ascend.npu_add_rms_norm_bias(
+                rms_norm_input, residual, rms_norm_weight, None, self.eps
+            )
+            out0 = output[0]
+            out1 = output[2]
+            out0 = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(out0, True)
+            quantized_output = torch.ops.vllm.quantize(out0, scale, scale_reciprocal, offset)
+            return quantized_output, out1
+
+        return pattern
+
+    def get_replacement(self):
+        def replacement(
+            rms_norm_input: torch.Tensor,
+            residual: torch.Tensor,
+            rms_norm_weight: torch.Tensor,
+            scale: torch.Tensor,
+            scale_reciprocal: torch.Tensor,
+            offset: torch.Tensor,
+        ):
+            """
+            Replacement for the AddRMSNormQuant fusion.
+            """
+            output = torch.ops.npu.npu_add_rms_norm_quant(
+                rms_norm_input, residual, rms_norm_weight, scale, offset, epsilon=self.eps
+            )
+            quantized_output = output[0]
+            out1 = output[2]
+            quantized_output = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(quantized_output, True)
+            return quantized_output, out1
+
+        return replacement
+
+
+class AddRMSNormQuantSPPatternWithBias(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 AddRMSNormQuant fusion pattern.
+        """
+        rms_norm_input = torch.randn(2, 4, device="npu", dtype=self.dtype)
+        residual = torch.randn(2, 4, device="npu", dtype=self.dtype)
+        rms_norm_weight = torch.randn(4, device="npu", dtype=self.dtype)
+        rmsnorm_bias = torch.randn(4, device="npu", dtype=self.dtype)
+        scale = torch.ones(4, device="npu", dtype=self.dtype)
+        scale_reciprocal = torch.ones(4, device="npu", dtype=self.dtype)
+        offset = torch.zeros(4, device="npu", dtype=self.dtype)
+        return [rms_norm_input, residual, rms_norm_weight, scale, scale_reciprocal, offset, rmsnorm_bias]
+
+    def get_pattern(self):
+        def pattern(
+            rms_norm_input: torch.Tensor,
+            residual: torch.Tensor,
+            rms_norm_weight: torch.Tensor,
+            scale: torch.Tensor,
+            scale_reciprocal: torch.Tensor,
+            offset: torch.Tensor,
+            bias: torch.Tensor,
+        ):
+            """
+            Pattern for AddRMSNormQuant fusion.
+            """
+            output = torch.ops._C_ascend.npu_add_rms_norm_bias(
+                rms_norm_input, residual, rms_norm_weight, bias, self.eps
+            )
+            out0 = output[0]
+            out1 = output[2]
+            out0 = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(out0, True)
+            quantized_output = torch.ops.vllm.quantize(out0, scale, scale_reciprocal, offset)
+            return quantized_output, out1
+
+        return pattern
+
+    def get_replacement(self):
+        def replacement(
+            rms_norm_input: torch.Tensor,
+            residual: torch.Tensor,
+            rms_norm_weight: torch.Tensor,
+            scale: torch.Tensor,
+            scale_reciprocal: torch.Tensor,
+            offset: torch.Tensor,
+            bias: torch.Tensor,
+        ):
+            """
+            Replacement for the AddRMSNormQuant fusion.
+            """
+            output = torch.ops.npu.npu_add_rms_norm_quant(
+                rms_norm_input, residual, rms_norm_weight, scale, offset, epsilon=self.eps, beta=bias
+            )
+            quantized_output = output[0]
+            out1 = output[2]
+            quantized_output = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(quantized_output, True)
+            return quantized_output, out1
+
+        return replacement
+
+
+class AddRMSNormDynamicQuantPattern(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 AddRMSNormQuant fusion pattern.
+        """
+        rms_norm_input = torch.randn(2, 4, device="npu", dtype=self.dtype)
+        residual = torch.randn(2, 4, device="npu", dtype=self.dtype)
+        rms_norm_weight = torch.randn(4, 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 AddRMSNormQuant 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_quant(out0)
+            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 AddRMSNormQuant fusion.
+            """
+            output = torch.ops.npu.npu_add_rms_norm_dynamic_quant(
+                rms_norm_input, residual, rms_norm_weight, epsilon=self.eps, output_mask=[True, False]
+            )
+            return (
+                output[0],
+                output[3],
+                output[2],
+            )
+
+        return replacement
+
+
+class AddRMSNormDynamicQuantPatternWithBias(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 AddRMSNormQuant fusion pattern.
+        """
+        rms_norm_input = torch.randn(2, 4, device="npu", dtype=self.dtype)
+        residual = torch.randn(2, 4, device="npu", dtype=self.dtype)
+        rms_norm_weight = torch.randn(4, device="npu", dtype=self.dtype)
+        rmsnorm_bias = torch.randn(4, device="npu", dtype=self.dtype)
+        return [rms_norm_input, residual, rms_norm_weight, rmsnorm_bias]
+
+    def get_pattern(self):
+        def pattern(
+            rms_norm_input: torch.Tensor,
+            residual: torch.Tensor,
+            rms_norm_weight: torch.Tensor,
+            bias: torch.Tensor,
+        ):
+            """
+            Pattern for AddRMSNormQuant fusion.
+            """
+            output = torch.ops._C_ascend.npu_add_rms_norm_bias(
+                rms_norm_input, residual, rms_norm_weight, bias, self.eps
+            )
+            out0 = output[0]
+            out1 = output[2]
+            quantized_output = torch.ops.npu.npu_dynamic_quant(out0)
+            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,
+            bias: torch.Tensor,
+        ):
+            """
+            Replacement for the AddRMSNormQuant fusion.
+            """
+            output = torch.ops.npu.npu_add_rms_norm_dynamic_quant(
+                rms_norm_input, residual, rms_norm_weight, epsilon=self.eps, output_mask=[True, False], beta=bias
+            )
+            return (
+                output[0],
+                output[3],
+                output[2],
+            )
+
+        return replacement
+
+
+class AddRMSNormDynamicQuantSPPattern(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 AddRMSNormQuant fusion pattern.
+        """
+        rms_norm_input = torch.randn(2, 4, device="npu", dtype=self.dtype)
+        residual = torch.randn(2, 4, device="npu", dtype=self.dtype)
+        rms_norm_weight = torch.randn(4, 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 AddRMSNormQuant 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_quant(out0)
+            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 AddRMSNormQuant fusion.
+            """
+            output = torch.ops.npu.npu_add_rms_norm_dynamic_quant(
+                rms_norm_input, residual, rms_norm_weight, epsilon=self.eps, output_mask=[True, False]
+            )
+            out3 = output[3]
+            quantized_output = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(output[0], True)
+            out3 = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(out3, True)
+            return quantized_output, out3, output[2]
+
+        return replacement
+
+
+class AddRMSNormDynamicQuantSPPatternWithBias(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 AddRMSNormQuant fusion pattern.
+        """
+        rms_norm_input = torch.randn(2, 4, device="npu", dtype=self.dtype)
+        residual = torch.randn(2, 4, device="npu", dtype=self.dtype)
+        rms_norm_weight = torch.randn(4, device="npu", dtype=self.dtype)
+        rmsnorm_bias = torch.randn(4, device="npu", dtype=self.dtype)
+        return [rms_norm_input, residual, rms_norm_weight, rmsnorm_bias]
+
+    def get_pattern(self):
+        def pattern(
+            rms_norm_input: torch.Tensor,
+            residual: torch.Tensor,
+            rms_norm_weight: torch.Tensor,
+            bias: torch.Tensor,
+        ):
+            """
+            Pattern for AddRMSNormQuant fusion.
+            """
+            output = torch.ops._C_ascend.npu_add_rms_norm_bias(
+                rms_norm_input, residual, rms_norm_weight, bias, 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_quant(out0)
+            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,
+            bias: torch.Tensor,
+        ):
+            """
+            Replacement for the AddRMSNormQuant fusion.
+            """
+            output = torch.ops.npu.npu_add_rms_norm_dynamic_quant(
+                rms_norm_input, residual, rms_norm_weight, epsilon=self.eps, output_mask=[True, False], beta=bias
+            )
+            out3 = output[3]
+            quantized_output = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(output[0], True)
+            out3 = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(out3, True)
+            return quantized_output, out3, output[2]
+
+        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
+
+
+def _model_uses_w4a4_quant(vllm_config: VllmConfig | None) -> bool:
+    """Check whether the model uses W4A4 int4 quantization for any layer.
+
+    W4A4 int4 schemes (e.g. W4A4_DYNAMIC, W4A4_FLATQUANT_DYNAMIC)
+    are incompatible with the fuse_norm_quant optimization,
+    so callers use this to disable that fusion.
+    """
+    if vllm_config is None:
+        return False
+    quant_config = getattr(vllm_config, "quant_config", None)
+    if quant_config is None:
+        return False
+    quant_description = getattr(quant_config, "quant_description", None)
+    if not quant_description:
+        return False
+    w4a4_int4_schemes = ["W4A4_DYNAMIC", "W4A4_FLATQUANT_DYNAMIC"]
+    return any(
+        isinstance(quant_type, str) and quant_type in w4a4_int4_schemes for quant_type in quant_description.values()
+    )
+
+
+class AddRMSNormQuantFusionPass(VllmInductorPass):
+    """
+    A pass for fusing AddRMSNorm and W8A8 quantization operations on Ascend.
+    """
+
+    def __init__(self, vllm_config: VllmConfig):
+        super().__init__(vllm_config)
+        self.pattern_match_passes: PatternMatcherPass = PatternMatcherPass(pass_name="rmsnorm_quant_fusion_pass")
+
+        dtype = vllm_config.model_config.dtype
+        if dtype not in (torch.bfloat16, torch.float16):
+            logger.debug("Quant fusion not enabled: unsupported dtype %s", dtype)
+            return
+
+        if _model_uses_w4a4_quant(vllm_config):
+            logger.debug(
+                "Quant fusion not enabled: the model contains "
+                "W4A4 quantized weights, which are incompatible with the "
+                "norm-quant fusion pass."
+            )
+            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)
+                AddRMSNormQuantPatternWithBias(vllm_config, eps=eps).register(self.pattern_match_passes)
+                AddRMSNormQuantSPPatternWithBias(vllm_config, eps=eps).register(self.pattern_match_passes)
+                AddRMSNormDynamicQuantPatternWithBias(vllm_config, eps=eps).register(self.pattern_match_passes)
+                AddRMSNormDynamicQuantSPPatternWithBias(vllm_config, eps=eps).register(self.pattern_match_passes)
+
+    def __call__(self, graph: torch.fx.Graph):
+        self.begin()
+        self.matched_count = self.pattern_match_passes.apply(graph)
+        logger.debug("Replaced %s 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.
+        """
+        return True

vllm_ascend/compilation/passes/qknorm_rope_fusion_pass.py

@@ -0,0 +1,244 @@
+#
+# 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.
+#
+import torch
+from torch._inductor.pattern_matcher import PatternMatcherPass, PatternPrettyPrinter
+from vllm.compilation.passes.vllm_inductor_pass import VllmInductorPass
+from vllm.config import VllmConfig, get_layers_from_vllm_config
+from vllm.config.compilation import Range
+from vllm.logger import logger
+from vllm.model_executor.layers.attention import Attention
+
+from vllm_ascend.compilation.passes.base_pattern import BasePattern
+from vllm_ascend.device.device_op import DeviceOperator
+from vllm_ascend.utils import get_rope_dim
+
+
+class QKNormRopeFusionPattern(BasePattern):
+    def __init__(self, vllm_config, head_dim, num_heads, num_kv_heads, eps=1e-6):
+        super().__init__(vllm_config, eps)
+        self.head_dim = head_dim
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.device = vllm_config.device_config.device if vllm_config.device_config else None
+        self.rope_dim = get_rope_dim(vllm_config)
+
+    def get_inputs(self):
+        T = 5
+        max_position_embeddings = 16384
+        qkv = torch.empty(T, self.q_size + 2 * self.kv_size, dtype=torch.bfloat16, device="npu")
+        q_weight = torch.empty(self.head_dim, dtype=torch.bfloat16, device="npu")
+        k_weight = torch.empty(self.head_dim, dtype=torch.bfloat16, device="npu")
+        cos_sin_cache = torch.empty(max_position_embeddings, self.head_dim, dtype=torch.bfloat16, device="npu")
+        positions = torch.ones(T, dtype=torch.int64, device="npu")
+        return [qkv, q_weight, k_weight, cos_sin_cache, positions]
+
+    def get_pattern(self):
+        def pattern(
+            qkv: torch.Tensor,
+            q_weight: torch.Tensor,
+            k_weight: torch.Tensor,
+            cos_sin_cache: torch.Tensor,
+            positions: torch.Tensor,
+        ):
+            q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+
+            q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim)
+            q_norm_out, _ = torch.ops.npu.npu_rms_norm(q_by_head, q_weight, self.eps)
+
+            k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim)
+            k_norm_out, _ = torch.ops.npu.npu_rms_norm(k_by_head, k_weight, self.eps)
+
+            q_flat = q_norm_out.view(q.shape)
+            k_flat = k_norm_out.view(k.shape)
+            q_rope, k_rope = torch.ops.vllm.npu_rotary_embedding(
+                positions, q_flat, k_flat, cos_sin_cache, self.head_dim, self.rope_dim, True
+            )
+
+            return q_rope, k_rope, v
+
+        return pattern
+
+    def get_replacement(self):
+        def replacement(
+            qkv: torch.Tensor,
+            q_weight: torch.Tensor,
+            k_weight: torch.Tensor,
+            cos_sin_cache: torch.Tensor,
+            positions: torch.Tensor,
+        ):
+            results = DeviceOperator.split_qkv_rmsnorm_rope(
+                input=qkv,
+                q_weight=q_weight,
+                k_weight=k_weight,
+                q_hidden_size=self.q_size,
+                kv_hidden_size=self.kv_size,
+                head_dim=self.head_dim,
+                eps=self.eps,
+                q_bias=None,
+                k_bias=None,
+                cos_sin_cache=cos_sin_cache,
+                positions=positions,
+            )
+
+            return results
+
+        return replacement
+
+
+class QKNormRopeFusionPatternWithBias(BasePattern):
+    def __init__(self, vllm_config, head_dim, num_heads, num_kv_heads, eps=1e-6):
+        super().__init__(vllm_config, eps)
+        self.head_dim = head_dim
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.device = vllm_config.device_config.device if vllm_config.device_config else None
+        self.rope_dim = get_rope_dim(vllm_config)
+
+    def get_inputs(self):
+        T = 5
+        max_position_embeddings = 16384
+        qkv = torch.empty(T, self.q_size + 2 * self.kv_size, dtype=torch.bfloat16, device="npu")
+        q_weight = torch.empty(self.head_dim, dtype=torch.bfloat16, device="npu")
+        k_weight = torch.empty(self.head_dim, dtype=torch.bfloat16, device="npu")
+        q_bias = torch.empty(self.head_dim, dtype=torch.bfloat16, device="npu")
+        k_bias = torch.empty(self.head_dim, dtype=torch.bfloat16, device="npu")
+        cos_sin_cache = torch.empty(max_position_embeddings, self.head_dim, dtype=torch.bfloat16, device="npu")
+        positions = torch.ones(T, dtype=torch.int64, device="npu")
+
+        return [qkv, q_weight, k_weight, q_bias, k_bias, cos_sin_cache, positions]
+
+    def get_pattern(self):
+        def pattern(
+            qkv: torch.Tensor,
+            q_weight: torch.Tensor,
+            k_weight: torch.Tensor,
+            q_bias: torch.Tensor,
+            k_bias: torch.Tensor,
+            cos_sin_cache: torch.Tensor,
+            positions: torch.Tensor,
+        ):
+            q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+
+            q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim)
+            q_norm_out, _ = torch.ops.npu.npu_rms_norm(q_by_head, q_weight, self.eps)
+            q_normed = q_norm_out + q_bias
+
+            k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim)
+            k_norm_out, _ = torch.ops.npu.npu_rms_norm(k_by_head, k_weight, self.eps)
+            k_normed = k_norm_out + k_bias
+
+            q_flat = q_normed.view(q.shape)
+            k_flat = k_normed.view(k.shape)
+            q_rope, k_rope = torch.ops.vllm.npu_rotary_embedding(
+                positions, q_flat, k_flat, cos_sin_cache, self.head_dim, self.rope_dim, True
+            )
+
+            return q_rope, k_rope, v
+
+        return pattern
+
+    def get_replacement(self):
+        def replacement(
+            qkv: torch.Tensor,
+            q_weight: torch.Tensor,
+            k_weight: torch.Tensor,
+            q_bias: torch.Tensor,
+            k_bias: torch.Tensor,
+            cos_sin_cache: torch.Tensor,
+            positions: torch.Tensor,
+        ):
+            results = DeviceOperator.split_qkv_rmsnorm_rope(
+                input=qkv,
+                q_weight=q_weight,
+                k_weight=k_weight,
+                q_hidden_size=self.q_size,
+                kv_hidden_size=self.kv_size,
+                head_dim=self.head_dim,
+                eps=self.eps,
+                q_bias=q_bias,
+                k_bias=k_bias,
+                cos_sin_cache=cos_sin_cache,
+                positions=positions,
+            )
+            return results
+
+        return replacement
+
+
+class QKNormRopeFusionPass(VllmInductorPass):
+    """
+    A pass for fusing QKV split and RMSNorm operations into a single qk_rmsnorm operator.
+    """
+
+    def __init__(self, vllm_config: VllmConfig):
+        super().__init__(vllm_config)
+        self.pattern_match_passes: PatternMatcherPass = PatternMatcherPass(pass_name="qknorm_rope_fusion_pass")
+
+        dtype = vllm_config.model_config.dtype
+        if dtype not in (torch.bfloat16,):
+            logger.debug("QKNorm and Rope fusion not enabled: unsupported dtype %s", dtype)
+            return
+
+        # use one attn layer to get meta (such as head_dim) for QKNormRopeFusionPattern
+        attn_layers: dict[str, Attention] = get_layers_from_vllm_config(vllm_config, Attention)
+        if len(attn_layers) == 0:
+            logger.debug("QKNorm and Rope fusion enabled, but no Attention layers were discovered.")
+            return
+        layer = next(iter(attn_layers.values()))
+        for epsilon in [1e-6, 1e-5]:
+            if layer.head_size != 128:
+                logger.debug("QKNorm and Rope fusion not enabled: head_dim %d is not equal of 128", layer.head_size)
+                continue
+            QKNormRopeFusionPattern(
+                vllm_config=vllm_config,
+                head_dim=layer.head_size,
+                num_heads=layer.num_heads,
+                num_kv_heads=layer.num_kv_heads,
+                eps=epsilon,
+            ).register(self.pattern_match_passes)
+
+            QKNormRopeFusionPatternWithBias(
+                vllm_config=vllm_config,
+                head_dim=layer.head_size,
+                num_heads=layer.num_heads,
+                num_kv_heads=layer.num_kv_heads,
+                eps=epsilon,
+            ).register(self.pattern_match_passes)
+
+    def __call__(self, graph: torch.fx.Graph):
+        self.begin()
+        self.matched_count = self.pattern_match_passes.apply(graph)
+        logger.debug("Fused %s QKNorm and Rope patterns", self.matched_count)
+        logger.debug("Patterns registered for replacement:")
+        pattern_idx = 0
+        for pattern_entry in self.pattern_match_passes.patterns.values():
+            for p in pattern_entry:
+                p_str = PatternPrettyPrinter.run(p.pattern)
+                logger.debug("Pattern %d: %s", pattern_idx, p_str)
+                pattern_idx += 1
+        self.end_and_log()
+
+    def is_applicable_for_range(self, compile_range: Range) -> bool:
+        """
+        Check if the pass is applicable for the current configuration.
+        """
+        return True

vllm_ascend/compilation/passes/sequence_parallelism.py

@@ -0,0 +1,234 @@
+import torch
+import torch._inductor.pattern_matcher as pm
+from torch._inductor.pattern_matcher import PatternMatcherPass
+from vllm.compilation.passes.vllm_inductor_pass import VllmInductorPass
+from vllm.config import VllmConfig
+from vllm.config.utils import Range
+from vllm.distributed import get_tensor_model_parallel_world_size, get_tp_group, tensor_model_parallel_all_reduce
+from vllm.logger import logger
+
+from vllm_ascend.compilation.passes.noop_elimination import NoOpEliminationPass
+from vllm_ascend.utils import is_moe_model
+
+SP_MIN_TOKEN_NUM_DEFAULT = 1000
+
+
+def get_sp_min_token_num(config: VllmConfig) -> int:
+    if is_moe_model(config):
+        return 1
+
+    return SP_MIN_TOKEN_NUM_DEFAULT
+
+
+class _SequenceParallelPatternHelper:
+    """Helper for sequence parallelism patterns.
+
+    Provides TP communication helper methods: _all_reduce, _reduce_scatter,
+    _all_gather, and tensor creation utilities.
+    """
+
+    def __init__(
+        self,
+        epsilon: float,
+        dtype: torch.dtype,
+        device: str,
+    ):
+        self.eps = epsilon
+        self.dtype = dtype
+        self.device = device
+        self.tp_group = get_tp_group()
+        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_rank = get_tp_group().rank_in_group
+
+    def _all_reduce(self, x: torch.Tensor) -> torch.Tensor:
+        return tensor_model_parallel_all_reduce(x)
+
+    def _reduce_scatter(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.ops.vllm.reduce_scatter(x, dim=0, world_size=self.tp_size, group_name=self.tp_group.unique_name)
+
+    def _all_gather(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.ops.vllm.all_gather(x, dim=0, world_size=self.tp_size, group_name=self.tp_group.unique_name)
+
+    def empty(self, *args, **kws):
+        return torch.empty(*args, dtype=self.dtype, device="npu", **kws)
+
+
+class MiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
+    """Replaces all_reduce + AddRMSNormBias with reduce_scatter + AddRMSNormBias
+    + all_gather for middle-layer sequence parallelism."""
+
+    def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6):
+        super().__init__(eps, vllm_config.model_config.dtype, torch.npu.current_device())
+
+    def empty(self, *args, **kws):
+        return torch.empty(*args, dtype=self.dtype, device="npu", **kws)
+
+    def get_inputs(self):
+        """
+        Generate example inputs.
+        """
+        input = self.empty(8, 16)
+        weight = self.empty(16)
+        residual = self.empty(8, 16)
+        return [input, weight, residual]
+
+    def register(self, pm_pass: PatternMatcherPass):
+        def pattern(
+            input: torch.Tensor,
+            weight: torch.Tensor,
+            residual: torch.Tensor,
+        ) -> tuple[torch.Tensor, torch.Tensor]:
+            x = self._all_reduce(input)
+            result, _, residual = torch.ops._C_ascend.npu_add_rms_norm_bias(x, residual, weight, None, self.eps)
+
+            return result, residual
+
+        def replacement(
+            input: torch.Tensor,
+            weight: torch.Tensor,
+            residual: torch.Tensor,
+        ) -> tuple[torch.Tensor, torch.Tensor]:
+            reduce_scatter = self._reduce_scatter(input)
+            residual = torch.ops.vllm.maybe_chunk_residual(reduce_scatter, residual)
+            result, _, residual = torch.ops._C_ascend.npu_add_rms_norm_bias(
+                reduce_scatter, residual, weight, None, self.eps
+            )
+            all_gather = self._all_gather(result)
+            return all_gather, residual
+
+        pm.register_replacement(pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass)
+
+
+class LastAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
+    """Same as MiddleAllReduceRMSNormPattern but for the last layer
+    (no residual backprop)."""
+
+    def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6):
+        super().__init__(eps, vllm_config.model_config.dtype, torch.npu.current_device())
+
+    def get_inputs(self):
+        input = self.empty(8, 16)
+        weight = self.empty(16)
+        residual = self.empty(8, 16)
+        return [input, weight, residual]
+
+    def register(self, pm_pass: PatternMatcherPass):
+        def pattern(
+            input: torch.Tensor,
+            weight: torch.Tensor,
+            residual: torch.Tensor,
+        ) -> torch.Tensor:
+            x = self._all_reduce(input)
+            result, _, _ = torch.ops._C_ascend.npu_add_rms_norm_bias(x, residual, weight, None, self.eps)
+
+            return result
+
+        def replacement(
+            input: torch.Tensor,
+            weight: torch.Tensor,
+            residual: torch.Tensor,
+        ) -> torch.Tensor:
+            reduce_scatter = self._reduce_scatter(input)
+            residual = torch.ops.vllm.maybe_chunk_residual(reduce_scatter, residual)
+            result, _, _ = torch.ops._C_ascend.npu_add_rms_norm_bias(reduce_scatter, residual, weight, None, self.eps)
+            all_gather = self._all_gather(result)
+            return all_gather
+
+        pm.register_replacement(pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass)
+
+
+class Qwen3VLMiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
+    """For Qwen3-VL middle layers with hidden_states + deepstack_input_embeds add.
+
+    Replaces all_reduce + add + AddRMSNormBias with reduce_scatter +
+    chunk(deepstack_input_embeds) + add + AddRMSNormBias + all_gather.
+    """
+
+    def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6):
+        super().__init__(eps, vllm_config.model_config.dtype, torch.npu.current_device())
+
+    def get_inputs(self):
+        input = self.empty(8, 16)
+        weight = self.empty(16)
+        residual = self.empty(8, 16)
+        deepstack_input_embeds = self.empty(8, 16)
+        return [input, weight, residual, deepstack_input_embeds]
+
+    def register(self, pm_pass: PatternMatcherPass):
+        def pattern(
+            input: torch.Tensor,
+            weight: torch.Tensor,
+            residual: torch.Tensor,
+            deepstack_input_embeds: torch.Tensor,
+        ) -> tuple[torch.Tensor, torch.Tensor]:
+            x = self._all_reduce(input)
+            add_ = x + deepstack_input_embeds
+            result, _, residual = torch.ops._C_ascend.npu_add_rms_norm_bias(add_, residual, weight, None, self.eps)
+
+            return result, residual
+
+        def replacement(
+            input: torch.Tensor,
+            weight: torch.Tensor,
+            residual: torch.Tensor,
+            deepstack_input_embeds: torch.Tensor,
+        ) -> tuple[torch.Tensor, torch.Tensor]:
+            reduce_scatter = self._reduce_scatter(input)
+            chunk = deepstack_input_embeds.chunk(self.tp_size)[self.tp_rank]
+            add_ = reduce_scatter + chunk
+            residual = torch.ops.vllm.maybe_chunk_residual(reduce_scatter, residual)
+            result, _, residual = torch.ops._C_ascend.npu_add_rms_norm_bias(add_, residual, weight, None, self.eps)
+            all_gather = self._all_gather(result)
+            return all_gather, residual
+
+        pm.register_replacement(pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass)
+
+
+class SequenceParallelismPass(VllmInductorPass):
+    """Sequence parallelism compilation pass.
+
+    Registers and applies the above patterns. Runs noop cleanup first, then
+    uses token range to determine whether to enable SP.
+    """
+
+    def __init__(self, config: VllmConfig):
+        super().__init__(config)
+
+        self.patterns: PatternMatcherPass = PatternMatcherPass(pass_name="npu_sequence_parallelism_pass")
+        self.noop_cleanup = NoOpEliminationPass(config)
+
+        for epsilon in [1e-5, 1e-6]:
+            MiddleAllReduceRMSNormPattern(config, epsilon).register(self.patterns)
+
+            LastAllReduceRMSNormPattern(config, epsilon).register(self.patterns)
+
+            Qwen3VLMiddleAllReduceRMSNormPattern(config, epsilon).register(self.patterns)
+
+        self.min_tokens = get_sp_min_token_num(config)
+
+    def __call__(self, graph: torch.fx.Graph):
+        self.begin()
+        self.noop_cleanup(graph)  # Eliminate redundant view-like operations
+        logger.debug("after noop_cleanup %s", graph.graph)
+        self.matched_count = self.patterns.apply(graph)
+        logger.debug("Replaced %s patterns", self.matched_count)
+        logger.debug("after apply replacement %s", graph.graph)
+
+        from torch._inductor.pattern_matcher import PatternPrettyPrinter
+
+        pattern_idx = 0
+        for pattern_entry in self.patterns.patterns.values():
+            for p in pattern_entry:
+                p_str = PatternPrettyPrinter.run(p.pattern)
+                logger.debug("Pattern %d: %s", pattern_idx, p_str)
+                pattern_idx += 1
+
+        self.end_and_log()
+
+    def is_applicable_for_range(self, compile_range: Range) -> bool:
+        """
+        Check if the pass is applicable for the current configuration.
+        """
+        applicable = compile_range.start >= self.min_tokens
+        logger.debug("SequenceParallelismPass compile_range=%r applicable=%r", compile_range, applicable)
+        return applicable

vllm_ascend/compilation/passes/sequence_parallelism_moe.py

@@ -0,0 +1,204 @@
+import torch
+import torch._inductor.pattern_matcher as pm
+from torch._inductor.pattern_matcher import PatternMatcherPass
+from vllm.compilation.passes.vllm_inductor_pass import PatternPrettyPrinter, VllmInductorPass
+from vllm.config import VllmConfig
+from vllm.config.utils import Range
+from vllm.logger import logger
+
+from vllm_ascend.compilation.passes.sequence_parallelism import (
+    _SequenceParallelPatternHelper,
+    get_sp_min_token_num,
+)
+
+
+class MiddleLayerAllgatherAddRMSNormPattern(_SequenceParallelPatternHelper):
+    """Replaces all_gather + slice + AddRMSNormBias with AddRMSNormBias +
+    all_gather to avoid middle-layer shape mismatch."""
+
+    def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6):
+        super().__init__(eps, vllm_config.model_config.dtype, torch.npu.current_device())
+
+    def get_inputs(self):
+        input = self.empty(5, 16)
+        weight = self.empty(16)
+        residual = self.empty(8, 16)
+        # num_tokens = 8
+        return [input, weight, residual]
+
+    def get_scalar_inputs(self):
+        return {"num_tokens": 8}
+
+    def register(self, pm_pass: PatternMatcherPass):
+        def pattern(
+            input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor, num_tokens
+        ) -> tuple[torch.Tensor, torch.Tensor]:
+            all_gather = self._all_gather(input)
+            x_sliced = all_gather[:num_tokens]
+            result, _, residual = torch.ops._C_ascend.npu_add_rms_norm_bias(x_sliced, residual, weight, None, self.eps)
+
+            return result, residual
+
+        def replacement(
+            input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor, num_tokens
+        ) -> tuple[torch.Tensor, torch.Tensor]:
+            residual = torch.ops.vllm.maybe_chunk_residual(input, residual)
+            result, _, residual = torch.ops._C_ascend.npu_add_rms_norm_bias(input, residual, weight, None, self.eps)
+            all_gather = self._all_gather(result)
+            return all_gather, residual
+
+        pm.register_replacement(
+            pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass, scalar_workaround=self.get_scalar_inputs()
+        )
+
+
+class LastLayerAllgatherRMSNormPattern(_SequenceParallelPatternHelper):
+    """Same as MiddleLayerAllgatherAddRMSNormPattern but for the last layer (no residual)
+    all_gather + RMSNorm fusion."""
+
+    def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6):
+        super().__init__(eps, vllm_config.model_config.dtype, torch.npu.current_device())
+
+    def get_inputs(self):
+        input = self.empty(5, 16)
+        weight = self.empty(16)
+        residual = self.empty(8, 16)
+        return [input, weight, residual]
+
+    def get_scalar_inputs(self):
+        return {"num_tokens": 8}
+
+    def register(self, pm_pass: PatternMatcherPass):
+        def pattern(
+            input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor, num_tokens
+        ) -> tuple[torch.Tensor, torch.Tensor]:
+            all_gather = self._all_gather(input)
+            x_sliced = all_gather[:num_tokens]
+            result, _, _ = torch.ops._C_ascend.npu_add_rms_norm_bias(x_sliced, residual, weight, None, self.eps)
+
+            return result
+
+        def replacement(
+            input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor, num_tokens
+        ) -> tuple[torch.Tensor, torch.Tensor]:
+            residual = torch.ops.vllm.maybe_chunk_residual(input, residual)
+            result, _, _ = torch.ops._C_ascend.npu_add_rms_norm_bias(input, residual, weight, None, self.eps)
+            all_gather = self._all_gather(result)
+            return all_gather
+
+        pm.register_replacement(
+            pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass, scalar_workaround=self.get_scalar_inputs()
+        )
+
+
+class Qwen3VLMiddleLayerAllgatherAddRMSNormPattern(_SequenceParallelPatternHelper):
+    """Replaces all_gather + slice + add + AddRMSNormBias with add(chunk) +
+    AddRMSNormBias + all_gather for Qwen3-VL-style all_gather path."""
+
+    def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6):
+        super().__init__(eps, vllm_config.model_config.dtype, torch.npu.current_device())
+
+    def get_inputs(self):
+        input = self.empty(5, 16)
+        weight = self.empty(16)
+        residual = self.empty(8, 16)
+        deepstack_input_embeds = self.empty(8, 16)
+        return [input, weight, residual, deepstack_input_embeds]
+
+    def get_scalar_inputs(self):
+        return {"num_tokens": 8}
+
+    def register(self, pm_pass: PatternMatcherPass):
+        def pattern(
+            input: torch.Tensor,
+            weight: torch.Tensor,
+            residual: torch.Tensor,
+            deepstack_input_embeds: torch.Tensor,
+            num_tokens,
+        ) -> tuple[torch.Tensor, torch.Tensor]:
+            all_gather = self._all_gather(input)
+            x_sliced = all_gather[:num_tokens]
+            add_ = x_sliced + deepstack_input_embeds
+            result, _, residual = torch.ops._C_ascend.npu_add_rms_norm_bias(add_, residual, weight, None, self.eps)
+
+            return result, residual
+
+        def replacement(
+            input: torch.Tensor,
+            weight: torch.Tensor,
+            residual: torch.Tensor,
+            deepstack_input_embeds: torch.Tensor,
+            num_tokens,
+        ) -> tuple[torch.Tensor, torch.Tensor]:
+            chunk = deepstack_input_embeds.chunk(self.tp_size)[self.tp_rank]
+            add_ = input + chunk
+            residual = torch.ops.vllm.maybe_chunk_residual(input, residual)
+            result, _, residual = torch.ops._C_ascend.npu_add_rms_norm_bias(add_, residual, weight, None, self.eps)
+            all_gather = self._all_gather(result)
+            return all_gather, residual
+
+        pm.register_replacement(
+            pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass, scalar_workaround=self.get_scalar_inputs()
+        )
+
+
+class AllGatherChunkNoOpPattern(_SequenceParallelPatternHelper):
+    """Folds all_gather + sequence_parallel_chunk_impl into identity (no-op)."""
+
+    def __init__(self, vllm_config: VllmConfig, eps: float = 1e-6):
+        super().__init__(eps, vllm_config.model_config.dtype, torch.npu.current_device())
+
+    def get_inputs(self):
+        return [self.empty(8, 16)]
+
+    def register(self, pm_pass: PatternMatcherPass):
+        def pattern(input: torch.Tensor) -> torch.Tensor:
+            gathered = self._all_gather(input)
+            return torch.ops.vllm.sequence_parallel_chunk_impl(gathered)
+
+        def replacement(input: torch.Tensor) -> torch.Tensor:
+            return input
+
+        pm.register_replacement(pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass)
+
+
+class SequenceParallelismMoePass(VllmInductorPass):
+    """Sequence parallelism AllGather epilogue pass.
+
+    Applies AllGather-based patterns: MiddleLayerAllgatherAddRMSNormPattern,
+    LastLayerAllgatherRMSNormPattern, Qwen3VLMiddleLayerAllgatherAddRMSNormPattern,
+    and AllGatherChunkNoOpPattern (all_gather + sequence_parallel_chunk_impl -> identity).
+    """
+
+    def __init__(self, config: VllmConfig):
+        super().__init__(config)
+
+        self.patterns: PatternMatcherPass = PatternMatcherPass(pass_name="npu_sequence_parallelism_allgather_ep_pass")
+
+        for epsilon in [1e-5, 1e-6]:
+            MiddleLayerAllgatherAddRMSNormPattern(config, epsilon).register(self.patterns)
+            LastLayerAllgatherRMSNormPattern(config, epsilon).register(self.patterns)
+            Qwen3VLMiddleLayerAllgatherAddRMSNormPattern(config, epsilon).register(self.patterns)
+
+        AllGatherChunkNoOpPattern(config).register(self.patterns)
+
+        self.min_tokens = get_sp_min_token_num(config)
+
+    def __call__(self, graph: torch.fx.Graph):
+        self.begin()
+        logger.debug("before apply replacement %s", str(graph))
+        self.matched_count = self.patterns.apply(graph)
+        logger.debug("after apply replacement %s", str(graph))
+        logger.debug("SequenceParallelismMoePass replaced %s patterns", self.matched_count)
+        pattern_idx = 0
+        for pattern_entry in self.patterns.patterns.values():
+            for p in pattern_entry:
+                p_str = PatternPrettyPrinter.run(p.pattern)
+                logger.debug("Pattern %d: %s", pattern_idx, p_str)
+                pattern_idx += 1
+        self.end_and_log()
+
+    def is_applicable_for_range(self, compile_range: Range) -> bool:
+        applicable = compile_range.start >= self.min_tokens
+        logger.debug("SequenceParallelismMoePass compile_range=%r applicable=%r", compile_range, applicable)
+        return applicable

vllm_ascend/compilation/passes/utils/npugraph_ex_utils_check.py

@@ -0,0 +1,75 @@
+#
+# 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 torch._inductor.pattern_matcher import Match
+from vllm.logger import logger
+
+
+def extra_stream_scope_check(match: Match) -> bool:
+    """
+    Checks if all nodes in the same stream.
+    """
+    non_default_streams = set()
+    has_default = False
+
+    for node in match.nodes:
+        if node.op == "call_function":
+            current_stream = node.meta.get("stream_label")
+            if current_stream is None:
+                has_default = True
+            else:
+                non_default_streams.add(current_stream)
+                if len(non_default_streams) > 1:
+                    logger.debug(
+                        "Cross-stream operation detected in pattern match for AddRMSNormQuant. "
+                        "Multiple streams found: %s. Fusion is not supported for cross-stream operations.",
+                        non_default_streams,
+                    )
+                    return False
+
+    if has_default and len(non_default_streams) > 0:
+        logger.debug(
+            "Cross-stream operation detected in pattern match for AddRMSNormQuant. "
+            "Multiple streams found: %s. Fusion is not supported for cross-stream operations.",
+            non_default_streams,
+        )
+        return False
+
+    return True
+
+
+_register_patterns = set()
+
+
+def check_and_register_fusion_pass(pattern_class: type, **kwargs):
+    global _register_patterns
+    eps = kwargs.get("eps", 1e-6)
+    pattern_key = str(pattern_class.__name__) + str(eps)
+    if pattern_key in _register_patterns:
+        return
+
+    pattern = pattern_class(**kwargs)
+    try:
+        pattern.register()
+        _register_patterns.add(pattern_key)
+    except RuntimeError as e:
+        if "Duplicate pattern" in str(e):
+            logger.warning("Pattern %s eps %s has been registered", pattern_class.__name__, eps)
+            _register_patterns.add(pattern_key)
+        else:
+            raise e