| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Iterable |
| 5 | from dataclasses import asdict, dataclass |
| 6 | from typing import Any |
| 7 | |
| 8 | from vllm.model_executor.kernels.mhc.tilelang_kernels import ( |
| 9 | compute_num_split, |
| 10 | mhc_pre_big_fuse_with_norm_tilelang, |
| 11 | ) |
| 12 | from vllm.model_executor.warmup.jit_warmup import VllmJitKernel |
| 13 | from vllm.platforms import current_platform |
| 14 | from vllm.utils.deep_gemm import is_deep_gemm_supported |
| 15 | from vllm.utils.math_utils import cdiv |
| 16 | |
| 17 | |
| 18 | def compute_mhc_pre_num_splits(input_size: int, num_tokens: int) -> int: |
| 19 | splits = compute_num_split(64, input_size, cdiv(num_tokens, 64)) |
| 20 | # Bound both GEMM and fused-normalization specializations during startup. |
| 21 | return 1 if splits == 1 else 4 if splits <= 4 else 16 |
| 22 | |
| 23 | |
| 24 | class MHCPreNormKernel(VllmJitKernel["MHCPreNormKernel.CompileKey"]): |
| 25 | @dataclass(frozen=True) |
| 26 | class CompileKey: |
| 27 | hidden_size: int |
| 28 | rms_eps: float |
| 29 | hc_pre_eps: float |
| 30 | hc_sinkhorn_eps: float |
| 31 | hc_post_mult_value: float |
| 32 | sinkhorn_repeat: int |
| 33 | norm_eps: float |
| 34 | n_splits: int |
| 35 | hc_mult: int |
| 36 | use_pre_mix_in: bool |
| 37 | rms_numel: int |
| 38 | save_pre_mix: bool = True |
| 39 | split_mode: str = "fused" |
| 40 | write_aux: bool = False |
| 41 | |
| 42 | kernel: Any = staticmethod(mhc_pre_big_fuse_with_norm_tilelang) |
| 43 | |
| 44 | def dispatch(self, *, n_splits, **fields) -> CompileKey: # type: ignore[override] |
| 45 | return self.CompileKey(n_splits=n_splits, **fields) |
| 46 | |
| 47 | def get_warmup_keys( |
| 48 | self, |
| 49 | *, |
| 50 | max_tokens: int, |
| 51 | extra_splits: Iterable[int] = (), |
| 52 | **fields, |
| 53 | ) -> list[CompileKey]: |
| 54 | # The split heuristic changes only at 64-token boundaries. Callers that |
| 55 | # reach this kernel through another GEMM (the fused post + pre-norm |
| 56 | # GEMM) pass the splits that GEMM picks in extra_splits. |
| 57 | splits = sorted( |
| 58 | set(extra_splits) |
| 59 | | ( |
| 60 | { |
| 61 | compute_mhc_pre_num_splits(fields["rms_numel"], num_tokens) |
| 62 | for num_tokens in range(1, max_tokens + 1, 64) |
| 63 | } |
| 64 | if is_deep_gemm_supported() |
| 65 | else {1} |
| 66 | ) |
| 67 | ) |
| 68 | return self._trace_dispatch(self.dispatch)(n_splits=splits, **fields) |
| 69 | |
| 70 | def compile(self, compile_key: CompileKey) -> None: |
| 71 | if compile_key not in self._compiled_cache: |
| 72 | self._compiled_cache[compile_key] = self.kernel.compile( |
| 73 | **asdict(compile_key) |
| 74 | ) |
| 75 | |
| 76 | def compile_many(self, compile_keys: Iterable[CompileKey]) -> None: |
| 77 | missing = list( |
| 78 | dict.fromkeys( |
| 79 | key for key in compile_keys if key not in self._compiled_cache |
| 80 | ) |
| 81 | ) |
| 82 | if len(missing) < 2: |
| 83 | return super().compile_many(missing) |
| 84 | |
| 85 | # TileLang elaborates serially before compiling independent TIR modules. |
| 86 | compiled = self.kernel.par_compile( |
| 87 | [asdict(key) for key in missing], num_workers=min(4, len(missing)) |
| 88 | ) |
| 89 | self._compiled_cache.update(zip(missing, compiled, strict=True)) |
| 90 | |
| 91 | def __call__(self, *tensors, **fields): |
| 92 | compile_key = self.dispatch( |
| 93 | n_splits=tensors[0].shape[0], |
| 94 | **fields, |
| 95 | ) |
| 96 | if not current_platform.is_cuda(): |
| 97 | return self.kernel(*tensors, **asdict(compile_key)) |
| 98 | return self._get_or_compile(compile_key)(*tensors) |
| 99 | |
| 100 | |
| 101 | MHC_PRE_NORM_KERNEL = MHCPreNormKernel() |