[0.11.0]Chery pick pta upgrade change (vllm-project#3940)

wangxiyuan · web-flow · commit 416943357166 · 2025-10-31T22:14:26.000+08:00
This PR cherry-pick two commit from main to upgrade torch-npu to 2.7.1
official release

---------

Signed-off-by: wangxiyuan &lt;wangxiyuan1007@gmail.com&gt;
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -20,6 +20,13 @@ set(VLLM_ASCEND_INSTALL_PATH "${CMAKE_INSTALL_PREFIX}")
 
 find_package(Torch REQUIRED)
 
+run_python(TORCH_VERSION
+  "import torch; print(torch.__version__)" "Failed to locate torch path")
+# check torch version is 2.7.1
+if(NOT ${TORCH_VERSION} VERSION_EQUAL "2.7.1")
+  message(FATAL_ERROR "Expected PyTorch version 2.7.1, but found ${TORCH_VERSION}")
+endif()
+
 set(RUN_MODE "npu" CACHE STRING "cpu/sim/npu")
 set(SOC_VERSION ${SOC_VERSION})
 message(STATUS "Detected SOC version: ${SOC_VERSION}")
diff --git a/README.md b/README.md
@@ -43,7 +43,7 @@ By using vLLM Ascend plugin, popular open-source models, including Transformer-l
 - Software:
   * Python >= 3.9, < 3.12
   * CANN >= 8.2.rc1 (Ascend HDK version refers to [here](https://www.hiascend.com/document/detail/zh/canncommercial/82RC1/releasenote/releasenote_0000.html))
-  * PyTorch >= 2.7.1, torch-npu >= 2.7.1.dev20250724
+  * PyTorch == 2.7.1, torch-npu == 2.7.1
   * vLLM (the same version as vllm-ascend)
 
 ## Getting Started
diff --git a/README.zh.md b/README.zh.md
@@ -44,7 +44,7 @@ vLLM 昇腾插件 (`vllm-ascend`) 是一个由社区维护的让vLLM在Ascend NP
 - 软件：
   * Python >= 3.9, < 3.12
   * CANN >= 8.2.rc1 (Ascend HDK 版本参考[这里](https://www.hiascend.com/document/detail/zh/canncommercial/82RC1/releasenote/releasenote_0000.html))
-  * PyTorch >= 2.7.1, torch-npu >= 2.7.1.dev20250724
+  * PyTorch == 2.7.1, torch-npu == 2.7.1
   * vLLM (与vllm-ascend版本一致)
 
 ## 开始使用
diff --git a/docs/source/installation.md b/docs/source/installation.md
@@ -13,8 +13,8 @@ This document describes how to install vllm-ascend manually.
     |---------------|----------------------------------|-------------------------------------------|
     | Ascend HDK    | Refer to [here](https://www.hiascend.com/document/detail/zh/canncommercial/82RC1/releasenote/releasenote_0000.html) | Required for CANN |
     | CANN          | >= 8.2.RC1                       | Required for vllm-ascend and torch-npu    |
-    | torch-npu     | >= 2.7.1.dev20250724             | Required for vllm-ascend, No need to install manually, it will be auto installed in below steps |
-    | torch         | >= 2.7.1                         | Required for torch-npu and vllm           |
+    | torch-npu     | == 2.7.1             | Required for vllm-ascend, No need to install manually, it will be auto installed in below steps |
+    | torch         | == 2.7.1                         | Required for torch-npu and vllm           |
 
 There are two installation methods:
 - **Using pip**: first prepare env manually or via CANN image, then install `vllm-ascend` using pip.
diff --git a/examples/disaggregated_prefill_v1/mooncake_connector_deployment_guide.md b/examples/disaggregated_prefill_v1/mooncake_connector_deployment_guide.md
@@ -5,7 +5,7 @@
  *  Software:
      *  Python >= 3.9, < 3.12
      *  CANN >= 8.2.rc1
-     *  PyTorch >= 2.7.1, torch-npu >= 2.7.1.dev20250724
+     *  PyTorch == 2.7.1, torch-npu == 2.7.1
      *  vLLM (same version as vllm-ascend)
      *  mooncake-transfer-engine reference documentation: https://github.com/kvcache-ai/Mooncake/blob/main/doc/zh/ascend_transport.md
 
diff --git a/examples/disaggregated_prefill_v1/mooncake_connector_store_deployment_guide.md b/examples/disaggregated_prefill_v1/mooncake_connector_store_deployment_guide.md
@@ -5,7 +5,7 @@
 * Software:
   * Python >= 3.9, < 3.12
   * CANN >= 8.2.rc1
-  * PyTorch >= 2.7.1, torch-npu >= 2.7.1.dev20250724
+  * PyTorch == 2.7.1, torch-npu == 2.7.1
   * vLLM：main branch
   * vLLM-Ascend：main branch
   * Mooncake：[AscendTransport/Mooncake at pooling-async-memcpy](https://github.com/AscendTransport/Mooncake/tree/pooling-async-memcpy)(Currently available branch code, continuously updated.)
diff --git a/pyproject.toml b/pyproject.toml
@@ -12,8 +12,8 @@ requires = [
     "scipy",
     "setuptools>=64",
     "setuptools-scm>=8",
-    "torch-npu==2.7.1.dev20250724",
-    "torch>=2.7.1",
+    "torch-npu==2.7.1",
+    "torch==2.7.1",
     "torchvision",
     "wheel",
     "msgpack",
diff --git a/requirements.txt b/requirements.txt
@@ -10,7 +10,7 @@ pyyaml
 scipy
 setuptools>=64
 setuptools-scm>=8
-torch>=2.7.1
+torch==2.7.1
 torchvision
 wheel
 opencv-python-headless<=4.11.0.86 # Required to avoid numpy version conflict with vllm
@@ -23,6 +23,6 @@ quart
 numba
 
 # Install torch_npu
---pre
---extra-index-url https://mirrors.huaweicloud.com/ascend/repos/pypi
-torch-npu==2.7.1.dev20250724
+#--pre
+#--extra-index-url https://mirrors.huaweicloud.com/ascend/repos/pypi
+torch-npu==2.7.1
diff --git a/tests/ut/ops/test_layernorm.py b/tests/ut/ops/test_layernorm.py
@@ -7,7 +7,6 @@
 
 from tests.ut.base import PytestBase
 from vllm_ascend.quantization.w8a8 import AscendW8A8LinearMethod
-from vllm_ascend.utils import version_check
 
 
 def mock_rms_norm(x, weight, eps):
@@ -18,15 +17,6 @@ def mock_add_rms_norm(x, residual, weight, eps):
     return 2 * x, None, 2 * residual
 
 
-def mock_add_rms_norm_quant(x, residual, weight, quant_scale, quant_offset,
-                            epsilon):
-    x_out = 2 * x
-    residual_out = 2 * residual
-    x_out_quant = x_out.to(torch.int8)
-    residual_out_quant = residual_out.to(torch.int8)
-    return x_out_quant, None, residual_out_quant
-
-
 def mock_add_rms_norm_quant_with_bias(x, residual, weight, quant_scale,
                                       quant_offset, beta, epsilon):
     x_out = 2 * x
@@ -43,10 +33,8 @@ def context(self, mocker: MockerFixture):
         mocker.patch("torch_npu.npu_rms_norm", side_effect=mock_rms_norm)
         mocker.patch("torch_npu.npu_add_rms_norm",
                      side_effect=mock_add_rms_norm)
-        torch_npu_check = version_check()
-        arnq_side_effect = mock_add_rms_norm_quant_with_bias if torch_npu_check else mock_add_rms_norm_quant
         mocker.patch("torch_npu.npu_add_rms_norm_quant",
-                     side_effect=arnq_side_effect)
+                     side_effect=mock_add_rms_norm_quant_with_bias)
         mocker.patch("torch.ops.vllm.maybe_wait_prefetch_done",
                      side_effect=lambda x: None)
 
@@ -82,8 +70,7 @@ def test_forward_oot_with_quant_fusion(self, mocker: MockerFixture):
 
         mock_model_instance = mocker.MagicMock()
         mock_forward_context.model_instance = mock_model_instance
-        torch_npu_check = version_check()
-        num_hidden_layers = 3 if torch_npu_check else 2
+        num_hidden_layers = 3
         mock_model_instance.model.layers = [
             mocker.MagicMock() for _ in range(num_hidden_layers)
         ]
@@ -136,34 +123,31 @@ def test_forward_oot_with_quant_fusion(self, mocker: MockerFixture):
         assert mock_forward_context.fusion_linear == "gate_up_dense"
         assert mock_forward_context.layer_idx == 1
 
-        if torch_npu_check:
-            mock_forward_context.fusion_linear = "gate_moe"
+        mock_forward_context.fusion_linear = "gate_moe"
         x_out, residual_out = layer.forward_oot(x, residual)
 
-        assert mock_get_forward_context.call_count == 6
-        fusion_linear_expected = "qkv_moe" if torch_npu_check else "qkv_dense"
+        assert mock_get_forward_context.call_count == 5
+        fusion_linear_expected = "qkv_moe"
         assert mock_forward_context.fusion_linear == fusion_linear_expected
         assert mock_forward_context.layer_idx == 2
 
         x_out, residual_out = layer.forward_oot(x, residual)
 
-        assert mock_get_forward_context.call_count == 7
-        fusion_linear_expected = "gate_moe" if torch_npu_check else "qkv_dense"
+        assert mock_get_forward_context.call_count == 6
+        fusion_linear_expected = "gate_moe"
         assert mock_forward_context.fusion_linear == fusion_linear_expected
         assert mock_forward_context.layer_idx == 2
 
-        if not torch_npu_check:
-            return
         # last layer returned directly
         x_out, residual_out = layer.forward_oot(x, residual)
 
-        assert mock_get_forward_context.call_count == 8
+        assert mock_get_forward_context.call_count == 7
         assert mock_forward_context.fusion_linear == "qkv_moe"
         assert mock_forward_context.layer_idx == 3
 
         x_out, residual_out = layer.forward_oot(x, residual)
 
-        assert mock_get_forward_context.call_count == 9
+        assert mock_get_forward_context.call_count == 8
         assert mock_forward_context.fusion_linear == "qkv_moe"
         assert mock_forward_context.layer_idx == 3
 
diff --git a/tests/ut/torchair/quantization/test_torchair_w8a8_dynamic.py b/tests/ut/torchair/quantization/test_torchair_w8a8_dynamic.py
@@ -23,9 +23,9 @@ def setUp(self):
     @patch("torch_npu.npu_swiglu")
     @patch("torch_npu.npu_dynamic_quant")
     @patch("torch_npu.npu_moe_finalize_routing")
-    @patch("torch_npu.npu_moe_init_routing")
+    @patch("torch_npu.npu_moe_init_routing_quant")
     def test_torchair_fused_experts_with_all2all(
-            self, mock_moe_init_routing, mock_moe_finalize_routing,
+            self, mock_npu_moe_init_routing_quant, mock_moe_finalize_routing,
             mock_dynamic_quant, mock_swiglu, mock_grouped_matmul,
             mock_moe_re_routing, mock_all_to_all_single):
 
@@ -38,11 +38,10 @@ def test_torchair_fused_experts_with_all2all(
         placeholder_ones = torch.ones(self.num_tokens, dtype=torch.int32)
         mock_all_to_all_single.side_effect = lambda output, input, *args, **kwargs: output.copy_(
             input)
-        mock_moe_init_routing.return_value = (
-            placeholder_int8,
-            placeholder_ones,
-            placeholder_ones,
-        )
+        mock_npu_moe_init_routing_quant.return_value = (
+            placeholder_int8, placeholder_ones, placeholder_ones,
+            torch.bincount(placeholder_ones, minlength=len(expert_map)),
+            torch.randn(self.num_tokens))
         mock_moe_re_routing.return_value = (placeholder_int8, self.placeholder,
                                             torch.randint(0,
                                                           100,
diff --git a/vllm_ascend/ascend_forward_context.py b/vllm_ascend/ascend_forward_context.py
@@ -11,8 +11,7 @@
                                   set_forward_context)
 
 import vllm_ascend.envs as envs_ascend
-from vllm_ascend.utils import (enable_sp, has_layer_idx, is_moe_model,
-                               version_check)
+from vllm_ascend.utils import enable_sp, has_layer_idx, is_moe_model
 
 if TYPE_CHECKING:
     from vllm_ascend.ops.weight_prefetch import WeightPrefetchMethod
@@ -162,9 +161,7 @@ def set_ascend_forward_context(
         # this optim now just support dense models due to the specific operators used.
         # Once the necessary conditions are met, support for MOE models will also be added.
         from vllm_ascend.quantization.quant_config import AscendQuantConfig
-        model_type_scope = ["llama", "qwen2", "qwen3"]
-        if version_check():
-            model_type_scope.append("qwen3_moe")
+        model_type_scope = ["llama", "qwen2", "qwen3", "qwen3_moe"]
         addrmsnorm_quant_fusion_enabled = isinstance(vllm_config.quant_config, AscendQuantConfig) and \
             vllm_config.model_config.hf_config.model_type in model_type_scope and \
             forward_context.layer_idx is not None
diff --git a/vllm_ascend/attention/attention_v1.py b/vllm_ascend/attention/attention_v1.py
@@ -38,7 +38,7 @@
                                                update_graph_params_workspaces)
 from vllm_ascend.ops.attention import vanilla_chunked_prefill
 from vllm_ascend.utils import (ACL_FORMAT_FRACTAL_NZ, aligned_16, is_310p,
-                               nd_to_nz_2d, nd_to_nz_spec, version_check)
+                               nd_to_nz_2d, nd_to_nz_spec)
 
 from ..utils import weak_ref_tensors
 
@@ -321,7 +321,6 @@ def __init__(
         self.num_queries_per_kv = self.num_heads // self.num_kv_heads
         self.key_cache = None
         self.value_cache = None
-        self.torch_npu_check = version_check()
 
     def _forward_prefill_no_cache(
         self,
@@ -429,22 +428,21 @@ def _forward_decode_only(
             forward_context: ForwardContext = get_forward_context()
             num_tokens = query.shape[0]
             if forward_context.capturing:
-                if self.torch_npu_check:
-                    # Get workspace from cache or calculate it if not present.
-                    workspace = graph_params.workspaces.get(num_tokens)
-                    if workspace is None:
-                        workspace = torch_npu._npu_paged_attention_get_workspace(
-                            query=query,
-                            key_cache=self.key_cache,
-                            value_cache=self.value_cache,
-                            num_kv_heads=self.num_kv_heads,
-                            num_heads=self.num_heads,
-                            scale_value=self.scale,
-                            block_table=attn_metadata.block_tables,
-                            context_lens=attn_metadata.seq_lens,
-                            out=output)
-                        update_graph_params_workspaces(
-                            num_tokens, weak_ref_tensors(workspace))
+                # Get workspace from cache or calculate it if not present.
+                workspace = graph_params.workspaces.get(num_tokens)
+                if workspace is None:
+                    workspace = torch_npu._npu_paged_attention_get_workspace(
+                        query=query,
+                        key_cache=self.key_cache,
+                        value_cache=self.value_cache,
+                        num_kv_heads=self.num_kv_heads,
+                        num_heads=self.num_heads,
+                        scale_value=self.scale,
+                        block_table=attn_metadata.block_tables,
+                        context_lens=attn_metadata.seq_lens,
+                        out=output)
+                    update_graph_params_workspaces(num_tokens,
+                                                   weak_ref_tensors(workspace))
 
                 # Handle graph capturing mode
                 stream = torch_npu.npu.current_stream()
@@ -466,30 +464,17 @@ def _forward_decode_only(
                 ))
 
                 torch.npu.graph_task_group_begin(stream)
-
-                if self.torch_npu_check:
-                    torch_npu._npu_paged_attention(
-                        query=query,
-                        key_cache=self.key_cache,
-                        value_cache=self.value_cache,
-                        num_kv_heads=self.num_kv_heads,
-                        num_heads=self.num_heads,
-                        scale_value=self.scale,
-                        block_table=attn_metadata.block_tables,
-                        context_lens=attn_metadata.seq_lens,
-                        out=output,
-                        workspace=workspace)
-                else:
-                    torch_npu._npu_paged_attention(
-                        query=query,
-                        key_cache=self.key_cache,
-                        value_cache=self.value_cache,
-                        num_kv_heads=self.num_kv_heads,
-                        num_heads=self.num_heads,
-                        scale_value=self.scale,
-                        block_table=attn_metadata.block_tables,
-                        context_lens=attn_metadata.seq_lens,
-                        out=output)
+                torch_npu._npu_paged_attention(
+                    query=query,
+                    key_cache=self.key_cache,
+                    value_cache=self.value_cache,
+                    num_kv_heads=self.num_kv_heads,
+                    num_heads=self.num_heads,
+                    scale_value=self.scale,
+                    block_table=attn_metadata.block_tables,
+                    context_lens=attn_metadata.seq_lens,
+                    out=output,
+                    workspace=workspace)
                 handle = torch.npu.graph_task_group_end(stream)
                 graph_params.handles[num_tokens].append(handle)
             else:
diff --git a/vllm_ascend/compilation/acl_graph.py b/vllm_ascend/compilation/acl_graph.py
@@ -18,8 +18,6 @@
 from vllm.logger import logger
 from vllm.platforms import current_platform
 
-from vllm_ascend.utils import version_check
-
 from ..utils import weak_ref_tensors
 
 
@@ -213,32 +211,20 @@ def update_attn_params(update_stream, forward_context, runtime_shape):
             output,
         ) = param
         seq_lens = forward_context.attn_metadata[key].seq_lens
-        torch_npu_check = version_check()
 
         with torch.npu.stream(update_stream):
             torch.npu.graph_task_update_begin(update_stream, handle)
-            if torch_npu_check:
-                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,
-                    workspace=graph_params.workspaces.get(runtime_shape))
-            else:
-                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._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,
+                workspace=graph_params.workspaces.get(runtime_shape))
             torch.npu.graph_task_update_end(update_stream)
 
             event.record(update_stream)
diff --git a/vllm_ascend/ops/layernorm.py b/vllm_ascend/ops/layernorm.py
diff --git a/vllm_ascend/ops/weight_prefetch.py b/vllm_ascend/ops/weight_prefetch.py
diff --git a/vllm_ascend/utils.py b/vllm_ascend/utils.py
