rocm.py 18.9 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
15
from vllm.logger import init_logger

16
from .interface import DeviceCapability, Platform, PlatformEnum, _Backend
17

18
if TYPE_CHECKING:
19
    from vllm.config import ModelConfig, VllmConfig
20

21
22
logger = init_logger(__name__)

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

30
31
32
33
34
35
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
zhuwenwen's avatar
zhuwenwen committed
36
37
38
39
# try:
#     import vllm._rocm_C  # noqa: F401
# except ImportError as e:
#     logger.warning("Failed to import from vllm._rocm_C with %r", e)
40

41

42
# Models not supported by ROCm.
43
_ROCM_UNSUPPORTED_MODELS: list[str] = []
44
45
46

# Models partially supported by ROCm.
# Architecture -> Reason.
zhuwenwen's avatar
zhuwenwen committed
47
48
49
50
# _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`")
51
_ROCM_PARTIALLY_SUPPORTED_MODELS: dict[str, str] = {
zhuwenwen's avatar
zhuwenwen committed
52
53
54
55
56
57
    # "Qwen2ForCausalLM":
    # _ROCM_SWA_REASON,
    # "MistralForCausalLM":
    # _ROCM_SWA_REASON,
    # "MixtralForCausalLM":
    # _ROCM_SWA_REASON,
zhuwenwen's avatar
zhuwenwen committed
58
59
60
    # "PaliGemmaForConditionalGeneration":
    # ("ROCm flash attention does not yet "
    #  "fully support 32-bit precision on PaliGemma"),
zhuwenwen's avatar
zhuwenwen committed
61
62
63
64
    # "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`")
65
}
66
_ROCM_DEVICE_ID_NAME_MAP: dict[str, str] = {
67
68
69
70
71
72
73
74
    "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",
}
75

76
# Prevent use of clashing `{CUDA/HIP}_VISIBLE_DEVICES``
77
78
79
80
81
82
# 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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

# 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


103
104
105
106
107
108
109
110
111
def device_id_to_physical_device_id(device_id: int) -> int:
    if "CUDA_VISIBLE_DEVICES" in os.environ:
        device_ids = os.environ["CUDA_VISIBLE_DEVICES"].split(",")
        physical_device_id = device_ids[device_id]
        return int(physical_device_id)
    else:
        return device_id


112
113
114
115
@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"])
116

117

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


124
@cache
125
126
127
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"])
128
129


130
@cache
131
132
133
134
135
136
137
138
139
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,
        alibi_slopes: Optional[torch.Tensor] = None) -> bool:
140

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

145
146
    # custom paged attn always supported on V0. On V1, requires sliding window
    # disabled due to observed numerical discrepancy.
zhuwenwen's avatar
zhuwenwen committed
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
    # 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)
    #             and max_seq_len <= 32768 and (envs.VLLM_ROCM_CUSTOM_PAGED_ATTN)
    #             and not (envs.VLLM_ROCM_USE_AITER_PAGED_ATTN
    #                      and envs.VLLM_ROCM_USE_AITER))

    # 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)
    #             and max_seq_len <= 32768 and alibi_slopes is None
    #             and kv_cache_dtype == "auto"
    #             and envs.VLLM_ROCM_CUSTOM_PAGED_ATTN)
167
    return False
168
169


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

179
    supported_quantization: list[str] = [
180
        "awq", "gptq", "fp8", "compressed-tensors", "fbgemm_fp8", "gguf",
zhuwenwen's avatar
zhuwenwen committed
181
        "quark", "ptpc_fp8", "moe_wna16", "blockwise_int8","w8a8_int8"
182
    ]
183

184
    @classmethod
185
    def get_attn_backend_cls(cls, selected_backend, head_size, dtype,
186
187
                             kv_cache_dtype, block_size, use_v1,
                             use_mla) -> str:
188
        if use_mla:
zhuwenwen's avatar
zhuwenwen committed
189
190
191
192
193
            if selected_backend == _Backend.TRITON_MLA or block_size != 64:
                if use_v1:
                    logger.info_once("Using Triton MLA backend on V1 engine.")
                    return ("vllm.v1.attention.backends.mla."
                            "triton_mla.TritonMLABackend")
194

zhuwenwen's avatar
zhuwenwen committed
195
196
                else:
                    logger.info("Using Triton MLA backend.")
