handle code conflict

Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>

Author: weijinqian_v1

tests/singlecard/test_aclgraph.py

@@ -103,3 +103,16 @@ def test_deepseek_raises_error(monkeypatch: pytest.MonkeyPatch) -> None:
                        max_model_len=1024,
                        enforce_eager=False)
         assert "ACL Graph does not support deepseek" in str(excinfo.value)
+
+
+@pytest.mark.skipif(os.getenv("VLLM_USE_V1") == "0",
+                    reason="aclgraph only support on v1")
+@pytest.mark.parametrize("model", MODELS)
+def test_ray_backend_sets_no_compilation(
+        model: str, monkeypatch: pytest.MonkeyPatch) -> None:
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_USE_V1", "1")
+        runner = VllmRunner(model,
+                            enforce_eager=False,
+                            distributed_executor_backend="ray")
+        assert runner.model.llm_engine.vllm_config.compilation_config.level == 0

vllm_ascend/ascend_forward_context.py

@@ -1,11 +1,13 @@
+import math
 from contextlib import contextmanager
 from enum import Enum
 from typing import Any, Optional
 
 import torch
 from vllm.config import VllmConfig
-from vllm.distributed import get_dp_group
+from vllm.distributed import get_dp_group, get_tp_group
 from vllm.forward_context import get_forward_context, set_forward_context
+from vllm.platforms import current_platform
 
 import vllm_ascend.envs as envs_ascend
 
@@ -15,16 +17,20 @@ class FusedMoEState(Enum):
     All2All = 1
     MC2 = 2
     All2AllSeq = 3
+    MC2_PREFILL = 4
 
 
 # TODO(zzzzwwjj): add soc_version to choose branch
 def get_fused_moe_state(ep_size: int, with_prefill: bool):
+    enable_chunk_mc2 = envs_ascend.VLLM_ASCEND_ENABLE_CHUNK_MC2
     if ep_size == 1:
         return FusedMoEState.AllGather
     elif envs_ascend.VLLM_ASCEND_ENABLE_MOE_ALL2ALL_SEQ:
         # MC2 Dispatch/Combine performs better than alltoall_seq in decoding stage.
         return FusedMoEState.All2AllSeq if (
             ep_size < 16 or with_prefill) else FusedMoEState.MC2
+    elif ep_size >= 16 and with_prefill and enable_chunk_mc2:
+        return FusedMoEState.MC2_PREFILL
     # NOTE: mc2 need ep_size >= 16 & all2all can't use in torchair graph.
     elif ep_size < 16 or with_prefill:
         return FusedMoEState.All2All
