@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9912b8ccb8
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 13, 00:45 UTC

[Build] Add OpenAI triton_kernels (#28788)

Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>

Varun Sundar Rabindranath committed Nov 18, 2025 at 19:45 UTC 9912b8ccb861593d76216afa583ac593faf5a309
6 files changed +119 -1
.gitignore
+3
@@ -4,6 +4,9 @@
4 # vllm-flash-attn built from source
5 vllm/vllm_flash_attn/*
6
7 +# OpenAI triton kernels copied from source
8 +vllm/third_party/triton_kernels/*
9 +
10 # triton jit
11 .triton
12
CMakeLists.txt
+5
@@ -1030,6 +1030,11 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
1030 WITH_SOABI)
1031 endif()
1032
1033 +# For CUDA and HIP builds also build the triton_kernels external package.
1034 +if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
1035 + include(cmake/external_projects/triton_kernels.cmake)
1036 +endif()
1037 +
1038 # For CUDA we also build and ship some external projects.
1039 if (VLLM_GPU_LANG STREQUAL "CUDA")
1040 include(cmake/external_projects/flashmla.cmake)
cmake/external_projects/triton_kernels.cmake new
+53
@@ -0,0 +1,53 @@
1 +# Install OpenAI triton_kernels from https://github.com/triton-lang/triton/tree/main/python/triton_kernels
2 +
3 +set(DEFAULT_TRITON_KERNELS_TAG "v3.5.0")
4 +
5 +# Set TRITON_KERNELS_SRC_DIR for use with local development with vLLM. We expect TRITON_KERNELS_SRC_DIR to
6 +# be directly set to the triton_kernels python directory.
7 +if (DEFINED ENV{TRITON_KERNELS_SRC_DIR})
8 + message(STATUS "[triton_kernels] Fetch from $ENV{TRITON_KERNELS_SRC_DIR}")
9 + FetchContent_Declare(
10 + triton_kernels
11 + SOURCE_DIR $ENV{TRITON_KERNELS_SRC_DIR}
12 + )
13 +
14 +else()
15 + set(TRITON_GIT "https://github.com/triton-lang/triton.git")
16 + message (STATUS "[triton_kernels] Fetch from ${TRITON_GIT}:${DEFAULT_TRITON_KERNELS_TAG}")
17 + FetchContent_Declare(
18 + triton_kernels
19 + # TODO (varun) : Fetch just the triton_kernels directory from Triton
20 + GIT_REPOSITORY https://github.com/triton-lang/triton.git
21 + GIT_TAG ${DEFAULT_TRITON_KERNELS_TAG}
22 + GIT_PROGRESS TRUE
23 + SOURCE_SUBDIR python/triton_kernels/triton_kernels
24 + )
25 +endif()
26 +
27 +# Fetch content
28 +FetchContent_MakeAvailable(triton_kernels)
29 +
30 +if (NOT triton_kernels_SOURCE_DIR)
31 + message (FATAL_ERROR "[triton_kernels] Cannot resolve triton_kernels_SOURCE_DIR")
32 +endif()
33 +
34 +if (DEFINED ENV{TRITON_KERNELS_SRC_DIR})
35 + set(TRITON_KERNELS_PYTHON_DIR "${triton_kernels_SOURCE_DIR}/")
36 +else()
37 + set(TRITON_KERNELS_PYTHON_DIR "${triton_kernels_SOURCE_DIR}/python/triton_kernels/triton_kernels/")
38 +endif()
39 +
40 +message (STATUS "[triton_kernels] triton_kernels is available at ${TRITON_KERNELS_PYTHON_DIR}")
41 +
42 +add_custom_target(triton_kernels)
43 +
44 +# Ensure the vllm/third_party directory exists before installation
45 +install(CODE "file(MAKE_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}/vllm/third_party/triton_kernels\")")
46 +
47 +## Copy .py files to install directory.
48 +install(DIRECTORY
49 + ${TRITON_KERNELS_PYTHON_DIR}
50 + DESTINATION
51 + vllm/third_party/triton_kernels/
52 + COMPONENT triton_kernels
53 + FILES_MATCHING PATTERN "*.py")
setup.py
+17
@@ -299,6 +299,20 @@ class cmake_build_ext(build_ext):
299 os.makedirs(os.path.dirname(dst_file), exist_ok=True)
300 self.copy_file(file, dst_file)
301
302 + if _is_cuda() or _is_hip():
303 + # copy vllm/third_party/triton_kernels/**/*.py from self.build_lib
304 + # to current directory so that they can be included in the editable
305 + # build
306 + print(
307 + f"Copying {self.build_lib}/vllm/third_party/triton_kernels "
308 + "to vllm/third_party/triton_kernels"
309 + )
310 + shutil.copytree(
311 + f"{self.build_lib}/vllm/third_party/triton_kernels",
312 + "vllm/third_party/triton_kernels",
313 + dirs_exist_ok=True,
314 + )
315 +
316
317 class precompiled_build_ext(build_ext):
318 """Disables extension building when using precompiled binaries."""
@@ -633,6 +647,9 @@ ext_modules = []
647 if _is_cuda() or _is_hip():
648 ext_modules.append(CMakeExtension(name="vllm._moe_C"))
649 ext_modules.append(CMakeExtension(name="vllm.cumem_allocator"))
650 + # Optional since this doesn't get built (produce an .so file). This is just
651 + # copying the relevant .py files from the source repository.
652 + ext_modules.append(CMakeExtension(name="vllm.triton_kernels", optional=True))
653
654 if _is_hip():
655 ext_modules.append(CMakeExtension(name="vllm._rocm_C"))
vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
+2
@@ -8,6 +8,7 @@ import torch
8 from vllm.logger import init_logger
9 from vllm.platforms import current_platform
10 from vllm.triton_utils import triton
11 +from vllm.utils.import_utils import has_triton_kernels
12 from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_newer
13
14 logger = init_logger(__name__)
@@ -15,6 +16,7 @@ logger = init_logger(__name__)
16
17 def _swizzle_mxfp4(quant_tensor, scale, num_warps):
18 """weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel"""
19 + assert has_triton_kernels()
20 import triton_kernels.matmul_ogs_details.opt_flags as opt_flags
21 from triton_kernels.numerics import InFlexData
22 from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
vllm/utils/import_utils.py
+39 -1
@@ -18,6 +18,10 @@ from typing import Any
18 import regex as re
19 from typing_extensions import Never
20
21 +from vllm.logger import init_logger
22 +
23 +logger = init_logger(__name__)
24 +
25
26 # TODO: This function can be removed if transformer_modules classes are
27 # serialized by value when communicating between processes
@@ -62,6 +66,35 @@ def import_pynvml():
66 return pynvml
67
68
69 +@cache
70 +def import_triton_kernels():
71 + """
72 + For convenience, prioritize triton_kernels that is available in
73 + `site-packages`. Use `vllm.third_party.triton_kernels` as a fall-back.
74 + """
75 + if _has_module("triton_kernels"):
76 + import triton_kernels
77 +
78 + logger.debug_once(
79 + f"Loading module triton_kernels from {triton_kernels.__file__}.",
80 + scope="local",
81 + )
82 + elif _has_module("vllm.third_party.triton_kernels"):
83 + import vllm.third_party.triton_kernels as triton_kernels
84 +
85 + logger.debug_once(
86 + f"Loading module triton_kernels from {triton_kernels.__file__}.",
87 + scope="local",
88 + )
89 + sys.modules["triton_kernels"] = triton_kernels
90 + else:
91 + logger.info_once(
92 + "triton_kernels unavailable in this build. "
93 + "Please consider installing triton_kernels from "
94 + "https://github.com/triton-lang/triton/tree/main/python/triton_kernels"
95 + )
96 +
97 +
98 def import_from_path(module_name: str, file_path: str | os.PathLike):
99 """
100 Import a Python file according to its file path.
@@ -397,7 +430,12 @@ def has_deep_gemm() -> bool:
430
431 def has_triton_kernels() -> bool:
432 """Whether the optional `triton_kernels` package is available."""
400 - return _has_module("triton_kernels")
433 + is_available = _has_module("triton_kernels") or _has_module(
434 + "vllm.third_party.triton_kernels"
435 + )
436 + if is_available:
437 + import_triton_kernels()
438 + return is_available
439
440
441 def has_tilelang() -> bool: