rocm.py 17.8 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 functools import cache, lru_cache, wraps
6
from typing import TYPE_CHECKING
7
8
9

import torch

10
import vllm.envs as envs
11
from vllm.logger import init_logger
12
from vllm.utils.torch_utils import cuda_device_count_stateless
13

14
from .interface import DeviceCapability, Platform, PlatformEnum
15

16
if TYPE_CHECKING:
17
    from vllm.attention.backends.registry import _Backend
18
    from vllm.config import VllmConfig
19
20
else:
    _Backend = None
21

22
23
logger = init_logger(__name__)

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

36
37
38
39
40
41
42
43
44
45
46
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)

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

# Models partially supported by ROCm.
# Architecture -> Reason.
52
53
54
55
56
57
_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`"
)
58
_ROCM_PARTIALLY_SUPPORTED_MODELS: dict[str, str] = {
59
60
61
    "Qwen2ForCausalLM": _ROCM_SWA_REASON,
    "MistralForCausalLM": _ROCM_SWA_REASON,
    "MixtralForCausalLM": _ROCM_SWA_REASON,
62
63
64
    "PaliGemmaForConditionalGeneration": (
        "ROCm flash attention does not yet fully support 32-bit precision on PaliGemma"
    ),
65
66
67
68
69
    "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`"
    ),
70
}
71
_ROCM_DEVICE_ID_NAME_MAP: dict[str, str] = {
72
73
74
    "0x74a0": "AMD_Instinct_MI300A",
    "0x74a1": "AMD_Instinct_MI300X",
    "0x74b5": "AMD_Instinct_MI300X",  # MI300X VF
75
    "0x74a2": "AMD_Instinct_MI308X",
76
77
78
79
80
    "0x74a5": "AMD_Instinct_MI325X",
    "0x74b9": "AMD_Instinct_MI325X",  # MI325X VF
    "0x74a9": "AMD_Instinct_MI300X_HF",
    "0x74bd": "AMD_Instinct_MI300X_HF",
}
81

82
# Prevent use of clashing `{CUDA/HIP}_VISIBLE_DEVICES`
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
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


108
109
110
111
112
113
@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"])


114
@cache
115
def on_mi3xx() -> bool:
116
    GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
117
118
119
120
121
122
123
    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"])
124
125


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


132
@cache
133
def use_rocm_custom_paged_attention(
134
135
136
137
138
139
140
    qtype: torch.dtype,
    head_size: int,
    block_size: int,
    gqa_ratio: int,
    max_seq_len: int,
    sliding_window: int,
    kv_cache_dtype: str,
141
142
    alibi_slopes: torch.Tensor | None = None,
    sinks: torch.Tensor | None = None,
143
) -> bool:
144
145
    GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
    ON_GFX9 = any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])
146
    ON_GFX11_GFX12 = any(arch in GPU_ARCH for arch in ["gfx11", "gfx12"])
147

148
149
    # custom paged attn always supported on V0. On V1, requires sliding window
    # disabled due to observed numerical discrepancy.
150
    if ON_GFX9:
151
152
153
154
155
156
157
158
159
160
161
        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)
            and max_seq_len <= 128 * 1024
            and (envs.VLLM_ROCM_CUSTOM_PAGED_ATTN)
            and not (envs.VLLM_ROCM_USE_AITER_PAGED_ATTN and envs.VLLM_ROCM_USE_AITER)
            and sinks is None
        )
162
163

    else:
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
        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)
            and max_seq_len <= 128 * 1024
            and alibi_slopes is None
            and kv_cache_dtype == "auto"
            and envs.VLLM_ROCM_CUSTOM_PAGED_ATTN
            and sinks is None
        )
181
182


183
184
class RocmPlatform(Platform):
    _enum = PlatformEnum.ROCM
185
    device_name: str = "rocm"
186
    device_type: str = "cuda"
187
    dispatch_key: str = "CUDA"
188
    ray_device_key: str = "GPU"
189
    dist_backend: str = "nccl"
190
191
    # rocm shares the same device control env var as CUDA
    device_control_env_var: str = "CUDA_VISIBLE_DEVICES"
192

193
    supported_quantization: list[str] = [
194
195
196
197
198
199
200
201
202
203
204
        "awq",
        "gptq",
        "fp8",
        "compressed-tensors",
        "fbgemm_fp8",
        "gguf",
        "quark",
        "ptpc_fp8",
        "mxfp4",
        "petit_nvfp4",
        "torchao",
205
    ]
206

207
    @classmethod
208
    def get_vit_attn_backend(cls, head_size: int, dtype: torch.dtype) -> "_Backend":
209
210
        from importlib.util import find_spec

211
        from vllm.attention.backends.registry import _Backend
