[Refactor][Device][4/N] Migrate worker and patch selection to hardware profiles

tests/ut/device/test_hardware_profile.py

@@ -24,14 +24,18 @@
 _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,
         HardwareCapability.MC2_HIERARCHY_COMM,
         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,
         }
     ),
 }

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:

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
 

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,14 +107,18 @@ 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,
         HardwareCapability.MC2_HIERARCHY_COMM,
         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,
                 }
             ),
         ),

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:

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

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()

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

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):

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(

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

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.