kanshan/TensorRT-LLMs
1
0
Fork 0
mirror of https://github.com/NVIDIA/TensorRT-LLM.git synced 2026-01-14 06:27:45 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

[None][fix] AutoDeploy: Use tmp folder for the load_moe_align (#9101)

Browse Source 
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
This commit is contained in:
Chenghao Zhang 2025-11-12 14:59:49 -08:00 committed by GitHub
parent 9241ccaf27
commit f1d637ec69
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/load_moe_align.py
View File

@ -4,6 +4,7 @@ Build moe_align CUDA extension eagerly with a persistent build directory
"""
import os
import tempfile
import torch
from torch.utils.cpp_extension import load
@ -12,21 +13,10 @@ from torch.utils.cpp_extension import load
os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "8.0;8.6;8.9;9.0")
THIS_DIR = os.path.dirname(os.path.abspath(__file__))
CACHE_ROOT = os.environ.get("AD_CACHE_DIR") or os.path.join(
os.environ.get("XDG_CACHE_HOME", os.path.join(os.path.expanduser("~"), ".cache")),
"ad_cache",
)
BUILD_DIR = os.path.join(CACHE_ROOT, "auto_deploy", "fused_moe", "moe_align")
try:
os.makedirs(BUILD_DIR, exist_ok=True)
except PermissionError:
import tempfile
# Fallback to the system temp dir while maintaining a stable subfolder layout
BUILD_DIR = os.path.join(
tempfile.gettempdir(), "ad_cache", "auto_deploy", "fused_moe", "moe_align"
)
os.makedirs(BUILD_DIR, exist_ok=True)
# Use system temp directory to avoid environment variable dependency
BUILD_DIR = os.path.join(tempfile.gettempdir(), "ad_cache", "auto_deploy", "fused_moe", "moe_align")
os.makedirs(BUILD_DIR, exist_ok=True)
moe_align_ext = load(
name="moe_align_ext",