Skip to content

GitLab

Commit 0d874a4e authored by wenjh's avatar wenjh
Browse files

Merge branch 'nv_main' of v2.12

parents a68e5f87 dfdd3820
Hide whitespace changes
Inline Side-by-side
Showing with 161 additions and 30 deletions
build_tools/__init__.py
View file @ 0d874a4e
# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
......
build_tools/build_ext.py
View file @ 0d874a4e
# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
......@@ -63,6 +63,12 @@ class CMakeExtension(setuptools.Extension):
f"-DCMAKE_BUILD_TYPE={build_type}",
f"-DCMAKE_INSTALL_PREFIX={install_dir}",
]
if bool(int(os.getenv("NVTE_USE_CCACHE", "0"))):
ccache_bin = os.getenv("NVTE_CCACHE_BIN", "ccache")
configure_command += [
f"-DCMAKE_CXX_COMPILER_LAUNCHER={ccache_bin}",
f"-DCMAKE_CUDA_COMPILER_LAUNCHER={ccache_bin}",
]
configure_command += self.cmake_flags
import pybind11
......
build_tools/jax.py
View file @ 0d874a4e
# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
......@@ -19,8 +19,29 @@ def install_requirements() -> List[str]:
def test_requirements() -> List[str]:
"""Test dependencies for TE/JAX extensions."""
return ["numpy"]
"""Test dependencies for TE/JAX extensions.
Triton Package Selection:
The triton package is selected based on NVTE_USE_PYTORCH_TRITON environment variable:
Default (NVTE_USE_PYTORCH_TRITON unset or "0"):
Returns 'triton' - OpenAI's standard package from PyPI.
Install with: pip install triton
NVTE_USE_PYTORCH_TRITON=1:
Returns 'pytorch-triton' - for mixed JAX+PyTorch environments.
Install with: pip install pytorch-triton --index-url https://download.pytorch.org/whl/cu121
Note: Do NOT install pytorch-triton from PyPI directly - that's a placeholder.
"""
use_pytorch_triton = bool(int(os.environ.get("NVTE_USE_PYTORCH_TRITON", "0")))
triton_package = "pytorch-triton" if use_pytorch_triton else "triton"
return [
"numpy",
triton_package,
]
def xla_path() -> str:
......
build_tools/pytorch.py
View file @ 0d874a4e
# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
......@@ -14,12 +14,17 @@ from typing import List
def install_requirements() -> List[str]:
"""Install dependencies for TE/PyTorch extensions."""
return ["torch>=2.1", "einops"] # "onnxscript==0.3.1", "onnx"]
return ["torch>=2.1", "einops", "packaging", "pydantic", "nvdlfw-inspect"]
def test_requirements() -> List[str]:
"""Test dependencies for TE/JAX extensions."""
return ["numpy", "torchvision", "transformers", "torchao==0.13"]
"""Test dependencies for TE/PyTorch extensions."""
return [
"numpy",
"torchvision",
"transformers",
"torchao==0.13",
]
def setup_pytorch_extension(
......
build_tools/te_version.py
View file @ 0d874a4e
# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
......
build_tools/utils.py
View file @ 0d874a4e
# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
......@@ -279,13 +279,9 @@ def get_cuda_include_dirs() -> Tuple[str, str]:
cuda_root = Path(nvidia.__file__).parent
return [
cuda_root / "cuda_nvcc" / "include",
cuda_root / "cublas" / "include",
cuda_root / "cuda_runtime" / "include",
cuda_root / "cudnn" / "include",
cuda_root / "cuda_cccl" / "include",
cuda_root / "nvtx" / "include",
cuda_root / "cuda_nvrtc" / "include",
subdir / "include"
for subdir in cuda_root.iterdir()
if subdir.is_dir() and (subdir / "include").is_dir()
]
......
build_tools/wheel_utils/Dockerfile.aarch
View file @ 0d874a4e
# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
......
build_tools/wheel_utils/Dockerfile.x86
View file @ 0d874a4e
# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
......
build_tools/wheel_utils/build_wheels.sh
View file @ 0d874a4e
# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
......
build_tools/wheel_utils/launch_aarch.sh
View file @ 0d874a4e
# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
......
build_tools/wheel_utils/launch_x86.sh
View file @ 0d874a4e
# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
......
docs/_static/css/output-style.css 0 → 100644
View file @ 0d874a4e
/* Custom styling for program output blocks */
.program-output {
background-color: #f8f9fa;
padding: 0; /* No padding at all */
margin: 0; /* No margins at all */
border-radius: 0; /* No rounded corners */
font-family: 'Courier New', monospace;
font-size: 14px;
line-height: 1.5;
width: 100%;
max-width: 100%;
}
.program-output pre {
margin: 0;
padding: 0;
background: transparent !important;
border: none !important;
color: #2c3e50;
width: 100%;
}
.program-output .highlight {
background: transparent !important;
margin: 0;
width: 100%;
}
/* Alternative lighter style */
.output-block {
background-color: #fafbfc;
border: 1px solid #e1e4e8;
padding: 10px 14px;
margin: 10px 0;
border-radius: 3px;
font-family: 'SF Mono', 'Consolas', monospace;
font-size: 13px;
color: #24292e;
}
/* Console-like output style */
.console-output {
background-color: #1e1e1e;
border-left: 3px solid #76b900;
padding: 14px 18px;
margin: 12px 0;
border-radius: 5px;
font-family: 'Fira Code', 'Consolas', monospace;
font-size: 13px;
color: #d4d4d4;
box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}
.console-output pre {
margin: 0;
color: #d4d4d4;
background: transparent !important;
}
docs/_static/css/rtabs.css 0 → 100644
View file @ 0d874a4e
/* Custom styling for sphinx-tabs */
.sphinx-tabs {
margin-bottom: 1rem;
}
.sphinx-tabs-tab {
background-color: #f4f4f4;
border: 1px solid #ccc;
border-bottom: none;
padding: 0.5rem 1rem;
margin-right: 0.5rem;
cursor: pointer;
font-weight: 500;
transition: background-color 0.2s;
}
.sphinx-tabs-tab:hover {
background-color: #e0e0e0;
}
.sphinx-tabs-tab[aria-selected="true"] {
background-color: #76b900; /* NVIDIA green */
color: white;
border-color: #76b900;
margin-right: 0.5rem;
}
.sphinx-tabs-panel {
border: 1px solid #ccc;
padding: 1rem;
background-color: #f9f9f9;
}
/* Dark mode support for RTD theme */
.rst-content .sphinx-tabs-tab {
color: #333;
}
.rst-content .sphinx-tabs-tab[aria-selected="true"] {
color: white;
}
docs/api/c/activation.rst
View file @ 0d874a4e
..
Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
See LICENSE for license information.
......
docs/api/c/cast.rst
View file @ 0d874a4e
..
Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
See LICENSE for license information.
......
docs/api/c/cast_transpose_noop.rst
View file @ 0d874a4e
..
Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
See LICENSE for license information.
......
docs/api/c/cudnn.rst
View file @ 0d874a4e
..
Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
See LICENSE for license information.
......
docs/api/c/fused_attn.rst
View file @ 0d874a4e
..
Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
See LICENSE for license information.
......
docs/api/c/fused_rope.rst
View file @ 0d874a4e
..
Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
See LICENSE for license information.
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment