cuda.py 23 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
6
"""Code inside this file can safely assume cuda platform, e.g. importing
pynvml. However, it should not initialize cuda context.
"""

7
import os
8
from collections.abc import Callable
9
from functools import cache, wraps
10
from typing import TYPE_CHECKING, Optional, TypeVar
11

12
import torch
13
from typing_extensions import ParamSpec
14

15
16
# import custom ops, trigger op registration
import vllm._C  # noqa
17
from vllm.logger import init_logger
18
from vllm.utils.import_utils import import_pynvml
19
from vllm.utils.torch_utils import cuda_device_count_stateless
20
from vllm.v1.attention.backends.registry import AttentionBackendEnum
21

22
from .interface import DeviceCapability, Platform, PlatformEnum
23

24
if TYPE_CHECKING:
25
    from vllm.config import VllmConfig
26
    from vllm.config.cache import CacheDType
27
    from vllm.v1.attention.selector import AttentionSelectorConfig
28
else:
29
30
    VllmConfig = None
    CacheDType = None
31

32
33
logger = init_logger(__name__)

34
35
36
_P = ParamSpec("_P")
_R = TypeVar("_R")

37
pynvml = import_pynvml()
38

39
40
41
42
# pytorch 2.5 uses cudnn sdpa by default, which will cause crash on some models
# see https://github.com/huggingface/diffusers/issues/9704 for details
torch.backends.cuda.enable_cudnn_sdp(False)

43

44
45
46
47
48
49
50
51
52
@cache
def _get_backend_priorities(
    use_mla: bool,
    device_capability: DeviceCapability,
) -> list[AttentionBackendEnum]:
    """Get backend priorities with lazy import to avoid circular dependency."""
    if use_mla:
        if device_capability.major == 10:
            return [
53
                AttentionBackendEnum.FLASHINFER_MLA,
54
                AttentionBackendEnum.CUTLASS_MLA,
55
                AttentionBackendEnum.FLASH_ATTN_MLA,
56
                AttentionBackendEnum.FLASHMLA,
57
58
59
60
61
62
                AttentionBackendEnum.TRITON_MLA,
                AttentionBackendEnum.FLASHMLA_SPARSE,
            ]
        else:
            return [
                AttentionBackendEnum.FLASH_ATTN_MLA,
63
                AttentionBackendEnum.FLASHMLA,
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
                AttentionBackendEnum.FLASHINFER_MLA,
                AttentionBackendEnum.TRITON_MLA,
                AttentionBackendEnum.FLASHMLA_SPARSE,
            ]
    else:
        if device_capability.major == 10:
            return [
                AttentionBackendEnum.FLASHINFER,
                AttentionBackendEnum.FLASH_ATTN,
                AttentionBackendEnum.TRITON_ATTN,
                AttentionBackendEnum.FLEX_ATTENTION,
            ]
        else:
            return [
                AttentionBackendEnum.FLASH_ATTN,
                AttentionBackendEnum.FLASHINFER,
                AttentionBackendEnum.TRITON_ATTN,
                AttentionBackendEnum.FLEX_ATTENTION,
            ]


85
def with_nvml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]:
86
    @wraps(fn)
87
    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
88
89
90
91
92
93
94
95
96
        pynvml.nvmlInit()
        try:
            return fn(*args, **kwargs)
        finally:
            pynvml.nvmlShutdown()

    return wrapper


97
98
class CudaPlatformBase(Platform):
    _enum = PlatformEnum.CUDA
99
    device_name: str = "cuda"
100
101
    device_type: str = "cuda"
    dispatch_key: str = "CUDA"
102
    ray_device_key: str = "GPU"
103
    dist_backend: str = "nccl"
104
    device_control_env_var: str = "CUDA_VISIBLE_DEVICES"
105

106
    @property
107
    def supported_dtypes(self) -> list[torch.dtype]:
108
109
110
        if self.has_device_capability(80):
            # Ampere and Hopper or later NVIDIA GPUs.
            return [torch.bfloat16, torch.float16, torch.float32]
111
        if self.has_device_capability(60):
112
113
114
115
116
117
            # Pascal, Volta and Turing NVIDIA GPUs, BF16 is not supported
            return [torch.float16, torch.float32]
        # Kepler and Maxwell NVIDIA GPUs, only FP32 is supported,
        # though vLLM doesn't support these GPUs.
        return [torch.float32]

118
119
120
121
122
    @classmethod
    def set_device(cls, device: torch.device) -> None:
        """
        Set the device for the current platform.
        """
