From b801633b6c4b9a95ab0435cf416dda9d6ef8ab05 Mon Sep 17 00:00:00 2001 From: frost_mourne <2906339855@qq.com> Date: Wed, 12 Aug 2026 03:04:09 +0000 Subject: [PATCH] [Refactor][Device][4/N] Migrate worker and patch selection to hardware profiles Signed-off-by: frost_mourne <2906339855@qq.com> --- tests/ut/device/test_hardware_profile.py | 10 ++++++++++ tests/ut/patch/worker/test_patch_kimi_k25.py | 8 ++++---- tests/ut/worker/a2/test_worker_v1.py | 20 +++++++++---------- vllm_ascend/device/hardware_profile.py | 18 +++++++++++++++++ vllm_ascend/patch/__init__.py | 2 +- vllm_ascend/patch/platform/__init__.py | 4 ++-- .../patch/platform/patch_distributed.py | 4 ++-- vllm_ascend/patch/worker/__init__.py | 4 ++-- vllm_ascend/patch/worker/patch_kimi_k25.py | 4 ++-- vllm_ascend/patch/worker/patch_mamba_utils.py | 4 ++-- vllm_ascend/patch/worker/patch_qwen3_5.py | 5 +++-- vllm_ascend/worker/worker.py | 9 ++++----- 12 files changed, 59 insertions(+), 33 deletions(-) diff --git a/tests/ut/device/test_hardware_profile.py b/tests/ut/device/test_hardware_profile.py index db1e528c2478..58c09d9914cd 100644 --- a/tests/ut/device/test_hardware_profile.py +++ b/tests/ut/device/test_hardware_profile.py @@ -24,6 +24,8 @@ _STANDARD_CAPABILITIES = frozenset( { HardwareCapability.AUTO_ENABLE_CUSTOM_OPS, + HardwareCapability.ATB_EXTENSIONS, + HardwareCapability.ATB_WARMUP, HardwareCapability.BGMV_SGMV_META_REGISTRATION, HardwareCapability.IRQ_CPU_RESERVATION, HardwareCapability.LORA_CUSTOM_OPS, @@ -31,7 +33,9 @@ HardwareCapability.NPUGRAPH_EX, HardwareCapability.RUNTIME_CUSTOM_OPS, HardwareCapability.SFA_DCP_REPLICATED_INDEXER, + HardwareCapability.STANDARD_MAMBA_PATCH, HardwareCapability.STANDARD_WORKER_PATCHES, + HardwareCapability.TRITON_BATCH_MEMCPY, } ) @@ -41,6 +45,8 @@ AscendDeviceType._310P: frozenset( { HardwareCapability.COMPATIBILITY_OP_IMPLEMENTATIONS, + HardwareCapability.DISTRIBUTED_COMMUNICATION_ADAPTATION, + HardwareCapability.GDN_COMPATIBILITY, HardwareCapability.IRQ_CPU_RESERVATION, HardwareCapability.RC_DEVICE_DISCOVERY, HardwareCapability.RUNTIME_CUSTOM_OPS, @@ -52,10 +58,14 @@ HardwareCapability.BGMV_SGMV_META_REGISTRATION, HardwareCapability.CLUSTER_CPU_TOPOLOGY, HardwareCapability.DYNAMIC_MX_QUANT_FUSION, + HardwareCapability.FP8_ATTENTION, + HardwareCapability.LOCAL_KV_COMM_RESOURCE, HardwareCapability.LORA_CUSTOM_OPS, HardwareCapability.NPUGRAPH_EX, HardwareCapability.REDUCED_CUDAGRAPH_CAPTURE_SIZES, + HardwareCapability.STANDARD_MAMBA_PATCH, HardwareCapability.STANDARD_WORKER_PATCHES, + HardwareCapability.TRITON_BATCH_MEMCPY, } ), } diff --git a/tests/ut/patch/worker/test_patch_kimi_k25.py b/tests/ut/patch/worker/test_patch_kimi_k25.py index cf30074839e7..227a35d99467 100644 --- a/tests/ut/patch/worker/test_patch_kimi_k25.py +++ b/tests/ut/patch/worker/test_patch_kimi_k25.py @@ -10,7 +10,8 @@ MoonViT3dPretrainedModel, ) -from vllm_ascend import utils as ascend_utils +from vllm_ascend.device.hardware import AscendDeviceType +from vllm_ascend.device.hardware_profile import get_hardware_profile from vllm_ascend.patch.worker import patch_kimi_k25 @@ -56,9 +57,8 @@ def test_a5_moonvit_to_patch_uses_current_vllm_contract(monkeypatch): original_to = MoonViT3dPretrainedModel.to original_forward = Learnable2DInterpPosEmbDivided_fixed.forward monkeypatch.setattr( - ascend_utils, - "get_ascend_device_type", - lambda: ascend_utils.AscendDeviceType.A5, + "vllm_ascend.device.hardware_profile.get_current_hardware_profile", + lambda: get_hardware_profile(AscendDeviceType.A5), ) try: diff --git a/tests/ut/worker/a2/test_worker_v1.py b/tests/ut/worker/a2/test_worker_v1.py index 0d3879fa4ee5..ae1744ca5b21 100644 --- a/tests/ut/worker/a2/test_worker_v1.py +++ b/tests/ut/worker/a2/test_worker_v1.py @@ -8,6 +8,8 @@ from vllm.v1.kv_cache_interface import FullAttentionSpec from tests.ut.base import TestBase +from vllm_ascend.device.hardware import AscendDeviceType +from vllm_ascend.device.hardware_profile import get_hardware_profile init_cached_hf_modules_path = "vllm.utils.import_utils.init_cached_hf_modules" kw_module = importlib.import_module("vllm_ascend.model_executor.warmup.kernel_warmup") @@ -481,7 +483,7 @@ def test_wake_up_does_not_transpose_moe_weights(self, mock_get_config, mock_allo @patch("vllm_ascend.worker.worker.MemorySnapshot") @patch("vllm_ascend.worker.worker.NPUWorker._init_worker_distributed_environment") @patch("vllm_ascend.worker.worker.init_device_properties_triton") - @patch("vllm_ascend.worker.worker.get_ascend_device_type") + @patch("vllm_ascend.worker.worker.get_current_hardware_profile") @patch("torch.npu.set_device") @patch("torch.npu.empty_cache") @patch("torch.npu.mem_get_info") @@ -497,11 +499,11 @@ def test_init_device( mock_current_platform, ): """Test _init_device method""" - from vllm_ascend.worker.worker import AscendDeviceType, NPUWorker + from vllm_ascend.worker.worker import NPUWorker # Setup mock mock_mem_get_info.return_value = (1000, 2000) - mock_get_device_type.return_value = AscendDeviceType.A2 + mock_get_device_type.return_value = get_hardware_profile(AscendDeviceType.A2) # Mock MemorySnapshot mock_snapshot = MagicMock() @@ -1388,8 +1390,7 @@ def test_load_model_sleep_mode_assertion_error(self, mock_allocator_class): self.assertIn("Sleep mode can only be", str(cm.exception)) @patch("vllm_ascend.worker.worker.set_random_seed") - @patch("vllm_ascend.worker.worker.get_ascend_device_type") - @patch("vllm_ascend.worker.worker.AscendDeviceType") + @patch("vllm_ascend.worker.worker.get_current_hardware_profile") @patch("vllm_ascend.worker.worker.get_ascend_config") @patch("vllm_ascend.worker.worker.logger") @patch("vllm_ascend.worker.worker.NPUWorker._warm_up_atb") @@ -1398,7 +1399,6 @@ def test_compile_or_warm_up_model_with_eager_mode( mock_warm_up_atb, mock_logger, mock_get_ascend_config, - mock_ascend_device_type, mock_get_ascend_device_type, mock_set_random_seed, ): @@ -1408,7 +1408,7 @@ def test_compile_or_warm_up_model_with_eager_mode( mock_ascend_config.ascend_compilation_config.enable_npugraph_ex = False mock_ascend_config.enable_cpu_binding = False mock_get_ascend_config.return_value = mock_ascend_config - mock_get_ascend_device_type.return_value = mock_ascend_device_type.A9B + mock_get_ascend_device_type.return_value = get_hardware_profile(AscendDeviceType.A2) from vllm_ascend.worker.worker import NPUWorker # Create worker mock @@ -1453,8 +1453,7 @@ def test_compile_or_warm_up_model_with_eager_mode( mock_kernel_warmup.assert_called_once_with(worker) @patch("vllm_ascend.worker.worker.set_random_seed") - @patch("vllm_ascend.worker.worker.get_ascend_device_type") - @patch("vllm_ascend.worker.worker.AscendDeviceType") + @patch("vllm_ascend.worker.worker.get_current_hardware_profile") @patch("vllm_ascend.worker.worker.CUDAGraphMode") @patch("vllm_ascend.worker.worker.get_ascend_config") @patch("vllm_ascend.worker.worker.logger") @@ -1465,7 +1464,6 @@ def test_compile_or_warm_up_model_with_graph_capture( mock_logger, mock_get_ascend_config, mock_cudagraph_mode, - mock_ascend_device_type, mock_get_ascend_device_type, mock_set_random_seed, ): @@ -1475,7 +1473,7 @@ def test_compile_or_warm_up_model_with_graph_capture( mock_ascend_config.ascend_compilation_config.enable_npugraph_ex = False mock_ascend_config.enable_cpu_binding = False mock_get_ascend_config.return_value = mock_ascend_config - mock_get_ascend_device_type.return_value = mock_ascend_device_type.A9B + mock_get_ascend_device_type.return_value = get_hardware_profile(AscendDeviceType.A2) mock_cudagraph_mode.NONE = mock_cudagraph_mode.NONE from vllm_ascend.worker.worker import NPUWorker diff --git a/vllm_ascend/device/hardware_profile.py b/vllm_ascend/device/hardware_profile.py index 804ade4ca5e9..862f24012a82 100644 --- a/vllm_ascend/device/hardware_profile.py +++ b/vllm_ascend/device/hardware_profile.py @@ -17,11 +17,17 @@ class HardwareCapability(Enum): """Independent SoC capabilities consumed by shared business logic.""" AUTO_ENABLE_CUSTOM_OPS = auto() + ATB_EXTENSIONS = auto() + ATB_WARMUP = auto() BGMV_SGMV_META_REGISTRATION = auto() CLUSTER_CPU_TOPOLOGY = auto() COMPATIBILITY_OP_IMPLEMENTATIONS = auto() + DISTRIBUTED_COMMUNICATION_ADAPTATION = auto() DYNAMIC_MX_QUANT_FUSION = auto() + FP8_ATTENTION = auto() + GDN_COMPATIBILITY = auto() IRQ_CPU_RESERVATION = auto() + LOCAL_KV_COMM_RESOURCE = auto() LORA_CUSTOM_OPS = auto() MC2_FULLMESH_V2_COMM = auto() MC2_HIERARCHY_COMM = auto() @@ -31,6 +37,8 @@ class HardwareCapability(Enum): RUNTIME_CUSTOM_OPS = auto() SFA_DCP_REPLICATED_INDEXER = auto() STANDARD_WORKER_PATCHES = auto() + STANDARD_MAMBA_PATCH = auto() + TRITON_BATCH_MEMCPY = auto() class AttentionBackendFamily(Enum): @@ -99,6 +107,8 @@ def supports(self, capability: HardwareCapability) -> bool: _STANDARD_CAPABILITIES = frozenset( { HardwareCapability.AUTO_ENABLE_CUSTOM_OPS, + HardwareCapability.ATB_EXTENSIONS, + HardwareCapability.ATB_WARMUP, HardwareCapability.BGMV_SGMV_META_REGISTRATION, HardwareCapability.IRQ_CPU_RESERVATION, HardwareCapability.LORA_CUSTOM_OPS, @@ -106,7 +116,9 @@ def supports(self, capability: HardwareCapability) -> bool: HardwareCapability.NPUGRAPH_EX, HardwareCapability.RUNTIME_CUSTOM_OPS, HardwareCapability.SFA_DCP_REPLICATED_INDEXER, + HardwareCapability.STANDARD_MAMBA_PATCH, HardwareCapability.STANDARD_WORKER_PATCHES, + HardwareCapability.TRITON_BATCH_MEMCPY, } ) _A3_CAPABILITIES = _STANDARD_CAPABILITIES | {HardwareCapability.MC2_FULLMESH_V2_COMM} @@ -147,6 +159,8 @@ def supports(self, capability: HardwareCapability) -> bool: capabilities=frozenset( { HardwareCapability.COMPATIBILITY_OP_IMPLEMENTATIONS, + HardwareCapability.DISTRIBUTED_COMMUNICATION_ADAPTATION, + HardwareCapability.GDN_COMPATIBILITY, HardwareCapability.IRQ_CPU_RESERVATION, HardwareCapability.RC_DEVICE_DISCOVERY, HardwareCapability.RUNTIME_CUSTOM_OPS, @@ -168,10 +182,14 @@ def supports(self, capability: HardwareCapability) -> bool: HardwareCapability.BGMV_SGMV_META_REGISTRATION, HardwareCapability.CLUSTER_CPU_TOPOLOGY, HardwareCapability.DYNAMIC_MX_QUANT_FUSION, + HardwareCapability.FP8_ATTENTION, + HardwareCapability.LOCAL_KV_COMM_RESOURCE, HardwareCapability.LORA_CUSTOM_OPS, HardwareCapability.NPUGRAPH_EX, HardwareCapability.REDUCED_CUDAGRAPH_CAPTURE_SIZES, + HardwareCapability.STANDARD_MAMBA_PATCH, HardwareCapability.STANDARD_WORKER_PATCHES, + HardwareCapability.TRITON_BATCH_MEMCPY, } ), ), diff --git a/vllm_ascend/patch/__init__.py b/vllm_ascend/patch/__init__.py index 785b3251ae86..3a33e611063c 100644 --- a/vllm_ascend/patch/__init__.py +++ b/vllm_ascend/patch/__init__.py @@ -214,7 +214,7 @@ # On 310P, override verify_and_update_config to align mamba_block_size and # attention block size to the 128-token kernel alignment, ensuring the # attention page size is >= mamba page size. This is the 310P counterpart -# of patch_mamba_config.py (loaded only when `is_310p()` is True). +# of patch_mamba_config.py (selected by the active hardware profile). # Related PR (if no, explain why): # No, 310P-specific kernel alignment requirement. # Future Plan: diff --git a/vllm_ascend/patch/platform/__init__.py b/vllm_ascend/patch/platform/__init__.py index cb77376726ad..d072c4e3d1c7 100644 --- a/vllm_ascend/patch/platform/__init__.py +++ b/vllm_ascend/patch/platform/__init__.py @@ -21,9 +21,9 @@ import vllm_ascend.patch.platform.patch_mla_prefill_backend # noqa import vllm_ascend.patch.platform.patch_pp_mtp # noqa import vllm_ascend.patch.platform.patch_use_v2_model_runner # noqa -from vllm_ascend.utils import is_310p +from vllm_ascend.device.hardware_profile import HardwareCapability, get_current_hardware_profile -if not is_310p(): +if get_current_hardware_profile().supports(HardwareCapability.STANDARD_MAMBA_PATCH): import vllm_ascend.patch.platform.patch_mamba_config # noqa else: import vllm_ascend.patch.platform.patch_mamba_config_310 # noqa diff --git a/vllm_ascend/patch/platform/patch_distributed.py b/vllm_ascend/patch/platform/patch_distributed.py index 8c085cce4f3d..fd463736665f 100644 --- a/vllm_ascend/patch/platform/patch_distributed.py +++ b/vllm_ascend/patch/platform/patch_distributed.py @@ -19,7 +19,7 @@ import torch -from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type +from vllm_ascend.device.hardware_profile import HardwareCapability, get_current_hardware_profile class NullHandle: @@ -85,5 +85,5 @@ def all_reduce( ) -if get_ascend_device_type() == AscendDeviceType._310P: +if get_current_hardware_profile().supports(HardwareCapability.DISTRIBUTED_COMMUNICATION_ADAPTATION): communication_adaptation_310p() diff --git a/vllm_ascend/patch/worker/__init__.py b/vllm_ascend/patch/worker/__init__.py index 9b1dd5d2ba84..ae05acf01809 100644 --- a/vllm_ascend/patch/worker/__init__.py +++ b/vllm_ascend/patch/worker/__init__.py @@ -17,7 +17,7 @@ from vllm.triton_utils import HAS_TRITON -from vllm_ascend.utils import is_310p +from vllm_ascend.device.hardware_profile import HardwareCapability, get_current_hardware_profile if HAS_TRITON: import vllm_ascend.patch.worker.patch_triton @@ -31,7 +31,7 @@ import vllm_ascend.patch.worker.patch_bind_kv_cache # noqa import vllm_ascend.patch.worker.patch_step3p5 # noqa -if not is_310p(): +if get_current_hardware_profile().supports(HardwareCapability.STANDARD_WORKER_PATCHES): import vllm_ascend.patch.worker.patch_qwen3_5 # noqa import vllm_ascend.patch.worker.patch_qwen3_dflash # noqa import vllm_ascend.patch.worker.patch_qwen3vl # noqa diff --git a/vllm_ascend/patch/worker/patch_kimi_k25.py b/vllm_ascend/patch/worker/patch_kimi_k25.py index bea4d288bc2d..1e13cd5df509 100644 --- a/vllm_ascend/patch/worker/patch_kimi_k25.py +++ b/vllm_ascend/patch/worker/patch_kimi_k25.py @@ -24,7 +24,7 @@ get_rope_shape_decorate, ) -from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type +from vllm_ascend.device.hardware_profile import HardwareCapability, get_current_hardware_profile @get_rope_shape_decorate @@ -82,7 +82,7 @@ def forward(self, x: torch.Tensor, grid_thws: torch.Tensor | list) -> torch.Tens # the `dtype=model_config.dtype` (e.g. bf16) would overwrite the fp8 parameters # created by the Ascend quantization scheme, causing a dtype mismatch later # in weight_loader when the checkpoint's fp8 weights are loaded. -if get_ascend_device_type() == AscendDeviceType.A5: +if get_current_hardware_profile().supports(HardwareCapability.FP8_ATTENTION): _original_moonvit_to = MoonViT3dPretrainedModel.to def _patched_moonvit_to(self, *args, **kwargs): diff --git a/vllm_ascend/patch/worker/patch_mamba_utils.py b/vllm_ascend/patch/worker/patch_mamba_utils.py index a16fa0a22a4c..b3f105293fcf 100644 --- a/vllm_ascend/patch/worker/patch_mamba_utils.py +++ b/vllm_ascend/patch/worker/patch_mamba_utils.py @@ -18,9 +18,9 @@ from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch from vllm.v1.worker.mamba_utils import MambaCopyBuffers +from vllm_ascend.device.hardware_profile import HardwareCapability, get_current_hardware_profile from vllm_ascend.ops.triton.batch_memcpy import batch_memcpy_kernel from vllm_ascend.ops.triton.mamba.postprocess import postprocess_mamba_fused_kernel -from vllm_ascend.utils import is_310p # Upstream uses 16 temporal-copy tiles to saturate H100/GB200. K3 already # exposes 138 independent state programs per request, while Triton-Ascend @@ -32,7 +32,7 @@ def _can_launch_triton_batch_memcpy() -> bool: - return not is_310p() + return get_current_hardware_profile().supports(HardwareCapability.TRITON_BATCH_MEMCPY) def _get_mamba_groups( diff --git a/vllm_ascend/patch/worker/patch_qwen3_5.py b/vllm_ascend/patch/worker/patch_qwen3_5.py index 8992dbf08c2c..ea2a5dd16b30 100644 --- a/vllm_ascend/patch/worker/patch_qwen3_5.py +++ b/vllm_ascend/patch/worker/patch_qwen3_5.py @@ -31,8 +31,9 @@ IntermediateTensors = None from vllm.model_executor.models.qwen3_next import Qwen3NextAttention +from vllm_ascend.device.hardware_profile import HardwareCapability, get_current_hardware_profile from vllm_ascend.ops.gdn import AscendGatedDeltaNetAttention -from vllm_ascend.utils import is_310p, vllm_version_is +from vllm_ascend.utils import vllm_version_is if vllm_version_is("0.27.1"): import vllm.model_executor.models.qwen3_next as qwen3_next_module @@ -231,7 +232,7 @@ def qwen3_5_mtp_forward( _GDN_PATCH_TARGET.get_state_shape = AscendGatedDeltaNetAttention.get_state_shape _GDN_PATCH_TARGET.get_attn_backend = AscendGatedDeltaNetAttention.get_attn_backend -if is_310p(): +if get_current_hardware_profile().supports(HardwareCapability.GDN_COMPATIBILITY): from vllm_ascend._310p.ops.fla.gdn_310 import AscendGatedDeltaNetAttention310 _GDN_PATCH_TARGET._forward_core = AscendGatedDeltaNetAttention310._forward_core diff --git a/vllm_ascend/worker/worker.py b/vllm_ascend/worker/worker.py index 3318e6fe5788..6429bbf6f121 100644 --- a/vllm_ascend/worker/worker.py +++ b/vllm_ascend/worker/worker.py @@ -63,6 +63,7 @@ _attach_profiling_chunk_execution_time, ) from vllm_ascend.cpu_binding import bind_cpus +from vllm_ascend.device.hardware_profile import HardwareCapability, get_current_hardware_profile from vllm_ascend.device_allocator.camem import CaMemAllocator from vllm_ascend.device_allocator.sleep_mem_optimized import SleepWakeupManager from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.layerwise_cache_layout import ( @@ -78,10 +79,8 @@ from vllm_ascend.ops.triton.triton_utils import init_device_properties_triton from vllm_ascend.profiler.torch_npu_profiler import TorchNPUProfilerWrapper from vllm_ascend.utils import ( - AscendDeviceType, check_ascend_device_type, enable_sp, - get_ascend_device_type, register_ascend_customop, setup_ascend_local_comm_res, ) @@ -135,7 +134,7 @@ def __init__( from vllm_ascend import ops ops.register_dummy_fusion_op() - if get_ascend_device_type() != AscendDeviceType.A5: + if get_current_hardware_profile().supports(HardwareCapability.ATB_EXTENSIONS): _register_atb_extensions() register_ascend_customop(vllm_config) # init ascend config and soc version @@ -435,7 +434,7 @@ def _init_device(self): gc.collect() torch.npu.empty_cache() - if get_ascend_device_type() == AscendDeviceType.A5: + if get_current_hardware_profile().supports(HardwareCapability.LOCAL_KV_COMM_RESOURCE): setup_ascend_local_comm_res(self.local_rank, self.vllm_config.kv_transfer_config) # take current memory snapshot @@ -822,7 +821,7 @@ def compile_or_warm_up_model(self) -> CompilationTimes: # Call ATB matmul to warm up; otherwise, the first operation (ReshapeAndCache) # may cause performance degradation at runtime. - if get_ascend_device_type() != AscendDeviceType.A5: + if get_current_hardware_profile().supports(HardwareCapability.ATB_WARMUP): self._warm_up_atb() # Bind after warmup so hot allocations are already materialized on the # worker process before migratepages/taskset run.