| 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 | import math |
| 8 | from collections.abc import Iterable |
| 9 | from dataclasses import dataclass |
| 10 | from typing import Any, Literal |
| 11 | |
| 12 | import torch |
| 13 | |
| 14 | logger = logging.getLogger(__name__) |
| 15 | |
| 16 | _cutedsl_available: bool | None = None |
| 17 | |
| 18 | |
| 19 | def is_available() -> bool: |
| 20 | global _cutedsl_available |
| 21 | if _cutedsl_available is not None: |
| 22 | return _cutedsl_available |
| 23 | try: |
| 24 | import cutlass # noqa: F401 |
| 25 | import cutlass.cute # noqa: F401 |
| 26 | |
| 27 | _cutedsl_available = True |
| 28 | except ImportError: |
| 29 | _cutedsl_available = False |
| 30 | logger.info("cuteDSL (CUTLASS Python) not available, ll_bf16_gemm disabled") |
| 31 | return _cutedsl_available |
| 32 | |
| 33 | |
| 34 | # Default configs |
| 35 | _DEFAULT_DOTPROD_BS = 128 |
| 36 | _DEFAULT_DOTPROD_MAX_M = 4 |
| 37 | _DEFAULT_SPLITK_CONFIG = (6, 4) |
| 38 | |
| 39 | # ll_bf16 router shapes covered by warmup/tuning: |
| 40 | # (4096, 256) # DSV4-Flash |
| 41 | # (6144, 256) # GLM5.2 |
| 42 | # (7168, 256) # DSV3.2 |
| 43 | # (7168, 384) # DSV4-Pro |
| 44 | |
| 45 | # SM100f-specific tuned configs |
| 46 | _SM100F_TUNED_DOTPROD_BS: dict[tuple[int, int], dict[int, int]] = { |
| 47 | (6144, 256): {M: 256 for M in (1, 3, 4)}, |
| 48 | } |
| 49 | _SM100F_TUNED_SPLITK_CONFIGS: dict[tuple[int, int], dict[int, tuple[int, int]]] = { |
| 50 | (4096, 256): { |
| 51 | **{M: (8, 5) for M in (5, 8)}, |
| 52 | 9: (8, 2), |
| 53 | }, |
| 54 | (7168, 256): {14: (8, 2)}, |
| 55 | (6144, 256): {M: (8, 2) for M in (9, 12, 16)}, |
| 56 | (7168, 384): {M: (7, 5) for M in (13, 16)}, |
| 57 | } |
| 58 | |
| 59 | # SM90-specific tuned configs |
| 60 | _SM90_TUNED_DOTPROD_BS: dict[tuple[int, int], dict[int, int]] = { |
| 61 | (4096, 256): {M: 256 for M in (1, 3)}, |
| 62 | (7168, 384): {M: 256 for M in (1, 2)}, |
| 63 | } |
| 64 | _SM90_TUNED_SPLITK_CONFIGS: dict[tuple[int, int], dict[int, tuple[int, int]]] = { |
| 65 | (4096, 256): { |
| 66 | **{M: (8, 2) for M in range(5, 8)}, |
| 67 | **{M: (8, 5) for M in range(10, 12)}, |
| 68 | **{M: (8, 2) for M in (13, 16)}, |
| 69 | **{M: (8, 5) for M in (14, 15)}, |
| 70 | }, |
| 71 | (7168, 256): {8: (6, 5)}, |
| 72 | (6144, 256): {M: (8, 2) for M in (9, 11)}, |
| 73 | (7168, 384): { |
| 74 | **{M: (8, 2) for M in (6, 8, 12)}, |
| 75 | **{M: (7, 5) for M in (7, 9, 10, 11, 13, 14, 15, 16)}, |
| 76 | }, |
| 77 | } |
| 78 | |
| 79 | |
| 80 | def _arch_tuned_configs() -> tuple[ |
| 81 | dict[tuple[int, int], dict[int, int]], |
| 82 | dict[tuple[int, int], dict[int, tuple[int, int]]], |
| 83 | ]: |
| 84 | from vllm.platforms import current_platform |
| 85 | |
| 86 | if current_platform.is_device_capability_family(100): |
| 87 | return ( |
| 88 | _SM100F_TUNED_DOTPROD_BS, |
| 89 | _SM100F_TUNED_SPLITK_CONFIGS, |
| 90 | ) |
| 91 | if current_platform.is_device_capability(90): |
| 92 | return ( |
| 93 | _SM90_TUNED_DOTPROD_BS, |
| 94 | _SM90_TUNED_SPLITK_CONFIGS, |
| 95 | ) |
| 96 | return {}, {} |
| 97 | |
| 98 | |
| 99 | _cute_ctx = None |
| 100 | |
| 101 | |
| 102 | def _cute(): |
| 103 | global _cute_ctx |
| 104 | if _cute_ctx is not None: |
| 105 | return _cute_ctx |
| 106 | import cutlass.cute as cute |
| 107 | from cuda.bindings.driver import CUstream |
| 108 | |
| 109 | _cute_ctx = (cute, CUstream) |
| 110 | return _cute_ctx |
| 111 | |
| 112 | |
| 113 | def _stream(): |
| 114 | _, CUstream = _cute() |
| 115 | from vllm.utils.torch_utils import current_stream |
| 116 | |
| 117 | return CUstream(current_stream().cuda_stream) |
| 118 | |
| 119 | |
| 120 | def _use_pdl() -> bool: |
| 121 | from vllm.platforms import current_platform |
| 122 | |
| 123 | return current_platform.is_arch_support_pdl() |
| 124 | |
| 125 | |
| 126 | class LLBf16Gemm: |
| 127 | @dataclass(frozen=True, slots=True) |
| 128 | class CompileKey: |
| 129 | backend: Literal["dotprod", "splitk"] |
| 130 | M: int = 0 |
| 131 | K: int = 0 |
| 132 | bs: int = 0 |
| 133 | split_k: int = 0 |
| 134 | num_stages: int = 0 |
| 135 | |
| 136 | def __init__(self, *, prefetch_pdl_weights: bool = False) -> None: |
| 137 | self._prefetch_pdl_weights = prefetch_pdl_weights |
| 138 | # Dot-prod: keyed on (M, K, bs), because M and K are Constexpr. |
| 139 | self._compiled_cache: dict[tuple[int, int, int], Any] = {} |
| 140 | # Split-K: keyed on (split_k, num_stages), fully shape-dynamic. |
| 141 | self._splitk_cache: dict[tuple[int, int], Any] = {} |
| 142 | |
| 143 | def dispatch(self, *, M: int, K: int, N: int) -> CompileKey: |
| 144 | tuned_bs, tuned_splitk = _arch_tuned_configs() |
| 145 | if M <= _DEFAULT_DOTPROD_MAX_M or K < 2048: |
| 146 | bs = tuned_bs.get((K, N), {}).get(M, _DEFAULT_DOTPROD_BS) |
| 147 | return self.CompileKey(backend="dotprod", M=M, K=K, bs=bs) |
| 148 | |
| 149 | split_k, num_stages = tuned_splitk.get((K, N), {}).get( |
| 150 | M, _DEFAULT_SPLITK_CONFIG |
| 151 | ) |
| 152 | return self.CompileKey(backend="splitk", split_k=split_k, num_stages=num_stages) |
| 153 | |
| 154 | def get_warmup_keys( |
| 155 | self, |
| 156 | *, |
| 157 | shapes: Iterable[tuple[int, int]], |
| 158 | m_values: Iterable[int], |
| 159 | ) -> list[CompileKey]: |
| 160 | return list( |
| 161 | dict.fromkeys( |
| 162 | self.dispatch(M=M, K=K, N=N) for K, N in shapes for M in m_values |
| 163 | ) |
| 164 | ) |
| 165 | |
| 166 | @staticmethod |
| 167 | def _fake_gemm_tensors(*, M, K, N, divisibility: int): |
| 168 | from cutlass import BFloat16, Float32 |
| 169 | from quack.compile_utils import make_fake_tensor |
| 170 | |
| 171 | hidden_states = make_fake_tensor(BFloat16, (M, K), divisibility=divisibility) |
| 172 | router_weight = make_fake_tensor(BFloat16, (N, K), divisibility=divisibility) |
| 173 | output = make_fake_tensor(Float32, (M, N), divisibility=1) |
| 174 | return hidden_states, router_weight, output |
| 175 | |
| 176 | def _compile_splitk(self, compile_key: CompileKey) -> None: |
| 177 | cute, _ = _cute() |
| 178 | from ._ll_bf16_splitk import LLBf16SplitK |
| 179 | |
| 180 | hidden_states, router_weight, output = self._fake_gemm_tensors( |
| 181 | M=cute.sym_int(), |
| 182 | K=cute.sym_int(), |
| 183 | N=cute.sym_int(), |
| 184 | divisibility=8, |
| 185 | ) |
| 186 | gemm = LLBf16SplitK( |
| 187 | tile_n=16, |
| 188 | tile_k=256, |
| 189 | num_stages=compile_key.num_stages, |
| 190 | num_dma_warps=4, |
| 191 | split_k=compile_key.split_k, |
| 192 | use_pdl=_use_pdl(), |
| 193 | ) |
| 194 | compiled = cute.compile( |
| 195 | gemm, |
| 196 | hidden_states, |
| 197 | router_weight, |
| 198 | output, |
| 199 | _stream(), |
| 200 | options="--enable-tvm-ffi", |
| 201 | ) |
| 202 | self._splitk_cache[(compile_key.split_k, compile_key.num_stages)] = compiled |
| 203 | logger.debug( |
| 204 | "Compiled ll_bf16_splitk: sk=%d ns=%d", |
| 205 | compile_key.split_k, |
| 206 | compile_key.num_stages, |
| 207 | ) |
| 208 | |
| 209 | def _compile_dotprod(self, compile_key: CompileKey) -> None: |
| 210 | cute, _ = _cute() |
| 211 | from ._ll_bf16_dotprod import LLBf16Dotprod |
| 212 | |
| 213 | N = cute.sym_int() |
| 214 | stride_divisibility = math.gcd(8, compile_key.K) |
| 215 | hidden_states, router_weight, output = self._fake_gemm_tensors( |
| 216 | M=compile_key.M, |
| 217 | K=compile_key.K, |
| 218 | N=N, |
| 219 | divisibility=stride_divisibility, |
| 220 | ) |
| 221 | gemm = LLBf16Dotprod( |
| 222 | k=compile_key.K, |
| 223 | bs=compile_key.bs, |
| 224 | use_pdl=_use_pdl(), |
| 225 | prefetch_pdl_weights=self._prefetch_pdl_weights, |
| 226 | ) |
| 227 | compiled = cute.compile( |
| 228 | gemm, |
| 229 | hidden_states, |
| 230 | router_weight, |
| 231 | output, |
| 232 | compile_key.M, |
| 233 | compile_key.K, |
| 234 | 1, # runtime N placeholder for fake-tensor compile |
| 235 | _stream(), |
| 236 | options="--enable-tvm-ffi --ptxas-options -maxrregcount=64", |
| 237 | ) |
| 238 | self._compiled_cache[(compile_key.M, compile_key.K, compile_key.bs)] = compiled |
| 239 | logger.debug( |
| 240 | "Compiled ll_bf16_dotprod: M=%d, K=%d, bs=%d", |
| 241 | compile_key.M, |
| 242 | compile_key.K, |
| 243 | compile_key.bs, |
| 244 | ) |
| 245 | |
| 246 | def compile(self, compile_key: CompileKey) -> None: |
| 247 | if compile_key.backend == "splitk": |
| 248 | splitk_cache_key = (compile_key.split_k, compile_key.num_stages) |
| 249 | if splitk_cache_key not in self._splitk_cache: |
| 250 | self._compile_splitk(compile_key) |
| 251 | return |
| 252 | |
| 253 | dotprod_cache_key = (compile_key.M, compile_key.K, compile_key.bs) |
| 254 | if dotprod_cache_key not in self._compiled_cache: |
| 255 | self._compile_dotprod(compile_key) |
| 256 | |
| 257 | def warmup( |
| 258 | self, |
| 259 | *, |
| 260 | shapes: Iterable[tuple[int, int]], |
| 261 | m_values: Iterable[int], |
| 262 | ) -> None: |
| 263 | for compile_key in self.get_warmup_keys(shapes=shapes, m_values=m_values): |
| 264 | self.compile(compile_key) |
| 265 | |
| 266 | @staticmethod |
| 267 | def _validate_inputs( |
| 268 | hidden_states: torch.Tensor, |
| 269 | router_weight: torch.Tensor, |
| 270 | output_dtype: torch.dtype, |
| 271 | ) -> None: |
| 272 | if hidden_states.dim() != 2 or router_weight.dim() != 2: |
| 273 | raise ValueError("hidden_states and router_weight must be 2D tensors") |
| 274 | if ( |
| 275 | hidden_states.dtype != torch.bfloat16 |
| 276 | or router_weight.dtype != torch.bfloat16 |
| 277 | ): |
| 278 | raise ValueError("hidden_states and router_weight must have dtype=bfloat16") |
| 279 | if hidden_states.device.type != "cuda" or router_weight.device.type != "cuda": |
| 280 | raise ValueError( |
| 281 | "hidden_states and router_weight must have device_type=cuda" |
| 282 | ) |
| 283 | if hidden_states.device != router_weight.device: |
| 284 | raise ValueError( |
| 285 | "hidden_states and router_weight must be on the same CUDA device" |
| 286 | ) |
| 287 | if output_dtype != torch.float32: |
| 288 | raise ValueError("ll_bf16_gemm only supports output_dtype=torch.float32") |
| 289 | if hidden_states.shape[1] != router_weight.shape[1]: |
| 290 | raise ValueError( |
| 291 | "hidden_states and router_weight must have matching K dimensions" |
| 292 | ) |
| 293 | # Kernels use vectorized bf16 loads and require 16-byte row alignment. |
| 294 | if hidden_states.shape[1] % 8 != 0: |
| 295 | raise ValueError("ll_bf16_gemm requires K to be divisible by 8") |
| 296 | if not hidden_states.is_contiguous() or not router_weight.is_contiguous(): |
| 297 | raise ValueError( |
| 298 | "hidden_states and router_weight must be contiguous row-major inputs" |
| 299 | ) |
| 300 | |
| 301 | def __call__( |
| 302 | self, |
| 303 | hidden_states: torch.Tensor, # [M, K] bf16 |
| 304 | router_weight: torch.Tensor, # [N, K] bf16 |
| 305 | output_dtype: torch.dtype = torch.float32, |
| 306 | ) -> torch.Tensor: # [M, N] fp32 |
| 307 | self._validate_inputs(hidden_states, router_weight, output_dtype) |
| 308 | |
| 309 | M, K = hidden_states.shape |
| 310 | N = router_weight.shape[0] |
| 311 | compile_key = self.dispatch(M=M, K=K, N=N) |
| 312 | if compile_key.backend == "splitk": |
| 313 | splitk_cache_key = (compile_key.split_k, compile_key.num_stages) |
| 314 | if splitk_cache_key not in self._splitk_cache: |
| 315 | self.compile(compile_key) |
| 316 | kernel = self._splitk_cache[splitk_cache_key] |
| 317 | else: |
| 318 | dotprod_cache_key = (compile_key.M, compile_key.K, compile_key.bs) |
| 319 | if dotprod_cache_key not in self._compiled_cache: |
| 320 | self.compile(compile_key) |
| 321 | kernel = self._compiled_cache[dotprod_cache_key] |
| 322 | |
| 323 | stream = _stream() |
| 324 | output = torch.empty(M, N, dtype=output_dtype, device=hidden_states.device) |
| 325 | if compile_key.backend == "splitk": |
| 326 | kernel(hidden_states, router_weight, output, stream, 1.0) |
| 327 | else: |
| 328 | kernel(hidden_states, router_weight, output, N, stream) |
| 329 | return output |
| 330 | |
| 331 | |
| 332 | ll_bf16_gemm_kernel = LLBf16Gemm() |
| 333 | ll_bf16_gemm_c1_pdl_kernel = LLBf16Gemm(prefetch_pdl_weights=True) |
| 334 | |
| 335 | |
| 336 | def ll_bf16_gemm( |
| 337 | hidden_states: torch.Tensor, |
| 338 | router_weight: torch.Tensor, |
| 339 | output_dtype: torch.dtype = torch.float32, |
| 340 | ) -> torch.Tensor: |
| 341 | kernel = ( |
| 342 | ll_bf16_gemm_c1_pdl_kernel |
| 343 | if hidden_states.shape[0] == 1 |
| 344 | else ll_bf16_gemm_kernel |
| 345 | ) |
| 346 | return kernel(hidden_states, router_weight, output_dtype) |