123
        torch.cuda.set_device(device)
124
125
126
127
128
        # With this trick we can force the device to be set eagerly
        # see https://github.com/pytorch/pytorch/issues/155668
        # for why and when it is needed
        _ = torch.zeros(1, device=device)

129
    @classmethod
130
    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
131
        raise NotImplementedError
132

133
134
135
    @classmethod
    def get_device_name(cls, device_id: int = 0) -> str:
        raise NotImplementedError
136

137
138
139
    @classmethod
    def get_device_total_memory(cls, device_id: int = 0) -> int:
        raise NotImplementedError
140

141
    @classmethod
142
    def is_fully_connected(cls, device_ids: list[int]) -> bool:
143
        raise NotImplementedError
144

145
146
147
    @classmethod
    def log_warnings(cls):
        pass
148

149
    @classmethod
150
    def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
151
        from vllm.v1.attention.backends.registry import AttentionBackendEnum
152

153
        parallel_config = vllm_config.parallel_config
154
        model_config = vllm_config.model_config
155

156
        if parallel_config.worker_cls == "auto":
157
            parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker"
158

159
160
161
        cache_config = vllm_config.cache_config
        if cache_config and cache_config.block_size is None:
            cache_config.block_size = 16
162

163
        # TODO(lucas): handle this more gracefully
164
        # Note: model_config may be None during testing
165
166
167
168
169
170
171
172
173
        # Note: block_size is initialized in
        # HybridAttentionMambaModelConfig.verify_and_update_config
        # for models with both attention and mamba,
        # and doesn't need to be reinitialized here
        if (
            model_config is not None
            and model_config.use_mla
            and cache_config.block_size is not None
        ):
174
            use_sparse = hasattr(vllm_config.model_config.hf_config, "index_topk")
175
            # If `--attention-config.backend` is not set and we are using MLA,
176
177
178
179
180
            # then we default to FlashMLA backend for non-blackwell GPUs,
            # else we default to CutlassMLA. For each case, we force the
            # required block_size.
            use_flashmla = False
            use_cutlass_mla = False
181
            use_flashinfer_mla = False
182

183
184
            from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported

185
            if vllm_config.attention_config.backend is None:
186
                # Default case
187
188
189
190
191
192
193
194
195
                hf_text_config = model_config.hf_text_config
                qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)
                if (
                    cls.is_device_capability_family(100)
                    and not use_sparse
                    and qk_nope_head_dim == 128
                ):
                    # Blackwell => Force FlashInfer MLA (unless sparse, i.e. DSv3.2)
                    # and only if qk_nope_head_dim == 128 (kernel constraint)
196
                    use_flashinfer_mla = True
197
198
199
                    # Set the backend in AttentionConfig so it's used during
                    # backend selection
                    vllm_config.attention_config.backend = (
200
                        AttentionBackendEnum.FLASHINFER_MLA
201
                    )
202
203
204
205
206
                elif cls.is_device_capability_family(100) and not use_sparse:
                    # Fall back to CUTLASS_MLA as 2nd priority on Blackwell
                    use_cutlass_mla = True
                elif is_flashmla_dense_supported()[0]:
                    # Non-Blackwell with FlashMLA support
207
                    use_flashmla = True
208
209
210
                else:
                    # Fallback: will use Triton MLA or other compatible backend
                    pass
211
212
            else:
                # Forced case
213
214
215
216
                backend = vllm_config.attention_config.backend
                use_flashmla = backend == AttentionBackendEnum.FLASHMLA
                use_cutlass_mla = backend == AttentionBackendEnum.CUTLASS_MLA
                use_flashinfer_mla = backend == AttentionBackendEnum.FLASHINFER_MLA
217

218
219
            if (
                use_flashmla
220
                and is_flashmla_dense_supported()[0]
221
                and cache_config.block_size % 64 != 0
222
            ):
223
                cache_config.block_size = 64
224
                logger.info("Forcing kv cache block size to 64 for FlashMLA backend.")
225

226
            if use_cutlass_mla and cache_config.block_size % 128 != 0:
227
                cache_config.block_size = 128
228
229
230
                logger.info(
                    "Forcing kv cache block size to 128 for CUTLASS_MLA backend."
                )
231

232
233
234
235
236
            if (
                use_flashinfer_mla
                and cache_config.block_size != 32
                and cache_config.block_size % 64 != 0
            ):
237
238
                cache_config.block_size = 64
                logger.info(
239
240
                    "Forcing kv cache block size to 64 for FlashInferMLA backend."
                )