@@ -40,7 +46,8 @@ def set_ascend_forward_context(
         num_tokens: Optional[int] = None,
         num_tokens_across_dp: Optional[torch.Tensor] = None,
         with_prefill: bool = True,
-        in_profile_run: bool = False):
+        in_profile_run: bool = False,
+        num_actual_tokens: Optional[int] = None):
     """A context manager that stores the current forward context,
     can be attention metadata, etc.
     We add some additional param into forward_context.
@@ -52,7 +59,6 @@ def set_ascend_forward_context(
                              num_tokens_across_dp=num_tokens_across_dp):
         forward_context = get_forward_context()
         forward_context.with_prefill = with_prefill
-
         ep_size = torch.distributed.get_world_size(
         ) if vllm_config.parallel_config.enable_expert_parallel else 1
 
@@ -66,19 +72,55 @@ def set_ascend_forward_context(
         # due to multiple warmups before actual capturing
         forward_context.capturing = False
 
+        if num_tokens is None and attn_metadata is not None:
+            if hasattr(attn_metadata, 'num_actual_tokens'):
+                # for v1 engine
+                num_tokens = attn_metadata.num_actual_tokens
+            else:
+                # for v0 engine
+                num_tokens = attn_metadata.num_prefill_tokens + attn_metadata.num_decode_tokens
+
+        if num_actual_tokens is None:
+            num_actual_tokens = num_tokens
+
         dp_world_size = get_dp_group().world_size
         if dp_world_size > 1 and forward_context.dp_metadata is not None:
-            forward_context.max_tokens_across_dp = forward_context.dp_metadata.max_tokens_across_dp_cpu.item(
+            max_tokens_across_dp = forward_context.dp_metadata.max_tokens_across_dp_cpu.item(
             )
-        elif num_tokens is not None:
-            forward_context.max_tokens_across_dp = num_tokens
-        elif attn_metadata is not None:
-            if hasattr(attn_metadata, 'num_actual_tokens'):
-                forward_context.max_tokens_across_dp = attn_metadata.num_actual_tokens
-            else:
-                forward_context.max_tokens_across_dp = attn_metadata.num_prefill_tokens + attn_metadata.num_decode_tokens
         else:
-            forward_context.max_tokens_across_dp = None
+            max_tokens_across_dp = num_tokens
+
+        forward_context.max_tokens_across_dp = max_tokens_across_dp
+
+        if num_tokens is not None:
+            tp_world_size = get_tp_group().world_size
+            world_size = torch.distributed.get_world_size()
+            # NOTE: token num which need to pad to when mc2
+            forward_context.padded_num_tokens = math.ceil(
+                max_tokens_across_dp / tp_world_size) * tp_world_size
+            # NOTE: mc2 op's param `global_bs`, add `world_size` to make `global_bs` absolutely larger than actual global_bs.
+            forward_context.global_bs = math.ceil(
+                max_tokens_across_dp / tp_world_size) * world_size
+
+            if fused_moe_state == FusedMoEState.MC2_PREFILL:
+                chunk_size = envs.VLLM_ASCEND_FUSED_MOE_MC2_CHUNK_SIZE
+                forward_context.max_num_chunks = math.ceil(
+                    math.ceil(max_tokens_across_dp / tp_world_size) /
+                    chunk_size)
+
+                forward_context.global_bs = math.ceil(
+                    math.ceil(max_tokens_across_dp / tp_world_size) /
+                    forward_context.max_num_chunks) * world_size
+
+                min_num_tokens = forward_context.max_num_chunks * tp_world_size
+                forward_context.padded_num_tokens = math.ceil(
+                    max_tokens_across_dp / min_num_tokens) * min_num_tokens
+
+            mc2_mask = torch.zeros(forward_context.padded_num_tokens,
+                                   dtype=torch.bool,
+                                   device=current_platform.device_type)
+            mc2_mask[:num_actual_tokens] = True
+            forward_context.mc2_mask = mc2_mask
 
         try:
             yield

vllm_ascend/attention/mla_v1.py

@@ -11,7 +11,6 @@
 from vllm.config import get_current_vllm_config
 from vllm.model_executor.layers.linear import (LinearBase,
                                                UnquantizedLinearMethod)
-from vllm.platforms import current_platform
 from vllm.utils import cdiv, round_down
 
 from vllm_ascend import envs
@@ -71,6 +70,7 @@ class ChunkedContextMetadata:
         max_seq_lens: list[int]
         workspace: torch.Tensor
         chunk_seq_lens: torch.Tensor
+        chunk_seq_lens_npu: torch.Tensor
 
     attn_mask: torch.Tensor
     query_lens: list[int]
@@ -99,7 +99,6 @@ class AscendMLADecodeMetadata:
     attn_mask: Optional[torch.Tensor] = None
     sin: torch.Tensor = None
     cos: torch.Tensor = None
-    mc2_mask: Optional[torch.Tensor] = None
 
 
 @dataclass
@@ -215,13 +214,6 @@ def __init__(self,
         self.cos_cache = None
         self.sin_cache = None
 
-    def generate_activate_mask(self, actual_seqs_num, batch_size):
-        mc2_mask = torch.zeros(batch_size,
-                               dtype=torch.bool,
-                               device=current_platform.device_type)
-        mc2_mask[:actual_seqs_num].fill_(True)
-        return mc2_mask
-
     def reorder_batch(self, input_batch: "InputBatch",
                       scheduler_output: "SchedulerOutput") -> bool:
         # We now want to reorder the batch so that the "decode" requests are at
@@ -364,7 +356,6 @@ def build_torchair_graph_dummy(
                          self.rope_dim,
                          dtype=self.runner.dtype,
                          device=device)
-        mc2_mask = self.generate_activate_mask(num_actual_tokens, num_reqs)
         decode_metadata = AscendMLADecodeMetadata(
             input_positions=input_positions,
             block_table=block_table,
@@ -374,8 +365,7 @@ def build_torchair_graph_dummy(
             attn_mask=self.runner.spec_attn_mask,
             actual_seq_q_lens=self.runner.actual_seq_q_lens[:num_reqs],
             sin=sin,
-            cos=cos,
-            mc2_mask=mc2_mask)
+            cos=cos)
         return self.metadata_cls(  # type: ignore
             num_input_tokens=num_actual_tokens,
             num_actual_tokens=num_actual_tokens,
@@ -481,6 +471,7 @@ def build(
                     seq_tot=chunk_seq_lens.sum(dim=1).tolist(),
                     max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),
                     chunk_seq_lens=chunk_seq_lens,
+                    chunk_seq_lens_npu=chunk_seq_lens.npu(),
                     workspace=self.chunked_prefill_workspace,
                 )
             prefill_input_positions = input_positions[tokens_start:]
@@ -547,15 +538,18 @@ def build(
                 actual_seq_q_lens = query_start_loc[1:].tolist(
                 ) + self.runner.actual_seq_q_lens[num_reqs:num_reqs +
                                                   num_reqs_pad_size]
+                # mtp torchair + PD scenario, last element of actual_seq_q_lens must equal to num_reqs_pad_size
+                num_padded_token_size = slot_mapping.size(0)
+                if actual_seq_q_lens[-1] != num_padded_token_size:
+                    actual_seq_q_lens.append(num_padded_token_size)
+                    seq_lens_list.append(0)
             else:
                 seq_lens_list = seq_lens.tolist()
 
             cos = self.cos_cache[input_positions].unsqueeze(  # type: ignore
                 1).unsqueeze(2)
             sin = self.sin_cache[input_positions].unsqueeze(  # type: ignore
                 1).unsqueeze(2)
-            mc2_mask = self.generate_activate_mask(
-                num_actual_tokens, num_reqs + num_reqs_pad_size)
 
             decode_metadata = AscendMLADecodeMetadata(
                 input_positions=input_positions,
@@ -566,8 +560,7 @@ def build(
                 attn_mask=self.runner.spec_attn_mask,
                 actual_seq_q_lens=actual_seq_q_lens,
                 sin=sin,
-                cos=cos,
-                mc2_mask=mc2_mask)
+                cos=cos)
 
         return self.metadata_cls(  # type: ignore
             num_actual_tokens=num_actual_tokens,
@@ -749,6 +742,8 @@ def _compute_prefill_context(
             toks = prefill_metadata.chunked_context.seq_tot[i]
 
             seq_len2 = prefill_metadata.chunked_context.chunk_seq_lens[i]
+            seq_len2_npu = prefill_metadata.chunked_context.chunk_seq_lens_npu[
+                i]
             seq_len = torch.stack([seq_len1, seq_len2])
             kv_c_normed = torch.empty(toks,
                                       num_heads,
@@ -765,7 +760,7 @@ def _compute_prefill_context(
                 cache_kv_c,
                 cache_k_pe,
                 prefill_metadata.block_table,
-                seq_len2.to(query.device),
+                seq_len2_npu,
                 seq_starts=prefill_metadata.chunked_context.starts[i],
                 key=kv_c_normed,
                 value=k_pe,

vllm_ascend/distributed/parallel_state.py

@@ -0,0 +1,49 @@
+from typing import Optional
+
+import torch
+from vllm.distributed.parallel_state import (GroupCoordinator, get_world_group,
+                                             init_model_parallel_group)
+
+# Currently, mc2 op need their own group coordinator.
+_MC2: Optional[GroupCoordinator] = None
+
+
+def get_mc2_group() -> GroupCoordinator:
+    assert _MC2 is not None, ("mc2 group is not initialized")
+    return _MC2
+
+
+def model_parallel_initialized():
+    return (_MC2 is not None)
+
+
+def init_ascend_model_parallel(
+    expert_parallel_size: int = 1,
+    world_size: Optional[int] = None,
+    backend: Optional[str] = None,
+):
+    if model_parallel_initialized():
+        return
+    assert torch.distributed.is_initialized()
+    world_size = world_size or torch.distributed.get_world_size()
+    backend = backend or torch.distributed.get_backend(
+        get_world_group().device_group)
+    num_expert_parallel_groups = world_size // expert_parallel_size
+
+    global _MC2
+    group_ranks = []
+    for i in range(num_expert_parallel_groups):
+        ranks = list(range(i, world_size, num_expert_parallel_groups))
+        group_ranks.append(ranks)
+
+    _MC2 = init_model_parallel_group(group_ranks,
+                                     get_world_group().local_rank,
+                                     backend,
+                                     group_name="mc2")
+
+
+def destroy_ascend_model_parallel():
+    global _MC2
+    if _MC2:
+        _MC2.destroy()
+    _MC2 = None

vllm_ascend/envs.py

@@ -142,6 +142,12 @@
     #   1: enable moe all2all seq.
     "VLLM_ASCEND_ENABLE_MOE_ALL2ALL_SEQ":
     lambda: bool(int(os.getenv('VLLM_ASCEND_ENABLE_MOE_ALL2ALL_SEQ', '0'))),
+# ENABLE chunk mc2
+    "VLLM_ASCEND_ENABLE_CHUNK_MC2":
+    lambda: bool(int(os.getenv("VLLM_ASCEND_ENABLE_CHUNK_MC2", "0"))),
+    # Batch MC2 in prefill: The number of tokens in each batch
+    "VLLM_ASCEND_FUSED_MOE_MC2_CHUNK_SIZE":
+    lambda: int(os.getenv("VLLM_ASCEND_FUSED_MOE_MC2_CHUNK_SIZE", "128")),
 }
 
 # end-env-vars-definition

vllm_ascend/models/qwen3.py

@@ -18,7 +18,7 @@
 from vllm.model_executor.sampling_metadata import SamplingMetadata
 from vllm.sequence import IntermediateTensors
 
-from vllm_ascend.ops.layernorm import AddRMSNormQuant
+from vllm_ascend.ops.layernorm import AddRMSNormW8A8Quant
 
 
 class CustomQwen3DecoderLayer(Qwen3DecoderLayer):
@@ -43,15 +43,15 @@ def __init__(
         assert isinstance(quant_config, AscendQuantConfig), \
             "Expected quant_config to be an instance of AscendQuantConfig"
 
-        if isinstance(self.self_attn.qkv_proj.quant_method,
+        if isinstance(self.self_attn.qkv_proj.quant_method.quant_method,
                       AscendW8A8LinearMethod):
-            self.input_layernorm = AddRMSNormQuant(
+            self.input_layernorm = AddRMSNormW8A8Quant(
                 config.hidden_size,
                 layer=self.self_attn.qkv_proj,
                 eps=config.rms_norm_eps)
-        if isinstance(self.mlp.gate_up_proj.quant_method,
+        if isinstance(self.mlp.gate_up_proj.quant_method.quant_method,
                       AscendW8A8LinearMethod):
-            self.post_attention_layernorm = AddRMSNormQuant(
+            self.post_attention_layernorm = AddRMSNormW8A8Quant(
                 config.hidden_size,
                 layer=self.mlp.gate_up_proj,
                 eps=config.rms_norm_eps)