| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from __future__ import annotations |
| 5 | |
| 6 | import logging |
| 7 | from collections.abc import Iterable |
| 8 | from dataclasses import dataclass |
| 9 | from functools import partial |
| 10 | from typing import Any |
| 11 | |
| 12 | import torch |
| 13 | |
| 14 | logger = logging.getLogger(__name__) |
| 15 | |
| 16 | _cutedsl_available: bool | None = None |
| 17 | |
| 18 | |
| 19 | @dataclass(frozen=True, slots=True) |
| 20 | class SkinnyGemmConfig: |
| 21 | num_rows: int |
| 22 | block_size: int |
| 23 | outputs_per_block: int |
| 24 | k_unroll: int = 1 |
| 25 | vector_width: int = 8 |
| 26 | static_k: int | None = None |
| 27 | |
| 28 | |
| 29 | class ShapeDynamicSkinnyGemm: |
| 30 | def __init__(self) -> None: |
| 31 | self._compiled: dict[tuple[torch.dtype, SkinnyGemmConfig, bool], Any] = {} |
| 32 | self._warmup_configs: set[tuple[torch.dtype, SkinnyGemmConfig, bool]] = set() |
| 33 | self._warmup_registered = False |
| 34 | |
| 35 | @staticmethod |
| 36 | def is_available() -> bool: |
| 37 | global _cutedsl_available |
| 38 | if _cutedsl_available is not None: |
| 39 | return _cutedsl_available |
| 40 | try: |
| 41 | import cutlass # noqa: F401 |
| 42 | import cutlass.cute # noqa: F401 |
| 43 | |
| 44 | _cutedsl_available = True |
| 45 | except ImportError: |
| 46 | _cutedsl_available = False |
| 47 | logger.info("cuteDSL is not available; skinny GEMM is disabled") |
| 48 | return _cutedsl_available |
| 49 | |
| 50 | @staticmethod |
| 51 | def _config(m: int, n: int, k: int) -> SkinnyGemmConfig: |
| 52 | num_rows = m |
| 53 | wide_block = 224 |
| 54 | if m == 1 and k >= 7168 and k % (wide_block * 8) == 0: |
| 55 | if n % 3 == 0: |
| 56 | k_unroll = 2 if n <= 2304 else 4 |
| 57 | return SkinnyGemmConfig(num_rows, wide_block, 3, k_unroll) |
| 58 | if 2304 < n < 4096 and n % 2 == 0: |
| 59 | return SkinnyGemmConfig(num_rows, wide_block, 2, k_unroll=4) |
| 60 | |
| 61 | if k <= 2048 or k % (128 * 8) != 0: |
| 62 | outputs_per_block = 2 if k <= 2048 else 4 |
| 63 | if n % outputs_per_block: |
| 64 | outputs_per_block = 1 |
| 65 | if k % (64 * 8) == 0: |
| 66 | return SkinnyGemmConfig(num_rows, 64, outputs_per_block, 2) |
| 67 | if k % (32 * 8) == 0: |
| 68 | return SkinnyGemmConfig(num_rows, 32, outputs_per_block, 2) |
| 69 | return SkinnyGemmConfig(num_rows, 32, outputs_per_block, 2, vector_width=4) |
| 70 | |
| 71 | block_size = 64 if 4096 <= n < 8192 else 128 |
| 72 | outputs_per_block = 1 if m == 1 and n <= 2304 else 2 |
| 73 | if n % outputs_per_block: |
| 74 | outputs_per_block = 1 |
| 75 | k_unroll = 2 if n <= 2304 or n >= 16384 else 1 |
| 76 | return SkinnyGemmConfig( |
| 77 | num_rows, |
| 78 | block_size, |
| 79 | outputs_per_block, |
| 80 | k_unroll=k_unroll, |
| 81 | ) |
| 82 | |
| 83 | @staticmethod |
| 84 | def _cutlass_dtype(dtype: torch.dtype): |
| 85 | from cutlass import BFloat16, Float16 |
| 86 | |
| 87 | return BFloat16 if dtype == torch.bfloat16 else Float16 |
| 88 | |
| 89 | @staticmethod |
| 90 | def _stream(): |
| 91 | from cuda.bindings.driver import CUstream |
| 92 | |
| 93 | from vllm.utils.torch_utils import current_stream |
| 94 | |
| 95 | return CUstream(current_stream().cuda_stream) |
| 96 | |
| 97 | @staticmethod |
| 98 | def _use_pdl() -> bool: |
| 99 | from vllm.platforms import current_platform |
| 100 | |
| 101 | return current_platform.is_arch_support_pdl() |
| 102 | |
| 103 | def _compile( |
| 104 | self, |
| 105 | dtype: torch.dtype, |
| 106 | config: SkinnyGemmConfig, |
| 107 | has_residual: bool, |
| 108 | ) -> None: |
| 109 | import cutlass.cute as cute |
| 110 | from quack.compile_utils import make_fake_tensor |
| 111 | |
| 112 | from ._skinny_gemm import CuteSkinnyGemm |
| 113 | |
| 114 | element_type = self._cutlass_dtype(dtype) |
| 115 | n = cute.sym_int(divisibility=config.outputs_per_block) |
| 116 | k = ( |
| 117 | config.static_k |
| 118 | if config.static_k is not None |
| 119 | else cute.sym_int(divisibility=config.block_size * config.vector_width) |
| 120 | ) |
| 121 | a = make_fake_tensor( |
| 122 | element_type, |
| 123 | (config.num_rows, k), |
| 124 | divisibility=config.vector_width, |
| 125 | ) |
| 126 | b = make_fake_tensor( |
| 127 | element_type, |
| 128 | (n, k), |
| 129 | divisibility=config.vector_width, |
| 130 | ) |
| 131 | c = make_fake_tensor(element_type, (config.num_rows, n), divisibility=1) |
| 132 | residual = make_fake_tensor(element_type, (config.num_rows, n), divisibility=1) |
| 133 | kernel = CuteSkinnyGemm( |
| 134 | element_type=element_type, |
| 135 | num_rows=config.num_rows, |
| 136 | block_size=config.block_size, |
| 137 | outputs_per_block=config.outputs_per_block, |
| 138 | vector_width=config.vector_width, |
| 139 | k_unroll=config.k_unroll, |
| 140 | has_residual=has_residual, |
| 141 | use_pdl=self._use_pdl(), |
| 142 | static_k=config.static_k, |
| 143 | ) |
| 144 | self._compiled[(dtype, config, has_residual)] = cute.compile( |
| 145 | kernel, |
| 146 | a, |
| 147 | b, |
| 148 | residual, |
| 149 | c, |
| 150 | self._stream(), |
| 151 | options="--enable-tvm-ffi --ptxas-options -maxrregcount=64", |
| 152 | ) |
| 153 | |
| 154 | def request_warmup_configs( |
| 155 | self, |
| 156 | dtype: torch.dtype, |
| 157 | configs: Iterable[SkinnyGemmConfig], |
| 158 | *, |
| 159 | has_residual: bool = False, |
| 160 | ) -> None: |
| 161 | """Request compilation of explicit measured configs before capture.""" |
| 162 | self._warmup_configs.update((dtype, config, has_residual) for config in configs) |
| 163 | if self._warmup_registered: |
| 164 | return |
| 165 | from vllm.model_executor.warmup.cutedsl_warmup import ( |
| 166 | register_cutedsl_warmup_provider, |
| 167 | ) |
| 168 | |
| 169 | register_cutedsl_warmup_provider(self) |
| 170 | self._warmup_registered = True |
| 171 | |
| 172 | def get_cutedsl_warmup_compile_units(self): |
| 173 | from vllm.model_executor.warmup.cutedsl_warmup import CuTeDSLCompileUnit |
| 174 | |
| 175 | return tuple( |
| 176 | CuTeDSLCompileUnit( |
| 177 | name=( |
| 178 | "shape-dynamic skinny GEMM with residual" |
| 179 | if has_residual |
| 180 | else "shape-dynamic skinny GEMM" |
| 181 | ), |
| 182 | key=("shape-dynamic-skinny-gemm", dtype, config, has_residual), |
| 183 | compile=partial(self._compile, dtype, config, has_residual), |
| 184 | ) |
| 185 | for dtype, config, has_residual in sorted( |
| 186 | self._warmup_configs, |
| 187 | key=lambda item: ( |
| 188 | str(item[0]), |
| 189 | item[1].num_rows, |
| 190 | item[1].block_size, |
| 191 | item[1].outputs_per_block, |
| 192 | item[1].k_unroll, |
| 193 | item[1].vector_width, |
| 194 | item[2], |
| 195 | ), |
| 196 | ) |
| 197 | ) |
| 198 | |
| 199 | def __call__( |
| 200 | self, |
| 201 | a: torch.Tensor, |
| 202 | b: torch.Tensor, |
| 203 | config: SkinnyGemmConfig | None = None, |
| 204 | residual: torch.Tensor | None = None, |
| 205 | ) -> torch.Tensor: |
| 206 | if a.dim() != 2 or b.dim() != 2: |
| 207 | raise ValueError("a and b must be 2D tensors") |
| 208 | if a.dtype not in (torch.bfloat16, torch.float16) or b.dtype != a.dtype: |
| 209 | raise ValueError("a and b must have the same BF16 or FP16 dtype") |
| 210 | if not a.is_cuda or not b.is_cuda or a.device != b.device: |
| 211 | raise ValueError("a and b must be CUDA tensors on the same device") |
| 212 | if not a.is_contiguous() or not b.is_contiguous(): |
| 213 | raise ValueError("a and b must be contiguous") |
| 214 | if a.shape[1] != b.shape[1]: |
| 215 | raise ValueError("a and b must have matching K dimensions") |
| 216 | if not 1 <= a.shape[0] <= 16: |
| 217 | raise ValueError("shape-dynamic skinny GEMM requires 1 <= M <= 16") |
| 218 | if residual is not None: |
| 219 | if residual.dim() != 2 or residual.shape != (a.shape[0], b.shape[0]): |
| 220 | raise ValueError("residual must have shape (M, N)") |
| 221 | if residual.dtype != a.dtype: |
| 222 | raise ValueError("residual must have the same dtype as a and b") |
| 223 | if residual.device != a.device or not residual.is_cuda: |
| 224 | raise ValueError("residual must be on the same CUDA device") |
| 225 | if not residual.is_contiguous(): |
| 226 | raise ValueError("residual must be contiguous") |
| 227 | |
| 228 | config = config or self._config(a.shape[0], b.shape[0], a.shape[1]) |
| 229 | if config.num_rows != a.shape[0]: |
| 230 | raise ValueError("config num_rows must match M") |
| 231 | if b.shape[0] % config.outputs_per_block != 0: |
| 232 | raise ValueError("N must be divisible by outputs_per_block") |
| 233 | if a.shape[1] % (config.block_size * config.vector_width) != 0: |
| 234 | raise ValueError( |
| 235 | "K must be divisible by block_size * vector_width for this config" |
| 236 | ) |
| 237 | if config.static_k is not None and a.shape[1] != config.static_k: |
| 238 | raise ValueError("input K must match config static_k") |
| 239 | has_residual = residual is not None |
| 240 | cache_key = (a.dtype, config, has_residual) |
| 241 | if cache_key not in self._compiled: |
| 242 | self._compile(a.dtype, config, has_residual) |
| 243 | output = torch.empty((a.shape[0], b.shape[0]), dtype=a.dtype, device=a.device) |
| 244 | residual_arg = output if residual is None else residual |
| 245 | self._compiled[cache_key](a, b, residual_arg, output, self._stream()) |
| 246 | return output |
| 247 | |
| 248 | |
| 249 | shape_dynamic_skinny_gemm = ShapeDynamicSkinnyGemm() |