197
                    return "vllm.attention.backends.triton_mla.TritonMLABackend"  
zhuwenwen's avatar
zhuwenwen committed
198
            else:
199
200
201
202
203
204
205
206
207
208
209
210
                if envs.VLLM_USE_FLASH_MLA:
                    from vllm.attention.backends.flashmla import (
                        is_flashmla_supported)
                    if not is_flashmla_supported()[0]:
                        logger.warning(
                            "FlashMLA backend is not supported due to %s",
                            is_flashmla_supported()[1])
                    elif block_size != 64:
                        logger.warning(
                            "FlashMLA backend is not supported for block size %d"
                            " (currently only supports block size 64).",
                            block_size)
zhuwenwen's avatar
zhuwenwen committed
211
                    else:
212
213
214
215
216
217
218
219
220
221
222
                        if use_v1:
                            logger.info_once(
                                "Using FlashMLA backend on V1 engine.")
                            return ("vllm.v1.attention.backends.mla."
                                    "flashmla.FlashMLABackend")
                        else:
                            logger.info("Using FlashMLA backend.")
                            return ("vllm.attention.backends."
                                    "flashmla.FlashMLABackend")
                else:
                    logger.info("Using Triton MLA backend (block size 64).")
zhuwenwen's avatar
zhuwenwen committed
223
                    return "vllm.attention.backends.triton_mla.TritonMLABackend"           
zhuwenwen's avatar
zhuwenwen committed
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240

            # from vllm.attention.backends.rocm_aiter_mla import (
            #     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:
            #         logger.info("Using Triton MLA backend.")
            #         return "vllm.attention.backends.triton_mla.TritonMLABackend"  # noqa: E501
            #     else:
            #         raise ValueError(
            #             f" The selected backend, {selected_backend.name},"
            #             f"does not support block size {block_size}.")
zhuwenwen's avatar
zhuwenwen committed
241
242
            # elif selected_backend == _Backend.ROCM_AITER_MLA \
            #     or selected_backend == _Backend.ROCM_AITER_MLA_VLLM_V1:
zhuwenwen's avatar
zhuwenwen committed
243
            #     if block_size == 1:
zhuwenwen's avatar
zhuwenwen committed
244
245
246
247
248
249
            #         if use_v1:
            #             logger.info("Using AITER MLA backend on V1 engine.")
            #             return "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend"  # noqa: E501
            #         else:
            #             logger.info("Using AITER MLA backend")
            #             return "vllm.attention.backends.rocm_aiter_mla.AiterMLABackend"  # noqa: E501
250
251
252
      
        if selected_backend is None or selected_backend == _Backend.FLASH_ATTN:
            selected_backend = _Backend.ROCM_FLASH
253

254
255
        if envs.VLLM_USE_V1:
            if envs.VLLM_FLASH_ATTN_V1 and block_size == 64:
256
                if cls.has_device_capability(80):
257
                    logger.info_once("Using Flash Attention backend on V1 engine. (only supports block size 64)")
258
259
                    return ("vllm.v1.attention.backends."
                            "flash_attn.FlashAttentionBackend")
260
            else:
261
262
263
                logger.info("Using Triton Attention backend on V1 engine.")
                return ("vllm.v1.attention.backends."
                        "triton_attn.TritonAttentionBackend")
264
265
266
267
268
269
270
271
        if selected_backend == _Backend.ROCM_FLASH:
            if not cls.has_device_capability(90):
                # not Instinct series GPUs.
                logger.info("flash_attn is not supported on NAVI GPUs.")
        else:
            logger.info("%s is not supported in AMD GPUs.", selected_backend)
        logger.info("Using ROCmFlashAttention backend.")
        return "vllm.attention.backends.rocm_flash_attn.ROCmFlashAttentionBackend"  # noqa: E501
272

273

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

282
    @classmethod
283
    @with_amdsmi_context
284
    def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:
