| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # SPDX-FileCopyrightText: Songlin Yang, Yu Zhang |
| 4 | # |
| 5 | # This file contains code copied from the flash-linear-attention project. |
| 6 | # The original source code was licensed under the MIT license and included |
| 7 | # the following copyright notice: |
| 8 | # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang |
| 9 | # ruff: noqa: E501 |
| 10 | import contextlib |
| 11 | import functools |
| 12 | import logging |
| 13 | import os |
| 14 | from collections.abc import Callable |
| 15 | from enum import Enum |
| 16 | from typing import Any, Literal |
| 17 | |
| 18 | import torch |
| 19 | |
| 20 | from vllm.platforms import current_platform |
| 21 | from vllm.triton_utils import triton |
| 22 | |
| 23 | logger = logging.getLogger(__name__) |
| 24 | |
| 25 | COMPILER_MODE = os.getenv("FLA_COMPILER_MODE") == "1" |
| 26 | FLA_CI_ENV = os.getenv("FLA_CI_ENV") == "1" |
| 27 | |
| 28 | SUPPRESS_LEVEL = int(os.getenv("GDN_RECOMPUTE_SUPPRESS_LEVEL", "0")) |
| 29 | |
| 30 | # Default chunk size used across FLA triton kernels (kda, chunk, chunk_o, etc.) |
| 31 | FLA_CHUNK_SIZE = 64 |
| 32 | |
| 33 | |
| 34 | def tensor_cache(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]: |
| 35 | """ |
| 36 | A decorator that caches the most recent results of a function with tensor inputs. |
| 37 | |
| 38 | This decorator will store the output of the decorated function for the most recent set of input tensors. |
| 39 | The cache is limited to a fixed size (default is 4). When the cache is full, the oldest entry will be removed. |
| 40 | |
| 41 | Args: |
| 42 | fn (Callable[..., torch.Tensor]): |
| 43 | The function to be decorated. It should take tensor inputs and return tensor outputs. |
| 44 | |
| 45 | Returns: |
| 46 | Callable[..., torch.Tensor]: |
| 47 | A wrapped version of the input function with single-entry caching. |
| 48 | """ |
| 49 | |
| 50 | cache_entries: tuple[tuple | None, dict | None, Any] = [] |
| 51 | cache_size = 8 |
| 52 | |
| 53 | @functools.wraps(fn) |
| 54 | def wrapper(*args: Any, **kwargs: Any) -> Any: |
| 55 | nonlocal cache_entries, cache_size |
| 56 | for i, entry in enumerate(cache_entries): |
| 57 | last_args, last_kwargs, last_result = entry |
| 58 | if ( |
| 59 | len(args) == len(last_args) |
| 60 | and len(kwargs) == len(last_kwargs) |
| 61 | and all(a is b for a, b in zip(args, last_args)) |
| 62 | and all( |
| 63 | k in last_kwargs and v is last_kwargs[k] for k, v in kwargs.items() |
| 64 | ) |
| 65 | ): |
| 66 | cache_entries = ( |
| 67 | cache_entries[:i] |
| 68 | + cache_entries[i + 1 :] |
| 69 | + [(args, kwargs, last_result)] |
| 70 | ) |
| 71 | return last_result |
| 72 | |
| 73 | result = fn(*args, **kwargs) |
| 74 | |
| 75 | if len(cache_entries) >= cache_size: |
| 76 | cache_entries = cache_entries[1:] |
| 77 | cache_entries.append((args, kwargs, result)) |
| 78 | return result |
| 79 | |
| 80 | return wrapper |
| 81 | |
| 82 | |
| 83 | def input_guard(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]: |
| 84 | """ |
| 85 | A decorator to make sure all input tensors are contiguous and set the device based on input tensors. |
| 86 | """ |
| 87 | |
| 88 | @functools.wraps(fn) |
| 89 | def wrapper(*args, **kwargs): |
| 90 | contiguous_args = ( |
| 91 | i if not isinstance(i, torch.Tensor) else i.contiguous() for i in args |
| 92 | ) |
| 93 | contiguous_kwargs = { |
| 94 | k: (v if not isinstance(v, torch.Tensor) else v.contiguous()) |
| 95 | for k, v in kwargs.items() |
| 96 | } |
| 97 | |
| 98 | tensor = None |
| 99 | for arg in args: |
| 100 | if isinstance(arg, torch.Tensor): |
| 101 | tensor = arg |
| 102 | break |
| 103 | if tensor is None: |
| 104 | for value in kwargs.values(): |
| 105 | if isinstance(value, torch.Tensor): |
| 106 | tensor = value |
| 107 | break |
| 108 | |
| 109 | if tensor is not None: |
| 110 | ctx = torch.accelerator.device_index(tensor.device.index) |
| 111 | else: |
| 112 | ctx = contextlib.nullcontext() |
| 113 | |
| 114 | with ctx: |
| 115 | return fn(*contiguous_args, **contiguous_kwargs) |
| 116 | |
| 117 | return wrapper |
| 118 | |
| 119 | |
| 120 | @functools.cache |
| 121 | def get_available_device() -> str: |
| 122 | try: |
| 123 | return triton.runtime.driver.active.get_current_target().backend |
| 124 | except (RuntimeError, AttributeError): |
| 125 | return "cpu" |
| 126 | |
| 127 | |
| 128 | @functools.cache |
| 129 | def _check_platform() -> Literal["nvidia", "amd", "intel", "musa"]: |
| 130 | device = get_available_device() |
| 131 | mapping = { |
| 132 | "cuda": "nvidia", |
| 133 | "hip": "amd", |
| 134 | "xpu": "intel", |
| 135 | } |
| 136 | # return the mapped value, or the original if not found |
| 137 | return mapping.get(device, device) |
| 138 | |
| 139 | |
| 140 | # For AMD GPUs, the triton backend is 'hip', while for Nvidia GPUs, the triton backend is 'cuda'. |
| 141 | # However, the torch backend is 'cuda' for both Nvidia and AMD GPUs. |
| 142 | # Therefore, we need to check the triton backend to determine the actual GPU vendor. |
| 143 | device = "cuda" if current_platform.is_cuda_alike() else get_available_device() |
| 144 | device_torch_lib = getattr(torch, device, None) |
| 145 | device_platform = _check_platform() |
| 146 | |
| 147 | is_amd = device_platform == "amd" |
| 148 | is_intel = device_platform == "intel" |
| 149 | is_nvidia = device_platform == "nvidia" |
| 150 | is_intel_alchemist = is_intel and "Intel(R) Arc(TM) A" in torch.xpu.get_device_name(0) |
| 151 | is_nvidia_hopper = is_nvidia and ( |
| 152 | "NVIDIA H" in torch.cuda.get_device_name(0) |
| 153 | or torch.cuda.get_device_capability()[0] >= 9 |
| 154 | ) |
| 155 | use_cuda_graph = is_nvidia and os.environ.get("FLA_USE_CUDA_GRAPH", "0") == "1" |
| 156 | is_gather_supported = hasattr(triton.language, "gather") |
| 157 | is_tma_supported = ( |
| 158 | is_nvidia_hopper |
| 159 | and os.getenv("FLA_USE_TMA", "0") == "1" |
| 160 | and ( |
| 161 | hasattr(triton.language, "_experimental_make_tensor_descriptor") |
| 162 | or hasattr(triton.language, "make_tensor_descriptor") |
| 163 | ) |
| 164 | ) |
| 165 | |
| 166 | |
| 167 | def get_all_max_shared_mem(): |
| 168 | try: |
| 169 | return [ |
| 170 | triton.runtime.driver.active.utils.get_device_properties(i)[ |
| 171 | "max_shared_mem" |
| 172 | ] |
| 173 | for i in range(device_torch_lib.device_count()) |
| 174 | ] |
| 175 | except BaseException: |
| 176 | return [-1] |
| 177 | |
| 178 | |
| 179 | class Backend(Enum): |
| 180 | ADA = 101376 # RTX 4090 |
| 181 | AMPERE = 166912 # A100 |
| 182 | HOPPER = 232448 # H100 |
| 183 | DEFAULT = 102400 # Default |
| 184 | |
| 185 | @classmethod |
| 186 | def get_shared_memory(cls, arch: str) -> int: |
| 187 | try: |
| 188 | return cls[arch.upper()].value |
| 189 | except KeyError: |
| 190 | return cls.DEFAULT.value |
| 191 | |
| 192 | |
| 193 | @functools.cache |
| 194 | def check_shared_mem(arch: str = "none", tensor_idx: int = 0) -> bool: |
| 195 | try: |
| 196 | device_shared_mem_list = get_all_max_shared_mem() |
| 197 | max_shared_memory = device_shared_mem_list[tensor_idx] |
| 198 | return max_shared_memory >= Backend.get_shared_memory(arch) |
| 199 | except Exception: |
| 200 | return False |