002f9879b2
Replaces cherry-picked upstream_ref with complete source trees. xllm/ — Iluvatar official C++ inference engine (15MB, 1470 files) Complete: kernels → layers → models → runtime → scheduler → api Excluded: .git, binary images, third_party submodule checkouts ds_vllm/ — Iluvatar official vllm fork (8MB, 703 files) Included: csrc/ (ALL CUDA kernels), fused_moe/, qwen3_5 model, _custom_ops Excluded: tests, benchmarks, docs, examples (not needed for reference) Critical call chains now fully traceable: MoE: moe_topk_softmax_kernels.cuh → ixformer.h → fused_moe.cpp → layer GDN: qwen3_gated_delta_net_base.cpp → qwen3_5_gated_delta_net.cpp Attention: ixformer.h → xllm_paged_attention → attention.cpp
216 lines
7.7 KiB
Python
216 lines
7.7 KiB
Python
import os
|
|
import platform
|
|
|
|
from typing import Optional
|
|
|
|
|
|
def get_cxx_abi() -> bool:
|
|
try:
|
|
import torch
|
|
return torch.compiled_with_cxx11_abi()
|
|
except ImportError:
|
|
return False
|
|
|
|
|
|
def get_python_include_path() -> Optional[str]:
|
|
try:
|
|
from sysconfig import get_paths
|
|
return get_paths()["include"]
|
|
except ImportError:
|
|
return None
|
|
|
|
|
|
def get_torch_root_path() -> Optional[str]:
|
|
try:
|
|
import torch
|
|
import os
|
|
return os.path.dirname(os.path.abspath(torch.__file__))
|
|
except ImportError:
|
|
return None
|
|
|
|
|
|
def get_torch_mlu_root_path() -> Optional[str]:
|
|
try:
|
|
import torch_mlu
|
|
import os
|
|
return os.path.dirname(os.path.abspath(torch_mlu.__file__))
|
|
except ImportError:
|
|
return None
|
|
|
|
|
|
def get_ixformer_root_path() -> Optional[str]:
|
|
try:
|
|
import ixformer
|
|
import os
|
|
return os.path.dirname(os.path.abspath(ixformer.__file__))
|
|
except ImportError:
|
|
return None
|
|
|
|
|
|
def get_cuda_root_path() -> Optional[str]:
|
|
try:
|
|
import torch
|
|
from torch.utils.cpp_extension import CUDA_HOME
|
|
if CUDA_HOME is None:
|
|
raise RuntimeError(
|
|
"PyTorch was not built with CUDA, or nvcc is not in PATH. "
|
|
"Please set CUDA_TOOLKIT_ROOT_DIR manually."
|
|
)
|
|
return CUDA_HOME
|
|
except ImportError:
|
|
return None
|
|
|
|
|
|
def get_torch_musa_root_path() -> Optional[str]:
|
|
try:
|
|
import torch_musa
|
|
import os
|
|
return os.path.dirname(os.path.abspath(torch_musa.__file__))
|
|
except ImportError:
|
|
return None
|
|
|
|
def prepend_path_env(var_name: str, path: str, sep: str = os.pathsep) -> None:
|
|
"""Prepend a path into a path env var without duplicates."""
|
|
if not path:
|
|
return
|
|
current = os.getenv(var_name, "")
|
|
entries = [item for item in current.split(sep) if item]
|
|
if path in entries:
|
|
entries = [item for item in entries if item != path]
|
|
entries.insert(0, path)
|
|
os.environ[var_name] = sep.join(entries)
|
|
|
|
|
|
def set_npu_torch_ld_library_path() -> None:
|
|
"""Only for NPU flow: ensure torch runtime libraries are discoverable."""
|
|
torch_root = os.getenv("PYTORCH_INSTALL_PATH") or get_torch_root_path() or ""
|
|
if not torch_root:
|
|
return
|
|
|
|
# Order keeps current behavior: torch.libs > torch > torch/lib
|
|
for path in (f"{torch_root}.libs", torch_root, os.path.join(torch_root, "lib")):
|
|
if os.path.isdir(path):
|
|
prepend_path_env("LD_LIBRARY_PATH", path)
|
|
|
|
|
|
def set_common_envs() -> None:
|
|
os.environ["PYTHON_INCLUDE_PATH"] = get_python_include_path() or ""
|
|
torch_root = get_torch_root_path() or ""
|
|
os.environ["PYTHON_LIB_PATH"] = torch_root
|
|
os.environ["LIBTORCH_ROOT"] = torch_root
|
|
os.environ["PYTORCH_INSTALL_PATH"] = torch_root
|
|
|
|
|
|
def set_npu_envs() -> None:
|
|
PYTORCH_NPU_INSTALL_PATH = os.getenv("PYTORCH_NPU_INSTALL_PATH")
|
|
if not PYTORCH_NPU_INSTALL_PATH:
|
|
os.environ["PYTORCH_NPU_INSTALL_PATH"] = "/usr/local/libtorch_npu"
|
|
|
|
set_common_envs()
|
|
set_npu_torch_ld_library_path()
|
|
NPU_TOOLKIT_HOME = os.getenv("NPU_TOOLKIT_HOME")
|
|
if not NPU_TOOLKIT_HOME:
|
|
os.environ["NPU_TOOLKIT_HOME"] = "/usr/local/Ascend/ascend-toolkit/latest"
|
|
NPU_TOOLKIT_HOME = "/usr/local/Ascend/ascend-toolkit/latest"
|
|
LD_LIBRARY_PATH = os.getenv("LD_LIBRARY_PATH", "")
|
|
arch = platform.machine()
|
|
LD_LIBRARY_PATH = NPU_TOOLKIT_HOME+"/lib64" + ":" + \
|
|
NPU_TOOLKIT_HOME+"/lib64/plugin/opskernel" + ":" + \
|
|
NPU_TOOLKIT_HOME+"/lib64/plugin/nnengine" + ":" + \
|
|
NPU_TOOLKIT_HOME+"/opp/built-in/op_impl/ai_core/tbe/op_tiling/lib/linux/"+arch + ":" + \
|
|
NPU_TOOLKIT_HOME+"/opp/vendors/xllm/op_api/lib" + ":" + \
|
|
NPU_TOOLKIT_HOME+"/tools/aml/lib64" + ":" + \
|
|
NPU_TOOLKIT_HOME+"/tools/aml/lib64/plugin" + ":" + \
|
|
LD_LIBRARY_PATH
|
|
os.environ["LD_LIBRARY_PATH"] = LD_LIBRARY_PATH
|
|
PYTHONPATH = os.getenv("PYTHONPATH", "")
|
|
PYTHONPATH = NPU_TOOLKIT_HOME+"/python/site-packages" + ":" + \
|
|
NPU_TOOLKIT_HOME+"/opp/built-in/op_impl/ai_core/tbe" + ":" + \
|
|
PYTHONPATH
|
|
os.environ["PYTHONPATH"] = PYTHONPATH
|
|
PATH = os.getenv("PATH", "")
|
|
PATH = NPU_TOOLKIT_HOME+"/bin" + ":" + \
|
|
NPU_TOOLKIT_HOME+"/compiler/ccec_compiler/bin" + ":" + \
|
|
NPU_TOOLKIT_HOME+"/tools/ccec_compiler/bin" + ":" + \
|
|
PATH
|
|
os.environ["PATH"] = PATH
|
|
os.environ["ASCEND_AICPU_PATH"] = NPU_TOOLKIT_HOME
|
|
os.environ["ASCEND_OPP_PATH"] = NPU_TOOLKIT_HOME+"/opp"
|
|
os.environ["TOOLCHAIN_HOME"] = NPU_TOOLKIT_HOME+"/toolkit"
|
|
os.environ["NPU_HOME_PATH"] = NPU_TOOLKIT_HOME
|
|
|
|
ATB_PATH = os.getenv("ATB_PATH")
|
|
if not ATB_PATH:
|
|
os.environ["ATB_PATH"] = "/usr/local/Ascend/nnal/atb"
|
|
ATB_PATH = "/usr/local/Ascend/nnal/atb"
|
|
|
|
|
|
cxx_abi = "1" if get_cxx_abi() else "0"
|
|
ATB_HOME_PATH = os.path.join(ATB_PATH, "latest", "atb", "cxx_abi_" + cxx_abi)
|
|
os.environ["ATB_HOME_PATH"] = ATB_HOME_PATH
|
|
LD_LIBRARY_PATH = os.getenv("LD_LIBRARY_PATH", "")
|
|
LD_LIBRARY_PATH = ATB_HOME_PATH+"/lib" + ":" + \
|
|
ATB_HOME_PATH+"/examples" + ":" + \
|
|
ATB_HOME_PATH+"/tests/atbopstest" + ":" + \
|
|
LD_LIBRARY_PATH
|
|
os.environ["LD_LIBRARY_PATH"] = LD_LIBRARY_PATH
|
|
PATH = os.getenv("PATH", "")
|
|
PATH = ATB_HOME_PATH+"/bin" + ":" + PATH
|
|
os.environ["PATH"] = PATH
|
|
|
|
os.environ["ATB_STREAM_SYNC_EVERY_KERNEL_ENABLE"] = "0"
|
|
os.environ["ATB_STREAM_SYNC_EVERY_RUNNER_ENABLE"] = "0"
|
|
os.environ["ATB_STREAM_SYNC_EVERY_OPERATION_ENABLE"] = "0"
|
|
os.environ["ATB_OPSRUNNER_SETUP_CACHE_ENABLE"] = "1"
|
|
os.environ["ATB_OPSRUNNER_KERNEL_CACHE_TYPE"] = "3"
|
|
os.environ["ATB_OPSRUNNER_KERNEL_CACHE_LOCAL_COUNT"] = "1"
|
|
os.environ["ATB_OPSRUNNER_KERNEL_CACHE_GLOABL_COUNT"] = "5"
|
|
os.environ["ATB_OPSRUNNER_KERNEL_CACHE_TILING_SIZE"] = "10240"
|
|
os.environ["ATB_WORKSPACE_MEM_ALLOC_ALG_TYPE"] = "1"
|
|
os.environ["ATB_WORKSPACE_MEM_ALLOC_GLOBAL"] = "0"
|
|
os.environ["ATB_COMPARE_TILING_EVERY_KERNEL"] = "0"
|
|
os.environ["ATB_HOST_TILING_BUFFER_BLOCK_NUM"] = "128"
|
|
os.environ["ATB_DEVICE_TILING_BUFFER_BLOCK_NUM"] = "32"
|
|
os.environ["ATB_SHARE_MEMORY_NAME_SUFFIX"] = ""
|
|
os.environ["ATB_LAUNCH_KERNEL_WITH_TILING"] = "1"
|
|
os.environ["ATB_MATMUL_SHUFFLE_K_ENABLE"] = "1"
|
|
os.environ["ATB_RUNNER_POOL_SIZE"] = "64"
|
|
os.environ["ASDOPS_HOME_PATH"] = ATB_HOME_PATH
|
|
os.environ["ASDOPS_MATMUL_PP_FLAG"] = "1"
|
|
os.environ["ASDOPS_LOG_LEVEL"] = "ERROR"
|
|
os.environ["ASDOPS_LOG_TO_STDOUT"] = "0"
|
|
os.environ["ASDOPS_LOG_TO_FILE"] = "1"
|
|
os.environ["ASDOPS_LOG_TO_FILE_FLUSH"] = "0"
|
|
os.environ["ASDOPS_LOG_TO_BOOST_TYPE"] = "atb"
|
|
os.environ["ASDOPS_LOG_PATH"] = "~"
|
|
os.environ["ASDOPS_TILING_PARSE_CACHE_DISABLE"] = "0"
|
|
os.environ["LCCL_DETERMINISTIC"] = "0"
|
|
os.environ["LCCL_PARALLEL"] = "0"
|
|
|
|
|
|
def set_mlu_envs() -> None:
|
|
set_common_envs()
|
|
os.environ["PYTORCH_MLU_INSTALL_PATH"] = get_torch_mlu_root_path() or ""
|
|
|
|
|
|
def set_cuda_envs() -> None:
|
|
set_common_envs()
|
|
os.environ["CUDA_TOOLKIT_ROOT_DIR"] = get_cuda_root_path() or ""
|
|
|
|
|
|
def set_ilu_envs() -> None:
|
|
set_common_envs()
|
|
os.environ["IXFORMER_INSTALL_PATH"] = get_ixformer_root_path() or ""
|
|
|
|
|
|
def set_musa_envs() -> None:
|
|
set_common_envs()
|
|
os.environ["PYTORCH_MUSA_INSTALL_PATH"] = get_torch_musa_root_path() or ""
|
|
import torch_musa
|
|
from torch_musa.utils.musa_extension import MUSA_HOME
|
|
os.environ["TORCH_MUSA_PYTHONPATH"] = torch_musa.core.cmake_prefix_path
|
|
os.environ["MUSA_TOOLKIT_ROOT_DIR"] = MUSA_HOME
|
|
os.environ["MKL_DIR"] = "/opt/intel/oneapi/mkl/lib/cmake/mkl"
|
|
os.environ["MKLROOT"] = "/opt/intel/oneapi/mkl"
|
|
os.environ["TorchMusa_DIR"] = torch_musa.core.cmake_prefix_path + "/TorchMusa"
|
|
os.environ["MUSAMAPPING_PATH"] = MUSA_HOME + "/tools/musamapping" |