| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from abc import abstractmethod |
| 5 | from collections.abc import Callable, Mapping |
| 6 | from typing import Any, ClassVar, Generic, TypeAlias, TypeVar |
| 7 | |
| 8 | from vllm.model_executor.warmup.jit_warmup import ( |
| 9 | VllmJitKernel, |
| 10 | ) |
| 11 | |
| 12 | DEFAULT_CUTEDSL_COMPILE_OPTIONS = "--enable-tvm-ffi" |
| 13 | |
| 14 | CompileKeyT = TypeVar("CompileKeyT") |
| 15 | CuTeDSLLaunchSpec: TypeAlias = ( |
| 16 | tuple[CompileKeyT, tuple[Any, ...]] |
| 17 | | tuple[CompileKeyT, tuple[Any, ...], Any] |
| 18 | | tuple[ |
| 19 | CompileKeyT, |
| 20 | tuple[Any, ...], |
| 21 | Any, |
| 22 | Callable[[], Any], |
| 23 | ] |
| 24 | ) |
| 25 | |
| 26 | |
| 27 | def cutedsl_fake_stream(*, use_tvm_ffi_env_stream: bool = True) -> Any: |
| 28 | import cutlass.cute as cute |
| 29 | |
| 30 | return cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=use_tvm_ffi_env_stream) |
| 31 | |
| 32 | |
| 33 | def compile_cutedsl( |
| 34 | entry: Callable[..., Any], |
| 35 | *args: Any, |
| 36 | options: str = DEFAULT_CUTEDSL_COMPILE_OPTIONS, |
| 37 | use_tvm_ffi_env_stream: bool = True, |
| 38 | ) -> Any: |
| 39 | import cutlass.cute as cute |
| 40 | |
| 41 | return cute.compile( |
| 42 | entry, |
| 43 | *args, |
| 44 | cutedsl_fake_stream(use_tvm_ffi_env_stream=use_tvm_ffi_env_stream), |
| 45 | options=options, |
| 46 | ) |
| 47 | |
| 48 | |
| 49 | class VllmCuTeDSLJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]): |
| 50 | """CuTeDSL owner whose compiled executor is shared by warmup and runtime.""" |
| 51 | |
| 52 | kernel: ClassVar[Any] |
| 53 | bind_launch_inputs = False |
| 54 | |
| 55 | @abstractmethod |
| 56 | def warmup_inputs(self, compile_key: CompileKeyT) -> tuple[Any, ...]: |
| 57 | """Return fake arguments that compile one executor specialization.""" |
| 58 | raise NotImplementedError |
| 59 | |
| 60 | def compile(self, compile_key: CompileKeyT) -> None: |
| 61 | if compile_key in self._compiled_cache: |
| 62 | return |
| 63 | self._compiled_cache[compile_key] = compile_cutedsl( |
| 64 | self.kernel(compile_key), |
| 65 | *self.warmup_inputs(compile_key), |
| 66 | ) |
| 67 | |
| 68 | def launch( |
| 69 | self, |
| 70 | launch_spec: CuTeDSLLaunchSpec[CompileKeyT], |
| 71 | _inputs: Mapping[str, Any], |
| 72 | ) -> Any: |
| 73 | # (compile_key, args), optionally followed by output and epilogue. |
| 74 | compile_key, launch_args = launch_spec[:2] |
| 75 | executor = self._get_or_compile(compile_key) |
| 76 | result = executor(*launch_args) |
| 77 | if len(launch_spec) == 4: |
| 78 | return launch_spec[3]() |
| 79 | return launch_spec[2] if len(launch_spec) == 3 else result |