[https://nvbugs/5761391][fix] Cherry-pick #10471: Include triton-kernels as a packaged dependency

.github/CODEOWNERS

@@ -219,6 +219,12 @@ docs/source/performance/perf-benchmarking.md @NVIDIA/trtllm-bench-reviewers
 ## Any changes to versions, additions, or removals of third-party libraries
 /3rdparty/** @NVIDIA/trt-llm-oss-compliance
 
+### Vendored Third-Party Code (triton-kernels)
+## This is a temporary vendored copy of triton-kernels from the Triton project (MIT License).
+## Do not accept contributions to this directory - it should only be updated via scripts/vendor_triton_kernels.py
+## This can be removed if and when triton-kernels is published as a separate wheel.
+/triton_kernels/** @NVIDIA/trt-llm-oss-compliance
+
 ### Docker & Installation Scripts
 ## These scripts install and pin dependency versions
 /docker/common/** @NVIDIA/trt-llm-setup-infra-devs @NVIDIA/trt-llm-infra-devs @NVIDIA/trt-llm-oss-compliance

.pre-commit-config.yaml

@@ -1369,6 +1369,9 @@ common-files: &common_files |
         triton_backend/tools/whisper/client.py |
     )$
 
+# Global exclude pattern for vendored third-party code
+exclude: '^triton_kernels/'
+
 default_install_hook_types: [pre-commit, commit-msg]
 repos:
 -   repo: https://github.com/pycqa/isort

ATTRIBUTIONS-Python.md

@@ -62379,7 +62379,7 @@ Copyright 2018- The Hugging Face team. All rights reserved.
   - `Homepage`: https://github.com/huggingface/transformers
 
 
-## triton (3.5.0)
+## triton (3.5.1)
 
 ### Licenses
 License: `MIT License`
@@ -62417,6 +62417,40 @@ License: `MIT License`
   - `Homepage`: https://github.com/triton-lang/triton/
 
 
+## triton-kernels (3.5.1)
+
+### Licenses
+License: `MIT License`
+
+  - `LICENSE` (from triton repository root):
+```
+Copyright 2018-2020 Philippe Tillet
+Copyright 2020-2022 OpenAI
+
+Permission is hereby granted, free of charge, to any person obtaining
+a copy of this software and associated documentation files
+(the "Software"), to deal in the Software without restriction,
+including without limitation the rights to use, copy, modify, merge,
+publish, distribute, sublicense, and/or sell copies of the Software,
+and to permit persons to whom the Software is furnished to do so,
+subject to the following conditions:
+
+The above copyright notice and this permission notice shall be
+included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
+IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
+TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
+SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+```
+
+### URLs
+  - `Source`: https://github.com/triton-lang/triton/tree/v3.5.1/python/triton_kernels
+
+
 ## tritonclient (2.63.0)
 
 ### Licenses

examples/models/core/gpt_oss/README.md

@@ -107,33 +107,10 @@ Once again, the function call works successfully, this time using a different fu
 
 ## Using OpenAI Triton Kernels for MoE
 
-OpenAI ships a set of Triton kernels optimized for its MoE models. TensorRT-LLM can leverage these kernels; enable them with the steps below:
 
-1. **Build and install Triton** (tested with the commit below):
+OpenAI ships a set of Triton kernels optimized for its MoE models.
 
-```bash
-git clone https://github.com/triton-lang/triton.git
-cd triton
-# Specific commit verified with TensorRT-LLM
-git checkout f3067cd3bd0c29065fa4ecdb724b6f29cbabea5f
-pip install -r python/requirements.txt          # build-time dependencies
-pip install wheel build
-python3 setup.py bdist_wheel
-pip install ./dist/*.whl
-```
-
-2. **Expose the Triton kernels to TensorRT-LLM**
-   The kernels are not packaged in the wheel, so set the environment variable `TRITON_ROOT` to your Triton clone:
-
-```bash
-export TRITON_ROOT=/local/user/triton
-# TensorRT-LLM expects the kernels at:
-#   $TRITON_ROOT/python/triton_kernels
-```
-
-3. **Select Triton as the MoE backend**
-
-• **trtllm-serve** (or other similar commands) — add this snippet to the YAML file passed via `--config`:
+To use the Triton MoE backend with **trtllm-serve** (or other similar commands), add this snippet to the YAML file passed via `--config`:
 
 ```yaml
 moe_config:

pyproject.toml

@@ -10,13 +10,15 @@ build-backend = "setuptools.build_meta"
 ####################################################################################################
 [tool.isort]
 line_length = 80
+known_first_party = ["tensorrt_llm"]
+known_third_party = ["triton_kernels"]
 
 [tool.yapf]
 based_on_style = "pep8"
 column_limit = 80
 
 [tool.codespell]
-skip = ".git,3rdparty,tests/integration/test_input_files**,**.jsonl,**.json"
+skip = ".git,3rdparty,triton_kernels,tests/integration/test_input_files**,**.jsonl,**.json"
 exclude-file = "examples/models/core/whisper/tokenizer.py"
 ignore-words-list = "rouge,inout,atleast,strat,nd,subtile,thrid,improbe,NotIn,te,iteract,anythin,tru,Tracin,vEw,dOut"
 
@@ -42,6 +44,7 @@ fix = true
 # orders of magnitude faster, so we should move to deprecate `yapf`.
 exclude = [
     "**3rdparty/**",
+    "triton_kernels/**",
     ".devcontainer/make_env.py",
     ".github/scripts/label_community_user.py",
     ".github/scripts/pr_checklist_check.py",
@@ -1461,6 +1464,7 @@ convention = "google"
 
 [tool.ruff.lint.isort]
 known-first-party = ["tensorrt_llm"]
+known-third-party = ["triton_kernels"]
 split-on-trailing-comma = false
 
 
@@ -1493,7 +1497,7 @@ disallow_untyped_calls = false
 disallow_incomplete_defs = false
 disallow_untyped_defs = false
 warn_return_any = false
-exclude = []
+exclude = ["triton_kernels"]
 
 
 [[tool.mypy.overrides]]

requirements.txt

@@ -65,7 +65,7 @@ meson
 ninja
 blake3
 soundfile
-triton==3.5.1
+triton==3.5.1 # NOTE: if you update this, you must also run scripts/vendor_triton_kernels.py to vendor the new version of triton_kernels
 tiktoken
 blobfile
 openai-harmony==0.0.4

scripts/vendor_triton_kernels.py

@@ -0,0 +1,200 @@
+#!/usr/bin/env python3
+"""Script to vendor triton-kernels into TensorRT-LLM.
+
+This script:
+1. Clones the Triton repo at a specific tag to a temp directory
+2. Copies the triton_kernels module to the repo root as a top-level package
+3. Adds attribution headers to all Python files
+4. Copies the LICENSE file from Triton
+5. Creates a VERSION file to track the vendored version
+6. Creates a README.md with clear copyright attribution
+
+To update to a new version:
+    python scripts/vendor_triton_kernels.py --tag v3.6.0
+"""
+
+import argparse
+import shutil
+import subprocess
+import tempfile
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).parent.parent.resolve()
+TRITON_REPO_URL = "https://github.com/triton-lang/triton.git"
+TRITON_KERNELS_MODULE_PATH = "python/triton_kernels/triton_kernels"
+DEST_PATH = REPO_ROOT / "triton_kernels"
+
+VENDORED_NOTICE = "# This file is vendored from the Triton project. DO NOT EDIT THIS FILE DIRECTLY."
+ATTRIBUTION_HEADER = f"""\
+{VENDORED_NOTICE}
+# Source: https://github.com/triton-lang/triton/tree/{{tag}}/{{original_file}}
+# Triton is licensed under the MIT License.
+"""
+
+
+def clone_triton(tag: str, dest_dir: str) -> tuple[Path, Path]:
+    """Clone the Triton repo at the specified tag. Returns (module_path, repo_root)."""
+    print(f"Cloning Triton repo at tag {tag}...")
+
+    subprocess.run(
+        ["git", "clone", "--depth", "1", "--branch", tag, TRITON_REPO_URL, dest_dir],
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+
+    repo_root = Path(dest_dir)
+    triton_kernels_module_path = repo_root / TRITON_KERNELS_MODULE_PATH
+    if not triton_kernels_module_path.exists():
+        raise RuntimeError(f"triton_kernels module not found at {triton_kernels_module_path}")
+
+    return triton_kernels_module_path, repo_root
+
+
+def add_attribution_header(file_path: Path, tag: str, original_rel_path: str) -> None:
+    content = file_path.read_text()
+
+    # Handle shebang and encoding declarations
+    lines = content.split("\n")
+    insert_pos = 0
+    preserved_lines = []
+
+    for i, line in enumerate(lines):
+        if line.startswith("#!") or line.startswith("# -*-") or line.startswith("# coding"):
+            preserved_lines.append(line)
+            insert_pos = i + 1
+        else:
+            break
+
+    header = ATTRIBUTION_HEADER.format(tag=tag, original_file=original_rel_path)
+
+    new_content = "\n".join(preserved_lines)
+    if preserved_lines:
+        new_content += "\n"
+    new_content += header
+
+    # Add blank line between header and content if file has content
+    remaining_content = "\n".join(lines[insert_pos:])
+    if remaining_content.strip():
+        new_content += "\n"
+    new_content += remaining_content
+
+    file_path.write_text(new_content)
+
+
+def copy_triton_kernels(src_path: Path, dest_path: Path, tag: str) -> list[str]:
+    """Copy triton_kernels module to destination and add attribution headers."""
+    print(f"Copying triton_kernels to {dest_path}...")
+
+    if dest_path.exists():
+        print(f"  Removing existing {dest_path}")
+        shutil.rmtree(dest_path)
+
+    shutil.copytree(src_path, dest_path)
+
+    # Add attribution headers to all existing Python files
+    python_files = []
+    for py_file in dest_path.rglob("*.py"):
+        rel_path = py_file.relative_to(dest_path)
+        original_rel_path = f"{TRITON_KERNELS_MODULE_PATH}/{rel_path}"
+        python_files.append(str(rel_path))
+        add_attribution_header(py_file, tag, original_rel_path)
+
+    # Create __init__.py files in subdirs that don't have them.
+    # Triton's upstream code relies on implicit namespace packages (PEP 420), but
+    # setuptools' find_packages() requires __init__.py to discover subpackages.
+    for subdir in dest_path.rglob("*"):
+        if subdir.is_dir():
+            init_file = subdir / "__init__.py"
+            if not init_file.exists():
+                print(f"  Creating {init_file.relative_to(dest_path)}")
+                init_file.write_text(f"{VENDORED_NOTICE}\n")
+
+    print(f"  Copied triton_kernels module to {dest_path}")
+    return python_files
+
+
+def copy_license(triton_repo_root: Path, dest_path: Path) -> None:
+    """Copy the Triton LICENSE file."""
+    print("Copying LICENSE file...")
+
+    license_src = triton_repo_root / "LICENSE"
+    license_dest = dest_path / "LICENSE"
+
+    shutil.copy2(license_src, license_dest)
+    print(f"  Copied LICENSE to {license_dest}")
+
+
+def create_version_file(dest_path: Path, tag: str) -> None:
+    """Create a VERSION file to track which version was vendored."""
+    print("Creating VERSION file...")
+
+    version_file = dest_path / "VERSION"
+    version_content = f"""{tag}
+# This file tracks the version of triton-kernels that was vendored.
+# To update, run: python scripts/vendor_triton_kernels.py --tag <new-tag>
+"""
+
+    version_file.write_text(version_content)
+    print(f"  Created {version_file}")
+
+
+def create_readme(dest_path: Path, tag: str) -> None:
+    """Create a README.md with clear copyright attribution."""
+    print("Creating README.md...")
+
+    readme_file = dest_path / "README.md"
+    readme_content = f"""# Vendored triton_kernels
+
+This directory contains code vendored from the [Triton](https://github.com/triton-lang/triton) project.
+
+| | |
+|---|---|
+| **Copyright** | The Triton Authors |
+| **License** | MIT (see [LICENSE](LICENSE) file in this directory) |
+| **Source** | https://github.com/triton-lang/triton/tree/{tag}/python/triton_kernels/triton_kernels |
+| **Version** | `{tag}` |
+
+## Attribution
+
+This code is the work of the Triton authors and is included here under the MIT License.
+Each Python file includes an attribution header indicating its origin.
+
+## Do Not Edit
+
+This code is vendored verbatim and should not be modified directly.
+To update to a newer version, run:
+
+```bash
+python scripts/vendor_triton_kernels.py --tag <new-tag>
+```
+"""
+
+    readme_file.write_text(readme_content)
+    print(f"  Created {readme_file}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Vendor triton-kernels into TensorRT-LLM")
+    parser.add_argument(
+        "--tag",
+        required=True,
+        help="Triton git tag to vendor from. See the list of tags at https://github.com/triton-lang/triton/tags",
+    )
+    args = parser.parse_args()
+
+    print(f"Vendoring triton-kernels from Triton {args.tag}")
+    print(f"Destination: {DEST_PATH}")
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        triton_kernels_src, triton_repo_root = clone_triton(args.tag, tmp_dir)
+        copy_triton_kernels(triton_kernels_src, DEST_PATH, args.tag)
+        copy_license(triton_repo_root, DEST_PATH)
+        create_version_file(DEST_PATH, args.tag)
+        create_readme(DEST_PATH, args.tag)
+
+    print("SUCCESS: triton-kernels has been vendored.")
+
+
+if __name__ == "__main__":
+    main()

setup.py

@@ -305,6 +305,12 @@ def extract_from_precompiled(precompiled_location: str, package_data: List[str],
 with open("README.md", "r", encoding="utf-8") as fh:
     long_description = fh.read()
 
+packages = find_packages()
+# Add vendored triton_kernels as an explicit top-level package.
+# This is vendored from the Triton project and kept at repo root so its
+# internal absolute imports (e.g., "from triton_kernels.foo import bar") work.
+packages += find_packages(include=["triton_kernels", "triton_kernels.*"])
+
 # https://setuptools.pypa.io/en/latest/references/keywords.html
 setup(
     name='tensorrt_llm',
@@ -318,7 +324,7 @@ def extract_from_precompiled(precompiled_location: str, package_data: List[str],
     author="NVIDIA Corporation",
     url="https://github.com/NVIDIA/TensorRT-LLM",
     download_url="https://github.com/NVIDIA/TensorRT-LLM/tags",
-    packages=find_packages(),
+    packages=packages,
     # TODO Add windows support for python bindings.
     classifiers=[
         "Development Status :: 4 - Beta",
@@ -331,6 +337,7 @@ def extract_from_precompiled(precompiled_location: str, package_data: List[str],
     keywords="nvidia tensorrt deeplearning inference",
     package_data={
         'tensorrt_llm': package_data,
+        'triton_kernels': ['LICENSE', 'VERSION', 'README.md'],
     },
     license_files=get_license(),
     entry_points={