Merge branch 'main' into async-sched-dp

Author: jesse996

tests/ut/core/test_scheduler.py

@@ -8,6 +8,7 @@
                          SchedulerConfig, SpeculativeConfig, VllmConfig)
 from vllm.multimodal.inputs import PlaceholderRange
 from vllm.sampling_params import SamplingParams
+from vllm.utils import sha256
 from vllm.v1.core.kv_cache_utils import (get_request_block_hasher,
                                          init_none_hash)
 from vllm.v1.core.sched.output import SchedulerOutput
@@ -38,7 +39,7 @@ def create_requests(
     max_tokens: int = 16,
     stop_token_ids: Optional[list[int]] = None,
     block_size: int = 3,
-    hash_fn=hash,
+    hash_fn=sha256,
 ):
     init_none_hash(hash_fn)
     prompt_logprobs = PROMPT_LOGPROBS

tests/ut/kv_connector/utils.py

@@ -10,6 +10,7 @@
 from vllm import SamplingParams
 from vllm.config import (CacheConfig, DeviceConfig, KVTransferConfig,
                          ModelConfig, SchedulerConfig, VllmConfig)
+from vllm.utils import sha256
 from vllm.v1.core.kv_cache_utils import (get_request_block_hasher,
                                          init_none_hash)
 from vllm.v1.core.sched.scheduler import Scheduler
