| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from __future__ import annotations |
| 5 | |
| 6 | from typing import Any |
| 7 | |
| 8 | import torch |
| 9 | |
| 10 | from vllm.kernels.helion.case_key import CaseKey |
| 11 | from vllm.logger import init_logger |
| 12 | from vllm.utils.import_utils import has_helion |
| 13 | |
| 14 | if not has_helion(): |
| 15 | raise ImportError( |
| 16 | "silu_mul_fp8 Helion kernel requires helion to be installed. " |
| 17 | "Install it with: pip install helion" |
| 18 | ) |
| 19 | |
| 20 | import helion.language as hl |
| 21 | |
| 22 | from vllm.kernels.helion.register import register_kernel |
| 23 | |
| 24 | logger = init_logger(__name__) |
| 25 | |
| 26 | |
| 27 | def generate_silu_mul_fp8_inputs() -> dict[CaseKey, tuple[Any, ...]]: |
| 28 | intermediate_sizes = [2048, 2880, 4096, 8192, 11008, 14336] |
| 29 | |
| 30 | # Use the same num_tokens values as vLLM's default cudagraph capture sizes. |
| 31 | # See vllm/config/vllm.py _set_cudagraph_sizes() for the canonical formula. |
| 32 | num_tokens_list = [1, 2, 4] + list(range(8, 256, 8)) + list(range(256, 513, 16)) |
| 33 | |
| 34 | inputs: dict[CaseKey, tuple[Any, ...]] = {} |
| 35 | for num_tokens in num_tokens_list: |
| 36 | for intermediate_size in intermediate_sizes: |
| 37 | input_tensor = torch.randn( |
| 38 | num_tokens, |
| 39 | 2 * intermediate_size, |
| 40 | device="cuda", |
| 41 | dtype=torch.bfloat16, |
| 42 | ) |
| 43 | scale = torch.tensor([1.0], device="cuda", dtype=torch.float32) |
| 44 | |
| 45 | key = CaseKey({"intermediate": intermediate_size, "numtokens": num_tokens}) |
| 46 | inputs[key] = (input_tensor, scale) |
| 47 | |
| 48 | return inputs |
| 49 | |
| 50 | |
| 51 | _pick_cache: dict[tuple[int, int], CaseKey | None] = {} |
| 52 | |
| 53 | |
| 54 | def pick_silu_mul_fp8_config( |
| 55 | args: tuple[Any, ...], config_keys: list[CaseKey] |
| 56 | ) -> CaseKey | None: |
| 57 | """Pick the best pre-tuned config for the given input shape. |
| 58 | |
| 59 | Selection strategy: |
| 60 | 1. Find the closest intermediate_size among available configs |
| 61 | (exact match preferred). |
| 62 | 2. Among the num_tokens values tuned for that intermediate_size, pick |
| 63 | the smallest num_tokens >= the input's num_tokens. If the input is |
| 64 | larger than all available num_tokens, fall back to the largest. |
| 65 | """ |
| 66 | if not config_keys: |
| 67 | return None |
| 68 | |
| 69 | input_tensor, _scale = args |
| 70 | intermediate_size = int(input_tensor.shape[-1]) // 2 |
| 71 | num_tokens = int(input_tensor.view(-1, input_tensor.shape[-1]).shape[0]) |
| 72 | |
| 73 | cache_key = (num_tokens, intermediate_size) |
| 74 | cached = _pick_cache.get(cache_key) |
| 75 | if cached is not None: |
| 76 | return cached |
| 77 | |
| 78 | by_isize: dict[int, list[int]] = {} |
| 79 | for k in config_keys: |
| 80 | if k.is_default(): |
| 81 | continue |
| 82 | by_isize.setdefault(k["intermediate"], []).append(k["numtokens"]) |
| 83 | |
| 84 | if not by_isize: |
| 85 | return None |
| 86 | |
| 87 | best_isize = min(by_isize, key=lambda s: abs(s - intermediate_size)) |
| 88 | available = sorted(by_isize[best_isize]) |
| 89 | best_ntokens = next((n for n in available if n >= num_tokens), available[-1]) |
| 90 | |
| 91 | result = CaseKey({"intermediate": best_isize, "numtokens": best_ntokens}) |
| 92 | _pick_cache[cache_key] = result |
| 93 | return result |
| 94 | |
| 95 | |
| 96 | @register_kernel( |
| 97 | config_picker=pick_silu_mul_fp8_config, |
| 98 | input_generator=generate_silu_mul_fp8_inputs, |
| 99 | ) |
| 100 | def silu_mul_fp8(input: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: |
| 101 | original_shape = input.shape |
| 102 | two_d = hl.specialize(original_shape[-1]) |
| 103 | d = two_d // 2 |
| 104 | output_shape = original_shape[:-1] + (d,) |
| 105 | |
| 106 | input_2d = input.view(-1, original_shape[-1]) |
| 107 | m = input_2d.shape[0] |
| 108 | |
| 109 | # TODO(gmagogsfm): Support for more float8 subtypes (e4m3fnuz, e5m2) coming |
| 110 | out = torch.empty((m, d), device=input.device, dtype=torch.float8_e4m3fn) |
| 111 | |
| 112 | input_part_a = input_2d[:, :d] |
| 113 | input_part_b = input_2d[:, d:] |
| 114 | |
| 115 | assert scale.numel() == 1, "Scale must be a scalar Tensor" |
| 116 | |
| 117 | for tile_m, tile_n in hl.tile([m, d]): |
| 118 | a_vals = input_part_a[tile_m, tile_n] |
| 119 | silu_result = torch.nn.functional.silu(a_vals) |
| 120 | b_vals = input_part_b[tile_m, tile_n] |
| 121 | result = silu_result * b_vals |
| 122 | result_f32 = result.to(torch.float32) |
| 123 | scale_val = hl.load(scale, [0]) |
| 124 | inv_scale = 1.0 / scale_val |
| 125 | result_scaled = result_f32 * inv_scale |
| 126 | out[tile_m, tile_n] = result_scaled.to(out.dtype) |
| 127 | |
| 128 | return out.view(output_shape) |
| 129 | |
| 130 | |
| 131 | def silu_mul_fp8_baseline(input: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: |
| 132 | output_shape = input.shape[:-1] + (input.shape[-1] // 2,) |
| 133 | out = torch.empty(output_shape, dtype=torch.float8_e4m3fn, device=input.device) |
| 134 | torch.ops._C.silu_and_mul_quant(out, input, scale) |
| 135 | return out |