241

242
243
244
245
            # TODO(Chen): remove this hacky code
            if use_sparse and cache_config.block_size != 64:
                cache_config.block_size = 64
                logger.info(
246
247
                    "Forcing kv cache block size to 64 for FlashMLASparse backend."
                )
248

249
250
251
252
253
254
255
256
257
258
259
260
261
262
        scheduler_config = vllm_config.scheduler_config
        # Note: model_config may be None during testing
        if (
            model_config is not None
            and model_config.is_mm_prefix_lm
            and scheduler_config.is_multimodal_model
            and not scheduler_config.disable_chunked_mm_input
        ):
            logger.warning(
                "Forcing --disable_chunked_mm_input for models "
                "with multimodal-bidirectional attention."
            )
            scheduler_config.disable_chunked_mm_input = True

263
    @classmethod
264
    def get_current_memory_usage(
265
        cls, device: torch.types.Device | None = None
266
    ) -> float:
267
        torch.cuda.empty_cache()
268
269
270
        torch.cuda.reset_peak_memory_stats(device)
        return torch.cuda.max_memory_allocated(device)

271
    @classmethod
272
    def get_valid_backends(
273
        cls,
274
275
        device_capability: DeviceCapability,
        attn_selector_config: "AttentionSelectorConfig",
276
277
278
279
280
281
282
    ) -> tuple[
        list[tuple["AttentionBackendEnum", int]],
        dict["AttentionBackendEnum", list[str]],
    ]:
        valid_backends_priorities = []
        invalid_reasons = {}

283
284
285
286
287
        # TurboQuant KV cache: route directly to TQ backend
        kv_cache_dtype = attn_selector_config.kv_cache_dtype
        if kv_cache_dtype is not None and kv_cache_dtype.startswith("turboquant_"):
            return [(AttentionBackendEnum.TURBOQUANT, 0)], {}

288
289
290
        backend_priorities = _get_backend_priorities(
            attn_selector_config.use_mla, device_capability
        )
291
292
293
294
        for priority, backend in enumerate(backend_priorities):
            try:
                backend_class = backend.get_class()
                invalid_reasons_i = backend_class.validate_configuration(
295
296
                    device_capability=device_capability,
                    **attn_selector_config._asdict(),
297
                )
298
299
300
301
302
303
            except ImportError:
                invalid_reasons_i = ["ImportError"]
            if invalid_reasons_i:
                invalid_reasons[backend] = invalid_reasons_i
            else:
                valid_backends_priorities.append((backend, priority))
304

305
        return valid_backends_priorities, invalid_reasons
306

307
308
309
310
    @classmethod
    def get_attn_backend_cls(
        cls,
        selected_backend: "AttentionBackendEnum",
311
        attn_selector_config: "AttentionSelectorConfig",
312
313
314
315
    ) -> str:
        device_capability = cls.get_device_capability()
        assert device_capability is not None

316
        attn_selector_config = attn_selector_config._replace(block_size=None)
317
318
319
320
321
        # First try checking just the selected backend, if there is one.
        if selected_backend is not None:
            try:
                backend_class = selected_backend.get_class()
                invalid_reasons = backend_class.validate_configuration(
322
323
                    device_capability=device_capability,
                    **attn_selector_config._asdict(),
324
                )
325
326
327
328
329
330
            except ImportError:
                invalid_reasons = ["ImportError"]
            if invalid_reasons:
                raise ValueError(
                    f"Selected backend {selected_backend} is not valid for "
                    f"this configuration. Reason: {invalid_reasons}"
331
                )
332
333
334
335
336
337
338
            else:
                logger.info("Using %s backend.", selected_backend)
                return selected_backend.get_path()

        # No selected backend or the selected backend is invalid,
        # so we try finding a valid backend.
        valid_backends_priorities, invalid_reasons = cls.get_valid_backends(
339
340
            device_capability=device_capability,
            attn_selector_config=attn_selector_config,
341
        )
342
343
344
345
346
347
348
        reasons_str = (
            "{"
            + ", ".join(
                f"{backend.name}: [{', '.join(reasons)}]"
                for backend, reasons in invalid_reasons.items()
            )
            + "}"
349
        )
350
        config_str = attn_selector_config.__repr__()
351
352
353
354
355
356
357
358
359
        logger.debug_once(
            f"Some attention backends are not valid for {cls.device_name} with "
            f"{config_str}. Reasons: {reasons_str}."
        )
        if len(valid_backends_priorities) == 0:
            raise ValueError(
                f"No valid attention backend found for {cls.device_name} "
                f"with {config_str}. Reasons: {reasons_str}."
            )