zhuwenwen's avatar
zhuwenwen committed
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
        """
        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
305

306
    @classmethod
307
    @with_amdsmi_context
308
    @lru_cache(maxsize=8)
309
    def get_device_name(cls, device_id: int = 0) -> str:
310
        physical_device_id = device_id_to_physical_device_id(device_id)
311
        handle = amdsmi_get_processor_handles()[physical_device_id]
312
313
        # return amdsmi_get_gpu_asic_info(handle)["market_name"]
        return torch.cuda.get_device_name(device_id)
314

zhuwenwen's avatar
zhuwenwen committed
315
316
317
318
    @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
319

320
321
322
323
324
325
326
327
328
329
    @classmethod
    def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:
        if enforce_eager:
            logger.warning(
                "To see benefits of async output processing, enable CUDA "
                "graph. Since, enforce-eager is enabled, async output "
                "processor cannot be used")
            return False
        return True

330
    @classmethod
331
    def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
332
333
334
335
        cache_config = vllm_config.cache_config
        if cache_config and cache_config.block_size is None:
            cache_config.block_size = 16

336
337
338
339
        parallel_config = vllm_config.parallel_config
        scheduler_config = vllm_config.scheduler_config
        if parallel_config.worker_cls == "auto":
            if scheduler_config.is_multi_step:
340
341
342
                if envs.VLLM_USE_V1:
                    raise NotImplementedError(
                        "Multi-step scheduling is not supported (and not "
343
                        "needed) on vLLM V1. Please launch without "
344
345
346
347
                        "--num-scheduler-steps.")
                else:
                    parallel_config.worker_cls = \
                        "vllm.worker.multi_step_worker.MultiStepWorker"
348
            elif vllm_config.speculative_config:
349
                if envs.VLLM_USE_V1:
zhuwenwen's avatar
zhuwenwen committed
350
351
352
353
354
                    # raise NotImplementedError(
                    #     "Speculative decoding is not yet supported on vLLM V1."
                    # )
                    parallel_config.worker_cls = \
                            "vllm.v1.worker.gpu_worker.Worker"
355
356
357
358
359
                else:
                    parallel_config.worker_cls = \
                        "vllm.spec_decode.spec_decode_worker.create_spec_worker"
                    parallel_config.sd_worker_cls = \
                        "vllm.worker.worker.Worker"
360
            else:
361
362
363
364
365
                if envs.VLLM_USE_V1:
                    parallel_config.worker_cls = \
                            "vllm.v1.worker.gpu_worker.Worker"
                else:
                    parallel_config.worker_cls = "vllm.worker.worker.Worker"
366

367
368
369
370
371
372
373
374
375
376
377
378
    @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)

379
380
381
382
383
384
    @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"
385
386
                " is not set, disabling VLLM_USE_TRITON_AWQ.")
            envs.VLLM_USE_TRITON_AWQ = False
387
388
389
390

    @classmethod
    def get_punica_wrapper(cls) -> str:
        return "vllm.lora.punica_wrapper.punica_gpu.PunicaWrapperGPU"
391
392
393
394
395
396

    @classmethod
    def get_current_memory_usage(cls,
                                 device: Optional[torch.types.Device] = None
                                 ) -> float:
        torch.cuda.reset_peak_memory_stats(device)
zhuwenwen's avatar
zhuwenwen committed
397
398
399
        # return torch.cuda.mem_get_info(device)[1] - torch.cuda.mem_get_info(
        #     device)[0]
        return torch.cuda.max_memory_allocated(device)
400
401
402
403

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

405
406
407
408
409
    @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"])

410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
    @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
426
427

    @classmethod
428
    def supports_v1(cls, model_config: "ModelConfig") -> bool:
429
430
        # V1 support on AMD gpus is experimental
        return True
431
432
433
434
435

    @classmethod
    def use_custom_allreduce(cls) -> bool:
        # We only enable custom allreduce for MI300 series
        gcn_arch = torch.cuda.get_device_properties(0).gcnArchName
436
        supported_archs = ['gfx94', 'gfx95']
437
        return any(gfx in gcn_arch for gfx in supported_archs)
438
439
440
441

    @classmethod
    def get_cu_count(cls, device_id: int = 0) -> int:
        return torch.cuda.get_device_properties(
442
            device_id).multi_processor_count
443
444
445
446

    @classmethod
    def is_navi(cls) -> bool:
        return 'gfx1' in torch.cuda.get_device_properties(0).gcnArchName
447
448
449
450

    @classmethod
    def get_piecewise_backend_cls(cls) -> str:
        return "vllm.compilation.cuda_piecewise_backend.CUDAPiecewiseBackend"  # noqa
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480

    @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