212
213

        if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA and on_gfx9():
214
            return _Backend.ROCM_AITER_FA
215
216

        if on_gfx9() and find_spec("flash_attn") is not None:
217
            return _Backend.FLASH_ATTN
218

219
220
        return _Backend.TORCH_SDPA

221
    @classmethod
222
223
224
225
226
227
228
229
230
231
232
233
    def get_attn_backend_cls(
        cls,
        selected_backend,
        head_size,
        dtype,
        kv_cache_dtype,
        block_size,
        use_v1,
        use_mla,
        has_sink,
        use_sparse,
    ) -> str:
234
        from vllm.attention.backends.registry import _Backend
235

236
        if use_sparse:
237
            raise NotImplementedError("Sparse Attention is not supported on ROCm.")
238
        if use_mla:
239
240
241
            if not use_v1:
                raise RuntimeError(
                    "MLA attention backends require the V1 engine. "
242
243
                    "Set VLLM_USE_V1=1 to enable them."
                )
244

245
            from vllm.v1.attention.backends.mla.rocm_aiter_mla import (
246
247
                is_aiter_mla_enabled,
            )
248
249

            if selected_backend is None:
250
251
252
253
254
                selected_backend = (
                    _Backend.ROCM_AITER_MLA
                    if is_aiter_mla_enabled() or block_size == 1
                    else _Backend.TRITON_MLA
                )
255
256
257

            if selected_backend == _Backend.TRITON_MLA:
                if block_size != 1:
258
                    logger.info_once("Using Triton MLA backend on V1 engine.")
259
                    return "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend"
260
261
                raise ValueError(
                    f" The selected backend, {selected_backend.name},"
262
263
                    f"does not support block size {block_size}."
                )
264
            if selected_backend == _Backend.ROCM_AITER_MLA:
265
                if block_size == 1:
266
                    logger.info("Using AITER MLA backend on V1 engine.")
267
268
269
                    return (
                        "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend"  # noqa: E501
                    )
270
271
                raise ValueError(
                    f" The selected backend, {selected_backend.name},"
272
                    f"does not support block size {block_size}."
273
274
                    "(currently only supports block size 1)"
                )
275
276
            raise ValueError(
                f" The selected backend, {selected_backend.name},"
277
278
                f"is not MLA type while requested for MLA backend."
            )
279

280
        if envs.VLLM_USE_V1:
281
282
283
            if selected_backend == _Backend.FLEX_ATTENTION:
                logger.info("Using FlexAttention backend on V1 engine.")
                return "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend"
284
285
286
287
            if (
                envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA and on_gfx9()
            ) or selected_backend == _Backend.ROCM_AITER_FA:
                logger.info("Using Aiter Flash Attention backend on V1 engine.")
288
289
290
291
                return (
                    "vllm.v1.attention.backends."
                    "rocm_aiter_fa.AiterFlashAttentionBackend"
                )
292
293
294
295
296
297
298
299
300
301
            if (
                envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION
            ) or selected_backend == _Backend.ROCM_AITER_UNIFIED_ATTN:
                logger.info("Using Aiter Unified Attention backend on V1 engine.")
                return (
                    "vllm.v1.attention.backends."
                    "rocm_aiter_unified_attn.RocmAiterUnifiedAttentionBackend"
                )
            if (
                envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION
302
303
                or selected_backend == _Backend.ROCM_ATTN
            ):
304
305
                # rocm specific backend, with aiter and/or
                #   triton prefix-prefill
306
                logger.info("Using Rocm Attention backend on V1 engine.")
307
                return "vllm.v1.attention.backends.rocm_attn.RocmAttentionBackend"
308
309
310
            # default case, using triton unified attention
            logger.info("Using Triton Attention backend on V1 engine.")
            return "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend"
311
312
        raise RuntimeError(
            "V0 attention backends have been removed. Set VLLM_USE_V1=1 "
313
314
            "to select a supported backend."
        )
315

316
317
318
319
320
321
322
    @classmethod
    def set_device(cls, device: torch.device) -> None:
        """
        Set the device for the current platform.
        """
        torch.cuda.set_device(device)

323
    @classmethod
324
    @lru_cache(maxsize=8)
325
    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
326
327
        major, minor = torch.cuda.get_device_capability(device_id)
        return DeviceCapability(major=major, minor=minor)
328

329
    @classmethod
330
    @with_amdsmi_context
331
    def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:
332
333
334
        """
        Query if the set of gpus are fully connected by xgmi (1 hop)
        """
335
        handles = [amdsmi_get_processor_handles()[i] for i in physical_device_ids]
336
337
338
339
        for i, handle in enumerate(handles):
            for j, peer_handle in enumerate(handles):
                if i < j:
                    try:
