| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Utility methods for model layers.""" |
| 4 | |
| 5 | import functools |
| 6 | from collections.abc import Callable |
| 7 | from dataclasses import dataclass |
| 8 | |
| 9 | import torch |
| 10 | |
| 11 | from vllm import _custom_ops as ops |
| 12 | from vllm import envs |
| 13 | from vllm._aiter_ops import rocm_aiter_ops |
| 14 | from vllm.logger import init_logger |
| 15 | from vllm.platforms import CpuArchEnum, current_platform |
| 16 | from vllm.utils.flashinfer import ( |
| 17 | flashinfer_bf16_mm, |
| 18 | is_flashinfer_cutedsl_bf16_gemm_supported, |
| 19 | ) |
| 20 | from vllm.utils.platform_utils import num_compute_units |
| 21 | from vllm.utils.torch_utils import direct_register_custom_op |
| 22 | |
| 23 | logger = init_logger(__name__) |
| 24 | |
| 25 | |
| 26 | def get_token_bin_counts_and_mask( |
| 27 | tokens: torch.Tensor, |
| 28 | vocab_size: int, |
| 29 | num_seqs: int, |
| 30 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 31 | # Compute the bin counts for the tokens. |
| 32 | # vocab_size + 1 for padding. |
| 33 | bin_counts = torch.zeros( |
| 34 | (num_seqs, vocab_size + 1), dtype=torch.long, device=tokens.device |
| 35 | ) |
| 36 | bin_counts.scatter_add_(1, tokens, torch.ones_like(tokens)) |
| 37 | bin_counts = bin_counts[:, :vocab_size] |
| 38 | mask = bin_counts > 0 |
| 39 | |
| 40 | return bin_counts, mask |
| 41 | |
| 42 | |
| 43 | def apply_penalties( |
| 44 | logits: torch.Tensor, |
| 45 | prompt_tokens_tensor: torch.Tensor, |
| 46 | output_tokens_tensor: torch.Tensor, |
| 47 | presence_penalties: torch.Tensor, |
| 48 | frequency_penalties: torch.Tensor, |
| 49 | repetition_penalties: torch.Tensor, |
| 50 | ) -> torch.Tensor: |
| 51 | """Applies penalties in place to the logits tensor |
| 52 | logits : The input logits tensor of shape [num_seqs, vocab_size] |
| 53 | prompt_tokens_tensor: A tensor containing the prompt tokens. The prompts |
| 54 | are padded to the maximum prompt length within the batch using |
| 55 | `vocab_size` as the padding value. The value `vocab_size` is used |
| 56 | for padding because it does not correspond to any valid token ID |
| 57 | in the vocabulary. |
| 58 | output_tokens_tensor: The output tokens tensor. |
| 59 | presence_penalties: The presence penalties of shape (num_seqs, ) |
| 60 | frequency_penalties: The frequency penalties of shape (num_seqs, ) |
| 61 | repetition_penalties: The repetition penalties of shape (num_seqs, ) |
| 62 | """ |
| 63 | num_seqs, vocab_size = logits.shape |
| 64 | _, prompt_mask = get_token_bin_counts_and_mask( |
| 65 | prompt_tokens_tensor, vocab_size, num_seqs |
| 66 | ) |
| 67 | output_bin_counts, output_mask = get_token_bin_counts_and_mask( |
| 68 | output_tokens_tensor, vocab_size, num_seqs |
| 69 | ) |
| 70 | |
| 71 | # Apply repetition penalties as a custom op |
| 72 | from vllm._custom_ops import apply_repetition_penalties |
| 73 | |
| 74 | apply_repetition_penalties(logits, prompt_mask, output_mask, repetition_penalties) |
| 75 | |
| 76 | # We follow the definition in OpenAI API. |
| 77 | # Refer to https://platform.openai.com/docs/api-reference/parameter-details |
| 78 | logits -= frequency_penalties.unsqueeze(dim=1) * output_bin_counts |
| 79 | logits -= presence_penalties.unsqueeze(dim=1) * output_mask |
| 80 | return logits |
| 81 | |
| 82 | |
| 83 | def default_unquantized_gemm( |
| 84 | layer: torch.nn.Module, |
| 85 | x: torch.Tensor, |
| 86 | weight: torch.Tensor, |
| 87 | bias: torch.Tensor | None = None, |
| 88 | ): |
| 89 | return torch.nn.functional.linear(x, weight, bias) |
| 90 | |
| 91 | |
| 92 | _FlashInferBf16RuntimeCheck = Callable[ |
| 93 | [torch.Tensor, torch.Tensor, torch.Tensor | None], bool |
| 94 | ] |
| 95 | |
| 96 | |
| 97 | @dataclass(frozen=True) |
| 98 | class _FlashInferBf16Backend: |
| 99 | flashinfer_backend: str |
| 100 | is_supported: Callable[[], bool] |
| 101 | can_implement: _FlashInferBf16RuntimeCheck |
| 102 | |
| 103 | |
| 104 | def _can_use_flashinfer_cutedsl_bf16( |
| 105 | x: torch.Tensor, |
| 106 | weight: torch.Tensor, |
| 107 | bias: torch.Tensor | None, |
| 108 | ) -> bool: |
| 109 | if not ( |
| 110 | current_platform.is_cuda() and current_platform.is_device_capability_family(100) |
| 111 | ): |
| 112 | return False |
| 113 | if x.ndim < 1 or weight.ndim != 2: |
| 114 | return False |
| 115 | if ( |
| 116 | not x.is_cuda |
| 117 | or not weight.is_cuda |
| 118 | or x.device != weight.device |
| 119 | or x.dtype != torch.bfloat16 |
| 120 | or weight.dtype != torch.bfloat16 |
| 121 | or not x.is_contiguous() |
| 122 | or not weight.is_contiguous() |
| 123 | ): |
| 124 | return False |
| 125 | |
| 126 | k = x.shape[-1] |
| 127 | n = weight.shape[0] |
| 128 | if ( |
| 129 | k <= 0 |
| 130 | or n <= 0 |
| 131 | or weight.shape[1] != k |
| 132 | or k % 128 != 0 |
| 133 | or x.data_ptr() % 32 != 0 |
| 134 | or weight.data_ptr() % 32 != 0 |
| 135 | ): |
| 136 | return False |
| 137 | |
| 138 | m = x.numel() // k |
| 139 | if not 1 <= m <= 32: |
| 140 | return False |
| 141 | return bias is None or ( |
| 142 | bias.is_cuda |
| 143 | and bias.device == x.device |
| 144 | and bias.dtype == torch.bfloat16 |
| 145 | and bias.ndim == 1 |
| 146 | and bias.shape[0] == n |
| 147 | and bias.is_contiguous() |
| 148 | ) |
| 149 | |
| 150 | |
| 151 | _FLASHINFER_BF16_BACKENDS = { |
| 152 | "flashinfer_cutedsl": _FlashInferBf16Backend( |
| 153 | flashinfer_backend="cute-dsl", |
| 154 | is_supported=is_flashinfer_cutedsl_bf16_gemm_supported, |
| 155 | can_implement=_can_use_flashinfer_cutedsl_bf16, |
| 156 | ), |
| 157 | } |
| 158 | |
| 159 | |
| 160 | def _get_flashinfer_bf16_backend(vllm_backend: str) -> _FlashInferBf16Backend: |
| 161 | backend_spec = _FLASHINFER_BF16_BACKENDS.get(vllm_backend) |
| 162 | if backend_spec is None: |
| 163 | supported = ", ".join(sorted(_FLASHINFER_BF16_BACKENDS)) |
| 164 | raise ValueError( |
| 165 | f"Unsupported vLLM FlashInfer BF16 backend {vllm_backend!r}; " |
| 166 | f"supported backends: {supported}" |
| 167 | ) |
| 168 | return backend_spec |
| 169 | |
| 170 | |
| 171 | def cuda_flashinfer_bf16_gemm_impl( |
| 172 | x: torch.Tensor, |
| 173 | weight: torch.Tensor, |
| 174 | bias: torch.Tensor | None, |
| 175 | pdl: bool, |
| 176 | vllm_backend: str, |
| 177 | ) -> torch.Tensor: |
| 178 | backend_spec = _get_flashinfer_bf16_backend(vllm_backend) |
| 179 | if not backend_spec.can_implement(x, weight, bias): |
| 180 | return torch.nn.functional.linear(x, weight, bias) |
| 181 | |
| 182 | k = x.shape[-1] |
| 183 | n = weight.shape[0] |
| 184 | x_2d = x.view(-1, k) |
| 185 | out_2d = flashinfer_bf16_mm( |
| 186 | x_2d, |
| 187 | weight.t(), |
| 188 | bias, |
| 189 | pdl, |
| 190 | backend_spec.flashinfer_backend, |
| 191 | ) |
| 192 | return out_2d.view(*x.shape[:-1], n) |
| 193 | |
| 194 | |
| 195 | def cuda_flashinfer_bf16_gemm_fake( |
| 196 | x: torch.Tensor, |
| 197 | weight: torch.Tensor, |
| 198 | bias: torch.Tensor | None, |
| 199 | pdl: bool, |
| 200 | vllm_backend: str, |
| 201 | ) -> torch.Tensor: |
| 202 | return x.new_empty((*x.shape[:-1], weight.shape[0])) |
| 203 | |
| 204 | |
| 205 | def cuda_flashinfer_bf16_gemm( |
| 206 | layer: torch.nn.Module, |
| 207 | x: torch.Tensor, |
| 208 | weight: torch.Tensor, |
| 209 | bias: torch.Tensor | None = None, |
| 210 | *, |
| 211 | vllm_backend: str, |
| 212 | pdl: bool, |
| 213 | ) -> torch.Tensor: |
| 214 | return torch.ops.vllm.cuda_flashinfer_bf16_gemm( |
| 215 | x, |
| 216 | weight, |
| 217 | bias, |
| 218 | pdl, |
| 219 | vllm_backend, |
| 220 | ) |
| 221 | |
| 222 | |
| 223 | direct_register_custom_op( |
| 224 | op_name="cuda_flashinfer_bf16_gemm", |
| 225 | op_func=cuda_flashinfer_bf16_gemm_impl, |
| 226 | fake_impl=cuda_flashinfer_bf16_gemm_fake, |
| 227 | ) |
| 228 | |
| 229 | |
| 230 | def use_aiter_triton_gemm(n, m, k, dtype): |
| 231 | if ( |
| 232 | not rocm_aiter_ops.is_triton_gemm_enabled() |
| 233 | # MI300's - fp8nuz=True |
| 234 | or current_platform.is_fp8_fnuz() |
| 235 | or dtype not in [torch.float16, torch.bfloat16] |
| 236 | ): |
| 237 | return False |
| 238 | |
| 239 | # use hipblaslt for the larger GEMMs |
| 240 | if n > 2048 and m > 512: |
| 241 | return False |
| 242 | return ( |
| 243 | (m == 5120 and k == 2880) |
| 244 | or (m == 2880 and k == 4096) |
| 245 | or (m == 128 and k == 2880) |
| 246 | or (m == 640 and k == 2880) |
| 247 | or (m == 2880 and k == 512) |
| 248 | ) |
| 249 | |
| 250 | |
| 251 | def wvsplitkrc_dispatch(n: int, k: int, m: int, cu_count: int) -> tuple[int, bool]: |
| 252 | """Pick the K-shard split for wvSplitKrc and say whether the shape fits. |
| 253 | |
| 254 | Mirrors wvSplitKrc() in csrc/rocm/skinny_gemms.cu, which is also where the |
| 255 | shard cap is explained. Both must pick the same chunkk or the workspace |
| 256 | check here bounds the wrong k_rnd. |
| 257 | |
| 258 | Returns: |
| 259 | The CHUNKK the kernel will dispatch with, and whether the CU budget and |
| 260 | split-K workspace admit the shape at all. |
| 261 | |
| 262 | """ |
| 263 | # Next ^2 of n |
| 264 | N_p2 = 1 << (n - 1).bit_length() |
| 265 | # How many of 4 waves in a group can work on same 16 Ms at same time? |
| 266 | # This reduces the Ms each group works on, i.e. increasing the CUs needed. |
| 267 | GrpsShrB = min(N_p2 // 16, 4) |
| 268 | # With 64 Ms per CU (each of 4 SIMDs working on a 16x16 tile), and each |
| 269 | # working on a 512-shard of K, how many CUs would we need? |
| 270 | CuNeeded = ((m + 64 - 1) // 64) * ((k + 512 - 1) // 512) * GrpsShrB |
| 271 | |
| 272 | CHUNKK2_MAX_SHARDS = 11 |
| 273 | shards_chunkk2 = (k + 256 - 1) // 256 # 256-wide shards |
| 274 | chunkk = ( |
| 275 | 2 |
| 276 | if ( |
| 277 | N_p2 != 16 |
| 278 | and CuNeeded * 2 <= cu_count |
| 279 | and shards_chunkk2 <= CHUNKK2_MAX_SHARDS |
| 280 | ) |
| 281 | else 1 |
| 282 | ) |
| 283 | |
| 284 | # Deterministic reduction stores one fp32 partial per (M, N, k-shard); all |
| 285 | # of them must fit the split-K workspace. |
| 286 | k_rnd = (k + 512 // chunkk - 1) // (512 // chunkk) |
| 287 | fits = N_p2 * m * k_rnd <= 128 * 1024 * 12 and CuNeeded <= cu_count |
| 288 | return chunkk, fits |
| 289 | |
| 290 | |
| 291 | def rocm_unquantized_gemm_impl( |
| 292 | x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None |
| 293 | ) -> torch.Tensor: |
| 294 | from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx950, on_gfx1250 |
| 295 | |
| 296 | n = x.numel() // x.size(-1) |
| 297 | m = weight.shape[0] |
| 298 | k = weight.shape[1] |
| 299 | |
| 300 | cu_count = num_compute_units() |
| 301 | |
| 302 | _, fits_wvsplitkrc = wvsplitkrc_dispatch(n, k, m, cu_count) |
| 303 | |
| 304 | skinny_operands_compatible = weight.is_contiguous() and ( |
| 305 | bias is None or bias.is_contiguous() |
| 306 | ) |
| 307 | |
| 308 | use_skinny_reduce_counting = ( |
| 309 | envs.VLLM_ROCM_USE_SKINNY_GEMM |
| 310 | and on_gfx950() |
| 311 | and x.dtype in [torch.float16, torch.bfloat16] |
| 312 | and x.dim() == 2 |
| 313 | and ( |
| 314 | 10 <= n <= 128 |
| 315 | and k % 8 == 0 |
| 316 | and k > 512 |
| 317 | and m % 16 == 0 |
| 318 | and fits_wvsplitkrc |
| 319 | and skinny_operands_compatible |
| 320 | ) |
| 321 | ) |
| 322 | |
| 323 | if use_skinny_reduce_counting: |
| 324 | x_view = x.reshape(-1, x.size(-1)).contiguous() |
| 325 | return ops.wvSplitKrc(x_view, weight, cu_count, bias) |
| 326 | |
| 327 | # gfx1250's aiter gemm_a16w16 uses the gluon backend, which requires |
| 328 | # K % 256 == 0 (it walks K with fixed-size descriptors and won't pad a |
| 329 | # partial last tile). Some whitelisted shapes have K=2880 (e.g. gpt-oss-120b |
| 330 | # hidden), so skip aiter there and fall back to the torch GEMM path below. |
| 331 | if use_aiter_triton_gemm(n, m, k, x.dtype) and not (on_gfx1250() and k % 256 != 0): |
| 332 | from aiter.ops.triton.gemm_a16w16 import gemm_a16w16 |
| 333 | |
| 334 | return gemm_a16w16(x, weight, bias) |
| 335 | |
| 336 | use_skinny = ( |
| 337 | envs.VLLM_ROCM_USE_SKINNY_GEMM |
| 338 | and (on_gfx9() or on_gfx1x()) |
| 339 | # build (gfx9/gfx11 ISA); fall back to torch GEMM there. |
| 340 | # TODO GFX1250: Include once skinny GEMM is supported on gfx1250 |
| 341 | and x.dtype in [torch.float16, torch.bfloat16] |
| 342 | and k % 8 == 0 |
| 343 | and skinny_operands_compatible |
| 344 | ) |
| 345 | |
| 346 | if use_skinny: |
| 347 | # The skinny kernels assume contiguous K elements. A shape-preserving |
| 348 | # reshape can retain a transposed activation's non-contiguous strides. |
| 349 | # Note: Only build that view inside the branches that consume it. |
| 350 | if (m == 1 or m > 8) and 0 < n <= 5: |
| 351 | x_view = x.reshape(-1, x.size(-1)).contiguous() |
| 352 | cu_count = num_compute_units() |
| 353 | out = ops.wvSplitK(weight, x_view, cu_count, bias) |
| 354 | return out.reshape(*x.shape[:-1], weight.shape[0]) |
| 355 | elif m % 4 == 0 and n == 1 and k <= 8192 and bias is None: |
| 356 | x_view = x.reshape(-1, x.size(-1)).contiguous() |
| 357 | out = ops.LLMM1(weight, x_view, 4) |
| 358 | return out.reshape(*x.shape[:-1], weight.shape[0]) |
| 359 | |
| 360 | if rocm_aiter_ops.is_tgemm_enabled(): |
| 361 | from aiter.tuned_gemm import tgemm |
| 362 | |
| 363 | return tgemm.mm(x, weight, bias) |
| 364 | |
| 365 | return torch.nn.functional.linear(x, weight, bias) |
| 366 | |
| 367 | |
| 368 | def rocm_unquantized_gemm_fake( |
| 369 | x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None |
| 370 | ) -> torch.Tensor: |
| 371 | return x.new_empty((*x.shape[:-1], weight.shape[0])) |
| 372 | |
| 373 | |
| 374 | def rocm_unquantized_gemm( |
| 375 | layer: torch.nn.Module, |
| 376 | x: torch.Tensor, |
| 377 | weight: torch.Tensor, |
| 378 | bias: torch.Tensor | None = None, |
| 379 | ) -> torch.Tensor: |
| 380 | # dispatch_unquantized_gemm() picks this backend for any ROCm build, |
| 381 | # regardless of the tensor's actual device; the custom op below is only |
| 382 | # registered for the CUDA/HIP dispatch key, so route CPU tensors (e.g. |
| 383 | # from unit tests that build layers without moving them to the GPU) |
| 384 | # through the device-agnostic fallback instead of hard-crashing. |
| 385 | if not x.is_cuda: |
| 386 | return default_unquantized_gemm(layer, x, weight, bias) |
| 387 | return torch.ops.vllm.rocm_unquantized_gemm(x, weight, bias) |
| 388 | |
| 389 | |
| 390 | direct_register_custom_op( |
| 391 | op_name="rocm_unquantized_gemm", |
| 392 | op_func=rocm_unquantized_gemm_impl, |
| 393 | fake_impl=rocm_unquantized_gemm_fake, |
| 394 | ) |
| 395 | |
| 396 | |
| 397 | @functools.cache |
| 398 | def warmup_rocm_skinny_gemm_workspaces(device: torch.device) -> None: |
| 399 | """Eagerly allocate wvSplitKrc's per-device split-K workspace pool. |
| 400 | |
| 401 | wvSplitKrc partitions one per-device allocation into ``kWvSlots`` slots |
| 402 | (csrc/rocm/skinny_gemms.cu) and hands each stream one on first use, so that |
| 403 | two streams never share the split-K partials and counters. |
| 404 | |
| 405 | The pool is otherwise created lazily on the first qualifying GEMM |
| 406 | (csrc/rocm/skinny_gemms.cu), which can be the first real request — after |
| 407 | the KV cache backing buffer exists. If it landed in that segment's rounding |
| 408 | tail, it would pin the entire segment at engine shutdown; it could also land |
| 409 | inside a cudagraph capture, where it would be taken from the graph's private |
| 410 | pool and its zero-fill would become a replayed graph node. |
| 411 | """ |
| 412 | from vllm.platforms.rocm import on_gfx950 |
| 413 | |
| 414 | if not on_gfx950(): |
| 415 | return |
| 416 | try: |
| 417 | x = torch.zeros(16, 1024, dtype=torch.bfloat16, device=device) |
| 418 | weight = torch.zeros(32, 1024, dtype=torch.bfloat16, device=device) |
| 419 | ops.wvSplitKrc(x, weight, num_compute_units()) |
| 420 | except Exception: |
| 421 | logger.debug("wvSplitKrc workspace warmup failed", exc_info=True) |
| 422 | |
| 423 | |
| 424 | # Above this weight size, oneDNN's onednn_mm consistently matches or beats |
| 425 | # the SGL AMX kernel once M grows past decode-sized batches, and is within |
| 426 | # noise of it at decode-sized M -- so larger weights default to oneDNN |
| 427 | # rather than SGL. 1 MiB comfortably covers MoE router/gate weights (e.g. |
| 428 | # (2048, 128) .. (2880, 32) bf16/fp16, 180-720 KiB) while staying well below |
| 429 | # any dense qkv/o_proj/gate_up/down/lm_head projection in practice. This |
| 430 | # threshold is derived from bf16/fp16 unquantized dense-GEMM benchmarks only, |
| 431 | # so it does not apply to the int8 scaled_mm path below. |
| 432 | _CPU_SGL_GEMM_MAX_WEIGHT_BYTES = 1 * 1024 * 1024 |
| 433 | |
| 434 | |
| 435 | def check_cpu_sgl_kernel(n: int, k: int, dtype: torch.dtype) -> bool: |
| 436 | if not torch.cpu._is_amx_tile_supported() or dtype not in ( |
| 437 | torch.bfloat16, |
| 438 | torch.float16, |
| 439 | torch.int8, |
| 440 | ): |
| 441 | return False |
| 442 | if dtype == torch.float16 and not torch.cpu._is_amx_fp16_supported(): |
| 443 | # AMX-BF16/INT8 (amx_tile) and AMX-FP16 are separate CPU ISA |
| 444 | # extensions -- e.g. Sapphire/Emerald Rapids expose the former but |
| 445 | # not the latter -- and can_use_brgemm<at::Half> (gemm.h) always |
| 446 | # attempts brgemm for fp16 regardless of M, so this needs its own |
| 447 | # capability check rather than piggybacking on amx_tile. |
| 448 | return False |
| 449 | if dtype == torch.int8: |
| 450 | # int8_scaled_mm_with_quant requires the packed weight to stay int8 |
| 451 | # (gemm_int8.cpp); convert_weight_packed's N < TILE_N fallback |
| 452 | # returns a float32 tensor instead (gemm.cpp), which would trip |
| 453 | # that check, so N must be a full TILE_N tile here. |
| 454 | return k % 32 == 0 and n % 16 == 0 |
| 455 | if n * k * dtype.itemsize > _CPU_SGL_GEMM_MAX_WEIGHT_BYTES: |
| 456 | return False |
| 457 | if n < 16: |
| 458 | # convert_weight_packed transposes to fp32 instead of VNNI-packing |
| 459 | # when N < TILE_N (gemm.cpp), and weight_packed_linear detects that |
| 460 | # (via the packed weight's dtype) and routes to its fp32/brgemm |
| 461 | # fallback kernel -- no N/K alignment required in that regime. |
| 462 | return True |
| 463 | return k % 32 == 0 and n % 16 == 0 |
| 464 | |
| 465 | |
| 466 | def dispatch_cpu_unquantized_gemm( |
| 467 | layer: torch.nn.Module, |
| 468 | remove_weight: bool, |
| 469 | ) -> None: |
| 470 | # skip for missing layers |
| 471 | if layer.weight.is_meta: |
| 472 | layer.cpu_linear = torch.nn.functional.linear |
| 473 | return |
| 474 | |
| 475 | # Skip CPU GEMM dispatch for non-2D weights (e.g. MoE 3D expert weights). |
| 476 | # These layers are handled by their own specialized methods. |
| 477 | if layer.weight.ndim != 2: |
| 478 | # This is not a linear layer. |
| 479 | # For now it should be a causal_conv1d op or MoE 3D expert weights. |
| 480 | # Causal-conv weights use [dim, 1, width]. |
| 481 | # The C++ causal_conv1d kernels support aarch64 with bf16 and |
| 482 | # x86 CPUs that use VDPBF16PS (no AMX tiles). So the |
| 483 | # weight prepack applies to aarch64 with bf16 and any AVX-512BF16 CPU, |
| 484 | # not just AMX (e.g. AMD Zen5/Turin). |
| 485 | is_causal_conv1d_weight = ( |
| 486 | layer.weight.ndim == 3 |
| 487 | and layer.weight.size(1) == 1 |
| 488 | and layer.weight.size(2) == 4 |
| 489 | and layer.weight.dtype == torch.bfloat16 |
| 490 | ) |
| 491 | is_arm_bf16 = ( |
| 492 | current_platform.get_cpu_architecture() == CpuArchEnum.ARM |
| 493 | and torch.cpu.get_capabilities().get("bf16", False) |
| 494 | ) |
| 495 | if ( |
| 496 | (is_arm_bf16 or torch.cpu._is_avx512_bf16_supported()) |
| 497 | and is_causal_conv1d_weight |
| 498 | and hasattr(ops, "causal_conv1d_weight_pack") |
| 499 | ): |
| 500 | # prepack conv weight |
| 501 | unpacked = ( |
| 502 | layer.weight.view( |
| 503 | layer.weight.size(0), |
| 504 | layer.weight.size(2), |
| 505 | ) |
| 506 | .contiguous() |
| 507 | .clone() |
| 508 | ) |
| 509 | # Stash the un-packed (dim, width) weight so the speculative-decode |
| 510 | # and unsupported-layout fallbacks can still use ordinary weights. |
| 511 | layer._cpu_unpacked_conv_weight = unpacked |
| 512 | layer.weight.data = ops.causal_conv1d_weight_pack(unpacked) |
| 513 | return |
| 514 | |
| 515 | N, K = layer.weight.size() |
| 516 | dtype = layer.weight.dtype |
| 517 | |
| 518 | # Zen CPU path: zentorch_linear_unary with optional eager weight prepacking. |
| 519 | if current_platform.is_zen_cpu() and hasattr( |
| 520 | torch.ops.zentorch, "zentorch_linear_unary" |
| 521 | ): |
| 522 | zen_weight = layer.weight.detach() |
| 523 | is_prepacked = False |
| 524 | |
| 525 | if envs.VLLM_ZENTORCH_WEIGHT_PREPACK and hasattr( |
| 526 | torch.ops.zentorch, "zentorch_weight_prepack_for_linear" |
| 527 | ): |
| 528 | zen_weight = torch.ops.zentorch.zentorch_weight_prepack_for_linear( |
| 529 | zen_weight |
| 530 | ) |
| 531 | is_prepacked = True |
| 532 | |
| 533 | layer.cpu_linear = lambda x, weight, bias, _p=is_prepacked: ( |
| 534 | torch.ops.zentorch.zentorch_linear_unary( |
| 535 | x, zen_weight, bias, is_weight_prepacked=_p |
| 536 | ) |
| 537 | ) |
| 538 | if remove_weight: |
| 539 | layer.weight = torch.nn.Parameter( |
| 540 | torch.empty(0, dtype=dtype, device=layer.weight.device), |
| 541 | requires_grad=False, |
| 542 | ) |
| 543 | logger.debug_once( |
| 544 | "CPU unquantized GEMM dispatch: using zentorch_linear_unary (prepacked=%s)", |
| 545 | is_prepacked, |
| 546 | ) |
| 547 | return |
| 548 | |
| 549 | # Small weights (e.g. MoE router/gate projections, where N is the expert |
| 550 | # count rather than a hidden-size-scaled dimension) never reach oneDNN's |
| 551 | # compute-bound regime, no matter how large the batch gets: SGL's lower |
| 552 | # per-call dispatch overhead wins consistently across the full measured |
| 553 | # M range. Larger dense projections (qkv/o_proj/gate_up/down/lm_head) |
| 554 | # cross over to favoring oneDNN once batch size grows past decode-sized |
| 555 | # M, so they keep using oneDNN below. |
| 556 | if check_cpu_sgl_kernel(N, K, dtype): |
| 557 | # For small size GEMM, packed_weight might be float32 |
| 558 | packed_weight = torch.ops._C.convert_weight_packed(layer.weight) |
| 559 | if getattr(layer, "bias", None) is not None: |
| 560 | layer.bias.data = layer.bias.to(torch.float32) |
| 561 | layer.cpu_linear = lambda x, weight, bias: ops.weight_packed_linear_cpu( |
| 562 | x, |
| 563 | packed_weight, |
| 564 | N, |
| 565 | bias, |
| 566 | ) |
| 567 | if remove_weight: |
| 568 | layer.weight = torch.nn.Parameter( |
| 569 | torch.empty(0, dtype=dtype, device=layer.weight.device), |
| 570 | requires_grad=False, |
| 571 | ) |
| 572 | logger.debug_once( |
| 573 | "CPU unquantized GEMM dispatch: using sgl-kernel weight_packed_linear" |
| 574 | ) |
| 575 | return |
| 576 | |
| 577 | if ( |
| 578 | ops._supports_onednn |
| 579 | and current_platform.get_cpu_architecture() != CpuArchEnum.POWERPC |
| 580 | ): |
| 581 | try: |
| 582 | origin_weight = layer.weight |
| 583 | handler = ops.create_onednn_mm(origin_weight.t(), 32) |
| 584 | layer.cpu_linear = lambda x, weight, bias: ops.onednn_mm(handler, x, bias) |
| 585 | if remove_weight: |
| 586 | layer.weight = torch.nn.Parameter( |
| 587 | torch.empty(0, dtype=dtype, device=layer.weight.device), |
| 588 | requires_grad=False, |
| 589 | ) |
| 590 | logger.debug_once("CPU unquantized GEMM dispatch: using oneDNN onednn_mm") |
| 591 | return |
| 592 | except RuntimeError as e: |
| 593 | logger.warning_once( |
| 594 | "Failed to create oneDNN linear, fallback to torch linear." |
| 595 | f" Exception: {e}" |
| 596 | ) |
| 597 | |
| 598 | # fallback case |
| 599 | layer.cpu_linear = lambda x, weight, bias: torch.nn.functional.linear( |
| 600 | x, weight, bias |
| 601 | ) |
| 602 | logger.debug_once( |
| 603 | "CPU unquantized GEMM dispatch: using torch.nn.functional.linear (fallback)" |
| 604 | ) |
| 605 | |
| 606 | |
| 607 | def cpu_unquantized_gemm( |
| 608 | layer: torch.nn.Module, |
| 609 | x: torch.Tensor, |
| 610 | weight: torch.Tensor, |
| 611 | bias: torch.Tensor | None = None, |
| 612 | ): |
| 613 | return layer.cpu_linear(x, weight, bias) |
| 614 | |
| 615 | |
| 616 | def dispatch_unquantized_gemm( |
| 617 | linear_backend: str = "auto", |
| 618 | ) -> Callable[..., torch.Tensor]: |
| 619 | if current_platform.is_rocm(): |
| 620 | return rocm_unquantized_gemm |
| 621 | elif current_platform.is_cpu(): |
| 622 | return cpu_unquantized_gemm |
| 623 | elif not current_platform.is_cuda(): |
| 624 | return default_unquantized_gemm |
| 625 | |
| 626 | backend_spec = _FLASHINFER_BF16_BACKENDS.get(linear_backend) |
| 627 | if backend_spec is None: |
| 628 | return default_unquantized_gemm |
| 629 | |
| 630 | if not backend_spec.is_supported(): |
| 631 | logger.warning_once( |
| 632 | "--linear-backend=%s requested FlashInfer mm_bf16 backend %r, " |
| 633 | "but it is unavailable on the current hardware or environment; " |
| 634 | "using automatic selection for unquantized linear layers.", |
| 635 | linear_backend, |
| 636 | backend_spec.flashinfer_backend, |
| 637 | ) |
| 638 | return default_unquantized_gemm |
| 639 | |
| 640 | logger.info_once( |
| 641 | "Using FlashInfer %s for eligible unquantized BF16 GEMMs.", |
| 642 | backend_spec.flashinfer_backend, |
| 643 | ) |
| 644 | return functools.partial( |
| 645 | cuda_flashinfer_bf16_gemm, |
| 646 | vllm_backend=linear_backend, |
| 647 | pdl=current_platform.is_arch_support_pdl(), |
| 648 | ) |