rocm.py 18.6 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
import os
5
from datetime import timedelta
6
from functools import cache, lru_cache, wraps
7
from typing import TYPE_CHECKING, Optional
8
9

import torch
10
11
from torch.distributed import PrefixStore, ProcessGroup
from torch.distributed.distributed_c10d import is_nccl_available
12

13
import vllm.envs as envs
14
from vllm.logger import init_logger
15
from vllm.utils import cuda_device_count_stateless
16

17
from .interface import DeviceCapability, Platform, PlatformEnum
18

19
if TYPE_CHECKING:
20
    from vllm.attention.backends.registry import _Backend
21
    from vllm.config import ModelConfig, VllmConfig
22
23
else:
    _Backend = None
24

25
26
logger = init_logger(__name__)

27
try:
28
29
30
    from amdsmi import (AmdSmiException, amdsmi_get_gpu_asic_info,
                        amdsmi_get_processor_handles, amdsmi_init,
                        amdsmi_shut_down, amdsmi_topo_get_link_type)
31
32
33
except ImportError as e:
    logger.warning("Failed to import from amdsmi with %r", e)

34
35
36
37
38
39
40
41
42
43
44
try:
    import vllm._C  # noqa: F401
except ImportError as e:
    logger.warning("Failed to import from vllm._C with %r", e)

# import custom ops, trigger op registration
try:
    import vllm._rocm_C  # noqa: F401
except ImportError as e:
    logger.warning("Failed to import from vllm._rocm_C with %r", e)

45
# Models not supported by ROCm.
46
_ROCM_UNSUPPORTED_MODELS: list[str] = []
47
48
49
50
51
52
53

# Models partially supported by ROCm.
# Architecture -> Reason.
_ROCM_SWA_REASON = ("Sliding window attention (SWA) is not yet supported in "
                    "Triton flash attention. For half-precision SWA support, "
                    "please use CK flash attention by setting "
                    "`VLLM_USE_TRITON_FLASH_ATTN=0`")
54
_ROCM_PARTIALLY_SUPPORTED_MODELS: dict[str, str] = {
55
56
57
58
59
60
61
62
63
64
65
66
67
68
    "Qwen2ForCausalLM":
    _ROCM_SWA_REASON,
    "MistralForCausalLM":
    _ROCM_SWA_REASON,
    "MixtralForCausalLM":
    _ROCM_SWA_REASON,
    "PaliGemmaForConditionalGeneration":
    ("ROCm flash attention does not yet "
     "fully support 32-bit precision on PaliGemma"),
    "Phi3VForCausalLM":
    ("ROCm Triton flash attention may run into compilation errors due to "
     "excessive use of shared memory. If this happens, disable Triton FA "
     "by setting `VLLM_USE_TRITON_FLASH_ATTN=0`")
}
69
_ROCM_DEVICE_ID_NAME_MAP: dict[str, str] = {
70
71
72
73
74
75
76
77
    "0x74a0": "AMD_Instinct_MI300A",
    "0x74a1": "AMD_Instinct_MI300X",
    "0x74b5": "AMD_Instinct_MI300X",  # MI300X VF
    "0x74a5": "AMD_Instinct_MI325X",
    "0x74b9": "AMD_Instinct_MI325X",  # MI325X VF
    "0x74a9": "AMD_Instinct_MI300X_HF",
    "0x74bd": "AMD_Instinct_MI300X_HF",
}
78

79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
# Prevent use of clashing `{CUDA/HIP}_VISIBLE_DEVICES``
if "HIP_VISIBLE_DEVICES" in os.environ:
    val = os.environ["HIP_VISIBLE_DEVICES"]
    if cuda_val := os.environ.get("CUDA_VISIBLE_DEVICES", None):
        assert val == cuda_val
    else:
        os.environ["CUDA_VISIBLE_DEVICES"] = val

# AMDSMI utils
# Note that NVML is not affected by `{CUDA/HIP}_VISIBLE_DEVICES`,
# all the related functions work on real physical device ids.
# the major benefit of using AMDSMI is that it will not initialize CUDA


def with_amdsmi_context(fn):

    @wraps(fn)
    def wrapper(*args, **kwargs):
        amdsmi_init()
        try:
            return fn(*args, **kwargs)
        finally:
            amdsmi_shut_down()

    return wrapper


106
107
108
109
110
111
@cache
def on_gfx1x() -> bool:
    GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
    return any(arch in GPU_ARCH for arch in ["gfx11", "gfx12"])


112
@cache
113
def on_mi3xx() -> bool:
114
    GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
115
116
117
118
119
120
121
    return any(arch in GPU_ARCH for arch in ["gfx942", "gfx950"])


@cache
def on_gfx9() -> bool:
    GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
    return any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])
122
123


124
125
126
127
128
129
@cache
def on_gfx950() -> bool:
    GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
    return any(arch in GPU_ARCH for arch in ["gfx950"])


130
@cache
131
132
133
134
135
136
137
138
def use_rocm_custom_paged_attention(
        qtype: torch.dtype,
        head_size: int,
        block_size: int,
        gqa_ratio: int,
        max_seq_len: int,
        sliding_window: int,
        kv_cache_dtype: str,
139
140
        alibi_slopes: Optional[torch.Tensor] = None,
        sinks: Optional[torch.Tensor] = None) -> bool:
141

142
143
    GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
    ON_GFX9 = any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])
144
    ON_GFX11_GFX12 = any(arch in GPU_ARCH for arch in ["gfx11", "gfx12"])
145

146
147
    # custom paged attn always supported on V0. On V1, requires sliding window
    # disabled due to observed numerical discrepancy.
148
149
150
151
152
153
154
    if ON_GFX9:
        return ((not envs.VLLM_USE_V1 or sliding_window == 0
                 or sliding_window == (-1, -1))
                and (qtype == torch.half or qtype == torch.bfloat16)
                and (head_size == 64 or head_size == 128)
                and (block_size == 16 or block_size == 32)
                and (gqa_ratio >= 1 and gqa_ratio <= 16)
155
156
                and max_seq_len <= 128 * 1024
                and (envs.VLLM_ROCM_CUSTOM_PAGED_ATTN)
157
                and not (envs.VLLM_ROCM_USE_AITER_PAGED_ATTN
158
                         and envs.VLLM_ROCM_USE_AITER) and sinks is None)
159
160
161
162
163
164
165

    else:
        return (ON_GFX11_GFX12 and (not envs.VLLM_USE_V1 or sliding_window == 0
                                    or sliding_window == (-1, -1))
                and (qtype == torch.half or qtype == torch.bfloat16)
                and head_size == 128 and block_size == 16
                and (gqa_ratio >= 3 and gqa_ratio <= 16)
166
                and max_seq_len <= 128 * 1024 and alibi_slopes is None
167
                and kv_cache_dtype == "auto"
168
                and envs.VLLM_ROCM_CUSTOM_PAGED_ATTN and sinks is None)
169
170


171
172
class RocmPlatform(Platform):
    _enum = PlatformEnum.ROCM
173
    device_name: str = "rocm"
174
    device_type: str = "cuda"
175
    dispatch_key: str = "CUDA"
176
    ray_device_key: str = "GPU"
177
    dist_backend: str = "nccl"
178
179
    # rocm shares the same device control env var as CUDA
    device_control_env_var: str = "CUDA_VISIBLE_DEVICES"
180

181
    supported_quantization: list[str] = [
182
        "awq", "gptq", "fp8", "compressed-tensors", "fbgemm_fp8", "gguf",
183
        "quark", "ptpc_fp8", "mxfp4", "petit_nvfp4", "torchao"
184
    ]
185

186
    @classmethod
187
    def get_vit_attn_backend(cls, head_size: int,
188
189
                             dtype: torch.dtype) -> "_Backend":
        from vllm.attention.backends.registry import _Backend
190
191
192
193
194
195
196
        if (envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA
                and on_gfx9()):
            # Note: AITER FA is only supported for Qwen-VL models.
            # TODO: Add support for other VL models in their model class.
            return _Backend.ROCM_AITER_FA
        if on_gfx9():
            return _Backend.FLASH_ATTN
