init v0.23.0

Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>
2026-08-27 15:11:51 +08:00
parent b582a8e7d1
commit 7f8a1b1f7a
2849 changed files with 712887 additions and 22001 deletions
--- a/vllm_ascend/compilation/passes/sequence_parallelism_moe.py
+++ b/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