From 33d73bf8d954566a0d167555a34ab553a7400a19 Mon Sep 17 00:00:00 2001 From: zyuli Date: Wed, 15 Jul 2026 00:43:11 +0000 Subject: [PATCH 01/16] [KMCompiler][TLERaw] Fix cuda device-side not supporting multiple functions --- .../experimental/tle/language/raw/core.py | 10 +++--- .../experimental/tle/raw/cuda/runtime.py | 4 ++- .../experimental/tle/raw/mlir/runtime.py | 4 ++- .../experimental/tle/raw/tops/runtime.py | 4 ++- third_party/tle/triton_tle.cc | 28 ++++++++------- third_party/tle/triton_tle_raw.cc | 36 ++++++++++++------- 6 files changed, 55 insertions(+), 31 deletions(-) diff --git a/python/triton/experimental/tle/language/raw/core.py b/python/triton/experimental/tle/language/raw/core.py index 3ae5e56e3..722feb47a 100644 --- a/python/triton/experimental/tle/language/raw/core.py +++ b/python/triton/experimental/tle/language/raw/core.py @@ -3,9 +3,9 @@ from triton.experimental.tle.language.gpu import buffered_tensor -def _resolve_alias_indices(func, llvm, handles, output_indices, _semantic): +def _resolve_alias_indices(func, llvm, handles, output_indices, extern_func_name, _semantic): if output_indices is None: - return _semantic.builder.compute_alias_operand_indices(llvm, handles) + return _semantic.builder.compute_alias_operand_indices(llvm, handles, extern_func_name) return output_indices @@ -52,8 +52,10 @@ def _tle_raw_call(func, args, *, output_indices, hint, smem, _semantic): else: context = _semantic.builder.get_context() llvm = func.make_llvm(context) - alias_indices = _resolve_alias_indices(func, llvm, handles, output_indices, _semantic) - dsl_region_op = func.create_region_by_llvm(_semantic.builder, llvm, handles, alias_indices, hint) + alias_indices = _resolve_alias_indices(func, llvm, handles, output_indices, func.extern_func_name or "", + _semantic) + dsl_region_op = func.create_region_by_llvm(_semantic.builder, llvm, handles, alias_indices, hint, + func.extern_func_name or "") return _wrap_results(args, alias_indices, dsl_region_op, smem=smem) diff --git a/python/triton/experimental/tle/raw/cuda/runtime.py b/python/triton/experimental/tle/raw/cuda/runtime.py index 2fc2ae69f..9e4a86b4d 100644 --- a/python/triton/experimental/tle/raw/cuda/runtime.py +++ b/python/triton/experimental/tle/raw/cuda/runtime.py @@ -72,7 +72,8 @@ def register_pending_source(self, *, hint: str = "") -> str: extra={"source_file": self.source_file}, ) - def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint: str = ""): + def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint: str = "", + extern_func_name: str = ""): return builder.create_tle_raw_region_by_llvm_func( llvm, self.region_dialect, @@ -80,6 +81,7 @@ def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint handles, alias_indices, hint, + extern_func_name, ) def create_region_deferred(self, builder, source_id: str, handles, alias_indices, hint: str = ""): diff --git a/python/triton/experimental/tle/raw/mlir/runtime.py b/python/triton/experimental/tle/raw/mlir/runtime.py index 2b3e633d3..743ef5511 100644 --- a/python/triton/experimental/tle/raw/mlir/runtime.py +++ b/python/triton/experimental/tle/raw/mlir/runtime.py @@ -99,7 +99,8 @@ def create_region_deferred(self, builder, source_id: str, handles, alias_indices hint, ) - def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint: str = ""): + def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint: str = "", + extern_func_name: str = ""): return builder.create_tle_raw_region_by_llvm_func( llvm, self.region_dialect, @@ -107,6 +108,7 @@ def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint handles, alias_indices, hint, + extern_func_name, ) @cached_property diff --git a/python/triton/experimental/tle/raw/tops/runtime.py b/python/triton/experimental/tle/raw/tops/runtime.py index 056015a99..ca62b9b72 100644 --- a/python/triton/experimental/tle/raw/tops/runtime.py +++ b/python/triton/experimental/tle/raw/tops/runtime.py @@ -173,7 +173,8 @@ def _rewrite_gcu_intrinsics(llvm_ir: str) -> str: print("// ---- end ----") return result - def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint: str = ""): + def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint: str = "", + extern_func_name: str = ""): return builder.create_tle_raw_region_by_llvm_func( llvm, self.region_dialect, @@ -181,6 +182,7 @@ def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint handles, alias_indices, hint, + extern_func_name, ) def make_llvm(self, mlir_context) -> str: diff --git a/third_party/tle/triton_tle.cc b/third_party/tle/triton_tle.cc index 7709e90b2..3aefb71ea 100644 --- a/third_party/tle/triton_tle.cc +++ b/third_party/tle/triton_tle.cc @@ -62,13 +62,16 @@ namespace tle = triton::tle; extern std::vector computeAliasOperandIndices(TritonOpBuilder &self, std::string_view text, - const std::vector &args); + const std::vector &args, + std::string_view funcName); -extern tle::DSLRegionOp createTLERawRegionByLLVMFunc( - TritonOpBuilder &self, std::string_view text, - std::string_view regionDialect, std::string_view argDialect, - const std::vector &args, - const std::vector &aliasOperandIndices, std::string_view hint); +extern tle::DSLRegionOp +createTLERawRegionByLLVMFunc(TritonOpBuilder &self, std::string_view text, + std::string_view regionDialect, + std::string_view argDialect, + const std::vector &args, + const std::vector &aliasOperandIndices, + std::string_view hint, std::string_view funcName); extern tle::DSLRegionOp createTLERawRegionDeferred( TritonOpBuilder &self, std::string_view sourceId, @@ -697,12 +700,13 @@ void init_tle_raw_ir(py::module &&m) { .def("dump", &tle::YieldOp::dump); auto *builder_cls = ir::getBuilderClass(); - builder_cls->def("compute_alias_operand_indices", - &computeAliasOperandIndices); - builder_cls->def( - "create_tle_raw_region_by_llvm_func", &createTLERawRegionByLLVMFunc, - py::arg("text"), py::arg("region_dialect"), py::arg("arg_dialect"), - py::arg("args"), py::arg("output_operand_indices"), py::arg("hint") = ""); + builder_cls->def("compute_alias_operand_indices", &computeAliasOperandIndices, + py::arg("text"), py::arg("args"), py::arg("func_name") = ""); + builder_cls->def("create_tle_raw_region_by_llvm_func", + &createTLERawRegionByLLVMFunc, py::arg("text"), + py::arg("region_dialect"), py::arg("arg_dialect"), + py::arg("args"), py::arg("output_operand_indices"), + py::arg("hint") = "", py::arg("func_name") = ""); builder_cls->def( "create_tle_raw_region_deferred", &createTLERawRegionDeferred, py::arg("source_id"), py::arg("region_dialect"), py::arg("arg_dialect"), diff --git a/third_party/tle/triton_tle_raw.cc b/third_party/tle/triton_tle_raw.cc index f790e9ab3..6409b3893 100644 --- a/third_party/tle/triton_tle_raw.cc +++ b/third_party/tle/triton_tle_raw.cc @@ -7,6 +7,7 @@ #include "tle/utils/include/AnalyzeReturnType.h" #include "tle/utils/include/TleRawMaterialize.h" #include "llvm/ADT/STLExtras.h" +#include using namespace mlir; namespace tle = triton::tle; @@ -18,6 +19,12 @@ StringAttr getOptionalStringAttr(OpBuilder &builder, std::string_view value) { return builder.getStringAttr(value); } +std::optional getOptionalFuncName(std::string_view value) { + if (value.empty()) + return std::nullopt; + return llvm::StringRef(value.data(), value.size()); +} + void setDeferredMetadataAttrs(tle::DSLRegionOp op, OpBuilder &builder, std::string_view sourceId) { if (!sourceId.empty()) @@ -37,13 +44,15 @@ tle::DSLRegionOp createDSLRegionOp( } } // namespace -std::vector -computeAliasOperandIndices(TritonOpBuilder &self, std::string_view text, - const std::vector &args) { +std::vector computeAliasOperandIndices(TritonOpBuilder &self, + std::string_view text, + const std::vector &args, + std::string_view funcName) { OwningOpRef module = tle::raw::parseLLVMModule(self.getContext(), text); assert(module && "Failed to parse LLVM IR text"); - LLVM::LLVMFuncOp func = tle::raw::findExternalLLVMFunc(module.get(), {}); + LLVM::LLVMFuncOp func = tle::raw::findExternalLLVMFunc( + module.get(), getOptionalFuncName(funcName)); assert(func && "No function found in LLVM IR text"); SmallVector funcArgToDslArg = @@ -61,15 +70,18 @@ computeAliasOperandIndices(TritonOpBuilder &self, std::string_view text, return std::vector(result.begin(), result.end()); } -tle::DSLRegionOp createTLERawRegionByLLVMFunc( - TritonOpBuilder &self, std::string_view text, - std::string_view regionDialect, std::string_view argDialect, - const std::vector &args, - const std::vector &aliasOperandIndices, std::string_view hint) { +tle::DSLRegionOp +createTLERawRegionByLLVMFunc(TritonOpBuilder &self, std::string_view text, + std::string_view regionDialect, + std::string_view argDialect, + const std::vector &args, + const std::vector &aliasOperandIndices, + std::string_view hint, std::string_view funcName) { OwningOpRef module = tle::raw::parseLLVMModule(self.getContext(), text); assert(module && "Failed to parse LLVM IR text"); - LLVM::LLVMFuncOp func = tle::raw::findExternalLLVMFunc(module.get(), {}); + LLVM::LLVMFuncOp func = tle::raw::findExternalLLVMFunc( + module.get(), getOptionalFuncName(funcName)); assert(func && "No function found in LLVM IR text"); OpBuilder &builder = self.getBuilder(); @@ -79,8 +91,8 @@ tle::DSLRegionOp createTLERawRegionByLLVMFunc( } ModuleOp curModule = cast(curOp); - auto funcOpOrErr = - tle::raw::cloneLLVMSymbolsAndLookupFunc(curModule, module.get(), {}); + auto funcOpOrErr = tle::raw::cloneLLVMSymbolsAndLookupFunc( + curModule, module.get(), getOptionalFuncName(funcName)); assert(succeeded(funcOpOrErr)); LLVM::LLVMFuncOp funcOp = *funcOpOrErr; From 23bc3513ff70204b8d199db132d14c0810f7805f Mon Sep 17 00:00:00 2001 From: zyuli Date: Wed, 15 Jul 2026 04:02:23 +0000 Subject: [PATCH 02/16] [KMCompiler][TLERaw] Support nvshmem --- .../triton/experimental/tle/raw/cache_key.py | 10 + .../experimental/tle/raw/cuda/runtime.py | 44 +- .../tle/raw/nvshmem/common-host.cu | 56 ++ .../experimental/tle/raw/nvshmem/utils.py | 256 ++++++++ .../01-simple-shift/simple-shift-device.cu | 11 + .../01-simple-shift/simple-shift-host.cu | 47 ++ .../nvshmem/01-simple-shift/simple-shift.py | 73 +++ .../02-allgather-gemm/ag-gemm-device.cu | 31 + .../nvshmem/02-allgather-gemm/ag-gemm-host.cu | 63 ++ .../raw/nvshmem/02-allgather-gemm/ag-gemm.py | 555 ++++++++++++++++++ third_party/nvidia/backend/compiler.py | 5 + 11 files changed, 1150 insertions(+), 1 deletion(-) create mode 100644 python/triton/experimental/tle/raw/nvshmem/common-host.cu create mode 100644 python/triton/experimental/tle/raw/nvshmem/utils.py create mode 100644 python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-device.cu create mode 100644 python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-host.cu create mode 100644 python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py create mode 100644 python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm-device.cu create mode 100644 python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm-host.cu create mode 100644 python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm.py diff --git a/python/triton/experimental/tle/raw/cache_key.py b/python/triton/experimental/tle/raw/cache_key.py index 7dfbc6140..6ede3e08c 100644 --- a/python/triton/experimental/tle/raw/cache_key.py +++ b/python/triton/experimental/tle/raw/cache_key.py @@ -94,6 +94,16 @@ def compute_tle_raw_source_cache_key( return hasher.hexdigest() +def compute_tle_raw_host_cache_key(source: Union[str, Path], arch: str) -> str: + """Hash CUDA host source files.""" + source_path = Path(source).resolve() + hasher = hashlib.sha256() + hasher.update(_read_source(source_path).encode()) + hasher.update(str(arch).encode()) + + return hasher.hexdigest() + + def bind_tle_raw_source_cache_key(edsl: Any, **dialect_kwargs) -> None: """Attach __triton_tle_raw_source_cache_key__ to a @dialect edsl object.""" if getattr(edsl, TLE_RAW_SOURCE_CACHE_KEY_ATTR, None) is not None: diff --git a/python/triton/experimental/tle/raw/cuda/runtime.py b/python/triton/experimental/tle/raw/cuda/runtime.py index 9e4a86b4d..6f243eb13 100644 --- a/python/triton/experimental/tle/raw/cuda/runtime.py +++ b/python/triton/experimental/tle/raw/cuda/runtime.py @@ -6,17 +6,23 @@ from pathlib import Path import subprocess from typing import Any, Final +import ctypes +from triton import knobs import torch from triton._C.libtriton import llvm # pyright: ignore[reportMissingImports] from triton._C.libtriton.tle.llvm import parse_llvm_ir # pyright: ignore[reportMissingImports] from triton.experimental.tle.raw.source_store import register_source +from triton.experimental.tle.raw.nvshmem.utils import get_nvshmem_home # TODO: We use cli tools to compile CUDA code temporarily, and plan to replace it with LLVM components Python bindings in the future. CLANG = os.getenv("CLANG", "clang") CLANG_FLAGS = shlex.split(os.getenv("CLANG_FLAGS", "")) +_cumodule_hook_installed = False +_nvshmemx_cumodule_init = None + def _sanitize_clang_ir(ir: str) -> str: # Newer clang emits attributes that this Triton branch's LLVM parser does @@ -46,10 +52,43 @@ def _get_cuda_gpu_arch() -> str: return f"--cuda-gpu-arch=sm_{major}{minor}" +def _get_nvshmemx_cumodule_init(): + global _nvshmemx_cumodule_init + if _nvshmemx_cumodule_init is not None: + return _nvshmemx_cumodule_init + + nvshmem_home = get_nvshmem_home() + library = ctypes.CDLL(str(Path(nvshmem_home) / "lib" / "libnvshmem_host.so")) + fn = library.nvshmemx_cumodule_init + fn.argtypes = [ctypes.c_void_p] + fn.restype = ctypes.c_int + _nvshmemx_cumodule_init = fn + return fn + + +def _install_cumodule_hook(): + global _cumodule_hook_installed + if _cumodule_hook_installed: + return + + def hook(*args, **kwargs): + key = kwargs["key"] + function = kwargs["fn"].jit_function + device = kwargs["compile"]["device"] + kernel = function.device_caches[device][0].get(key) + assert kernel is not None + kernel._init_handles() + result = _get_nvshmemx_cumodule_init()(ctypes.c_void_p(kernel.module)) + assert result == 0, f"nvshmemx_cumodule_init failed: {result}" + + knobs.runtime.jit_post_compile_hook = hook + _cumodule_hook_installed = True + + class CUDAJITFunction(object): def __init__(self, fn: Any, file: Path, *args, **kwargs) -> None: - super().__init__(*args, **{k: v for k, v in kwargs.items() if k not in ("extern_func_name", "deferred")}) + super().__init__() self.fn: Final[Any] = fn self.code: Final[str] = file.read_text() self.region_dialect: Final[str] = "cuda" @@ -60,6 +99,9 @@ def __init__(self, fn: Any, file: Path, *args, **kwargs) -> None: self.deferred: Final[bool] = kwargs.get("deferred", False) self.__triton_builtin__: Final[bool] = True + if "nvshmem" in self.code: + _install_cumodule_hook() + def register_pending_source(self, *, hint: str = "") -> str: if not self.extern_func_name: raise RuntimeError("deferred tle_raw CUDA source requires extern_func_name= " diff --git a/python/triton/experimental/tle/raw/nvshmem/common-host.cu b/python/triton/experimental/tle/raw/nvshmem/common-host.cu new file mode 100644 index 000000000..5f1382b7f --- /dev/null +++ b/python/triton/experimental/tle/raw/nvshmem/common-host.cu @@ -0,0 +1,56 @@ +#include +#include +#include +#include +#include +#include + +#undef CUDA_CHECK +#define CUDA_CHECK(stmt) \ + do { \ + cudaError_t result = (stmt); \ + if (cudaSuccess != result) { \ + fprintf(stderr, "[%s:%d] cuda failed with %s \n", __FILE__, __LINE__, \ + cudaGetErrorString(result)); \ + exit(-1); \ + } \ + } while (0) + +extern "C" int nvshmem_get_unique_id_bytes(void *uid_buffer, + size_t uid_buffer_size) { + if (uid_buffer_size < sizeof(nvshmemx_uniqueid_t)) { + return -1; + } + + nvshmemx_uniqueid_t uid; + nvshmemx_get_uniqueid(&uid); + memcpy(uid_buffer, &uid, sizeof(uid)); + return 0; +} + +extern "C" int nvshmem_init_from_torch_distributed(int rank, int nranks, + int cuda_device, + void *uid_buffer, + size_t uid_buffer_size) { + if (uid_buffer_size < sizeof(nvshmemx_uniqueid_t)) { + return -1; + } + + CUDA_CHECK(cudaSetDevice(cuda_device)); + + nvshmemx_uniqueid_t uid; + memcpy(&uid, uid_buffer, sizeof(uid)); + + nvshmemx_init_attr_t attr; + memset(&attr, 0, sizeof(attr)); + nvshmemx_set_attr_uniqueid_args(rank, nranks, &uid, &attr); + nvshmemx_init_attr(NVSHMEMX_INIT_WITH_UNIQUEID, &attr); + + return 0; +} + +extern "C" int nvshmem_finalize_from_torch_distributed() { + nvshmem_finalize(); + + return 0; +} diff --git a/python/triton/experimental/tle/raw/nvshmem/utils.py b/python/triton/experimental/tle/raw/nvshmem/utils.py new file mode 100644 index 000000000..e6dd25c9d --- /dev/null +++ b/python/triton/experimental/tle/raw/nvshmem/utils.py @@ -0,0 +1,256 @@ +import ctypes +import datetime +import os +import warnings +from pathlib import Path + +import torch +import functools +import sysconfig +import subprocess +import re +import tempfile +import shlex +from triton.runtime.cache import get_cache_manager +from triton.experimental.tle.raw.cache_key import compute_tle_raw_host_cache_key + +try: + from cuda.bindings import driver as cuda + from cuda.bindings import runtime as cudart +except ImportError: + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message=r"The cuda\.(cuda|cudart) module is deprecated.*", + category=FutureWarning, + ) + from cuda import cuda as cuda + from cuda import cudart as cudart + + +def CUDA_CHECK(err): + if isinstance(err, cuda.CUresult): + if err != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"Cuda Error: {err}: {cuda.cuGetErrorName(err)}") + elif isinstance(err, cudart.cudaError_t): + if err != cudart.cudaError_t.cudaSuccess: + raise RuntimeError(f"Cuda Error: {err}: {cudart.cudaGetErrorString(err)}") + else: + raise RuntimeError(f"Unknown error type: {err}") + + +@functools.lru_cache() +def get_nvshmem_home() -> Path: + if (nvshmem_home := os.getenv("NVSHMEM_HOME")) is not None: + return Path(nvshmem_home) + + try: + import nvidia.nvshmem + return Path(nvidia.nvshmem.__path__[0]) + except Exception: + raise RuntimeError("Cannot resolve NVSHMEM_HOME: environment variable not set and nvidia.nvshmem not found") + + +@functools.lru_cache() +def get_nvcc(): + return _path_to_binary("nvcc") + + +@functools.lru_cache() +def _path_to_binary(binary: str): + binary += sysconfig.get_config_var("EXE") + paths = [os.environ.get(f"TRITON_{binary.upper()}_PATH", "")] + + cuda_home = os.getenv("CUDA_HOME", "/usr/local/cuda") + paths += [f"{cuda_home}/bin/{binary}"] + + for path in paths: + if os.path.exists(path) and os.path.isfile(path): + result = subprocess.check_output([path, "--version"], stderr=subprocess.STDOUT) + if result is not None: + version = re.search(r".*release (\d+\.\d+).*", result.decode("utf-8"), flags=re.MULTILINE) + if version is not None: + return path, version.group(1) + raise RuntimeError(f"Cannot find {binary}") + + +def _compile_cuda_host_to_cache( + source, + nvshmem_home, + arch: str = None, + force: bool = False, +) -> tuple[Path, str, bool]: + source_path = Path(source).expanduser().resolve() + if not arch: + from triton.experimental.tle.raw.cuda.runtime import _get_cuda_gpu_arch + arch = _get_cuda_gpu_arch().split('=')[1] + nvshmem_home = get_nvshmem_home() + + host_cache_key = compute_tle_raw_host_cache_key(source_path, arch) + output_name = source_path.with_suffix(".so").name + cache = get_cache_manager(host_cache_key) + + cached = None if force else cache.get_file(output_name) + if cached is not None: + return Path(cached), host_cache_key, True + + temporary = tempfile.NamedTemporaryFile( + prefix=f".{output_name}.", + suffix=".tmp", + delete=False, + ) + temporary_path = Path(temporary.name) + temporary.close() + nvcc, _ = get_nvcc() + command = [ + nvcc, + "-shared", + "-Xcompiler", + "-fPIC", + "-rdc=true", + f"-arch={arch}", + f"-I{nvshmem_home / 'include'}", + f"-L{nvshmem_home / 'lib'}", + "-lnvshmem_host", + "-lnvshmem_device", + "-o", + str(temporary_path), + str(source_path), + ] + try: + build = subprocess.run(command, capture_output=True) + if build.returncode != 0: + raise RuntimeError("nvcc failed while compiling CUDA host library\n" + f"command: {shlex.join(command)}\n" + f"stderr:\n{build.stderr.decode()}") + cached_path = cache.put(temporary_path.read_bytes(), output_name, binary=True) + return Path(cached_path), host_cache_key, False + finally: + temporary_path.unlink(missing_ok=True) + + +class CudaHostLibrary: + + def __init__(self, library_path): + self.path = Path(library_path).expanduser().resolve() + self.library = ctypes.CDLL(str(self.path)) + + def __getattr__(self, name): + return getattr(self.library, name) + + +def get_common_host_source() -> Path: + return Path(__file__).resolve().parents[0] / "common-host.cu" + + +def load_host(source, nvshmem_home=None, arch=None, force: bool = False) -> CudaHostLibrary: + path, _, _ = _compile_cuda_host_to_cache(source, nvshmem_home, arch, force=force) + return CudaHostLibrary(path) + + +def load_common_host(source=None, nvshmem_home=None, arch=None, force: bool = False) -> CudaHostLibrary: + return load_host(source or get_common_host_source(), nvshmem_home, arch, force) + + +def init_torch_distributed(): + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(local_rank) + torch.distributed.init_process_group( + backend="cpu:gloo,cuda:nccl", + rank=rank, + world_size=world_size, + device_id=torch.device("cuda", local_rank), + timeout=datetime.timedelta(seconds=1800), + ) + group = torch.distributed.new_group(ranks=list(range(world_size)), backend="nccl") + torch.distributed.barrier(group=group) + return group + + +def init_nvshmem_by_torch_pg(common, group): + rank = group.rank() + world_size = group.size() + uid_size = 1024 + + if rank == 0: + temp_buffer = ctypes.create_string_buffer(uid_size) + result = common.nvshmem_get_unique_id_bytes(temp_buffer, uid_size) + assert result == 0, f"nvshmemx_get_uniqueid failed: {result}" + uid = bytes(temp_buffer.raw) + else: + uid = bytes(uid_size) + + objects = [uid] + torch.distributed.broadcast_object_list( + objects, + src=torch.distributed.get_global_rank(group, 0), + group=group, + ) + uid_buffer = ctypes.create_string_buffer(objects[0], uid_size) + result = common.nvshmem_init_from_torch_distributed( + rank, + world_size, + int(os.environ["LOCAL_RANK"]), + uid_buffer, + uid_size, + ) + assert result == 0, f"NVSHMEM init failed: {result}" + torch.distributed.barrier(group=group) + + +def tensor_from_pointer(pointer, shape, dtype, device): + elements = 1 + for extent in shape: + elements *= extent + storage = torch._C._construct_storage_from_data_pointer( + pointer.value, + device, + elements * dtype.itemsize, + ) + return torch.empty(0, dtype=dtype, device=device).set_(storage).view(shape) + + +def set_signal_cuda_ptr(signal_ptr, signal, stream): + (err, ) = cuda.cuStreamWriteValue64( + stream.cuda_stream, + signal_ptr, + signal, + cuda.CUstreamWriteValue_flags.CU_STREAM_WRITE_VALUE_DEFAULT, + ) + CUDA_CHECK(err) + + +def print_perf( + name: str, + value: float, + group, + rank: int, + world_size: int, + unit: str = "ms", +): + for index in range(world_size): + torch.distributed.barrier(group=group) + if rank == index: + print(f"{name} #{rank}: {value:.3f} {unit}", flush=True) + + +def print_perf_mean( + name: str, + value: float, + group, + rank: int, + world_size: int, + unit: str = "ms", +): + device = torch.device("cuda", int(os.environ["LOCAL_RANK"])) + value_tensor = torch.tensor(float(value), device=device, dtype=torch.float64) + torch.distributed.all_reduce( + value_tensor, + op=torch.distributed.ReduceOp.SUM, + group=group, + ) + mean_value = (value_tensor / world_size).item() + if rank == 0: + print(f"{name} mean: {mean_value:.3f} {unit}", flush=True) diff --git a/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-device.cu b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-device.cu new file mode 100644 index 000000000..f2e4d780c --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-device.cu @@ -0,0 +1,11 @@ +extern "C" __device__ int nvshmem_my_pe(); +extern "C" __device__ int nvshmem_n_pes(); +extern "C" __device__ void nvshmem_int_p(int *dest, int value, int pe); + +extern "C" __device__ __attribute__((always_inline)) void +simple_shift(__attribute__((address_space(1))) int *destination) { + int mype = nvshmem_my_pe(); + int npes = nvshmem_n_pes(); + int peer = (mype + 1) % npes; + nvshmem_int_p(destination, mype, peer); +} diff --git a/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-host.cu b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-host.cu new file mode 100644 index 000000000..610cf7d66 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-host.cu @@ -0,0 +1,47 @@ +#include +#include +#include +#include + +#undef CUDA_CHECK +#define CUDA_CHECK(stmt) \ + do { \ + cudaError_t result = (stmt); \ + if (cudaSuccess != result) { \ + fprintf(stderr, "[%s:%d] cuda failed with %s \n", __FILE__, __LINE__, \ + cudaGetErrorString(result)); \ + exit(-1); \ + } \ + } while (0) + +extern "C" void simple_shift_before_launch(int *mype, int *npes, + int *mype_in_node, int *npes_in_node, + cudaStream_t *stream, int **dst, + int **data_h) { + nvshmem_init(); + *mype = nvshmem_my_pe(); + *npes = nvshmem_n_pes(); + *mype_in_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE); + *npes_in_node = nvshmem_team_n_pes(NVSHMEMX_TEAM_NODE); + + CUDA_CHECK(cudaSetDevice(*mype_in_node)); + CUDA_CHECK(cudaStreamCreate(stream)); + + *dst = (int *)nvshmem_malloc(sizeof(int)); + *data_h = (int *)malloc(sizeof(int)); +} + +extern "C" void simple_shift_after_launch(cudaStream_t stream, void *dst, + void *data_h, int mype, int npes) { + nvshmemx_barrier_all_on_stream(stream); + cudaMemcpyAsync(data_h, dst, sizeof(int), cudaMemcpyDeviceToHost, stream); + cudaStreamSynchronize(stream); + + int *data = (int *)data_h; + printf("%d: received message %d\n", mype, data[0]); + + nvshmem_free(dst); + free(data_h); + + nvshmem_finalize(); +} diff --git a/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py new file mode 100644 index 000000000..e9a79a82d --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py @@ -0,0 +1,73 @@ +import ctypes +import torch +import triton +import triton.experimental.tle.language.raw as tle_raw + +from pathlib import Path +from triton.experimental.tle.raw import dialect + +from triton.experimental.tle.raw.nvshmem.utils import ( + load_host, + tensor_from_pointer, +) + + +@dialect( + name="cuda", + compiler="clang", + file=(Path(__file__).parent / "simple-shift-device.cu"), + extern_func_name="simple_shift", +) +def simple_shift(*args, **kwargs): + ... + + +@triton.jit +def simple_shift_kernel(destination_ptr, ): + tle_raw.call(simple_shift, [destination_ptr]) + + +def simpe_shift(): + host_source = Path(__file__).with_name("simple-shift-host.cu") + host_lib = load_host(source=host_source) + + mype = ctypes.c_int() + npes = ctypes.c_int() + mype_in_node = ctypes.c_int() + npes_in_node = ctypes.c_int() + stream_ptr = ctypes.c_void_p() + destination_ptr = ctypes.c_void_p() + host_data_ptr = ctypes.c_void_p() + host_lib.simple_shift_before_launch( + ctypes.byref(mype), + ctypes.byref(npes), + ctypes.byref(mype_in_node), + ctypes.byref(npes_in_node), + ctypes.byref(stream_ptr), + ctypes.byref(destination_ptr), + ctypes.byref(host_data_ptr), + ) + + device = triton.runtime.driver.active.get_active_torch_device() + destination = tensor_from_pointer( + destination_ptr, + (1, ), + torch.int32, + device, + ) + stream = torch.cuda.ExternalStream(stream_ptr.value, device=device) + + with torch.cuda.stream(stream): + simple_shift_kernel[(1, )](destination) + + host_lib.simple_shift_after_launch( + stream_ptr, + destination_ptr, + host_data_ptr, + mype_in_node.value, + npes_in_node.value, + ) + + +if __name__ == "__main__": + simpe_shift() diff --git a/python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm-device.cu b/python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm-device.cu new file mode 100644 index 000000000..15eafbb33 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm-device.cu @@ -0,0 +1,31 @@ +#include +#include +#include + +extern "C" __device__ uint64_t nvshmem_signal_wait_until(uint64_t *sig_addr, + int cmp, + uint64_t cmp_value); + +enum { + NVSHMEM_CMP_GE = 5, + NVSHMEM_SIGNAL_SET = 9, +}; + +extern "C" __device__ __attribute__((always_inline)) void +ag_mark_local_ready(__attribute__((address_space(1))) uint64_t *ready, + int rank) { + if (threadIdx.x == 0) { + __threadfence_system(); + ready[(size_t)rank] = 1; + } + __syncthreads(); +} + +extern "C" __device__ __attribute__((always_inline)) void +ag_wait_ready(__attribute__((address_space(1))) uint64_t *ready, + int signal_index) { + if (threadIdx.x == 0) { + nvshmem_signal_wait_until(ready + signal_index, NVSHMEM_CMP_GE, 1); + } + __syncthreads(); +} diff --git a/python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm-host.cu b/python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm-host.cu new file mode 100644 index 000000000..72d0898f9 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm-host.cu @@ -0,0 +1,63 @@ +#include +#include +#include +#include +#include +#include +#include + +#define CUDA_CHECK(stmt) \ + do { \ + cudaError_t result = (stmt); \ + if (result != cudaSuccess) { \ + fprintf(stderr, "[%s:%d] CUDA failed: %s\n", __FILE__, __LINE__, \ + cudaGetErrorString(result)); \ + return -1; \ + } \ + } while (0) + +extern "C" int ag_gemm_workspace_create(int elements_per_rank, void **workspace, + uint64_t **ready, int *mype, int *npes, + int *mype_in_node, int *npes_in_node) { + if (elements_per_rank <= 0 || workspace == nullptr || ready == nullptr) { + return -1; + } + + *mype = nvshmem_my_pe(); + *npes = nvshmem_n_pes(); + *mype_in_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE); + *npes_in_node = nvshmem_team_n_pes(NVSHMEMX_TEAM_NODE); + CUDA_CHECK(cudaSetDevice(*mype_in_node)); + + size_t workspace_bytes = (size_t)(*npes) * elements_per_rank * sizeof(__half); + *workspace = nvshmem_malloc(workspace_bytes); + *ready = (uint64_t *)nvshmem_calloc((size_t)(*npes), sizeof(uint64_t)); + if (*workspace == nullptr || *ready == nullptr) { + if (*ready != nullptr) { + nvshmem_free(*ready); + } + if (*workspace != nullptr) { + nvshmem_free(*workspace); + } + return -2; + } + + nvshmem_barrier_all(); + return 0; +} + +extern "C" int ag_gemm_workspace_destroy(void *workspace, void *ready) { + CUDA_CHECK(cudaDeviceSynchronize()); + nvshmem_barrier_all(); + nvshmem_free(ready); + nvshmem_free(workspace); + return 0; +} + +extern "C" void *ag_gemm_peer_workspace_ptr(void *workspace, int peer) { + return nvshmem_ptr(workspace, peer); +} + +extern "C" uint64_t *ag_gemm_peer_ready_ptr(uint64_t *ready, int peer) { + return (uint64_t *)nvshmem_ptr(ready, peer); +} diff --git a/python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm.py b/python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm.py new file mode 100644 index 000000000..108e36275 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm.py @@ -0,0 +1,555 @@ +import argparse +import ctypes +import os +from pathlib import Path +from typing import Optional + +import torch +import triton +import triton.language as tl +import triton.experimental.tle.language.raw as tle_raw +from triton.experimental.tle.raw import dialect + +from triton.experimental.tle.raw.nvshmem.utils import ( + print_perf_mean, + load_common_host, + load_host, + init_torch_distributed, + init_nvshmem_by_torch_pg, + tensor_from_pointer, + set_signal_cuda_ptr, + cudart, +) + + +def _device_dialect(function_name): + return dialect( + name="cuda", + compiler="clang", + file=Path(__file__).parent / "ag-gemm-device.cu", + extern_func_name=function_name, + ) + + +@_device_dialect("ag_mark_local_ready") +def mark_local_ready(*args, **kwargs): + ... + + +@_device_dialect("ag_wait_ready") +def wait_ready(*args, **kwargs): + ... + + +@triton.jit(do_not_specialize=["rank"]) +def set_local_ready(ready, rank): + tle_raw.call(mark_local_ready, [ready, rank]) + + +def cp_engine_producer_all_gather_put(host, local_tensor, ag_buffer, signal_buffer, M_per_rank, N, signal_target, rank, + local_world_size, world_size, intranode_ag_stream): + local_rank = rank % local_world_size + nbytes = M_per_rank * N * local_tensor.element_size() + for i in range(1, local_world_size): + segment = rank * M_per_rank * N + local_dst_rank = (local_rank + local_world_size - i) % local_world_size + peer_workspace = host.ag_gemm_peer_workspace_ptr( + ctypes.c_void_p(ag_buffer.data_ptr()), + local_dst_rank, + ) + peer_ready = host.ag_gemm_peer_ready_ptr( + ctypes.c_void_p(signal_buffer.data_ptr()), + local_dst_rank, + ) + assert peer_workspace, f"failed to get peer workspace pointer for PE {local_dst_rank}" + assert peer_ready, f"failed to get peer ready pointer for PE {local_dst_rank}" + + src_ptr = ag_buffer.data_ptr() + segment * local_tensor.element_size() + dst_ptr = peer_workspace + segment * local_tensor.element_size() + (err, ) = cudart.cudaMemcpyAsync( + dst_ptr, + src_ptr, + nbytes, + cudart.cudaMemcpyKind.cudaMemcpyDefault, + intranode_ag_stream.cuda_stream, + ) + set_signal_cuda_ptr( + peer_ready + rank * ctypes.sizeof(ctypes.c_uint64), + signal_target, + intranode_ag_stream, + ) + + +@triton.jit +def ag_gemm_consumer( + a_ptr, + b_ptr, + c_ptr, + ready, + M, + N, + K, + RANK: tl.constexpr, + NUM_RANKS: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + LOCAL_WORLD_SIZE: tl.constexpr, + NUM_SMS: tl.constexpr, + EPILOGUE_SUBTILE: tl.constexpr, +): + dtype = c_ptr.dtype.element_ty + start_pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + k_tiles = tl.cdiv(K, BLOCK_SIZE_K) + num_tiles = num_pid_m * num_pid_n + node_id = RANK // LOCAL_WORLD_SIZE + nnodes = NUM_RANKS // LOCAL_WORLD_SIZE + + a_desc = tl.make_tensor_descriptor( + a_ptr, + shape=[M, K], + strides=[K, 1], + block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_K], + ) + b_desc = tl.make_tensor_descriptor( + b_ptr, + shape=[N, K], + strides=[K, 1], + block_shape=[BLOCK_SIZE_N, BLOCK_SIZE_K], + ) + c_desc = tl.make_tensor_descriptor( + c_ptr, + shape=[M, N], + strides=[N, 1], + block_shape=[ + BLOCK_SIZE_M, + BLOCK_SIZE_N if not EPILOGUE_SUBTILE else BLOCK_SIZE_N // 2, + ], + ) + + tiles_per_SM = num_tiles // NUM_SMS + if start_pid < num_tiles % NUM_SMS: + tiles_per_SM += 1 + + tile_id = start_pid - NUM_SMS + ki = -1 + + pid_m = 0 + pid_n = 0 + offs_am = 0 + offs_bn = 0 + + M_per_rank = M // NUM_RANKS + pid_ms_per_rank = tl.cdiv(M_per_rank, BLOCK_SIZE_M) + + num_pid_in_group = GROUP_SIZE_M * num_pid_n + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + for _ in range(0, k_tiles * tiles_per_SM): + ki = tl.where(ki == k_tiles - 1, 0, ki + 1) + if ki == 0: + tile_id += NUM_SMS + group_id = tile_id // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (tile_id % group_size_m) + pid_n = (tile_id % num_pid_in_group) // group_size_m + + # swizzle m + if nnodes == 1: + alpha = 0 + beta = 0 + pid_m = (pid_m + ((((RANK ^ alpha) + beta) % NUM_RANKS) * pid_ms_per_rank)) % num_pid_m + else: + m_rank = pid_m // pid_ms_per_rank + pid_m_intra_rank = pid_m - m_rank * pid_ms_per_rank + m_node_id = m_rank // LOCAL_WORLD_SIZE + m_local_rank = m_rank % LOCAL_WORLD_SIZE + swizzle_m_node_id = (m_node_id + node_id) % nnodes + swizzle_m_local_rank = (m_local_rank + RANK) % LOCAL_WORLD_SIZE + swizzle_m_rank = swizzle_m_node_id * LOCAL_WORLD_SIZE + swizzle_m_local_rank + + pid_m = swizzle_m_rank * pid_ms_per_rank + pid_m_intra_rank + + offs_am = pid_m * BLOCK_SIZE_M + offs_bn = pid_n * BLOCK_SIZE_N + + source_rank = offs_am // M_per_rank + tle_raw.call(wait_ready, [ready, source_rank]) + + offs_k = ki * BLOCK_SIZE_K + a = a_desc.load([offs_am, offs_k]) + b = b_desc.load([offs_bn, offs_k]) + accumulator = tl.dot(a, b.T, accumulator) + + if ki == k_tiles - 1: + if EPILOGUE_SUBTILE: + acc = tl.reshape(accumulator, (BLOCK_SIZE_M, 2, BLOCK_SIZE_N // 2)) + acc = tl.permute(acc, (0, 2, 1)) + acc0, acc1 = tl.split(acc) + c0 = acc0.to(dtype) + c_desc.store([offs_am, offs_bn], c0) + c1 = acc1.to(dtype) + c_desc.store([offs_am, offs_bn + BLOCK_SIZE_N // 2], c1) + else: + c = accumulator.to(dtype) + c_desc.store([offs_am, offs_bn], c) + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + +def ag_gemm_op( + a, + b, + c, + rank, + num_ranks, + workspace, + ready, + comm_stream, + compute_stream, + host, + local_world_size, + block_m=128, + block_n=256, + block_k=64, + group_m=8, + num_warps=8, + num_stages=3, +): + assert a.shape[1] == b.shape[1], "incompatible GEMM dimensions" + assert a.dtype == b.dtype == c.dtype, "incompatible GEMM dtypes" + + m, k = workspace.shape + m_per_rank = m // num_ranks + n_per_rank = b.shape[0] + total_tiles = triton.cdiv(m, block_m) * triton.cdiv(n_per_rank, block_n) + + num_ag_sms = 0 # The copy engine does not occupy SM resources. + num_gemm_sms = torch.cuda.get_device_properties("cuda").multi_processor_count - num_ag_sms + + grid = (min(total_tiles, num_gemm_sms), ) + + current_stream = torch.cuda.current_stream(b.device) + comm_stream.wait_stream(current_stream) + cp_engine_producer_all_gather_put(host=host, local_tensor=a, ag_buffer=workspace, signal_buffer=ready, + M_per_rank=m_per_rank, N=k, signal_target=1, rank=rank, + local_world_size=local_world_size, world_size=num_ranks, + intranode_ag_stream=comm_stream) + + def alloc_fn(size: int, alignment: int, stream: Optional[int]): + return torch.empty(size, device="cuda", dtype=torch.int8) + + triton.set_allocator(alloc_fn) + + compiled = None + with torch.cuda.stream(compute_stream): + compute_stream.wait_stream(current_stream) + compiled = ag_gemm_consumer[grid]( + workspace, + b, + c, + ready, + m, + n_per_rank, + k, + rank, + num_ranks, + block_m, + block_n, + block_k, + group_m, + local_world_size, + NUM_SMS=num_gemm_sms, + EPILOGUE_SUBTILE=False, + num_warps=num_warps, + num_stages=num_stages, + ) + + current_stream.wait_stream(comm_stream) + current_stream.wait_stream(compute_stream) + + return compiled + + +def triton_prepare(a_local, workspace, ready, rank): + ready.zero_() + m_per_rank = a_local.shape[0] + workspace[rank * m_per_rank:(rank + 1) * m_per_rank].copy_(a_local) + set_local_ready[(1, )](ready, rank, num_warps=1) + + +def torch_ag_gemm(group, a_local, b, gathered): + torch.distributed.all_gather_into_tensor(gathered, a_local, group=group) + return torch.matmul(gathered, b.T) + + +def configure_host_library(host): + host.ag_gemm_workspace_create.argtypes = [ + ctypes.c_int, + ctypes.POINTER(ctypes.c_void_p), + ctypes.POINTER(ctypes.c_void_p), + ctypes.POINTER(ctypes.c_int), + ctypes.POINTER(ctypes.c_int), + ctypes.POINTER(ctypes.c_int), + ctypes.POINTER(ctypes.c_int), + ] + host.ag_gemm_workspace_create.restype = ctypes.c_int + host.ag_gemm_workspace_destroy.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + host.ag_gemm_workspace_destroy.restype = ctypes.c_int + + # The return value is a 64-bit address. + # Without declaring it as c_void_p, it will be truncated to 32 bits. + host.ag_gemm_peer_workspace_ptr.argtypes = [ctypes.c_void_p, ctypes.c_int] + host.ag_gemm_peer_workspace_ptr.restype = ctypes.c_void_p + host.ag_gemm_peer_ready_ptr.argtypes = [ctypes.c_void_p, ctypes.c_int] + host.ag_gemm_peer_ready_ptr.restype = ctypes.c_void_p + + +def create_workspace(host, world_size, m_per_rank, k, device): + workspace_ptr = ctypes.c_void_p() + ready_ptr = ctypes.c_void_p() + mype = ctypes.c_int() + npes = ctypes.c_int() + local_pe = ctypes.c_int() + local_npes = ctypes.c_int() + result = host.ag_gemm_workspace_create( + m_per_rank * k, + ctypes.byref(workspace_ptr), + ctypes.byref(ready_ptr), + ctypes.byref(mype), + ctypes.byref(npes), + ctypes.byref(local_pe), + ctypes.byref(local_npes), + ) + assert result == 0, f"workspace allocation failed: {result}" + assert npes.value == world_size + workspace = tensor_from_pointer( + workspace_ptr, + (world_size * m_per_rank, k), + torch.float16, + device, + ) + ready = tensor_from_pointer( + ready_ptr, + (world_size, ), + torch.uint64, + device, + ) + return workspace_ptr, ready_ptr, workspace, ready, mype, local_pe, local_npes + + +def perf(fn, group, warmup, iters, prepare_fn=None): + assert warmup >= 0 and iters > 0 + + def prepare(): + torch.cuda.synchronize() + torch.distributed.barrier(group=group) + + if prepare_fn is not None: + prepare_fn() + + torch.cuda.synchronize() + torch.distributed.barrier(group=group) + + for _ in range(warmup): + prepare() + fn() + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + + start.record() + for _ in range(iters): + prepare() + fn() + end.record() + torch.cuda.synchronize() + + profile = {"total": start.elapsed_time(end) / iters} + return profile + + +def parse_args(): + parser = argparse.ArgumentParser(description="NVSHMEM allgather GEMM") + parser.add_argument("--m-per-rank", type=int, default=1024) + parser.add_argument("--n-per-rank", type=int, default=1024) + parser.add_argument("--k", type=int, default=1024) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--iters", type=int, default=10) + parser.add_argument("--case", choices=("check", "perf"), default="check") + + return parser.parse_args() + + +def main(): + args = parse_args() + rank = int(os.environ.get("RANK", 0)) + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + world_size = int(os.environ.get("WORLD_SIZE", 1)) + local_world_size = int(os.environ["LOCAL_WORLD_SIZE"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + + assert world_size >= 2, "AG-GEMM requires at least two GPUs" + assert world_size == local_world_size, ("This example is designed for single-node testing: " + "WORLD_SIZE must equal LOCAL_WORLD_SIZE") + + if torch.cuda.get_device_capability()[0] < 9: + print("Skip the test because the device is not sm90 or higher") + import sys + sys.exit() + + group = init_torch_distributed() + host_source = Path(__file__).with_name("ag-gemm-host.cu") + host = load_host(host_source) + common = load_common_host() + configure_host_library(host) + init_nvshmem_by_torch_pg(common, group) + + dtype = torch.float16 + a_local = torch.randn((args.m_per_rank, args.k), dtype=dtype, device=device) + b = torch.randn((args.n_per_rank, args.k), dtype=dtype, device=device) + c = torch.empty( + (world_size * args.m_per_rank, args.n_per_rank), + dtype=a_local.dtype, + device=device, + ) + gathered = torch.empty( + (world_size * args.m_per_rank, args.k), + dtype=a_local.dtype, + device=device, + ) + + ( + workspace_ptr, + ready_ptr, + workspace, + ready, + mype, + local_pe, + local_npes, + ) = create_workspace( + host, + world_size, + args.m_per_rank, + args.k, + device, + ) + assert mype.value == rank + assert local_pe.value == local_rank + + comm_stream = torch.cuda.Stream(device=device, priority=-1) + compute_stream = torch.cuda.Stream(device=device, priority=-1) + + def triton_prapare_func(): + triton_prepare(a_local, workspace, ready, rank) + + def triton_func(): + return ag_gemm_op(a_local, b, c, rank, world_size, workspace, ready, comm_stream, compute_stream, host, + local_world_size) + + def torch_func(): + return torch_ag_gemm(group, a_local, b, gathered) + + def run_correctness(): + if rank == 0: + print("[check] start correctness validation", flush=True) + + triton_prapare_func() + triton_func() + result = host.ag_gemm_workspace_destroy(workspace_ptr, ready_ptr) + assert result == 0 + result = common.nvshmem_finalize_from_torch_distributed() + assert result == 0 + torch.cuda.synchronize() + torch.distributed.barrier(group=group) + + golden = torch_func() + torch.cuda.synchronize() + torch.distributed.barrier(group=group) + torch.distributed.destroy_process_group() + + torch.testing.assert_close(c, golden, atol=1e-3, rtol=1e-3) + if rank == 0: + print("[check] Pass!", flush=True) + print(f"configuration: GPUs={world_size}, " + f"A_local=({args.m_per_rank}, {args.k}), " + f"B_local=({args.n_per_rank}, {args.k}), ") + + # For performance testing + # It's best to test them separately + # For example, when testing Triton, comment out the Torch part. + def run_benchmark(): + if rank == 0: + print( + f"[bench] start benchmark: warmup={args.warmup}, " + f"iters={args.iters}", + flush=True, + ) + + triton_profile = perf( + triton_func, + group, + args.warmup, + args.iters, + prepare_fn=triton_prapare_func, + ) + result = host.ag_gemm_workspace_destroy(workspace_ptr, ready_ptr) + assert result == 0 + result = common.nvshmem_finalize_from_torch_distributed() + assert result == 0 + torch.cuda.synchronize() + torch.distributed.barrier(group=group) + + torch_profile = perf( + torch_func, + group, + args.warmup, + args.iters, + ) + torch.cuda.synchronize() + torch.distributed.barrier(group=group) + + print_perf_mean( + "triton ag-gemm", + triton_profile["total"], + group, + rank, + world_size, + ) + print_perf_mean( + "torch ag-gemm", + torch_profile["total"], + group, + rank, + world_size, + ) + print_perf_mean( + "speedup", + torch_profile["total"] / triton_profile["total"], + group, + rank, + world_size, + unit="x", + ) + + if rank == 0: + print(f"configuration: GPUs={world_size}, " + f"A_local=({args.m_per_rank}, {args.k}), " + f"B_local=({args.n_per_rank}, {args.k}).") + + torch.distributed.destroy_process_group() + + if args.case == "check": + run_correctness() + if args.case == "perf": + run_benchmark() + + +if __name__ == "__main__": + main() diff --git a/third_party/nvidia/backend/compiler.py b/third_party/nvidia/backend/compiler.py index fe4ecde0c..af0ec4345 100644 --- a/third_party/nvidia/backend/compiler.py +++ b/third_party/nvidia/backend/compiler.py @@ -139,6 +139,11 @@ def __post_init__(self): if not extern_libs.get('libdevice', None): extern_libs['libdevice'] = knobs.nvidia.libdevice_path or str(default_libdir / 'libdevice.10.bc') + # TODO: change it to use @dialect library=nvshmem as the condition for loading libnvshmem_device + nvshmem_home = os.getenv("NVSHMEM_HOME") + if nvshmem_home and (not extern_libs.get('libnvshmem_device', None)): + extern_libs['libnvshmem_device'] = str(Path(nvshmem_home) / 'lib' / 'libnvshmem_device.bc') + # flagtree tle distributed: Add distributed bitcode library(libflagcx_device.bc) if distributed features are enabled. extern_libs.update(Distributed().get_extern_libs()) object.__setattr__(self, 'extern_libs', tuple(extern_libs.items())) From ac3948a999caac5813f1256c919ed148f9c5b66a Mon Sep 17 00:00:00 2001 From: zyuli Date: Thu, 16 Jul 2026 09:02:56 +0000 Subject: [PATCH 03/16] [KMCompiler][TLERaw] Unify launcher to torchrun --- .../nvshmem/01-simple-shift/simple-shift-host.cu | 3 --- .../raw/nvshmem/01-simple-shift/simple-shift.py | 14 +++++++++++++- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-host.cu b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-host.cu index 610cf7d66..da769f89c 100644 --- a/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-host.cu +++ b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-host.cu @@ -18,7 +18,6 @@ extern "C" void simple_shift_before_launch(int *mype, int *npes, int *mype_in_node, int *npes_in_node, cudaStream_t *stream, int **dst, int **data_h) { - nvshmem_init(); *mype = nvshmem_my_pe(); *npes = nvshmem_n_pes(); *mype_in_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE); @@ -42,6 +41,4 @@ extern "C" void simple_shift_after_launch(cudaStream_t stream, void *dst, nvshmem_free(dst); free(data_h); - - nvshmem_finalize(); } diff --git a/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py index e9a79a82d..059483e15 100644 --- a/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py +++ b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py @@ -1,3 +1,4 @@ +import os import ctypes import torch import triton @@ -7,7 +8,10 @@ from triton.experimental.tle.raw import dialect from triton.experimental.tle.raw.nvshmem.utils import ( + load_common_host, load_host, + init_torch_distributed, + init_nvshmem_by_torch_pg, tensor_from_pointer, ) @@ -28,8 +32,14 @@ def simple_shift_kernel(destination_ptr, ): def simpe_shift(): + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + device = torch.device("cuda", local_rank) + + group = init_torch_distributed() host_source = Path(__file__).with_name("simple-shift-host.cu") host_lib = load_host(source=host_source) + common = load_common_host() + init_nvshmem_by_torch_pg(common, group) mype = ctypes.c_int() npes = ctypes.c_int() @@ -48,7 +58,6 @@ def simpe_shift(): ctypes.byref(host_data_ptr), ) - device = triton.runtime.driver.active.get_active_torch_device() destination = tensor_from_pointer( destination_ptr, (1, ), @@ -68,6 +77,9 @@ def simpe_shift(): npes_in_node.value, ) + common.nvshmem_finalize_from_torch_distributed() + torch.distributed.destroy_process_group() + if __name__ == "__main__": simpe_shift() From cb159433eb6a40376df4e3eb37a3085c9b368859 Mon Sep 17 00:00:00 2001 From: i3wanna2 <15910307812@163.com> Date: Thu, 16 Jul 2026 16:57:25 +0000 Subject: [PATCH 04/16] [KMCompiler][TLERaw] Lazy-load nvshmem/cuda imports and harden dialect registry Avoid pulling cuda-python/nvshmem on import dialect or ordinary CUDA compile paths (fixes nvidia unit/cuda and enflame collect failures). Register cuda backend with try/except like other dialects, and align tops_mlir with extern_func_name API. --- .../experimental/tle/raw/cuda/runtime.py | 3 +- .../experimental/tle/raw/nvshmem/utils.py | 46 ++++++++++++++----- python/triton/experimental/tle/raw/runtime.py | 19 ++++++-- .../experimental/tle/raw/tops/mlir_runtime.py | 6 ++- 4 files changed, 55 insertions(+), 19 deletions(-) diff --git a/python/triton/experimental/tle/raw/cuda/runtime.py b/python/triton/experimental/tle/raw/cuda/runtime.py index 6f243eb13..2c686a241 100644 --- a/python/triton/experimental/tle/raw/cuda/runtime.py +++ b/python/triton/experimental/tle/raw/cuda/runtime.py @@ -14,7 +14,6 @@ from triton._C.libtriton import llvm # pyright: ignore[reportMissingImports] from triton._C.libtriton.tle.llvm import parse_llvm_ir # pyright: ignore[reportMissingImports] from triton.experimental.tle.raw.source_store import register_source -from triton.experimental.tle.raw.nvshmem.utils import get_nvshmem_home # TODO: We use cli tools to compile CUDA code temporarily, and plan to replace it with LLVM components Python bindings in the future. CLANG = os.getenv("CLANG", "clang") @@ -57,6 +56,8 @@ def _get_nvshmemx_cumodule_init(): if _nvshmemx_cumodule_init is not None: return _nvshmemx_cumodule_init + from triton.experimental.tle.raw.nvshmem.utils import get_nvshmem_home + nvshmem_home = get_nvshmem_home() library = ctypes.CDLL(str(Path(nvshmem_home) / "lib" / "libnvshmem_host.so")) fn = library.nvshmemx_cumodule_init diff --git a/python/triton/experimental/tle/raw/nvshmem/utils.py b/python/triton/experimental/tle/raw/nvshmem/utils.py index e6dd25c9d..97ec264c0 100644 --- a/python/triton/experimental/tle/raw/nvshmem/utils.py +++ b/python/triton/experimental/tle/raw/nvshmem/utils.py @@ -14,21 +14,42 @@ from triton.runtime.cache import get_cache_manager from triton.experimental.tle.raw.cache_key import compute_tle_raw_host_cache_key -try: - from cuda.bindings import driver as cuda - from cuda.bindings import runtime as cudart -except ImportError: - with warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", - message=r"The cuda\.(cuda|cudart) module is deprecated.*", - category=FutureWarning, - ) - from cuda import cuda as cuda - from cuda import cudart as cudart +_cuda = None +_cudart = None + + +def _get_cuda_modules(): + global _cuda, _cudart + if _cuda is not None and _cudart is not None: + return _cuda, _cudart + try: + from cuda.bindings import driver as cuda + from cuda.bindings import runtime as cudart + except ImportError: + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message=r"The cuda\.(cuda|cudart) module is deprecated.*", + category=FutureWarning, + ) + from cuda import cuda as cuda + from cuda import cudart as cudart + _cuda = cuda + _cudart = cudart + return _cuda, _cudart + + + +def __getattr__(name): + # Preserve `from ...utils import cuda/cudart` after lazy-loading. + if name in ("cuda", "cudart"): + cuda, cudart = _get_cuda_modules() + return cuda if name == "cuda" else cudart + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") def CUDA_CHECK(err): + cuda, cudart = _get_cuda_modules() if isinstance(err, cuda.CUresult): if err != cuda.CUresult.CUDA_SUCCESS: raise RuntimeError(f"Cuda Error: {err}: {cuda.cuGetErrorName(err)}") @@ -213,6 +234,7 @@ def tensor_from_pointer(pointer, shape, dtype, device): def set_signal_cuda_ptr(signal_ptr, signal, stream): + cuda, _ = _get_cuda_modules() (err, ) = cuda.cuStreamWriteValue64( stream.cuda_stream, signal_ptr, diff --git a/python/triton/experimental/tle/raw/runtime.py b/python/triton/experimental/tle/raw/runtime.py index 577b9d813..8f7729c44 100644 --- a/python/triton/experimental/tle/raw/runtime.py +++ b/python/triton/experimental/tle/raw/runtime.py @@ -1,7 +1,12 @@ from .cache_key import bind_tle_raw_source_cache_key -from .cuda import CUDAJITFunction -registry = {"cuda": CUDAJITFunction} +registry = {} + +try: + from .cuda import CUDAJITFunction + registry["cuda"] = CUDAJITFunction +except ImportError: + pass try: from .mlir import MLIRJITFunction @@ -25,9 +30,13 @@ def dialect( ): def decorator(fn): - if name == "mlir" and name not in registry: - from .mlir import MLIRJITFunction - registry[name] = MLIRJITFunction + if name not in registry: + if name == "cuda": + from .cuda import CUDAJITFunction + registry[name] = CUDAJITFunction + elif name == "mlir": + from .mlir import MLIRJITFunction + registry[name] = MLIRJITFunction edsl = registry[name](fn, **kwargs) bind_tle_raw_source_cache_key(edsl, name=name, **kwargs) return edsl diff --git a/python/triton/experimental/tle/raw/tops/mlir_runtime.py b/python/triton/experimental/tle/raw/tops/mlir_runtime.py index 47833d5fb..3c6d5e063 100644 --- a/python/triton/experimental/tle/raw/tops/mlir_runtime.py +++ b/python/triton/experimental/tle/raw/tops/mlir_runtime.py @@ -53,6 +53,8 @@ def __init__(self, fn: Any, pipeline: Optional[List[str]] = None, context: Optio ctx = ir.Context() if context is None else context ctx.allow_unregistered_dialects = True self.context: Final[ir.Context] = ctx + self.extern_func_name: Final[Optional[str]] = kwargs.get("extern_func_name") + self.deferred: Final[bool] = kwargs.get("deferred", False) self.__triton_builtin__: Final[bool] = True def __deepcopy__(self, memo: Dict[int, Any]) -> TOPSMLIRJITFunction: @@ -345,7 +347,8 @@ def _unwrap_gpu_module(text: str) -> str: body = "\n".join(lines[start:end]) return f"module {{\n{body}\n}}\n" - def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint: str = ""): + def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint: str = "", + extern_func_name: str = ""): return builder.create_tle_raw_region_by_llvm_func( llvm, self.region_dialect, @@ -353,6 +356,7 @@ def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint handles, alias_indices, hint, + extern_func_name, ) def make_llvm(self, context=None) -> str: From 459f691ba7afc47992a48bd0d27305f492a8a042 Mon Sep 17 00:00:00 2001 From: i3wanna2 <15910307812@163.com> Date: Thu, 16 Jul 2026 17:13:24 +0000 Subject: [PATCH 05/16] [KMCompiler][TLERaw] Fix extra blank line in nvshmem utils --- python/triton/experimental/tle/raw/nvshmem/utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/triton/experimental/tle/raw/nvshmem/utils.py b/python/triton/experimental/tle/raw/nvshmem/utils.py index 97ec264c0..c2d77b61b 100644 --- a/python/triton/experimental/tle/raw/nvshmem/utils.py +++ b/python/triton/experimental/tle/raw/nvshmem/utils.py @@ -39,7 +39,6 @@ def _get_cuda_modules(): return _cuda, _cudart - def __getattr__(name): # Preserve `from ...utils import cuda/cudart` after lazy-loading. if name in ("cuda", "cudart"): From 63af2490bf589747cd5976c90ef86247595eb397 Mon Sep 17 00:00:00 2001 From: i3wanna2 <15910307812@163.com> Date: Thu, 16 Jul 2026 17:34:53 +0000 Subject: [PATCH 06/16] [KMCompiler][TLERaw] Add extern_func_name to TOPSJITFunction Align tops dialect with the multi-function raw call API so tle_raw.call does not AttributeError on Enflame tops tutorials. --- python/triton/experimental/tle/raw/tops/runtime.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/triton/experimental/tle/raw/tops/runtime.py b/python/triton/experimental/tle/raw/tops/runtime.py index ca62b9b72..9984b54e5 100644 --- a/python/triton/experimental/tle/raw/tops/runtime.py +++ b/python/triton/experimental/tle/raw/tops/runtime.py @@ -52,12 +52,14 @@ def edsl(*args, **kwargs): def __init__(self, fn: Any, file: Optional[Path] = None, arch: Optional[str] = None, extra_flags: Optional[List[str]] = None, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) + super().__init__(*args, **{k: v for k, v in kwargs.items() if k not in ("extern_func_name", "deferred")}) self.fn: Final[Any] = fn self.arch: Final[str] = arch or _get_gcu_arch() self.extra_flags: Final[List[str]] = extra_flags or [] self.region_dialect: Final[str] = "tops" self.arg_dialect: Final[str] = "llvm" + self.extern_func_name: Final[Optional[str]] = kwargs.get("extern_func_name") + self.deferred: Final[bool] = kwargs.get("deferred", False) self.__triton_builtin__: Final[bool] = True if file is not None: From e915385b1053af7781ecd190649d1d05f5dbb4c8 Mon Sep 17 00:00:00 2001 From: i3wanna2 <15910307812@163.com> Date: Thu, 16 Jul 2026 17:45:57 +0000 Subject: [PATCH 07/16] [KMCompiler][TLERaw] Use getattr for extern_func_name in raw call Match deferred handling so backends without the attribute still work. Drop redundant extern_func_name fields from tops/tops_mlir; keep the create_region_by_llvm parameter for API compatibility. --- python/triton/experimental/tle/language/raw/core.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/triton/experimental/tle/language/raw/core.py b/python/triton/experimental/tle/language/raw/core.py index 722feb47a..8eb97870e 100644 --- a/python/triton/experimental/tle/language/raw/core.py +++ b/python/triton/experimental/tle/language/raw/core.py @@ -52,10 +52,10 @@ def _tle_raw_call(func, args, *, output_indices, hint, smem, _semantic): else: context = _semantic.builder.get_context() llvm = func.make_llvm(context) - alias_indices = _resolve_alias_indices(func, llvm, handles, output_indices, func.extern_func_name or "", - _semantic) + extern_func_name = getattr(func, "extern_func_name", None) or "" + alias_indices = _resolve_alias_indices(func, llvm, handles, output_indices, extern_func_name, _semantic) dsl_region_op = func.create_region_by_llvm(_semantic.builder, llvm, handles, alias_indices, hint, - func.extern_func_name or "") + extern_func_name) return _wrap_results(args, alias_indices, dsl_region_op, smem=smem) From 5af081a5b45a617bd5a4443511ec343ee0503cac Mon Sep 17 00:00:00 2001 From: i3wanna2 <15910307812@163.com> Date: Thu, 16 Jul 2026 18:25:29 +0000 Subject: [PATCH 08/16] [KMCompiler][TLERaw] Extract RawJITFunction base for dialect backends Share extern_func_name/deferred and default create_region_by_llvm so cuda/mlir/tops stay aligned without duplicating dialect kwargs. --- .../experimental/tle/raw/cuda/runtime.py | 19 +++---------- .../experimental/tle/raw/mlir/runtime.py | 19 +++---------- python/triton/experimental/tle/raw/runtime.py | 27 +++++++++++++++++++ .../experimental/tle/raw/tops/mlir_runtime.py | 19 +++---------- .../experimental/tle/raw/tops/runtime.py | 19 +++---------- 5 files changed, 43 insertions(+), 60 deletions(-) diff --git a/python/triton/experimental/tle/raw/cuda/runtime.py b/python/triton/experimental/tle/raw/cuda/runtime.py index 2c686a241..bbb52f65a 100644 --- a/python/triton/experimental/tle/raw/cuda/runtime.py +++ b/python/triton/experimental/tle/raw/cuda/runtime.py @@ -14,6 +14,7 @@ from triton._C.libtriton import llvm # pyright: ignore[reportMissingImports] from triton._C.libtriton.tle.llvm import parse_llvm_ir # pyright: ignore[reportMissingImports] from triton.experimental.tle.raw.source_store import register_source +from triton.experimental.tle.raw.runtime import RawJITFunction # TODO: We use cli tools to compile CUDA code temporarily, and plan to replace it with LLVM components Python bindings in the future. CLANG = os.getenv("CLANG", "clang") @@ -86,19 +87,15 @@ def hook(*args, **kwargs): _cumodule_hook_installed = True -class CUDAJITFunction(object): +class CUDAJITFunction(RawJITFunction): def __init__(self, fn: Any, file: Path, *args, **kwargs) -> None: - super().__init__() - self.fn: Final[Any] = fn + super().__init__(fn, **kwargs) self.code: Final[str] = file.read_text() self.region_dialect: Final[str] = "cuda" self.lowered_region_dialect: Final[str] = "llvm" self.arg_dialect: Final[str] = "llvm" self.source_file: Final[str] = str(file) - self.extern_func_name = kwargs.get("extern_func_name", None) - self.deferred: Final[bool] = kwargs.get("deferred", False) - self.__triton_builtin__: Final[bool] = True if "nvshmem" in self.code: _install_cumodule_hook() @@ -117,15 +114,7 @@ def register_pending_source(self, *, hint: str = "") -> str: def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint: str = "", extern_func_name: str = ""): - return builder.create_tle_raw_region_by_llvm_func( - llvm, - self.region_dialect, - self.arg_dialect, - handles, - alias_indices, - hint, - extern_func_name, - ) + return super().create_region_by_llvm(builder, llvm, handles, alias_indices, hint, extern_func_name) def create_region_deferred(self, builder, source_id: str, handles, alias_indices, hint: str = ""): return builder.create_tle_raw_region_deferred( diff --git a/python/triton/experimental/tle/raw/mlir/runtime.py b/python/triton/experimental/tle/raw/mlir/runtime.py index 743ef5511..7065ff67c 100644 --- a/python/triton/experimental/tle/raw/mlir/runtime.py +++ b/python/triton/experimental/tle/raw/mlir/runtime.py @@ -10,16 +10,16 @@ from .codegen import MLIRCodeGenerator from triton.experimental.tle.raw.source_store import register_source +from triton.experimental.tle.raw.runtime import RawJITFunction _pending_jit_fn_key = "mlir_jit_fn" -class MLIRJITFunction(object): +class MLIRJITFunction(RawJITFunction): def __init__(self, fn: Any, pipeline: Optional[List[str]] = None, context: Optional[ir.Context] = None, *args, **kwargs) -> None: - super().__init__(*args, **{k: v for k, v in kwargs.items() if k not in ("extern_func_name", "deferred")}) - self.fn: Final[Any] = fn + super().__init__(fn, **kwargs) self.pipeline: Final[List[str]] = ([*pipeline] if pipeline is not None else [ "convert-scf-to-cf", "finalize-memref-to-llvm", @@ -33,9 +33,6 @@ def __init__(self, fn: Any, pipeline: Optional[List[str]] = None, context: Optio self.context: Final[ir.Context] = ir.Context() if context is None else context self.region_dialect: Final[str] = "mlir" self.arg_dialect: Final[str] = "llvm" - self.extern_func_name: Final[Optional[str]] = kwargs.get("extern_func_name") - self.deferred: Final[bool] = kwargs.get("deferred", False) - self.__triton_builtin__: Final[bool] = True def __deepcopy__(self, memo: Dict[int, Any]) -> MLIRJITFunction: return self.__class__(copy.deepcopy(self.fn, memo), copy.deepcopy(self.pipeline, memo), self.context) @@ -101,15 +98,7 @@ def create_region_deferred(self, builder, source_id: str, handles, alias_indices def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint: str = "", extern_func_name: str = ""): - return builder.create_tle_raw_region_by_llvm_func( - llvm, - self.region_dialect, - self.arg_dialect, - handles, - alias_indices, - hint, - extern_func_name, - ) + return super().create_region_by_llvm(builder, llvm, handles, alias_indices, hint, extern_func_name) @cached_property def src(self) -> str: diff --git a/python/triton/experimental/tle/raw/runtime.py b/python/triton/experimental/tle/raw/runtime.py index 8f7729c44..7326cd993 100644 --- a/python/triton/experimental/tle/raw/runtime.py +++ b/python/triton/experimental/tle/raw/runtime.py @@ -1,5 +1,32 @@ +from __future__ import annotations + +from typing import Any + from .cache_key import bind_tle_raw_source_cache_key + +class RawJITFunction: + """Shared @dialect state and default LLVM region materialization.""" + + def __init__(self, fn: Any, **kwargs) -> None: + self.fn = fn + self.extern_func_name = kwargs.get("extern_func_name", "") + self.deferred = kwargs.get("deferred", False) + self.__triton_builtin__ = True + + def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint: str = "", + extern_func_name: str = ""): + return builder.create_tle_raw_region_by_llvm_func( + llvm, + self.region_dialect, + self.arg_dialect, + handles, + alias_indices, + hint, + extern_func_name, + ) + + registry = {} try: diff --git a/python/triton/experimental/tle/raw/tops/mlir_runtime.py b/python/triton/experimental/tle/raw/tops/mlir_runtime.py index 3c6d5e063..c74c70339 100644 --- a/python/triton/experimental/tle/raw/tops/mlir_runtime.py +++ b/python/triton/experimental/tle/raw/tops/mlir_runtime.py @@ -13,11 +13,12 @@ from mlir.passmanager import PassManager from ..mlir.codegen import MLIRCodeGenerator +from triton.experimental.tle.raw.runtime import RawJITFunction _GCU_COMPILER_OPT = "/opt/triton_gcu/bin/gcu-compiler-opt" -class TOPSMLIRJITFunction(object): +class TOPSMLIRJITFunction(RawJITFunction): """TLE-Raw dialect for TOPS MLIR EDSL: writes MLIR that lowers to GCU-compatible LLVM IR. Usage: @@ -32,8 +33,7 @@ def edsl(x: Input[L["!llvm.ptr<1>"]], ...): def __init__(self, fn: Any, pipeline: Optional[List[str]] = None, context: Optional[ir.Context] = None, arch: str = "gcu400", use_gcu_opt: bool = True, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - self.fn: Final[Any] = fn + super().__init__(fn, **kwargs) self.arch: Final[str] = arch self.use_gcu_opt: Final[bool] = use_gcu_opt self.region_dialect: Final[str] = "tops" @@ -53,9 +53,6 @@ def __init__(self, fn: Any, pipeline: Optional[List[str]] = None, context: Optio ctx = ir.Context() if context is None else context ctx.allow_unregistered_dialects = True self.context: Final[ir.Context] = ctx - self.extern_func_name: Final[Optional[str]] = kwargs.get("extern_func_name") - self.deferred: Final[bool] = kwargs.get("deferred", False) - self.__triton_builtin__: Final[bool] = True def __deepcopy__(self, memo: Dict[int, Any]) -> TOPSMLIRJITFunction: return self.__class__( @@ -349,15 +346,7 @@ def _unwrap_gpu_module(text: str) -> str: def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint: str = "", extern_func_name: str = ""): - return builder.create_tle_raw_region_by_llvm_func( - llvm, - self.region_dialect, - self.arg_dialect, - handles, - alias_indices, - hint, - extern_func_name, - ) + return super().create_region_by_llvm(builder, llvm, handles, alias_indices, hint, extern_func_name) def make_llvm(self, context=None) -> str: if self.use_gcu_opt and os.path.isfile(_GCU_COMPILER_OPT): diff --git a/python/triton/experimental/tle/raw/tops/runtime.py b/python/triton/experimental/tle/raw/tops/runtime.py index 9984b54e5..d0627da3e 100644 --- a/python/triton/experimental/tle/raw/tops/runtime.py +++ b/python/triton/experimental/tle/raw/tops/runtime.py @@ -8,6 +8,7 @@ from triton._C.libtriton import llvm from triton._C.libtriton.tle.llvm import parse_llvm_ir from triton.backends.enflame.gcu_intrinsics import rewrite_intrinsics_to_placeholders +from triton.experimental.tle.raw.runtime import RawJITFunction def _find_tops_include_dir() -> str: @@ -41,7 +42,7 @@ def _get_gcu_arch() -> str: return os.getenv("GCU_ARCH", "gcu400") -class TOPSJITFunction(object): +class TOPSJITFunction(RawJITFunction): """TLE-Raw dialect for TOPS C++ (.tops) files compiled via topscc. Usage: @@ -52,15 +53,11 @@ def edsl(*args, **kwargs): def __init__(self, fn: Any, file: Optional[Path] = None, arch: Optional[str] = None, extra_flags: Optional[List[str]] = None, *args, **kwargs) -> None: - super().__init__(*args, **{k: v for k, v in kwargs.items() if k not in ("extern_func_name", "deferred")}) - self.fn: Final[Any] = fn + super().__init__(fn, **kwargs) self.arch: Final[str] = arch or _get_gcu_arch() self.extra_flags: Final[List[str]] = extra_flags or [] self.region_dialect: Final[str] = "tops" self.arg_dialect: Final[str] = "llvm" - self.extern_func_name: Final[Optional[str]] = kwargs.get("extern_func_name") - self.deferred: Final[bool] = kwargs.get("deferred", False) - self.__triton_builtin__: Final[bool] = True if file is not None: self.code: Final[str] = Path(file).read_text() @@ -177,15 +174,7 @@ def _rewrite_gcu_intrinsics(llvm_ir: str) -> str: def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint: str = "", extern_func_name: str = ""): - return builder.create_tle_raw_region_by_llvm_func( - llvm, - self.region_dialect, - self.arg_dialect, - handles, - alias_indices, - hint, - extern_func_name, - ) + return super().create_region_by_llvm(builder, llvm, handles, alias_indices, hint, extern_func_name) def make_llvm(self, mlir_context) -> str: llvm_ir_text = self._compile_tops_to_llvm_ir() From c5f67bd5c234267b95525a4ed50569e1b2697f5b Mon Sep 17 00:00:00 2001 From: zyuli Date: Fri, 17 Jul 2026 02:46:51 +0000 Subject: [PATCH 09/16] [KMCompiler][TLERaw] Add tle cuda and nvshmem test --- .github/workflows/nv3.6-build-and-test.yml | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/.github/workflows/nv3.6-build-and-test.yml b/.github/workflows/nv3.6-build-and-test.yml index 7babf1f48..1aebc5d45 100644 --- a/.github/workflows/nv3.6-build-and-test.yml +++ b/.github/workflows/nv3.6-build-and-test.yml @@ -114,7 +114,9 @@ jobs: python3 python/tutorials/tle/raw/mlir/05-topk.py --kernel tle python3 python/tutorials/tle/raw/mlir/06-test-vassert.py ## flagtree tle cuda - # TODO: These tests are currently skipped because the CLANG environment variable cannot be set. - # python3 python/tutorials/tle/raw/cuda/01-vector-add.py - # python3 python/tutorials/tle/raw/cuda/02-fused-softmax.py - # python3 python/tutorials/tle/raw/cuda/03-matrix-multiplication.py + python3 python/tutorials/tle/raw/cuda/01-vector-add.py + python3 python/tutorials/tle/raw/cuda/02-fused-softmax.py + python3 python/tutorials/tle/raw/cuda/03-matrix-multiplication.py + ## flagtree tle nvshmem + torchrun --nproc_per_node=2 python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py + torchrun --nproc_per_node=2 python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm.py From 289ef3df5a8702eb98d725b810d15792ef28d1e8 Mon Sep 17 00:00:00 2001 From: zyuli Date: Fri, 17 Jul 2026 02:57:13 +0000 Subject: [PATCH 10/16] [KMCompiler][TLERaw] Replace os.getenv with knobs --- python/triton/knobs.py | 2 ++ third_party/nvidia/backend/compiler.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/python/triton/knobs.py b/python/triton/knobs.py index 1ddd4020e..2611beb8c 100644 --- a/python/triton/knobs.py +++ b/python/triton/knobs.py @@ -501,6 +501,8 @@ class nvidia_knobs(base_knobs): libdevice_path: env_opt_str = env_opt_str("TRITON_LIBDEVICE_PATH") libcuda_path: env_opt_str = env_opt_str("TRITON_LIBCUDA_PATH") + nvshmem_home: env_opt_str = env_opt_str("NVSHMEM_HOME") + class amd_knobs(base_knobs): use_buffer_ops: env_bool = env_bool("AMDGCN_USE_BUFFER_OPS", True) diff --git a/third_party/nvidia/backend/compiler.py b/third_party/nvidia/backend/compiler.py index af0ec4345..e0c5ca03e 100644 --- a/third_party/nvidia/backend/compiler.py +++ b/third_party/nvidia/backend/compiler.py @@ -140,7 +140,7 @@ def __post_init__(self): extern_libs['libdevice'] = knobs.nvidia.libdevice_path or str(default_libdir / 'libdevice.10.bc') # TODO: change it to use @dialect library=nvshmem as the condition for loading libnvshmem_device - nvshmem_home = os.getenv("NVSHMEM_HOME") + nvshmem_home = knobs.nvidia.nvshmem_home if nvshmem_home and (not extern_libs.get('libnvshmem_device', None)): extern_libs['libnvshmem_device'] = str(Path(nvshmem_home) / 'lib' / 'libnvshmem_device.bc') From 51fa18ca91dc2ca7259d1bd30f9c1ea210bb92cd Mon Sep 17 00:00:00 2001 From: i3wanna2 <15910307812@163.com> Date: Fri, 17 Jul 2026 03:25:01 +0000 Subject: [PATCH 11/16] [KMCompiler][TLERaw] Auto-resolve clang and NVSHMEM for raw CUDA CI Discover clang>=20 and pip NVSHMEM when unset, wire knobs for optional overrides, and extend nv3.6 workflow with cuda/nvshmem tutorials including smem. --- .github/workflows/nv3.6-build-and-test.yml | 1 + .../experimental/tle/raw/cuda/runtime.py | 157 +++++++++++++++--- .../experimental/tle/raw/nvshmem/utils.py | 40 ++++- python/triton/knobs.py | 3 + third_party/nvidia/backend/compiler.py | 13 +- 5 files changed, 184 insertions(+), 30 deletions(-) diff --git a/.github/workflows/nv3.6-build-and-test.yml b/.github/workflows/nv3.6-build-and-test.yml index 1aebc5d45..52b8d4b24 100644 --- a/.github/workflows/nv3.6-build-and-test.yml +++ b/.github/workflows/nv3.6-build-and-test.yml @@ -117,6 +117,7 @@ jobs: python3 python/tutorials/tle/raw/cuda/01-vector-add.py python3 python/tutorials/tle/raw/cuda/02-fused-softmax.py python3 python/tutorials/tle/raw/cuda/03-matrix-multiplication.py + python3 python/tutorials/tle/raw/cuda/03-matrix-multiplication-smem.py ## flagtree tle nvshmem torchrun --nproc_per_node=2 python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py torchrun --nproc_per_node=2 python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm.py diff --git a/python/triton/experimental/tle/raw/cuda/runtime.py b/python/triton/experimental/tle/raw/cuda/runtime.py index bbb52f65a..e576ab0f7 100644 --- a/python/triton/experimental/tle/raw/cuda/runtime.py +++ b/python/triton/experimental/tle/raw/cuda/runtime.py @@ -1,27 +1,134 @@ from __future__ import annotations + +import ctypes +import functools import os import re import shlex +import shutil import struct -from pathlib import Path import subprocess +from pathlib import Path from typing import Any, Final -import ctypes -from triton import knobs import torch - +from triton import knobs from triton._C.libtriton import llvm # pyright: ignore[reportMissingImports] from triton._C.libtriton.tle.llvm import parse_llvm_ir # pyright: ignore[reportMissingImports] -from triton.experimental.tle.raw.source_store import register_source from triton.experimental.tle.raw.runtime import RawJITFunction +from triton.experimental.tle.raw.source_store import register_source -# TODO: We use cli tools to compile CUDA code temporarily, and plan to replace it with LLVM components Python bindings in the future. -CLANG = os.getenv("CLANG", "clang") -CLANG_FLAGS = shlex.split(os.getenv("CLANG_FLAGS", "")) +# TODO: Temporarily shell out to clang; replace with LLVM Python bindings later. +_MIN_CLANG_MAJOR = 20 -_cumodule_hook_installed = False -_nvshmemx_cumodule_init = None +# --------------------------------------------------------------------------- +# Clang toolchain: knobs override -> discover, always version-check (>= 20) +# --------------------------------------------------------------------------- + + +def _parse_clang_major(clang: str) -> int | None: + try: + out = subprocess.check_output([clang, "--version"], text=True, stderr=subprocess.STDOUT) + except (OSError, subprocess.CalledProcessError): + return None + match = re.search(r"clang version (\d+)\.", out) + if match is None: + match = re.search(r"version (\d+)\.", out) + return int(match.group(1)) if match else None + + +def _clang_meets_min_version(clang: str) -> bool: + major = _parse_clang_major(clang) + return major is not None and major >= _MIN_CLANG_MAJOR + + +def _discover_clang_binaries() -> list[str]: + """Prefer newer versioned binaries, then plain ``clang``.""" + found: list[str] = [] + seen: set[str] = set() + for name in ("clang-22", "clang-21", "clang-20", "clang"): + path = shutil.which(name) + if path is None: + continue + resolved = str(Path(path).resolve()) + if resolved in seen: + continue + seen.add(resolved) + found.append(path) + return found + + +@functools.lru_cache() +def _resolve_clang() -> str: + """Use user CLANG if usable; otherwise discover. Every candidate is version-checked.""" + tried: list[str] = [] + user_clang = knobs.nvidia.tle_raw_clang + if user_clang: + tried.append(user_clang) + if _clang_meets_min_version(user_clang): + return user_clang + + for candidate in _discover_clang_binaries(): + if candidate in tried: + continue + tried.append(candidate) + if _clang_meets_min_version(candidate): + return candidate + + detail = ", ".join(tried) if tried else "" + raise RuntimeError(f"TLE raw CUDA requires clang >= {_MIN_CLANG_MAJOR}. " + f"Tried: {detail}. Install clang-20+ or set CLANG to a suitable binary.") + + +# --------------------------------------------------------------------------- +# Clang compile flags (--cuda-path, includes, optional CLANG_FLAGS) +# --------------------------------------------------------------------------- + + +def _cuda_home() -> Path: + return Path(os.getenv("CUDA_HOME", "/usr/local/cuda")) + + +def _nvidia_backend_include() -> Path | None: + for parent in Path(__file__).resolve().parents: + candidate = parent / "third_party" / "nvidia" / "backend" / "include" + if candidate.is_dir(): + return candidate + return None + + +def _default_clang_flags() -> list[str]: + cuda_home = _cuda_home() + flags = [f"--cuda-path={cuda_home}", f"-I{cuda_home / 'include'}"] + backend_include = _nvidia_backend_include() + if backend_include is not None: + flags.append(f"-I{backend_include}") + try: + from triton.experimental.tle.raw.nvshmem.utils import try_get_nvshmem_home + nvshmem_home = try_get_nvshmem_home() + if nvshmem_home is not None: + flags.append(f"-I{nvshmem_home / 'include'}") + except Exception: + pass + return flags + + +def _clang_flags() -> list[str]: + extra = knobs.nvidia.tle_raw_clang_flags or "" + return [*_default_clang_flags(), *shlex.split(extra)] + + +def _get_cuda_gpu_arch() -> str: + arch = os.getenv("TLE_CUDA_ARCH") + if arch: + return f"--cuda-gpu-arch={arch}" + major, minor = torch.cuda.get_device_capability() + return f"--cuda-gpu-arch=sm_{major}{minor}" + + +# --------------------------------------------------------------------------- +# Sanitize clang LLVM IR for this Triton's parser +# --------------------------------------------------------------------------- def _sanitize_clang_ir(ir: str) -> str: @@ -44,12 +151,12 @@ def _replace_hex_float(match: re.Match[str]) -> str: return re.sub(r"f0x([0-9A-Fa-f]+)", _replace_hex_float, ir) -def _get_cuda_gpu_arch() -> str: - arch = os.getenv("TLE_CUDA_ARCH") - if arch: - return f"--cuda-gpu-arch={arch}" - major, minor = torch.cuda.get_device_capability() - return f"--cuda-gpu-arch=sm_{major}{minor}" +# --------------------------------------------------------------------------- +# NVSHMEM: post-compile cumodule init hook +# --------------------------------------------------------------------------- + +_cumodule_hook_installed = False +_nvshmemx_cumodule_init = None def _get_nvshmemx_cumodule_init(): @@ -57,10 +164,11 @@ def _get_nvshmemx_cumodule_init(): if _nvshmemx_cumodule_init is not None: return _nvshmemx_cumodule_init - from triton.experimental.tle.raw.nvshmem.utils import get_nvshmem_home - - nvshmem_home = get_nvshmem_home() - library = ctypes.CDLL(str(Path(nvshmem_home) / "lib" / "libnvshmem_host.so")) + from triton.experimental.tle.raw.nvshmem.utils import ( + get_nvshmem_home, + resolve_nvshmem_host_library, + ) + library = ctypes.CDLL(str(resolve_nvshmem_host_library(get_nvshmem_home()))) fn = library.nvshmemx_cumodule_init fn.argtypes = [ctypes.c_void_p] fn.restype = ctypes.c_int @@ -87,6 +195,11 @@ def hook(*args, **kwargs): _cumodule_hook_installed = True +# --------------------------------------------------------------------------- +# Dialect runtime +# --------------------------------------------------------------------------- + + class CUDAJITFunction(RawJITFunction): def __init__(self, fn: Any, file: Path, *args, **kwargs) -> None: @@ -129,7 +242,7 @@ def create_region_deferred(self, builder, source_id: str, handles, alias_indices def make_llvm(self, mlir_context) -> str: build = subprocess.run( [ - CLANG, + _resolve_clang(), "-x", "cuda", "--cuda-device-only", @@ -140,7 +253,7 @@ def make_llvm(self, mlir_context) -> str: "-", "-o", "-", - *CLANG_FLAGS, + *_clang_flags(), ], input=self.code.encode(), capture_output=True, diff --git a/python/triton/experimental/tle/raw/nvshmem/utils.py b/python/triton/experimental/tle/raw/nvshmem/utils.py index c2d77b61b..65df76800 100644 --- a/python/triton/experimental/tle/raw/nvshmem/utils.py +++ b/python/triton/experimental/tle/raw/nvshmem/utils.py @@ -61,8 +61,9 @@ def CUDA_CHECK(err): @functools.lru_cache() def get_nvshmem_home() -> Path: - if (nvshmem_home := os.getenv("NVSHMEM_HOME")) is not None: - return Path(nvshmem_home) + from triton import knobs + if knobs.nvidia.nvshmem_home is not None: + return Path(knobs.nvidia.nvshmem_home) try: import nvidia.nvshmem @@ -71,6 +72,32 @@ def get_nvshmem_home() -> Path: raise RuntimeError("Cannot resolve NVSHMEM_HOME: environment variable not set and nvidia.nvshmem not found") +def try_get_nvshmem_home() -> Path | None: + try: + return get_nvshmem_home() + except RuntimeError: + return None + + +def resolve_nvshmem_host_library(nvshmem_home: Path | None = None) -> Path: + """Prefer unversioned .so; fall back to .so.N (pip wheels often omit the symlink).""" + home = Path(nvshmem_home) if nvshmem_home is not None else get_nvshmem_home() + lib_dir = home / "lib" + for name in ("libnvshmem_host.so", "libnvshmem_host.so.3"): + path = lib_dir / name + if path.exists(): + return path + raise RuntimeError(f"Cannot find libnvshmem_host.so[.3] under {lib_dir}") + + +def resolve_nvshmem_device_bitcode(nvshmem_home: Path | None = None) -> Path | None: + home = Path(nvshmem_home) if nvshmem_home is not None else try_get_nvshmem_home() + if home is None: + return None + path = home / "lib" / "libnvshmem_device.bc" + return path if path.is_file() else None + + @functools.lru_cache() def get_nvcc(): return _path_to_binary("nvcc") @@ -122,6 +149,10 @@ def _compile_cuda_host_to_cache( temporary_path = Path(temporary.name) temporary.close() nvcc, _ = get_nvcc() + host_lib = resolve_nvshmem_host_library(nvshmem_home) + device_lib = nvshmem_home / "lib" / "libnvshmem_device.a" + if not device_lib.is_file(): + raise RuntimeError(f"Cannot find libnvshmem_device.a under {nvshmem_home / 'lib'}") command = [ nvcc, "-shared", @@ -130,9 +161,8 @@ def _compile_cuda_host_to_cache( "-rdc=true", f"-arch={arch}", f"-I{nvshmem_home / 'include'}", - f"-L{nvshmem_home / 'lib'}", - "-lnvshmem_host", - "-lnvshmem_device", + str(host_lib), + str(device_lib), "-o", str(temporary_path), str(source_path), diff --git a/python/triton/knobs.py b/python/triton/knobs.py index 2611beb8c..981094f34 100644 --- a/python/triton/knobs.py +++ b/python/triton/knobs.py @@ -502,6 +502,9 @@ class nvidia_knobs(base_knobs): libcuda_path: env_opt_str = env_opt_str("TRITON_LIBCUDA_PATH") nvshmem_home: env_opt_str = env_opt_str("NVSHMEM_HOME") + # TLE-raw CUDA: optional overrides; version / discovery live in tle/raw/cuda/runtime.py + tle_raw_clang: env_opt_str = env_opt_str("CLANG") + tle_raw_clang_flags: env_opt_str = env_opt_str("CLANG_FLAGS") class amd_knobs(base_knobs): diff --git a/third_party/nvidia/backend/compiler.py b/third_party/nvidia/backend/compiler.py index e0c5ca03e..680e2f63e 100644 --- a/third_party/nvidia/backend/compiler.py +++ b/third_party/nvidia/backend/compiler.py @@ -140,9 +140,16 @@ def __post_init__(self): extern_libs['libdevice'] = knobs.nvidia.libdevice_path or str(default_libdir / 'libdevice.10.bc') # TODO: change it to use @dialect library=nvshmem as the condition for loading libnvshmem_device - nvshmem_home = knobs.nvidia.nvshmem_home - if nvshmem_home and (not extern_libs.get('libnvshmem_device', None)): - extern_libs['libnvshmem_device'] = str(Path(nvshmem_home) / 'lib' / 'libnvshmem_device.bc') + # knobs.nvidia.nvshmem_home (NVSHMEM_HOME) is honored inside get_nvshmem_home(); + # otherwise fall back to pip nvidia.nvshmem auto-discovery. + if not extern_libs.get('libnvshmem_device', None): + try: + from triton.experimental.tle.raw.nvshmem.utils import resolve_nvshmem_device_bitcode + nvshmem_bc = resolve_nvshmem_device_bitcode() + except Exception: + nvshmem_bc = None + if nvshmem_bc is not None: + extern_libs['libnvshmem_device'] = str(nvshmem_bc) # flagtree tle distributed: Add distributed bitcode library(libflagcx_device.bc) if distributed features are enabled. extern_libs.update(Distributed().get_extern_libs()) From 5afb5e21f70225d0b9bed817e272e747e7ee3331 Mon Sep 17 00:00:00 2001 From: i3wanna2 <15910307812@163.com> Date: Fri, 17 Jul 2026 06:00:41 +0000 Subject: [PATCH 12/16] [KMCompiler][TLERaw] Address review: knobs comment and CI clang-22 path --- .github/workflows/nv3.6-build-and-test.yml | 1 + python/triton/knobs.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/nv3.6-build-and-test.yml b/.github/workflows/nv3.6-build-and-test.yml index 52b8d4b24..eec697287 100644 --- a/.github/workflows/nv3.6-build-and-test.yml +++ b/.github/workflows/nv3.6-build-and-test.yml @@ -114,6 +114,7 @@ jobs: python3 python/tutorials/tle/raw/mlir/05-topk.py --kernel tle python3 python/tutorials/tle/raw/mlir/06-test-vassert.py ## flagtree tle cuda + export CLANG=/usr/local/lib/python3.12/dist-packages/mlir/llvm_artifact/bin/clang-22 python3 python/tutorials/tle/raw/cuda/01-vector-add.py python3 python/tutorials/tle/raw/cuda/02-fused-softmax.py python3 python/tutorials/tle/raw/cuda/03-matrix-multiplication.py diff --git a/python/triton/knobs.py b/python/triton/knobs.py index 981094f34..8c744abf8 100644 --- a/python/triton/knobs.py +++ b/python/triton/knobs.py @@ -501,8 +501,8 @@ class nvidia_knobs(base_knobs): libdevice_path: env_opt_str = env_opt_str("TRITON_LIBDEVICE_PATH") libcuda_path: env_opt_str = env_opt_str("TRITON_LIBCUDA_PATH") + # flagtree tle-raw cuda nvshmem: optional overrides; discovery lives in tle/raw nvshmem_home: env_opt_str = env_opt_str("NVSHMEM_HOME") - # TLE-raw CUDA: optional overrides; version / discovery live in tle/raw/cuda/runtime.py tle_raw_clang: env_opt_str = env_opt_str("CLANG") tle_raw_clang_flags: env_opt_str = env_opt_str("CLANG_FLAGS") From f4a63892653f00d27d88387100da8b2d2ec39441 Mon Sep 17 00:00:00 2001 From: i3wanna2 <15910307812@163.com> Date: Fri, 17 Jul 2026 06:08:43 +0000 Subject: [PATCH 13/16] [KMCompiler][TLERaw] Shorten flagtree knobs comment --- python/triton/knobs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/triton/knobs.py b/python/triton/knobs.py index 8c744abf8..ef942090a 100644 --- a/python/triton/knobs.py +++ b/python/triton/knobs.py @@ -501,7 +501,7 @@ class nvidia_knobs(base_knobs): libdevice_path: env_opt_str = env_opt_str("TRITON_LIBDEVICE_PATH") libcuda_path: env_opt_str = env_opt_str("TRITON_LIBCUDA_PATH") - # flagtree tle-raw cuda nvshmem: optional overrides; discovery lives in tle/raw + # flagtree tle-raw cuda nvshmem nvshmem_home: env_opt_str = env_opt_str("NVSHMEM_HOME") tle_raw_clang: env_opt_str = env_opt_str("CLANG") tle_raw_clang_flags: env_opt_str = env_opt_str("CLANG_FLAGS") From 81e56784abe47bbf8f0bdf870001ad80ed949d73 Mon Sep 17 00:00:00 2001 From: i3wanna2 <15910307812@163.com> Date: Fri, 17 Jul 2026 09:57:18 +0000 Subject: [PATCH 14/16] [KMCompiler][TLERaw] Link nvshmem host via -L/-l and add CUDA requirements nvcc rejects versioned .so as a positional input; use -L/-l: and rpath. Add optional nvidia-nvshmem-cu12/cu13 requirement files for tle raw. --- python/requirements-nvshmem-cu12.txt | 2 ++ python/requirements-nvshmem-cu13.txt | 2 ++ .../triton/experimental/tle/raw/nvshmem/utils.py | 15 +++++++++++---- 3 files changed, 15 insertions(+), 4 deletions(-) create mode 100644 python/requirements-nvshmem-cu12.txt create mode 100644 python/requirements-nvshmem-cu13.txt diff --git a/python/requirements-nvshmem-cu12.txt b/python/requirements-nvshmem-cu12.txt new file mode 100644 index 000000000..720391de9 --- /dev/null +++ b/python/requirements-nvshmem-cu12.txt @@ -0,0 +1,2 @@ +#flagtree tle raw nvshmem +nvidia-nvshmem-cu12 diff --git a/python/requirements-nvshmem-cu13.txt b/python/requirements-nvshmem-cu13.txt new file mode 100644 index 000000000..147178646 --- /dev/null +++ b/python/requirements-nvshmem-cu13.txt @@ -0,0 +1,2 @@ +#flagtree tle raw nvshmem +nvidia-nvshmem-cu13 diff --git a/python/triton/experimental/tle/raw/nvshmem/utils.py b/python/triton/experimental/tle/raw/nvshmem/utils.py index 65df76800..ef0cf7256 100644 --- a/python/triton/experimental/tle/raw/nvshmem/utils.py +++ b/python/triton/experimental/tle/raw/nvshmem/utils.py @@ -149,10 +149,12 @@ def _compile_cuda_host_to_cache( temporary_path = Path(temporary.name) temporary.close() nvcc, _ = get_nvcc() + lib_dir = nvshmem_home / "lib" host_lib = resolve_nvshmem_host_library(nvshmem_home) - device_lib = nvshmem_home / "lib" / "libnvshmem_device.a" + device_lib = lib_dir / "libnvshmem_device.a" if not device_lib.is_file(): - raise RuntimeError(f"Cannot find libnvshmem_device.a under {nvshmem_home / 'lib'}") + raise RuntimeError(f"Cannot find libnvshmem_device.a under {lib_dir}") + # nvcc cannot take .so as a positional input; use -L/-l (-l: for versioned SONAME). command = [ nvcc, "-shared", @@ -161,8 +163,13 @@ def _compile_cuda_host_to_cache( "-rdc=true", f"-arch={arch}", f"-I{nvshmem_home / 'include'}", - str(host_lib), - str(device_lib), + f"-L{lib_dir}", + f"-l:{host_lib.name}", + "-lnvshmem_device", + "-Xlinker", + "-rpath", + "-Xlinker", + str(lib_dir), "-o", str(temporary_path), str(source_path), From 5e7c49903d286a45b02f8473d408703c2dedf9f4 Mon Sep 17 00:00:00 2001 From: i3wanna2 <15910307812@163.com> Date: Fri, 17 Jul 2026 11:38:45 +0000 Subject: [PATCH 15/16] [KMCompiler][TLERaw] Resolve per-SM NVSHMEM device bitcode NVSHMEM 3.7+ wheels drop unified libnvshmem_device.bc; fall back to libnvshmem_device_sm_XX.bc using the compile arch. --- .../experimental/tle/raw/nvshmem/utils.py | 21 ++++++++++++++++--- third_party/nvidia/backend/compiler.py | 2 +- 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/python/triton/experimental/tle/raw/nvshmem/utils.py b/python/triton/experimental/tle/raw/nvshmem/utils.py index ef0cf7256..8a42faed7 100644 --- a/python/triton/experimental/tle/raw/nvshmem/utils.py +++ b/python/triton/experimental/tle/raw/nvshmem/utils.py @@ -90,12 +90,27 @@ def resolve_nvshmem_host_library(nvshmem_home: Path | None = None) -> Path: raise RuntimeError(f"Cannot find libnvshmem_host.so[.3] under {lib_dir}") -def resolve_nvshmem_device_bitcode(nvshmem_home: Path | None = None) -> Path | None: +def resolve_nvshmem_device_bitcode(nvshmem_home: Path | None = None, arch: str | int | None = None) -> Path | None: + """Resolve device bitcode: unified .bc, else per-SM .bc (NVSHMEM >= ~3.7).""" home = Path(nvshmem_home) if nvshmem_home is not None else try_get_nvshmem_home() if home is None: return None - path = home / "lib" / "libnvshmem_device.bc" - return path if path.is_file() else None + lib_dir = home / "lib" + unified = lib_dir / "libnvshmem_device.bc" + if unified.is_file(): + return unified + # arch: sm90 / sm_90a / 90 -> try exact then next-lower known ships + if isinstance(arch, int): + sm = arch + else: + s = str(arch or "").lower().replace("sm_", "").replace("sm", "").rstrip("a") + sm = int(s) if s.isdigit() else None + sms = (90, 89, 80, 75, 70) + for n in ([sm] + [x for x in sms if sm is not None and x < sm]) if sm is not None else sms: + path = lib_dir / f"libnvshmem_device_sm_{n}.bc" + if path.is_file(): + return path + return None @functools.lru_cache() diff --git a/third_party/nvidia/backend/compiler.py b/third_party/nvidia/backend/compiler.py index 680e2f63e..764526c2e 100644 --- a/third_party/nvidia/backend/compiler.py +++ b/third_party/nvidia/backend/compiler.py @@ -145,7 +145,7 @@ def __post_init__(self): if not extern_libs.get('libnvshmem_device', None): try: from triton.experimental.tle.raw.nvshmem.utils import resolve_nvshmem_device_bitcode - nvshmem_bc = resolve_nvshmem_device_bitcode() + nvshmem_bc = resolve_nvshmem_device_bitcode(arch=self.arch) except Exception: nvshmem_bc = None if nvshmem_bc is not None: From acfda823c32fc1633b6ec2f64ee5d9aa79e4877b Mon Sep 17 00:00:00 2001 From: i3wanna2 <15910307812@163.com> Date: Fri, 17 Jul 2026 16:30:53 +0000 Subject: [PATCH 16/16] [KMCompiler][TLERaw] Gate nvshmem device bc on @dialect library Only link libnvshmem_device when library=nvshmem enables a utils flag, so ordinary CUDA compiles do not inject nvshmem into extern_libs. --- .../experimental/tle/raw/cuda/runtime.py | 5 +++- .../experimental/tle/raw/nvshmem/utils.py | 21 +++++++++++++++++ python/triton/experimental/tle/raw/runtime.py | 1 + .../nvshmem/01-simple-shift/simple-shift.py | 1 + .../raw/nvshmem/02-allgather-gemm/ag-gemm.py | 1 + third_party/nvidia/backend/compiler.py | 23 ++++++++++--------- 6 files changed, 40 insertions(+), 12 deletions(-) diff --git a/python/triton/experimental/tle/raw/cuda/runtime.py b/python/triton/experimental/tle/raw/cuda/runtime.py index e576ab0f7..2f550e323 100644 --- a/python/triton/experimental/tle/raw/cuda/runtime.py +++ b/python/triton/experimental/tle/raw/cuda/runtime.py @@ -210,7 +210,10 @@ def __init__(self, fn: Any, file: Path, *args, **kwargs) -> None: self.arg_dialect: Final[str] = "llvm" self.source_file: Final[str] = str(file) - if "nvshmem" in self.code: + if self.library == "nvshmem": + from triton.experimental.tle.raw.nvshmem.utils import enable_nvshmem_device_bc + enable_nvshmem_device_bc(True) + if self.library == "nvshmem" or "nvshmem" in self.code: _install_cumodule_hook() def register_pending_source(self, *, hint: str = "") -> str: diff --git a/python/triton/experimental/tle/raw/nvshmem/utils.py b/python/triton/experimental/tle/raw/nvshmem/utils.py index 8a42faed7..c6e76fba2 100644 --- a/python/triton/experimental/tle/raw/nvshmem/utils.py +++ b/python/triton/experimental/tle/raw/nvshmem/utils.py @@ -90,6 +90,27 @@ def resolve_nvshmem_host_library(nvshmem_home: Path | None = None) -> Path: raise RuntimeError(f"Cannot find libnvshmem_host.so[.3] under {lib_dir}") +# Set by @dialect(..., library="nvshmem"); CUDA backend links device .bc only when True. +_nvshmem_device_bc_enabled: bool = False + + +def enable_nvshmem_device_bc(enabled: bool = True) -> None: + global _nvshmem_device_bc_enabled + _nvshmem_device_bc_enabled = bool(enabled) + + +def is_nvshmem_device_bc_enabled() -> bool: + return _nvshmem_device_bc_enabled + + +def get_nvshmem_extern_libs(arch: str | int | None = None) -> dict[str, str]: + """Return {libnvshmem_device: path} when enabled; else {}.""" + if not is_nvshmem_device_bc_enabled(): + return {} + bc = resolve_nvshmem_device_bitcode(arch=arch) + return {"libnvshmem_device": str(bc)} if bc is not None else {} + + def resolve_nvshmem_device_bitcode(nvshmem_home: Path | None = None, arch: str | int | None = None) -> Path | None: """Resolve device bitcode: unified .bc, else per-SM .bc (NVSHMEM >= ~3.7).""" home = Path(nvshmem_home) if nvshmem_home is not None else try_get_nvshmem_home() diff --git a/python/triton/experimental/tle/raw/runtime.py b/python/triton/experimental/tle/raw/runtime.py index 7326cd993..e1f5cc7d9 100644 --- a/python/triton/experimental/tle/raw/runtime.py +++ b/python/triton/experimental/tle/raw/runtime.py @@ -12,6 +12,7 @@ def __init__(self, fn: Any, **kwargs) -> None: self.fn = fn self.extern_func_name = kwargs.get("extern_func_name", "") self.deferred = kwargs.get("deferred", False) + self.library = kwargs.get("library", "") or "" self.__triton_builtin__ = True def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint: str = "", diff --git a/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py index 059483e15..925b02807 100644 --- a/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py +++ b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py @@ -18,6 +18,7 @@ @dialect( name="cuda", + library="nvshmem", compiler="clang", file=(Path(__file__).parent / "simple-shift-device.cu"), extern_func_name="simple_shift", diff --git a/python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm.py b/python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm.py index 108e36275..88a792b21 100644 --- a/python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm.py +++ b/python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm.py @@ -25,6 +25,7 @@ def _device_dialect(function_name): return dialect( name="cuda", + library="nvshmem", compiler="clang", file=Path(__file__).parent / "ag-gemm-device.cu", extern_func_name=function_name, diff --git a/third_party/nvidia/backend/compiler.py b/third_party/nvidia/backend/compiler.py index 764526c2e..7070ccffc 100644 --- a/third_party/nvidia/backend/compiler.py +++ b/third_party/nvidia/backend/compiler.py @@ -139,17 +139,18 @@ def __post_init__(self): if not extern_libs.get('libdevice', None): extern_libs['libdevice'] = knobs.nvidia.libdevice_path or str(default_libdir / 'libdevice.10.bc') - # TODO: change it to use @dialect library=nvshmem as the condition for loading libnvshmem_device - # knobs.nvidia.nvshmem_home (NVSHMEM_HOME) is honored inside get_nvshmem_home(); - # otherwise fall back to pip nvidia.nvshmem auto-discovery. - if not extern_libs.get('libnvshmem_device', None): - try: - from triton.experimental.tle.raw.nvshmem.utils import resolve_nvshmem_device_bitcode + # flagtree tle raw: libnvshmem_device when @dialect(library="nvshmem") enabled in utils. + try: + from triton.experimental.tle.raw.nvshmem.utils import ( + is_nvshmem_device_bc_enabled, + resolve_nvshmem_device_bitcode, + ) + if is_nvshmem_device_bc_enabled(): nvshmem_bc = resolve_nvshmem_device_bitcode(arch=self.arch) - except Exception: - nvshmem_bc = None - if nvshmem_bc is not None: - extern_libs['libnvshmem_device'] = str(nvshmem_bc) + if nvshmem_bc is not None: + extern_libs.update({"libnvshmem_device": str(nvshmem_bc)}) + except Exception: + pass # flagtree tle distributed: Add distributed bitcode library(libflagcx_device.bc) if distributed features are enabled. extern_libs.update(Distributed().get_extern_libs()) @@ -440,7 +441,7 @@ def make_llir(self, src, metadata, options, capability): passes.ttgpuir.add_concurrency_sanitizer(pm) passes.ttgpuir.add_allocate_global_scratch_memory(pm) nvidia.passes.ttnvgpuir.add_proxy_fence_insertion(pm, capability) - # Materialize deferred tle_raw sources before inlining DSL regions. + # flagtree tle raw: Materialize deferred tle_raw sources before inlining DSL regions. from .deferred_raw import ( finish_deferred_raw_materialize, deferred_raw_materialize,