197
198
        return _Backend.TORCH_SDPA

199
    @classmethod
200
    def get_attn_backend_cls(cls, selected_backend, head_size, dtype,
201
                             kv_cache_dtype, block_size, use_v1, use_mla,
202
                             has_sink, use_sparse) -> str:
203
        from vllm.attention.backends.registry import _Backend
204
205
206
        if use_sparse:
            raise NotImplementedError(
                "Sparse Attention is not supported on ROCm.")
207
        if use_mla:
208
209
210
211
212
            if not use_v1:
                raise RuntimeError(
                    "MLA attention backends require the V1 engine. "
                    "Set VLLM_USE_V1=1 to enable them.")

213
            from vllm.v1.attention.backends.mla.rocm_aiter_mla import (
214
215
216
217
218
219
220
221
222
                is_aiter_mla_enabled)

            if selected_backend is None:
                selected_backend = (_Backend.ROCM_AITER_MLA if
                                    is_aiter_mla_enabled() or block_size == 1
                                    else _Backend.TRITON_MLA)

            if selected_backend == _Backend.TRITON_MLA:
                if block_size != 1:
223
224
225
226
227
228
                    logger.info_once("Using Triton MLA backend on V1 engine.")
                    return ("vllm.v1.attention.backends.mla."
                            "triton_mla.TritonMLABackend")
                raise ValueError(
                    f" The selected backend, {selected_backend.name},"
                    f"does not support block size {block_size}.")
229
            if selected_backend == _Backend.ROCM_AITER_MLA:
230
                if block_size == 1:
231
232
                    logger.info("Using AITER MLA backend on V1 engine.")
                    return "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend"  # noqa: E501
233
234
                raise ValueError(
                    f" The selected backend, {selected_backend.name},"
235
236
237
238
239
                    f"does not support block size {block_size}."
                    "(currently only supports block size 1)")
            raise ValueError(
                f" The selected backend, {selected_backend.name},"
                f"is not MLA type while requested for MLA backend.")
240

241
        if envs.VLLM_USE_V1:
242
243
244
245
246
            if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA \
                and on_gfx9():
                logger.info("Using Flash Attention backend on V1 engine.")
                return ("vllm.v1.attention.backends."
                        "rocm_aiter_fa.AiterFlashAttentionBackend")
247
248
249
            elif (envs.VLLM_ROCM_USE_AITER and
                envs.VLLM_USE_AITER_UNIFIED_ATTENTION) or \
                    envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION or \
250
                        selected_backend == _Backend.ROCM_ATTN:
251
252
253
254
255
                # rocm specific backend, with aiter and/or
                #   triton prefix-prefill
                logger.info("Using Rocm/Aiter Attention backend on V1 engine.")
                return ("vllm.v1.attention.backends."
                        "rocm_attn.RocmAttentionBackend")
256
            else:
257
                # default case, using triton unified attention
258
259
260
                logger.info("Using Triton Attention backend on V1 engine.")
                return ("vllm.v1.attention.backends."
                        "triton_attn.TritonAttentionBackend")
261
262
263
        raise RuntimeError(
            "V0 attention backends have been removed. Set VLLM_USE_V1=1 "
            "to select a supported backend.")
264

265
266
267
268
269
270
271
    @classmethod
    def set_device(cls, device: torch.device) -> None:
        """
        Set the device for the current platform.
        """
        torch.cuda.set_device(device)

272
    @classmethod
273
    @lru_cache(maxsize=8)
274
275
276
    def get_device_capability(cls,
                              device_id: int = 0
                              ) -> Optional[DeviceCapability]:
277
278
        major, minor = torch.cuda.get_device_capability(device_id)
        return DeviceCapability(major=major, minor=minor)
279

280
    @classmethod
281
    @with_amdsmi_context