360

361
362
363
364
365
366
367
368
        # We have found some valid backends. Select the one with the
        # highest priority.
        sorted_indices = sorted(
            range(len(valid_backends_priorities)),
            key=lambda i: valid_backends_priorities[i][1],
        )
        selected_index = sorted_indices[0]
        selected_backend = valid_backends_priorities[selected_index][0]
369
        logger.info_once(
370
            "Using %s attention backend out of potential backends: %s",
371
            selected_backend.name,
372
373
            tuple(b[0].name for b in valid_backends_priorities),
            scope="local",
374
375
376
        )

        return selected_backend.get_path()
377

378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
    @classmethod
    def get_supported_vit_attn_backends(cls) -> list["AttentionBackendEnum"]:
        return [
            AttentionBackendEnum.TORCH_SDPA,
            AttentionBackendEnum.FLASH_ATTN,
        ]

    @classmethod
    def get_vit_attn_backend(
        cls,
        head_size: int,
        dtype: torch.dtype,
        backend: Optional["AttentionBackendEnum"] = None,
    ) -> "AttentionBackendEnum":
        if backend is not None:
            assert backend in cls.get_supported_vit_attn_backends(), (
                f"Backend {backend} is not supported for vit attention. "
                f"Supported backends are: {cls.get_supported_vit_attn_backends()}"
            )
            logger.info_once(f"Using backend {backend} for vit attention")
            return backend

        # Try FlashAttention first
        if (cc := cls.get_device_capability()) and cc.major >= 8:
            try:
                backend_class = AttentionBackendEnum.FLASH_ATTN.get_class()
                if backend_class.supports_head_size(
                    head_size
                ) and backend_class.supports_dtype(dtype):
                    return AttentionBackendEnum.FLASH_ATTN
            except ImportError:
                pass

        return AttentionBackendEnum.TORCH_SDPA

413
414
415
416
    @classmethod
    def get_punica_wrapper(cls) -> str:
        return "vllm.lora.punica_wrapper.punica_gpu.PunicaWrapperGPU"

417
418
    @classmethod
    def get_device_communicator_cls(cls) -> str:
419
420
421
        return (
            "vllm.distributed.device_communicators.cuda_communicator.CudaCommunicator"  # noqa
        )
422

423
424
425
426
    @classmethod
    def supports_fp8(cls) -> bool:
        return cls.has_device_capability(89)

427
428
429
430
    @classmethod
    def use_custom_allreduce(cls) -> bool:
        return True

431
432
433
434
    @classmethod
    def opaque_attention_op(cls) -> bool:
        return True

435
    @classmethod
436
437
    def get_static_graph_wrapper_cls(cls) -> str:
        return "vllm.compilation.cuda_graph.CUDAGraphWrapper"
438

439
440
441
442
    @classmethod
    def device_count(cls) -> int:
        return cuda_device_count_stateless()

443
    @classmethod
444
445
    def check_if_supports_dtype(cls, dtype: torch.dtype):
        if dtype == torch.bfloat16:  # noqa: SIM102
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
            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 "
461
462
                    "`dtype` flag in CLI, for example: --dtype=half."
                )
463

464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
    @classmethod
    def insert_blocks_to_device(
        cls,
        src_cache: torch.Tensor,
        dst_cache: torch.Tensor,
        src_block_indices: torch.Tensor,
        dst_block_indices: torch.Tensor,
    ) -> None:
        """Copy blocks from src_cache to dst_cache on GPU."""
        _src_cache = src_cache[:, src_block_indices]
        dst_cache[:, dst_block_indices] = _src_cache.to(dst_cache.device)

    @classmethod
    def swap_out_blocks_to_host(
        cls,
        src_cache: torch.Tensor,
        dst_cache: torch.Tensor,
        src_block_indices: torch.Tensor,
        dst_block_indices: torch.Tensor,
    ) -> None:
        """Copy blocks from GPU to host (CPU)."""
        _src_cache = src_cache[:, src_block_indices]
        dst_cache[:, dst_block_indices] = _src_cache.cpu()

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

492
493
494
495
    @classmethod
    def support_static_graph_mode(cls) -> bool:
        return True

496

497
498
499
500
501
# NVML utils
# Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`,
# all the related functions work on real physical device ids.
# the major benefit of using NVML is that it will not initialize CUDA
class NvmlCudaPlatform(CudaPlatformBase):
502
    @classmethod
503
    @cache
504
    @with_nvml_context
505
    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
506
        try:
507
            physical_device_id = cls.device_id_to_physical_device_id(device_id)
508
509
510
511
512
513
514
515
516
517
            handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
            major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle)
            return DeviceCapability(major=major, minor=minor)
        except RuntimeError:
            return None

    @classmethod
    @with_nvml_context
    def has_device_capability(
        cls,
518
        capability: tuple[int, int] | int,
519
520
521
522
523
524
        device_id: int = 0,
    ) -> bool:
        try:
            return super().has_device_capability(capability, device_id)
        except RuntimeError:
            return False
525

526
    @classmethod
527
    @with_nvml_context
528
    def get_device_name(cls, device_id: int = 0) -> str:
529
        physical_device_id = cls.device_id_to_physical_device_id(device_id)
530
        return cls._get_physical_device_name(physical_device_id)
531

532
533
534
    @classmethod
    @with_nvml_context
    def get_device_uuid(cls, device_id: int = 0) -> str:
535
        physical_device_id = cls.device_id_to_physical_device_id(device_id)
536
537
538
        handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
        return pynvml.nvmlDeviceGetUUID(handle)

539
    @classmethod
540
    @with_nvml_context
541
    def get_device_total_memory(cls, device_id: int = 0) -> int:
542
        physical_device_id = cls.device_id_to_physical_device_id(device_id)
543
544
        handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
        return int(pynvml.nvmlDeviceGetMemoryInfo(handle).total)
545

546
    @classmethod
547
    @with_nvml_context
548
    def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:
549
550
551
        """
        query if the set of gpus are fully connected by nvlink (1 hop)
        """
552
        handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_ids]
553
554
555
556
557
        for i, handle in enumerate(handles):
            for j, peer_handle in enumerate(handles):
                if i < j:
                    try:
                        p2p_status = pynvml.nvmlDeviceGetP2PStatus(
558
559
560
561
                            handle,
                            peer_handle,
                            pynvml.NVML_P2P_CAPS_INDEX_NVLINK,
                        )
562
563
                        if p2p_status != pynvml.NVML_P2P_STATUS_OK:
                            return False
564
565
                    except pynvml.NVMLError:
                        logger.exception(
566
                            "NVLink detection failed. This is normal if"
567
568
                            " your machine has no NVLink equipped."
                        )
569
570
                        return False
        return True
571
572

    @classmethod
573
574
575
576
577
578
579
580
581
    def _get_physical_device_name(cls, device_id: int = 0) -> str:
        handle = pynvml.nvmlDeviceGetHandleByIndex(device_id)
        return pynvml.nvmlDeviceGetName(handle)

    @classmethod
    @with_nvml_context
    def log_warnings(cls):
        device_ids: int = pynvml.nvmlDeviceGetCount()
        if device_ids > 1:
582
583
584
585
586
            device_names = [cls._get_physical_device_name(i) for i in range(device_ids)]
            if (
                len(set(device_names)) > 1
                and os.environ.get("CUDA_DEVICE_ORDER") != "PCI_BUS_ID"
            ):
587
                logger.warning(
588
                    "Detected different devices in the system: %s. Please"
589
590
                    " make sure to set `CUDA_DEVICE_ORDER=PCI_BUS_ID` to "
                    "avoid unexpected behavior.",
591
                    ", ".join(device_names),
592
593
594
595
596
                )


class NonNvmlCudaPlatform(CudaPlatformBase):
    @classmethod
597
    @cache
598
599
600
601
602
603
604
605
606
607
608
609
610
611
    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability:
        major, minor = torch.cuda.get_device_capability(device_id)
        return DeviceCapability(major=major, minor=minor)

    @classmethod
    def get_device_name(cls, device_id: int = 0) -> str:
        return torch.cuda.get_device_name(device_id)

    @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

    @classmethod
612
    def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:
613
614
        logger.exception(
            "NVLink detection not possible, as context support was"
615
616
            " not found. Assuming no NVLink available."
        )
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
        return False


# Autodetect either NVML-enabled or non-NVML platform
# based on whether NVML is available.
nvml_available = False
try:
    try:
        pynvml.nvmlInit()
        nvml_available = True
    except Exception:
        # On Jetson, NVML is not supported.
        nvml_available = False
finally:
    if nvml_available:
        pynvml.nvmlShutdown()

CudaPlatform = NvmlCudaPlatform if nvml_available else NonNvmlCudaPlatform

636
CudaPlatform.log_warnings()