@@ -129,10 +130,10 @@ def create_request(
     """Make dummy request for testing."""
     global _none_hash_initialized
     if not _none_hash_initialized:
-        init_none_hash(hash)
+        init_none_hash(sha256)
         _none_hash_initialized = True
 
-    block_hasher = get_request_block_hasher(block_size, hash)
+    block_hasher = get_request_block_hasher(block_size, sha256)
 
     kv_transfer_params: Optional[dict[str, Any]] = None
 

tests/ut/worker/test_worker_v1.py

@@ -1008,7 +1008,10 @@ def test_load_model_sleep_mode_assertion_error(self, mock_allocator_class):
 
     @patch("vllm_ascend.worker.worker_v1.NPUPlatform.seed_everything")
     @patch("vllm_ascend.worker.worker_v1.logger")
-    def test_compile_or_warm_up_model_with_eager_mode(self, mock_logger,
+    @patch("torch_npu._npu_matmul_add_fp32")
+    def test_compile_or_warm_up_model_with_eager_mode(self,
+                                                      mock_npu_matmul_add,
+                                                      mock_logger,
                                                       mock_seed_everything):
         """Test compile_or_warm_up_model method - eager mode"""
         from vllm_ascend.worker.worker_v1 import NPUWorker
@@ -1050,10 +1053,14 @@ def test_compile_or_warm_up_model_with_eager_mode(self, mock_logger,
             # Verify seed setting
             mock_seed_everything.assert_called_once_with(12345)
 
+            # Verify calls
+            mock_npu_matmul_add.assert_called_once()
+
     @patch("vllm_ascend.worker.worker_v1.NPUPlatform.seed_everything")
     @patch("vllm_ascend.worker.worker_v1.logger")
+    @patch("torch_npu._npu_matmul_add_fp32")
     def test_compile_or_warm_up_model_with_graph_capture(
-            self, mock_logger, mock_seed_everything):
+            self, mock_npu_matmul_add, mock_logger, mock_seed_everything):
         """Test compile_or_warm_up_model method - with graph capture enabled"""
         from vllm_ascend.worker.worker_v1 import NPUWorker
 
@@ -1086,6 +1093,9 @@ def test_compile_or_warm_up_model_with_graph_capture(
             # Verify seed setting
             mock_seed_everything.assert_called_once_with(67890)
 
+            # Verify calls
+            mock_npu_matmul_add.assert_called_once()
+
     @patch("vllm_ascend.worker.worker_v1.CaMemAllocator")
     def test_initialize_from_config_with_sleep_mode(self,
                                                     mock_allocator_class):

vllm_ascend/models/qwen3_moe.py

@@ -20,7 +20,6 @@
 from typing import Optional, Union
 
 import torch
-import torch_npu
 from torch import nn
 from transformers import PretrainedConfig
 from vllm.compilation.decorators import support_torch_compile
@@ -277,11 +276,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.make_empty_intermediate_tensors = (
             make_empty_intermediate_tensors_factory(
                 ["hidden_states", "residual"], config.hidden_size))
-        # Call ATB matmul to warm up; otherwise, the first operation (ReshapeAndCache) may cause performance degradation at runtime.
-        x = torch.rand((2, 4), dtype=torch.float16).npu()
-        weight = torch.rand((2, 4), dtype=torch.float16).npu()
-        c = torch.rand((4, 4), dtype=torch.float32).npu()
-        torch_npu._npu_matmul_add_fp32(x, weight, c)
 
     def forward(
         self,

vllm_ascend/ops/linear.py

@@ -62,6 +62,7 @@ def __init__(
         prefix: str = "",
         *,
         return_bias: bool = True,
+        disable_tp: bool = False,
     ):
         self.comm_group = None
         if prefix.find("gate_up_proj") != -1 and mlp_tp_enable():
@@ -88,7 +89,8 @@ def __init__(
                                   params_dtype,
                                   quant_config,
                                   prefix,
-                                  return_bias=return_bias)
+                                  return_bias=return_bias,
+                                  disable_tp=disable_tp)
 
         self.gather_output = gather_output
 
@@ -137,6 +139,7 @@ def __init__(
         prefix: str = "",
         *,
         return_bias: bool = True,
+        disable_tp: bool = False,
     ):
         if prefix.find("down_proj") != -1 and mlp_tp_enable():
             comm_group = get_mlp_tp_group()
@@ -156,6 +159,7 @@ def __init__(
             self.forward_type = "normal"
         self.comm_group = comm_group
 
+        # TODO: check for disable_tp
         self.tp_size = self.comm_group.world_size
         self.tp_rank = self.comm_group.rank_in_group
 
@@ -171,7 +175,8 @@ def __init__(
                                   params_dtype,
                                   quant_config,
                                   prefix,
-                                  return_bias=return_bias)
+                                  return_bias=return_bias,
+                                  disable_tp=disable_tp)
 
         self.input_is_parallel = input_is_parallel
         self.reduce_results = reduce_results
@@ -392,6 +397,7 @@ def __init__(
         prefix: str = "",
         *,
         return_bias: bool = True,
+        disable_tp: bool = False,
     ):
         if prefix.find("gate_up_proj") != -1 and mlp_tp_enable():
             comm_group = get_mlp_tp_group()
@@ -403,6 +409,7 @@ def __init__(
             comm_group = get_tp_group()
             self.forward_type = "normal_tp"
         self.comm_group = comm_group
+        # TODO: check for disable_tp
         self.tp_rank = comm_group.rank_in_group
         self.tp_size = comm_group.world_size
 
@@ -418,7 +425,8 @@ def __init__(
                                             params_dtype=params_dtype,
                                             quant_config=quant_config,
                                             prefix=prefix,
-                                            return_bias=return_bias)
+                                            return_bias=return_bias,
+                                            disable_tp=disable_tp)
 
     def forward(
         self,
@@ -498,6 +506,7 @@ def __init__(
         prefix: str = "",
         *,
         return_bias: bool = True,
+        disable_tp: bool = False,
     ):
         if dense_optim_enable():
             self.forward_type = "dense_optim"
@@ -511,6 +520,7 @@ def __init__(
             total_num_kv_heads = total_num_heads
         self.total_num_kv_heads = total_num_kv_heads
         # Divide the weight matrix along the last dimension.
+        # TODO: check for disable_tp
         tp_size = self.comm_group.world_size
         self.num_heads = divide(self.total_num_heads, tp_size)
         if tp_size >= self.total_num_kv_heads:
@@ -537,7 +547,8 @@ def __init__(
                                             params_dtype=params_dtype,
                                             quant_config=quant_config,
                                             prefix=prefix,
-                                            return_bias=return_bias)
+                                            return_bias=return_bias,
+                                            disable_tp=disable_tp)
 
     def forward(
         self,
@@ -611,4 +622,4 @@ def __init__(
             self.quant_method = quant_config.get_quant_method(self,
                                                               prefix=prefix)
         self.return_bias = return_bias
-        self.disable_tp = disable_tp
+        self.disable_tp = disable_tp

vllm_ascend/worker/worker_v1.py

@@ -250,10 +250,19 @@ def compile_or_warm_up_model(self) -> None:
             self.model_runner._dummy_run(size)
         if not self.model_config.enforce_eager:
             self.model_runner.capture_model()
+        # Call ATB matmul to warm up; otherwise, the first operation (ReshapeAndCache)
+        # may cause performance degradation at runtime.
+        self._warm_up_atb()
         # Reset the seed to ensure that the random state is not affected by
         # the model initialization and profiling.
         NPUPlatform.seed_everything(self.model_config.seed)
 
+    def _warm_up_atb(self):
+        x = torch.rand((2, 4), dtype=torch.float16).npu()
+        weight = torch.rand((2, 4), dtype=torch.float16).npu()
+        c = torch.rand((4, 4), dtype=torch.float32).npu()
+        torch_npu._npu_matmul_add_fp32(x, weight, c)
+
     def get_model(self) -> nn.Module:
         return self.model_runner.get_model()