282
    def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
        """
        Query if the set of gpus are fully connected by xgmi (1 hop)
        """
        handles = [
            amdsmi_get_processor_handles()[i] for i in physical_device_ids
        ]
        for i, handle in enumerate(handles):
            for j, peer_handle in enumerate(handles):
                if i < j:
                    try:
                        link_type = amdsmi_topo_get_link_type(
                            handle, peer_handle)
                        # type is 2 for XGMI
                        if link_type["hops"] != 1 or link_type["type"] != 2:
                            return False
                    except AmdSmiException as error:
                        logger.error("AMD 1 hop XGMI detection failed.",
                                     exc_info=error)
                        return False
        return True

304
    @classmethod
305
    @with_amdsmi_context
306
    @lru_cache(maxsize=8)
307
    def get_device_name(cls, device_id: int = 0) -> str:
308
        physical_device_id = cls.device_id_to_physical_device_id(device_id)
309
        handle = amdsmi_get_processor_handles()[physical_device_id]
310
311
312
313
314
        asic_info = amdsmi_get_gpu_asic_info(handle)
        device_name: str = asic_info["device_id"]
        if device_name in _ROCM_DEVICE_ID_NAME_MAP:
            return _ROCM_DEVICE_ID_NAME_MAP[device_name]
        return asic_info["market_name"]
315
316
317
318
319

    @classmethod
    def get_device_total_memory(cls, device_id: int = 0) -> int:
        device_props = torch.cuda.get_device_properties(device_id)
        return device_props.total_memory
320
321

    @classmethod
322
    def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
323
324
        from vllm.config.compilation import CUDAGraphMode

325
        cache_config = vllm_config.cache_config
326
327
328
329
330
331
332
333
        compilation_config = vllm_config.compilation_config
        parallel_config = vllm_config.parallel_config
        is_eager_execution = compilation_config == CUDAGraphMode.NONE

        use_v1 = envs.VLLM_USE_V1
        use_aiter_rms_norm = envs.VLLM_ROCM_USE_AITER and \
             envs.VLLM_ROCM_USE_AITER_RMSNORM

334
335
336
        if cache_config and cache_config.block_size is None:
            cache_config.block_size = 16

337
        if parallel_config.worker_cls == "auto":
338
            parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker"
339
        #  Aiter rms norm perform best when CUDA Graph capture is enabled.
340
341
        if (use_v1 and use_aiter_rms_norm and not is_eager_execution
                and "-rms_norm" not in compilation_config.custom_ops):
342
            compilation_config.custom_ops.append("+rms_norm")
343

344
345
346
347
348
349
350
351
352
353
354
355
    @classmethod
    def verify_model_arch(cls, model_arch: str) -> None:
        if model_arch in _ROCM_UNSUPPORTED_MODELS:
            raise ValueError(f"Model architecture '{model_arch}' is not "
                             "supported by ROCm for now.")

        if model_arch in _ROCM_PARTIALLY_SUPPORTED_MODELS:
            msg = _ROCM_PARTIALLY_SUPPORTED_MODELS[model_arch]
            logger.warning(
                "Model architecture '%s' is partially "
                "supported by ROCm: %s", model_arch, msg)

356
357
358
359
360
361
362
363
    @classmethod
    def verify_quantization(cls, quant: str) -> None:
        super().verify_quantization(quant)
        if quant == "awq" and not envs.VLLM_USE_TRITON_AWQ:
            logger.warning(
                "Using AWQ quantization with ROCm, but VLLM_USE_TRITON_AWQ"
                " is not set, enabling VLLM_USE_TRITON_AWQ.")
        envs.VLLM_USE_TRITON_AWQ = True
364
365
366
367

    @classmethod
    def get_punica_wrapper(cls) -> str:
        return "vllm.lora.punica_wrapper.punica_gpu.PunicaWrapperGPU"
368
369
370
371
372
373

    @classmethod
    def get_current_memory_usage(cls,
                                 device: Optional[torch.types.Device] = None
                                 ) -> float:
        torch.cuda.reset_peak_memory_stats(device)
374
375
        return torch.cuda.mem_get_info(device)[1] - torch.cuda.mem_get_info(
            device)[0]
376
377
378
379

    @classmethod
    def get_device_communicator_cls(cls) -> str:
        return "vllm.distributed.device_communicators.cuda_communicator.CudaCommunicator"  # noqa
380