340
                        link_type = amdsmi_topo_get_link_type(handle, peer_handle)
341
342
343
344
                        # type is 2 for XGMI
                        if link_type["hops"] != 1 or link_type["type"] != 2:
                            return False
                    except AmdSmiException as error:
345
                        logger.error("AMD 1 hop XGMI detection failed.", exc_info=error)
346
347
348
                        return False
        return True

349
    @classmethod
350
    @with_amdsmi_context
351
    @lru_cache(maxsize=8)
352
    def get_device_name(cls, device_id: int = 0) -> str:
353
        physical_device_id = cls.device_id_to_physical_device_id(device_id)
354
        handle = amdsmi_get_processor_handles()[physical_device_id]
355
356
357
358
359
        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"]
360
361
362
363
364

    @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
365
366

    @classmethod
367
    def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
368
369
        from vllm.config.compilation import CUDAGraphMode

370
        cache_config = vllm_config.cache_config
371
372
373
374
375
        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
376
377
378
        use_aiter_rms_norm = (
            envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_RMSNORM
        )
379

380
381
382
        if cache_config and cache_config.block_size is None:
            cache_config.block_size = 16

383
        if parallel_config.worker_cls == "auto":
384
            parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker"
385
        #  Aiter rms norm perform best when CUDA Graph capture is enabled.
386
387
388
389
390
391
        if (
            use_v1
            and use_aiter_rms_norm
            and not is_eager_execution
            and "-rms_norm" not in compilation_config.custom_ops
        ):
392
            compilation_config.custom_ops.append("+rms_norm")
393

394
395
396
    @classmethod
    def verify_model_arch(cls, model_arch: str) -> None:
        if model_arch in _ROCM_UNSUPPORTED_MODELS:
397
398
399
            raise ValueError(
                f"Model architecture '{model_arch}' is not supported by ROCm for now."
            )
400
401
402
403

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

409
410
411
412
413
414
    @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"
415
416
                " is not set, enabling VLLM_USE_TRITON_AWQ."
            )
417
        os.environ["VLLM_USE_TRITON_AWQ"] = "1"
418
419
420
421

    @classmethod
    def get_punica_wrapper(cls) -> str:
        return "vllm.lora.punica_wrapper.punica_gpu.PunicaWrapperGPU"
422
423

    @classmethod
424
    def get_current_memory_usage(
425
        cls, device: torch.types.Device | None = None
426
    ) -> float:
427
        torch.cuda.reset_peak_memory_stats(device)
428
        return torch.cuda.mem_get_info(device)[1] - torch.cuda.mem_get_info(device)[0]
429
430
431

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

436
437
438
439
440
    @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"])

441
442
443
    @classmethod
    def supports_fp8(cls) -> bool:
        gcn_arch = torch.cuda.get_device_properties(0).gcnArchName
444
        return any(gfx in gcn_arch for gfx in ["gfx94", "gfx95", "gfx12"])
445
446
447
448

    @classmethod
    def is_fp8_fnuz(cls) -> bool:
        # only device 0 is checked, this assumes MI300 platforms are homogeneous
449
        return "gfx94" in torch.cuda.get_device_properties(0).gcnArchName
450
451
452
453
454
455
456

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

458
459
460
461
    @classmethod
    def use_custom_allreduce(cls) -> bool:
        # We only enable custom allreduce for MI300 series
        gcn_arch = torch.cuda.get_device_properties(0).gcnArchName
462
        supported_archs = ["gfx94", "gfx95"]
463
        return any(gfx in gcn_arch for gfx in supported_archs)
464

465
466
467
468
    @classmethod
    def opaque_attention_op(cls) -> bool:
        return True

469
470
    @classmethod
    def get_cu_count(cls, device_id: int = 0) -> int:
471
        return torch.cuda.get_device_properties(device_id).multi_processor_count
472
473
474

    @classmethod
    def is_navi(cls) -> bool:
475
        return "gfx1" in torch.cuda.get_device_properties(0).gcnArchName
476
477

    @classmethod
478
479
    def get_static_graph_wrapper_cls(cls) -> str:
        return "vllm.compilation.cuda_graph.CUDAGraphWrapper"
480

481
482
483
    @classmethod
    def device_count(cls) -> int:
        return cuda_device_count_stateless()
484

485
    @classmethod
486
487
    def check_if_supports_dtype(cls, dtype: torch.dtype):
        if dtype == torch.bfloat16:  # noqa: SIM102
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
            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 "
503
504
                    "`dtype` flag in CLI, for example: --dtype=half."
                )
505
506
507
508

    @classmethod
    def support_hybrid_kv_cache(cls) -> bool:
        return True
509
510
511
512

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