381
382
383
384
385
    @classmethod
    def supports_mx(cls) -> bool:
        gcn_arch = torch.cuda.get_device_properties(0).gcnArchName
        return any(gfx in gcn_arch for gfx in ["gfx95"])

386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
    @classmethod
    def supports_fp8(cls) -> bool:
        gcn_arch = torch.cuda.get_device_properties(0).gcnArchName
        return any(gfx in gcn_arch for gfx in ['gfx94', 'gfx95', 'gfx12'])

    @classmethod
    def is_fp8_fnuz(cls) -> bool:
        # only device 0 is checked, this assumes MI300 platforms are homogeneous
        return 'gfx94' in torch.cuda.get_device_properties(0).gcnArchName

    @classmethod
    def fp8_dtype(cls) -> torch.dtype:
        if cls.is_fp8_fnuz():
            return torch.float8_e4m3fnuz
        else:
            return torch.float8_e4m3fn
402

403
404
405
406
    @classmethod
    def use_custom_allreduce(cls) -> bool:
        # We only enable custom allreduce for MI300 series
        gcn_arch = torch.cuda.get_device_properties(0).gcnArchName
407
        supported_archs = ['gfx94', 'gfx95']
408
        return any(gfx in gcn_arch for gfx in supported_archs)
409

410
411
412
413
    @classmethod
    def opaque_attention_op(cls) -> bool:
        return True

414
415
416
    @classmethod
    def get_cu_count(cls, device_id: int = 0) -> int:
        return torch.cuda.get_device_properties(
417
            device_id).multi_processor_count
418
419
420
421

    @classmethod
    def is_navi(cls) -> bool:
        return 'gfx1' in torch.cuda.get_device_properties(0).gcnArchName
422
423

    @classmethod
424
425
    def get_static_graph_wrapper_cls(cls) -> str:
        return "vllm.compilation.cuda_graph.CUDAGraphWrapper"
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455

    @classmethod
    def stateless_init_device_torch_dist_pg(
        cls,
        backend: str,
        prefix_store: PrefixStore,
        group_rank: int,
        group_size: int,
        timeout: timedelta,
    ) -> ProcessGroup:
        assert is_nccl_available()
        pg: ProcessGroup = ProcessGroup(
            prefix_store,
            group_rank,
            group_size,
        )
        from torch.distributed.distributed_c10d import ProcessGroupNCCL

        backend_options = ProcessGroupNCCL.Options()
        backend_options._timeout = timeout

        backend_class = ProcessGroupNCCL(prefix_store, group_rank, group_size,
                                         backend_options)
        backend_type = ProcessGroup.BackendType.NCCL
        device = torch.device("cuda")
        pg._set_default_backend(backend_type)
        backend_class._set_sequence_number_for_group()

        pg._register_backend(device, backend_type, backend_class)
        return pg
456
457
458
459

    @classmethod
    def device_count(cls) -> int:
        return cuda_device_count_stateless()
460
461

    @classmethod
462
463
    def is_kv_cache_dtype_supported(cls, kv_cache_dtype: str,
                                    model_config: "ModelConfig") -> bool:
464
        return True
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484

    @classmethod
    def check_if_supports_dtype(cls, torch_dtype: torch.dtype):
        if torch_dtype == torch.bfloat16:  # noqa: SIM102
            if not cls.has_device_capability(80):
                capability = cls.get_device_capability()
                gpu_name = cls.get_device_name()

                if capability is None:
                    compute_str = "does not have a compute capability"
                else:
                    version_str = capability.as_version_str()
                    compute_str = f"has compute capability {version_str}"

                raise ValueError(
                    "Bfloat16 is only supported on GPUs "
                    "with compute capability of at least 8.0. "
                    f"Your {gpu_name} GPU {compute_str}. "
                    "You can use float16 instead by explicitly setting the "
                    "`dtype` flag in CLI, for example: --dtype=half.")
485
486
487
488

    @classmethod
    def support_hybrid_kv_cache(cls) -> bool:
        return True
489
490
491
492

    @classmethod
    def support_static_graph_mode(cls) -> bool:
        return True