From 011ab66dd1bc74b0724a5d38757039391642751f Mon Sep 17 00:00:00 2001 From: zyuli Date: Mon, 11 May 2026 10:09:24 +0800 Subject: [PATCH 1/8] support nvshmem --- clean_cache.sh | 1 + env.sh | 29 + .../PatternTritonGPUOpToLLVM.h | 5 + include/triton/Dialect/Triton/IR/TritonOps.td | 26 + lib/Conversion/TritonGPUToLLVM/CMakeLists.txt | 1 + .../TritonGPUToLLVM/ExternCallOpToLLVM.cpp | 61 ++ .../TritonToTritonGPUPass.cpp | 1 + lib/Dialect/Triton/IR/Ops.cpp | 18 + python/src/ir.cc | 8 + .../experimental/tle/language/raw/__init__.py | 4 +- .../experimental/tle/language/raw/core.py | 5 + .../experimental/tle/raw/cuda/runtime.py | 25 +- python/triton/language/core.py | 88 +++ .../01-simple-shift/simple-shift-device.cu | 8 + .../01-simple-shift/simple-shift-host.cu | 31 + .../nvshmem/01-simple-shift/simple-shift.py | 122 +++ support_nvshmem.patch | 697 ++++++++++++++++++ third_party/nvidia/backend/compiler.py | 45 +- third_party/nvidia/language/cuda/__init__.py | 2 + .../nvidia/language/cuda/libnvshmem_device.py | 31 + .../TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp | 2 + 21 files changed, 1201 insertions(+), 9 deletions(-) create mode 100755 clean_cache.sh create mode 100644 env.sh create mode 100644 lib/Conversion/TritonGPUToLLVM/ExternCallOpToLLVM.cpp 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 support_nvshmem.patch create mode 100644 third_party/nvidia/language/cuda/libnvshmem_device.py diff --git a/clean_cache.sh b/clean_cache.sh new file mode 100755 index 000000000..8cdaf2438 --- /dev/null +++ b/clean_cache.sh @@ -0,0 +1 @@ +rm -rf /home/zyuli/.triton/cache/* diff --git a/env.sh b/env.sh new file mode 100644 index 000000000..1603032d6 --- /dev/null +++ b/env.sh @@ -0,0 +1,29 @@ +# libz.so +# ln -s /usr/lib/x86_64-linux-gnu/libz.so.1 /home/zyuli/local/lib/libz.so +export LIBRARY_PATH=/home/zyuli/local/lib:$LIBRARY_PATH +export LD_LIBRARY_PATH=/home/zyuli/local/lib:$LD_LIBRARY_PATH + +# highest priority +# set clang include path +export CPLUS_INCLUDE_PATH=/usr/include/c++/11:/usr/include/x86_64-linux-gnu/c++/11:$CPLUS_INCLUDE_PATH +export C_INCLUDE_PATH=/usr/include/x86_64-linux-gnu:$C_INCLUDE_PATH + +# cmake reads and stores into CMAKE_CXX_FLAGS +export CXXFLAGS="-Wno-gcc-install-dir-libstdcxx --gcc-install-dir=/usr/lib/gcc/x86_64-linux-gnu/11" + +# clang and lld +export PATH=/home/zyuli/local/clang-lld/bin:$PATH + +# mlir +export PYTHONPATH=/home/zyuli/local/flagtree-mlir/python_packages/mlir_core:$PYTHONPATH + +# llvm +export LLVM_INCLUDE_DIRS=/home/zyuli/llvm-project/build-mlir/include +export LLVM_LIBRARY_DIR=/home/zyuli/llvm-project/build-mlir/lib +export LLVM_SYSPATH=/home/zyuli/llvm-project/build-mlir + +# flagtree +export MAX_JOBS=32 +export TRITON_BUILD_WITH_CLANG_LLD=1 +export TRITON_BUILD_DIR=/home/zyuli/build/flagtree_extern_call + diff --git a/include/triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h b/include/triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h index 680bf0e04..69102ed61 100644 --- a/include/triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h +++ b/include/triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h @@ -102,6 +102,11 @@ void populatePrintOpToLLVMPattern(LLVMTypeConverter &typeConverter, const TargetInfoBase &targetInfo, PatternBenefit benefit); +void populateExternCallOpToLLVMPattern(LLVMTypeConverter &typeConverter, + RewritePatternSet &patterns, + const TargetInfoBase &targetInfo, + PatternBenefit benefit); + void populateInstrumentationToLLVMPatterns(LLVMTypeConverter &typeConverter, const TargetInfoBase &targetInfo, RewritePatternSet &patterns, diff --git a/include/triton/Dialect/Triton/IR/TritonOps.td b/include/triton/Dialect/Triton/IR/TritonOps.td index 7fd215f9d..95239b3cb 100644 --- a/include/triton/Dialect/Triton/IR/TritonOps.td +++ b/include/triton/Dialect/Triton/IR/TritonOps.td @@ -845,6 +845,32 @@ def TT_MapElementwiseReturnOp: TT_Op<"map_elementwise.return", let assemblyFormat = "attr-dict ($result^ `:` type($result))?"; } +// +// External Call op +// +def TT_ExternCallOp : TT_Op<"extern_call", [ + DeclareOpInterfaceMethods, + ConditionallySpeculatable, +]> { + + let description = [{ + call an external function $symbol implemented in $libpath/$libname with $args + return $libpath/$libname:$symbol($args...) + }]; + + let arguments = (ins Variadic:$srcs, StrAttr:$libname, StrAttr:$libpath, StrAttr:$symbol, BoolAttr:$pure); + + let results = (outs Variadic:$result); + + let assemblyFormat = "operands attr-dict `:` functional-type(operands, $result)"; + + let extraClassDeclaration = [{ + // Interface method for ConditionallySpeculatable. + Speculation::Speculatability getSpeculatability(); + }]; + +} + // // External Elementwise op // diff --git a/lib/Conversion/TritonGPUToLLVM/CMakeLists.txt b/lib/Conversion/TritonGPUToLLVM/CMakeLists.txt index d4f49c8d1..44e52cf37 100644 --- a/lib/Conversion/TritonGPUToLLVM/CMakeLists.txt +++ b/lib/Conversion/TritonGPUToLLVM/CMakeLists.txt @@ -7,6 +7,7 @@ add_triton_library(TritonGPUToLLVM AssertOpToLLVM.cpp ControlFlowOpToLLVM.cpp ConvertLayoutOpToLLVM.cpp + ExternCallOpToLLVM.cpp ElementwiseOpToLLVM.cpp FuncOpToLLVM.cpp GatherOpToLLVM.cpp diff --git a/lib/Conversion/TritonGPUToLLVM/ExternCallOpToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/ExternCallOpToLLVM.cpp new file mode 100644 index 000000000..81b6bf256 --- /dev/null +++ b/lib/Conversion/TritonGPUToLLVM/ExternCallOpToLLVM.cpp @@ -0,0 +1,61 @@ +#include "mlir/Conversion/LLVMCommon/Pattern.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Support/LLVM.h" +#include "triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h" +#include "triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h" +#include "triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h" +#include "triton/Conversion/TritonGPUToLLVM/Utility.h" +#include "triton/Dialect/TritonGPU/IR/Dialect.h" + +namespace { + +class ExternCallOpConversion + : public ConvertOpToLLVMPattern { +public: + ExternCallOpConversion(const LLVMTypeConverter &converter, + const PatternBenefit &benefit) + : ConvertOpToLLVMPattern(converter, benefit) {} + + LogicalResult + matchAndRewrite(triton::ExternCallOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto loc = op->getLoc(); + + if (op->getNumResults() > 1) { + llvm::errs() << "ExternCallConversion does not support multi outs."; + return failure(); + } + + LLVM::LLVMVoidType voidTy = void_ty(op->getContext()); + auto newOperands = adaptor.getOperands(); + Type retType = + op->getNumResults() == 0 + ? voidTy + : this->getTypeConverter()->convertType(op->getResult(0).getType()); + std::string funcName = op.getSymbol().str(); + StringRef libname = op.getLibname(); + StringRef libpath = op.getLibpath(); + + Operation *externCallOp; + Type funcType = mlir::triton::gpu::getFunctionType(retType, newOperands); + LLVM::LLVMFuncOp funcOp = mlir::triton::gpu::appendOrGetExternFuncOp( + rewriter, op, funcName, funcType, libname, libpath); + externCallOp = LLVM::createLLVMCallOp(rewriter, loc, funcOp, newOperands); + + if (op->getNumResults() == 0) { + rewriter.eraseOp(op); + } else { + rewriter.replaceOp(op, externCallOp->getResult(0)); + } + + return success(); + } +}; + +} // namespace + +void mlir::triton::populateExternCallOpToLLVMPattern( + LLVMTypeConverter &typeConverter, RewritePatternSet &patterns, + const TargetInfoBase &targetInfo, PatternBenefit benefit) { + patterns.add(typeConverter, benefit); +} diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp index bf509ffe5..512ab61f6 100644 --- a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp +++ b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp @@ -618,6 +618,7 @@ void populateTritonPatterns(TritonGPUTypeConverter &typeConverter, GenericOpPattern, GenericOpPattern, GenericOpPattern, + GenericOpPattern, GenericOpPattern, GenericOpPattern, GenericOpPattern, diff --git a/lib/Dialect/Triton/IR/Ops.cpp b/lib/Dialect/Triton/IR/Ops.cpp index 27fa26554..854ae1f2b 100644 --- a/lib/Dialect/Triton/IR/Ops.cpp +++ b/lib/Dialect/Triton/IR/Ops.cpp @@ -1313,6 +1313,24 @@ Speculation::Speculatability ExternElementwiseOp::getSpeculatability() { return Speculation::NotSpeculatable; } +// -- ExternCallOp -- +void ExternCallOp::getEffects( + SmallVectorImpl> + &effects) { + if (getPure()) + return; + effects.emplace_back(MemoryEffects::Write::get(), + SideEffects::DefaultResource::get()); + effects.emplace_back(MemoryEffects::Read::get(), + SideEffects::DefaultResource::get()); +} + +Speculation::Speculatability ExternCallOp::getSpeculatability() { + if (getPure()) + return Speculation::Speculatable; + return Speculation::NotSpeculatable; +} + // -- GatherOp -- LogicalResult GatherOp::verify() { RankedTensorType indicesTy = getIndices().getType(); diff --git a/python/src/ir.cc b/python/src/ir.cc index e5591dcb0..9abfe8472 100644 --- a/python/src/ir.cc +++ b/python/src/ir.cc @@ -1673,6 +1673,14 @@ void init_triton_ir(py::module &&m) { return self.create(retType, argList, libName, libPath, symbol, isPure); }) + .def("create_extern_call", + [](TritonOpBuilder &self, const std::string &libName, + const std::string &libPath, const std::string &symbol, + std::vector &argList, const std::vector &retTypes, + bool isPure) -> OpState { + return self.create(retTypes, argList, libName, + libPath, symbol, isPure); + }) // Built-in instruction .def("create_get_program_id", [](TritonOpBuilder &self, int axis) -> Value { diff --git a/python/triton/experimental/tle/language/raw/__init__.py b/python/triton/experimental/tle/language/raw/__init__.py index d66f6cd38..790c65cf3 100644 --- a/python/triton/experimental/tle/language/raw/__init__.py +++ b/python/triton/experimental/tle/language/raw/__init__.py @@ -1,3 +1,3 @@ -from .core import call, call_smem +from .core import call, call_smem, call_nvshmem -__all__ = ["call", "call_smem"] +__all__ = ["call", "call_smem", "call_nvshmem"] diff --git a/python/triton/experimental/tle/language/raw/core.py b/python/triton/experimental/tle/language/raw/core.py index d7029380e..7e2538531 100644 --- a/python/triton/experimental/tle/language/raw/core.py +++ b/python/triton/experimental/tle/language/raw/core.py @@ -46,3 +46,8 @@ def call_smem(func, args, _semantic=None): return buffer_tensors[0] else: return tl.tuple(buffer_tensors) + + +@builtin +def call_nvshmem(func, outputs, inputs, _semantic=None): + func.make_cubin() diff --git a/python/triton/experimental/tle/raw/cuda/runtime.py b/python/triton/experimental/tle/raw/cuda/runtime.py index 9015f2bf7..955a366e6 100644 --- a/python/triton/experimental/tle/raw/cuda/runtime.py +++ b/python/triton/experimental/tle/raw/cuda/runtime.py @@ -3,20 +3,24 @@ from pathlib import Path import subprocess from typing import Any, Final +import torch from triton._C.libtriton import llvm # pyright: ignore[reportMissingImports] from triton._C.libtriton.tle.llvm import parse_llvm_ir # pyright: ignore[reportMissingImports] # 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") +NVCC = os.getenv("NVCC", "nvcc") class CUDAJITFunction(object): def __init__(self, fn: Any, file: Path, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) + super().__init__() self.fn: Final[Any] = fn self.code: Final[str] = file.read_text() + self.file: Final[Path] = file + self.libs = kwargs.get("library", {}) self.__triton_builtin__: Final[bool] = True def make_llvm(self, mlir_context) -> str: @@ -40,3 +44,22 @@ def make_llvm(self, mlir_context) -> str: llvm_context = llvm.context() module = parse_llvm_ir(build.stdout.decode(), llvm_context, mlir_context) return f"{module}" + + def make_cubin(self): + src = self.file + dst = Path(src).with_suffix('.o') + include_dirs = [] + for lib_name, lib_path in self.libs.items(): + # TODO: Remove the method of passing information by setting environment variables. + os.environ[(lib_name + "_home").upper()] = lib_path + include_dirs.append(os.path.join(lib_path, "include")) + include_flags = [f"-I{inc_dir}" for inc_dir in include_dirs] + prop = torch.cuda.get_device_properties(torch.cuda.current_device()) + arch = f"-arch=sm_{prop.major}{prop.minor}" + build = subprocess.run([NVCC, "-rdc=true", arch, *include_flags, "--extended-lambda", "-c", "-o", dst, src], + capture_output=True) + assert build.returncode == 0, (f"nvcc failed\nstderr:\n{build.stderr.decode()}") + # TODO: Remove the method of passing information by setting environment variables. + os.environ["USE_NVCC"] = 'True' + os.environ["CUDA_CUBIN"] = str(dst) + return diff --git a/python/triton/language/core.py b/python/triton/language/core.py index 439671c3a..5a9c56b12 100644 --- a/python/triton/language/core.py +++ b/python/triton/language/core.py @@ -3433,6 +3433,94 @@ def binary_op_type_legalization(lhs, rhs, semantic): return semantic.binary_op_type_checking_impl(lhs, rhs) +def dispatch_ec(func, lib_name: str, lib_path: str, args: list, arg_type_symbol_dict: dict, is_pure: bool, + _semantic=None): + ''' + Dispatch a function to a library + :param func: the function to dispatch + :param lib_name: the name of the library + :param lib_path: the path of the library + :param args: the arguments of the function + :param arg_type_symbol_dict: the type of the arguments + :param ret_shape: the shape of the return value + :param _semantic: the builder + :return: the return value of the function + ''' + if len(arg_type_symbol_dict) == 0: + raise ValueError("arg_type_symbol_dict is empty") + + num_args = len(list(arg_type_symbol_dict.keys())[0]) + if len(args) != num_args: + raise ValueError(f"length of input args does not match." + f"Expect {len(args)}, got {num_args}") + + arg_types = [] + arg_list = [] + for arg in args: + if isinstance(arg, tensor): + arg_types.append(arg.dtype) + arg_list.append(arg.handle) + else: + arg_types.append(type(arg)) + arg_list.append(arg) + arg_types = tuple(arg_types) + + if arg_types not in arg_type_symbol_dict: + raise ValueError(f"input arg type does not match." + f"Expect one of {arg_type_symbol_dict.keys()}, got {arg_types}") + else: + symbol = arg_type_symbol_dict[arg_types][0] + ret_types = arg_type_symbol_dict[arg_types][1] + if not isinstance(ret_types, (builtins.list, builtins.tuple)): + ret_types = [ret_types] + + if symbol == "": + raise ValueError("Symbol can not be empty") + call = func(lib_name, lib_path, symbol, arg_list, [ret_type.to_ir(_semantic.builder) for ret_type in ret_types], + is_pure) + + if len(ret_types) == 0: + return tensor(call, void) + if len(ret_types) == 1: + return tensor(call.get_result(0), ret_types[0]) + return tuple(tensor(call.get_result(i), ty) for i, ty in enumerate(ret_types)) + + +@builtin +def extern_call(lib_name: str, lib_path: str, args: list, arg_type_symbol_dict: dict, is_pure: bool, _semantic=None): + ''' + Dispatch an function to a library + :param lib_name: the name of the library + :param lib_path: the path of the library + :param args: the arguments of the function + :param arg_type_symbol_dict: the type of the arguments + :param is_pure: whether the function is pure + :param _semantic: the semantic + :return: the return value of the function + ''' + dispatch_args = args.copy() + all_scalar = True + arg_types = [] + for i in builtins.range(len(dispatch_args)): + dispatch_args[i] = _semantic.to_tensor(dispatch_args[i]) + arg_types.append(dispatch_args[i].dtype) + if dispatch_args[i].type.is_block(): + all_scalar = False + if not all_scalar: + raise ValueError("extern call only support inputs with scalr type") + + if len(arg_type_symbol_dict) == 0: + raise ValueError("arg_type_symbol_dict is empty") + + num_args = len(list(arg_type_symbol_dict.keys())[0]) + if len(args) != num_args: + raise ValueError(f"length of input args does not match." + f"Expect {len(args)}, got {num_args}") + + func = _semantic.builder.create_extern_call + return dispatch_ec(func, lib_name, lib_path, dispatch_args, arg_type_symbol_dict, is_pure, _semantic) + + def extern(fn): """A decorator for external functions.""" return builtin(fn) 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..f2910883b --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-device.cu @@ -0,0 +1,8 @@ +#include + +extern "C" __device__ void simple_shift(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..be9ea236a --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-host.cu @@ -0,0 +1,31 @@ +#include +#include +#include +#include +#include + +extern "C" void nvshmem_init_wrapper() { nvshmem_init(); } + +extern "C" int nvshmemx_cumodule_init_wrapper(CUmodule module) { + return nvshmemx_cumodule_init(module); +} + +extern "C" int nvshmem_team_mype_wrapper() { + int mype_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE); + return mype_node; +} + +// TODO: Adapt to different data types +extern "C" int *nvshmem_alloc_wrapper(int size) { + int *destination = (int *)nvshmem_malloc(sizeof(int) * size); + return destination; +} + +extern "C" void nvshmemx_barrier_warpper(cudaStream_t stream) { + nvshmemx_barrier_all_on_stream(stream); +} + +extern "C" void nvshmem_finalize_wrapper(int *dest) { + nvshmem_free(dest); + 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..2f87d70f6 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py @@ -0,0 +1,122 @@ +import os +import subprocess +import ctypes +import torch +import triton +import triton.knobs as knobs +import triton.experimental.tle.language.raw as tle_raw + +from pathlib import Path +from triton.experimental.tle.raw import dialect +from triton.language.extra.cuda import libnvshmem_device + + +@dialect( + name="cuda", + file=(Path(__file__).parent / "simple-shift-device.cu").resolve(), + library={"nvshmem": "/home/zyl/zyuli/envs/nvshmem/lib/python3.12/site-packages/nvidia/nvshmem"}, +) +def edsl(*args, **kwargs): + ... + + +@triton.jit +def simple_shift_kernel(destination_ptr, ): + # TODO: Combine with tle_raw.call, then dispatch + tle_raw.call_nvshmem(edsl, [], [destination_ptr]) + libnvshmem_device.simple_shift(destination_ptr) + + +def cuda_host_compile(cuda_host_path, cuda_host_lib): + NVCC = os.getenv("NVCC", "nvcc") + NVSHMEM_HOME = "/home/zyl/zyuli/envs/nvshmem/lib/python3.12/site-packages/nvidia/nvshmem" + include_path = f"-I{os.path.join(NVSHMEM_HOME, 'include')}" + lib_path = f"-L{os.path.join(NVSHMEM_HOME, 'lib')}" + + prop = torch.cuda.get_device_properties(torch.cuda.current_device()) + arch = f"-arch=sm_{prop.major}{prop.minor}" + tmp_file = Path(cuda_host_lib).with_suffix('.so.tmp') + build = [ + NVCC, "-shared", "-Xcompiler", "-fPIC", "-rdc=true", arch, include_path, lib_path, "-lnvshmem_host", + "-lnvshmem_device", "-o", tmp_file, cuda_host_path + ] + build = subprocess.run(build, capture_output=True) + assert build.returncode == 0, (f"NVCC host failed\nstderr:\n{build.stderr.decode()}") + tmp_file.rename(cuda_host_lib) + + +def simpe_shift(): + cu_file = (Path(__file__).parent / "simple-shift-host.cu").resolve() + lib_file = Path(cu_file).with_suffix('.so') + + rank = int(os.getenv("OMPI_COMM_WORLD_RANK", "0")) + if rank == 0: + cuda_host_compile(cu_file, lib_file) + + import time + timeout = 60 + start = time.time() + while True: + if lib_file.exists(): + try: + ctypes.CDLL(str(lib_file)) + break + except OSError: + pass + if time.time() - start > timeout: + raise RuntimeError(f"Timeout waiting for {lib_file}") + time.sleep(0.1) + + lib = ctypes.CDLL(lib_file) + lib.nvshmem_init_wrapper.argtypes = [] + lib.nvshmem_init_wrapper.restype = None + lib.nvshmemx_cumodule_init_wrapper.argtypes = [ctypes.c_void_p] + lib.nvshmemx_cumodule_init_wrapper.restype = ctypes.c_int + lib.nvshmem_team_mype_wrapper.argtypes = [] + lib.nvshmem_team_mype_wrapper.restype = ctypes.c_int + lib.nvshmem_alloc_wrapper.argtypes = [ctypes.c_int] + lib.nvshmem_alloc_wrapper.restype = ctypes.POINTER(ctypes.c_int) + lib.nvshmemx_barrier_warpper.argtypes = [ctypes.c_void_p] + lib.nvshmemx_barrier_warpper.restype = None + lib.nvshmem_finalize_wrapper.argtypes = [ctypes.POINTER(ctypes.c_int)] + lib.nvshmem_finalize_wrapper.restype = None + + lib.nvshmem_init_wrapper() + mype_node = lib.nvshmem_team_mype_wrapper() + torch.cuda.set_device(mype_node) + device = triton.runtime.driver.active.get_active_torch_device() + stream = torch.cuda.Stream() + + dest = lib.nvshmem_alloc_wrapper(1) + dest_addr = ctypes.cast(dest, ctypes.c_void_p).value + storage = torch._C._construct_storage_from_data_pointer(dest_addr, device, 4) + dest_tensor = torch.empty(0, dtype=torch.int32, device=device).set_(storage).view(1) + msg = torch.empty((1, ), dtype=torch.int32, pin_memory=True) + + def cumodule_init_hook(*args, **kwargs): + key = kwargs["key"] + jit_function = kwargs["fn"].jit_function + device = kwargs["compile"]["device"] + kernel_cache = jit_function.device_caches[device][0] + kernel = kernel_cache.get(key, None) + assert kernel is not None + kernel._init_handles() + ret = lib.nvshmemx_cumodule_init_wrapper(ctypes.c_void_p(kernel.module)) + assert ret == 0, f"nvshmemx_cumodule_init_wrapper failed: {ret}" + + knobs.runtime.jit_post_compile_hook = cumodule_init_hook + + simple_shift_kernel[(1, )](dest_tensor) + + stream_ptr = stream.cuda_stream + lib.nvshmemx_barrier_warpper(ctypes.c_void_p(stream_ptr)) + with torch.cuda.stream(stream): + msg.copy_(dest_tensor, non_blocking=True) + stream.synchronize() + + lib.nvshmem_finalize_wrapper(dest) + print(f"Rank {mype_node}: {msg}") + + +if __name__ == "__main__": + simpe_shift() diff --git a/support_nvshmem.patch b/support_nvshmem.patch new file mode 100644 index 000000000..f6d472e0d --- /dev/null +++ b/support_nvshmem.patch @@ -0,0 +1,697 @@ +diff --git a/include/triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h b/include/triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h +index 680bf0e04..69102ed61 100644 +--- a/include/triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h ++++ b/include/triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h +@@ -102,6 +102,11 @@ void populatePrintOpToLLVMPattern(LLVMTypeConverter &typeConverter, + const TargetInfoBase &targetInfo, + PatternBenefit benefit); + ++void populateExternCallOpToLLVMPattern(LLVMTypeConverter &typeConverter, ++ RewritePatternSet &patterns, ++ const TargetInfoBase &targetInfo, ++ PatternBenefit benefit); ++ + void populateInstrumentationToLLVMPatterns(LLVMTypeConverter &typeConverter, + const TargetInfoBase &targetInfo, + RewritePatternSet &patterns, +diff --git a/include/triton/Dialect/Triton/IR/TritonOps.td b/include/triton/Dialect/Triton/IR/TritonOps.td +index 7fd215f9d..95239b3cb 100644 +--- a/include/triton/Dialect/Triton/IR/TritonOps.td ++++ b/include/triton/Dialect/Triton/IR/TritonOps.td +@@ -845,6 +845,32 @@ def TT_MapElementwiseReturnOp: TT_Op<"map_elementwise.return", + let assemblyFormat = "attr-dict ($result^ `:` type($result))?"; + } + ++// ++// External Call op ++// ++def TT_ExternCallOp : TT_Op<"extern_call", [ ++ DeclareOpInterfaceMethods, ++ ConditionallySpeculatable, ++]> { ++ ++ let description = [{ ++ call an external function $symbol implemented in $libpath/$libname with $args ++ return $libpath/$libname:$symbol($args...) ++ }]; ++ ++ let arguments = (ins Variadic:$srcs, StrAttr:$libname, StrAttr:$libpath, StrAttr:$symbol, BoolAttr:$pure); ++ ++ let results = (outs Variadic:$result); ++ ++ let assemblyFormat = "operands attr-dict `:` functional-type(operands, $result)"; ++ ++ let extraClassDeclaration = [{ ++ // Interface method for ConditionallySpeculatable. ++ Speculation::Speculatability getSpeculatability(); ++ }]; ++ ++} ++ + // + // External Elementwise op + // +diff --git a/lib/Conversion/TritonGPUToLLVM/CMakeLists.txt b/lib/Conversion/TritonGPUToLLVM/CMakeLists.txt +index d4f49c8d1..44e52cf37 100644 +--- a/lib/Conversion/TritonGPUToLLVM/CMakeLists.txt ++++ b/lib/Conversion/TritonGPUToLLVM/CMakeLists.txt +@@ -7,6 +7,7 @@ add_triton_library(TritonGPUToLLVM + AssertOpToLLVM.cpp + ControlFlowOpToLLVM.cpp + ConvertLayoutOpToLLVM.cpp ++ ExternCallOpToLLVM.cpp + ElementwiseOpToLLVM.cpp + FuncOpToLLVM.cpp + GatherOpToLLVM.cpp +diff --git a/lib/Conversion/TritonGPUToLLVM/ExternCallOpToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/ExternCallOpToLLVM.cpp +new file mode 100644 +index 000000000..81b6bf256 +--- /dev/null ++++ b/lib/Conversion/TritonGPUToLLVM/ExternCallOpToLLVM.cpp +@@ -0,0 +1,61 @@ ++#include "mlir/Conversion/LLVMCommon/Pattern.h" ++#include "mlir/Dialect/GPU/IR/GPUDialect.h" ++#include "mlir/Support/LLVM.h" ++#include "triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h" ++#include "triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h" ++#include "triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h" ++#include "triton/Conversion/TritonGPUToLLVM/Utility.h" ++#include "triton/Dialect/TritonGPU/IR/Dialect.h" ++ ++namespace { ++ ++class ExternCallOpConversion ++ : public ConvertOpToLLVMPattern { ++public: ++ ExternCallOpConversion(const LLVMTypeConverter &converter, ++ const PatternBenefit &benefit) ++ : ConvertOpToLLVMPattern(converter, benefit) {} ++ ++ LogicalResult ++ matchAndRewrite(triton::ExternCallOp op, OpAdaptor adaptor, ++ ConversionPatternRewriter &rewriter) const override { ++ auto loc = op->getLoc(); ++ ++ if (op->getNumResults() > 1) { ++ llvm::errs() << "ExternCallConversion does not support multi outs."; ++ return failure(); ++ } ++ ++ LLVM::LLVMVoidType voidTy = void_ty(op->getContext()); ++ auto newOperands = adaptor.getOperands(); ++ Type retType = ++ op->getNumResults() == 0 ++ ? voidTy ++ : this->getTypeConverter()->convertType(op->getResult(0).getType()); ++ std::string funcName = op.getSymbol().str(); ++ StringRef libname = op.getLibname(); ++ StringRef libpath = op.getLibpath(); ++ ++ Operation *externCallOp; ++ Type funcType = mlir::triton::gpu::getFunctionType(retType, newOperands); ++ LLVM::LLVMFuncOp funcOp = mlir::triton::gpu::appendOrGetExternFuncOp( ++ rewriter, op, funcName, funcType, libname, libpath); ++ externCallOp = LLVM::createLLVMCallOp(rewriter, loc, funcOp, newOperands); ++ ++ if (op->getNumResults() == 0) { ++ rewriter.eraseOp(op); ++ } else { ++ rewriter.replaceOp(op, externCallOp->getResult(0)); ++ } ++ ++ return success(); ++ } ++}; ++ ++} // namespace ++ ++void mlir::triton::populateExternCallOpToLLVMPattern( ++ LLVMTypeConverter &typeConverter, RewritePatternSet &patterns, ++ const TargetInfoBase &targetInfo, PatternBenefit benefit) { ++ patterns.add(typeConverter, benefit); ++} +diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp +index b49871486..45c613465 100644 +--- a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp ++++ b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp +@@ -600,6 +600,7 @@ void populateTritonPatterns(TritonGPUTypeConverter &typeConverter, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, ++ GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, +diff --git a/lib/Dialect/Triton/IR/Ops.cpp b/lib/Dialect/Triton/IR/Ops.cpp +index 27fa26554..854ae1f2b 100644 +--- a/lib/Dialect/Triton/IR/Ops.cpp ++++ b/lib/Dialect/Triton/IR/Ops.cpp +@@ -1313,6 +1313,24 @@ Speculation::Speculatability ExternElementwiseOp::getSpeculatability() { + return Speculation::NotSpeculatable; + } + ++// -- ExternCallOp -- ++void ExternCallOp::getEffects( ++ SmallVectorImpl> ++ &effects) { ++ if (getPure()) ++ return; ++ effects.emplace_back(MemoryEffects::Write::get(), ++ SideEffects::DefaultResource::get()); ++ effects.emplace_back(MemoryEffects::Read::get(), ++ SideEffects::DefaultResource::get()); ++} ++ ++Speculation::Speculatability ExternCallOp::getSpeculatability() { ++ if (getPure()) ++ return Speculation::Speculatable; ++ return Speculation::NotSpeculatable; ++} ++ + // -- GatherOp -- + LogicalResult GatherOp::verify() { + RankedTensorType indicesTy = getIndices().getType(); +diff --git a/python/src/ir.cc b/python/src/ir.cc +index e5591dcb0..9abfe8472 100644 +--- a/python/src/ir.cc ++++ b/python/src/ir.cc +@@ -1673,6 +1673,14 @@ void init_triton_ir(py::module &&m) { + return self.create(retType, argList, libName, + libPath, symbol, isPure); + }) ++ .def("create_extern_call", ++ [](TritonOpBuilder &self, const std::string &libName, ++ const std::string &libPath, const std::string &symbol, ++ std::vector &argList, const std::vector &retTypes, ++ bool isPure) -> OpState { ++ return self.create(retTypes, argList, libName, ++ libPath, symbol, isPure); ++ }) + // Built-in instruction + .def("create_get_program_id", + [](TritonOpBuilder &self, int axis) -> Value { +diff --git a/python/triton/experimental/tle/language/raw/__init__.py b/python/triton/experimental/tle/language/raw/__init__.py +index d66f6cd38..790c65cf3 100644 +--- a/python/triton/experimental/tle/language/raw/__init__.py ++++ b/python/triton/experimental/tle/language/raw/__init__.py +@@ -1,3 +1,3 @@ +-from .core import call, call_smem ++from .core import call, call_smem, call_nvshmem + +-__all__ = ["call", "call_smem"] ++__all__ = ["call", "call_smem", "call_nvshmem"] +diff --git a/python/triton/experimental/tle/language/raw/core.py b/python/triton/experimental/tle/language/raw/core.py +index d7029380e..7e2538531 100644 +--- a/python/triton/experimental/tle/language/raw/core.py ++++ b/python/triton/experimental/tle/language/raw/core.py +@@ -46,3 +46,8 @@ def call_smem(func, args, _semantic=None): + return buffer_tensors[0] + else: + return tl.tuple(buffer_tensors) ++ ++ ++@builtin ++def call_nvshmem(func, outputs, inputs, _semantic=None): ++ func.make_cubin() +diff --git a/python/triton/experimental/tle/raw/cuda/runtime.py b/python/triton/experimental/tle/raw/cuda/runtime.py +index 9015f2bf7..955a366e6 100644 +--- a/python/triton/experimental/tle/raw/cuda/runtime.py ++++ b/python/triton/experimental/tle/raw/cuda/runtime.py +@@ -3,20 +3,24 @@ import os + from pathlib import Path + import subprocess + from typing import Any, Final ++import torch + + from triton._C.libtriton import llvm # pyright: ignore[reportMissingImports] + from triton._C.libtriton.tle.llvm import parse_llvm_ir # pyright: ignore[reportMissingImports] + + # 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") ++NVCC = os.getenv("NVCC", "nvcc") + + + class CUDAJITFunction(object): + + def __init__(self, fn: Any, file: Path, *args, **kwargs) -> None: +- super().__init__(*args, **kwargs) ++ super().__init__() + self.fn: Final[Any] = fn + self.code: Final[str] = file.read_text() ++ self.file: Final[Path] = file ++ self.libs = kwargs.get("library", {}) + self.__triton_builtin__: Final[bool] = True + + def make_llvm(self, mlir_context) -> str: +@@ -40,3 +44,22 @@ class CUDAJITFunction(object): + llvm_context = llvm.context() + module = parse_llvm_ir(build.stdout.decode(), llvm_context, mlir_context) + return f"{module}" ++ ++ def make_cubin(self): ++ src = self.file ++ dst = Path(src).with_suffix('.o') ++ include_dirs = [] ++ for lib_name, lib_path in self.libs.items(): ++ # TODO: Remove the method of passing information by setting environment variables. ++ os.environ[(lib_name + "_home").upper()] = lib_path ++ include_dirs.append(os.path.join(lib_path, "include")) ++ include_flags = [f"-I{inc_dir}" for inc_dir in include_dirs] ++ prop = torch.cuda.get_device_properties(torch.cuda.current_device()) ++ arch = f"-arch=sm_{prop.major}{prop.minor}" ++ build = subprocess.run([NVCC, "-rdc=true", arch, *include_flags, "--extended-lambda", "-c", "-o", dst, src], ++ capture_output=True) ++ assert build.returncode == 0, (f"nvcc failed\nstderr:\n{build.stderr.decode()}") ++ # TODO: Remove the method of passing information by setting environment variables. ++ os.environ["USE_NVCC"] = 'True' ++ os.environ["CUDA_CUBIN"] = str(dst) ++ return +diff --git a/python/triton/language/core.py b/python/triton/language/core.py +index 439671c3a..5a9c56b12 100644 +--- a/python/triton/language/core.py ++++ b/python/triton/language/core.py +@@ -3433,6 +3433,94 @@ def binary_op_type_legalization(lhs, rhs, semantic): + return semantic.binary_op_type_checking_impl(lhs, rhs) + + ++def dispatch_ec(func, lib_name: str, lib_path: str, args: list, arg_type_symbol_dict: dict, is_pure: bool, ++ _semantic=None): ++ ''' ++ Dispatch a function to a library ++ :param func: the function to dispatch ++ :param lib_name: the name of the library ++ :param lib_path: the path of the library ++ :param args: the arguments of the function ++ :param arg_type_symbol_dict: the type of the arguments ++ :param ret_shape: the shape of the return value ++ :param _semantic: the builder ++ :return: the return value of the function ++ ''' ++ if len(arg_type_symbol_dict) == 0: ++ raise ValueError("arg_type_symbol_dict is empty") ++ ++ num_args = len(list(arg_type_symbol_dict.keys())[0]) ++ if len(args) != num_args: ++ raise ValueError(f"length of input args does not match." ++ f"Expect {len(args)}, got {num_args}") ++ ++ arg_types = [] ++ arg_list = [] ++ for arg in args: ++ if isinstance(arg, tensor): ++ arg_types.append(arg.dtype) ++ arg_list.append(arg.handle) ++ else: ++ arg_types.append(type(arg)) ++ arg_list.append(arg) ++ arg_types = tuple(arg_types) ++ ++ if arg_types not in arg_type_symbol_dict: ++ raise ValueError(f"input arg type does not match." ++ f"Expect one of {arg_type_symbol_dict.keys()}, got {arg_types}") ++ else: ++ symbol = arg_type_symbol_dict[arg_types][0] ++ ret_types = arg_type_symbol_dict[arg_types][1] ++ if not isinstance(ret_types, (builtins.list, builtins.tuple)): ++ ret_types = [ret_types] ++ ++ if symbol == "": ++ raise ValueError("Symbol can not be empty") ++ call = func(lib_name, lib_path, symbol, arg_list, [ret_type.to_ir(_semantic.builder) for ret_type in ret_types], ++ is_pure) ++ ++ if len(ret_types) == 0: ++ return tensor(call, void) ++ if len(ret_types) == 1: ++ return tensor(call.get_result(0), ret_types[0]) ++ return tuple(tensor(call.get_result(i), ty) for i, ty in enumerate(ret_types)) ++ ++ ++@builtin ++def extern_call(lib_name: str, lib_path: str, args: list, arg_type_symbol_dict: dict, is_pure: bool, _semantic=None): ++ ''' ++ Dispatch an function to a library ++ :param lib_name: the name of the library ++ :param lib_path: the path of the library ++ :param args: the arguments of the function ++ :param arg_type_symbol_dict: the type of the arguments ++ :param is_pure: whether the function is pure ++ :param _semantic: the semantic ++ :return: the return value of the function ++ ''' ++ dispatch_args = args.copy() ++ all_scalar = True ++ arg_types = [] ++ for i in builtins.range(len(dispatch_args)): ++ dispatch_args[i] = _semantic.to_tensor(dispatch_args[i]) ++ arg_types.append(dispatch_args[i].dtype) ++ if dispatch_args[i].type.is_block(): ++ all_scalar = False ++ if not all_scalar: ++ raise ValueError("extern call only support inputs with scalr type") ++ ++ if len(arg_type_symbol_dict) == 0: ++ raise ValueError("arg_type_symbol_dict is empty") ++ ++ num_args = len(list(arg_type_symbol_dict.keys())[0]) ++ if len(args) != num_args: ++ raise ValueError(f"length of input args does not match." ++ f"Expect {len(args)}, got {num_args}") ++ ++ func = _semantic.builder.create_extern_call ++ return dispatch_ec(func, lib_name, lib_path, dispatch_args, arg_type_symbol_dict, is_pure, _semantic) ++ ++ + def extern(fn): + """A decorator for external functions.""" + return builtin(fn) +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..f2910883b +--- /dev/null ++++ b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-device.cu +@@ -0,0 +1,8 @@ ++#include ++ ++extern "C" __device__ void simple_shift(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..be9ea236a +--- /dev/null ++++ b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-host.cu +@@ -0,0 +1,31 @@ ++#include ++#include ++#include ++#include ++#include ++ ++extern "C" void nvshmem_init_wrapper() { nvshmem_init(); } ++ ++extern "C" int nvshmemx_cumodule_init_wrapper(CUmodule module) { ++ return nvshmemx_cumodule_init(module); ++} ++ ++extern "C" int nvshmem_team_mype_wrapper() { ++ int mype_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE); ++ return mype_node; ++} ++ ++// TODO: Adapt to different data types ++extern "C" int *nvshmem_alloc_wrapper(int size) { ++ int *destination = (int *)nvshmem_malloc(sizeof(int) * size); ++ return destination; ++} ++ ++extern "C" void nvshmemx_barrier_warpper(cudaStream_t stream) { ++ nvshmemx_barrier_all_on_stream(stream); ++} ++ ++extern "C" void nvshmem_finalize_wrapper(int *dest) { ++ nvshmem_free(dest); ++ 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..2f87d70f6 +--- /dev/null ++++ b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py +@@ -0,0 +1,122 @@ ++import os ++import subprocess ++import ctypes ++import torch ++import triton ++import triton.knobs as knobs ++import triton.experimental.tle.language.raw as tle_raw ++ ++from pathlib import Path ++from triton.experimental.tle.raw import dialect ++from triton.language.extra.cuda import libnvshmem_device ++ ++ ++@dialect( ++ name="cuda", ++ file=(Path(__file__).parent / "simple-shift-device.cu").resolve(), ++ library={"nvshmem": "/home/zyl/zyuli/envs/nvshmem/lib/python3.12/site-packages/nvidia/nvshmem"}, ++) ++def edsl(*args, **kwargs): ++ ... ++ ++ ++@triton.jit ++def simple_shift_kernel(destination_ptr, ): ++ # TODO: Combine with tle_raw.call, then dispatch ++ tle_raw.call_nvshmem(edsl, [], [destination_ptr]) ++ libnvshmem_device.simple_shift(destination_ptr) ++ ++ ++def cuda_host_compile(cuda_host_path, cuda_host_lib): ++ NVCC = os.getenv("NVCC", "nvcc") ++ NVSHMEM_HOME = "/home/zyl/zyuli/envs/nvshmem/lib/python3.12/site-packages/nvidia/nvshmem" ++ include_path = f"-I{os.path.join(NVSHMEM_HOME, 'include')}" ++ lib_path = f"-L{os.path.join(NVSHMEM_HOME, 'lib')}" ++ ++ prop = torch.cuda.get_device_properties(torch.cuda.current_device()) ++ arch = f"-arch=sm_{prop.major}{prop.minor}" ++ tmp_file = Path(cuda_host_lib).with_suffix('.so.tmp') ++ build = [ ++ NVCC, "-shared", "-Xcompiler", "-fPIC", "-rdc=true", arch, include_path, lib_path, "-lnvshmem_host", ++ "-lnvshmem_device", "-o", tmp_file, cuda_host_path ++ ] ++ build = subprocess.run(build, capture_output=True) ++ assert build.returncode == 0, (f"NVCC host failed\nstderr:\n{build.stderr.decode()}") ++ tmp_file.rename(cuda_host_lib) ++ ++ ++def simpe_shift(): ++ cu_file = (Path(__file__).parent / "simple-shift-host.cu").resolve() ++ lib_file = Path(cu_file).with_suffix('.so') ++ ++ rank = int(os.getenv("OMPI_COMM_WORLD_RANK", "0")) ++ if rank == 0: ++ cuda_host_compile(cu_file, lib_file) ++ ++ import time ++ timeout = 60 ++ start = time.time() ++ while True: ++ if lib_file.exists(): ++ try: ++ ctypes.CDLL(str(lib_file)) ++ break ++ except OSError: ++ pass ++ if time.time() - start > timeout: ++ raise RuntimeError(f"Timeout waiting for {lib_file}") ++ time.sleep(0.1) ++ ++ lib = ctypes.CDLL(lib_file) ++ lib.nvshmem_init_wrapper.argtypes = [] ++ lib.nvshmem_init_wrapper.restype = None ++ lib.nvshmemx_cumodule_init_wrapper.argtypes = [ctypes.c_void_p] ++ lib.nvshmemx_cumodule_init_wrapper.restype = ctypes.c_int ++ lib.nvshmem_team_mype_wrapper.argtypes = [] ++ lib.nvshmem_team_mype_wrapper.restype = ctypes.c_int ++ lib.nvshmem_alloc_wrapper.argtypes = [ctypes.c_int] ++ lib.nvshmem_alloc_wrapper.restype = ctypes.POINTER(ctypes.c_int) ++ lib.nvshmemx_barrier_warpper.argtypes = [ctypes.c_void_p] ++ lib.nvshmemx_barrier_warpper.restype = None ++ lib.nvshmem_finalize_wrapper.argtypes = [ctypes.POINTER(ctypes.c_int)] ++ lib.nvshmem_finalize_wrapper.restype = None ++ ++ lib.nvshmem_init_wrapper() ++ mype_node = lib.nvshmem_team_mype_wrapper() ++ torch.cuda.set_device(mype_node) ++ device = triton.runtime.driver.active.get_active_torch_device() ++ stream = torch.cuda.Stream() ++ ++ dest = lib.nvshmem_alloc_wrapper(1) ++ dest_addr = ctypes.cast(dest, ctypes.c_void_p).value ++ storage = torch._C._construct_storage_from_data_pointer(dest_addr, device, 4) ++ dest_tensor = torch.empty(0, dtype=torch.int32, device=device).set_(storage).view(1) ++ msg = torch.empty((1, ), dtype=torch.int32, pin_memory=True) ++ ++ def cumodule_init_hook(*args, **kwargs): ++ key = kwargs["key"] ++ jit_function = kwargs["fn"].jit_function ++ device = kwargs["compile"]["device"] ++ kernel_cache = jit_function.device_caches[device][0] ++ kernel = kernel_cache.get(key, None) ++ assert kernel is not None ++ kernel._init_handles() ++ ret = lib.nvshmemx_cumodule_init_wrapper(ctypes.c_void_p(kernel.module)) ++ assert ret == 0, f"nvshmemx_cumodule_init_wrapper failed: {ret}" ++ ++ knobs.runtime.jit_post_compile_hook = cumodule_init_hook ++ ++ simple_shift_kernel[(1, )](dest_tensor) ++ ++ stream_ptr = stream.cuda_stream ++ lib.nvshmemx_barrier_warpper(ctypes.c_void_p(stream_ptr)) ++ with torch.cuda.stream(stream): ++ msg.copy_(dest_tensor, non_blocking=True) ++ stream.synchronize() ++ ++ lib.nvshmem_finalize_wrapper(dest) ++ print(f"Rank {mype_node}: {msg}") ++ ++ ++if __name__ == "__main__": ++ simpe_shift() +diff --git a/third_party/nvidia/backend/compiler.py b/third_party/nvidia/backend/compiler.py +index 632b79075..482af7b32 100644 +--- a/third_party/nvidia/backend/compiler.py ++++ b/third_party/nvidia/backend/compiler.py +@@ -498,7 +498,7 @@ class CUDABackend(BaseBackend): + assert len(names) == 1 + metadata["name"] = names[0] + # post-process +- ptx_version = f'{ptx_version//10}.{ptx_version%10}' ++ ptx_version = '8.5' + ret = re.sub(r'\.version \d+\.\d+', f'.version {ptx_version}', ret, flags=re.MULTILINE) + ret = re.sub(r'\.target sm_\d+', f'.target sm_{capability}', ret, flags=re.MULTILINE) + if not knobs.compilation.dump_ir_extract_di_local_variables: +@@ -512,13 +512,15 @@ class CUDABackend(BaseBackend): + return ret + + def make_cubin(self, src, metadata, opt, capability): +- ptxas = get_ptxas(self.target.arch).path ++ ptxas = os.getenv("ptxas", "ptxas") + with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.ptx') as fsrc, \ + tempfile.NamedTemporaryFile(delete=False, mode='r', suffix='.log') as flog: + fsrc.write(src) + fsrc.flush() + fbin = fsrc.name + '.o' + ++ use_nvcc = os.getenv("USE_NVCC", '').lower() in ('1', 'true') ++ os.environ.pop("USE_NVCC", None) + debug_info = [] + if knobs.compilation.disable_line_info: + # This option is ignored if used without -lineinfo +@@ -539,9 +541,11 @@ class CUDABackend(BaseBackend): + # Accept more ptxas options if provided + ptx_extra_options = opt.ptx_options.split(" ") if opt.ptx_options else [] + ++ # If use nvshmem, we need to compile the ptx file into a relocatable object file and then link it with nvshmem library ++ compile_only = ["-c"] if use_nvcc else [] + ptxas_cmd = [ +- ptxas, *debug_info, *fmad, '-v', *disable_opt, *ptx_extra_options, f'--gpu-name={arch}', fsrc.name, +- '-o', fbin ++ ptxas, *compile_only, *debug_info, *fmad, '-v', *disable_opt, *ptx_extra_options, f'--gpu-name={arch}', ++ fsrc.name, '-o', fbin + ] + try: + subprocess.run(ptxas_cmd, check=True, close_fds=False, stderr=flog) +@@ -581,8 +585,37 @@ please share the reproducer above with Triton project. + """) + raise PTXASError(error) + +- with open(fbin, 'rb') as f: +- cubin = f.read() ++ if use_nvcc: ++ NVLINK = os.getenv("NVLINK", "nvlink") ++ NVSHMEM_HOME = os.getenv("NVSHMEM_HOME") ++ fbin_combined = fbin + ".combined.cubin" ++ cuda_cubin = os.getenv("CUDA_CUBIN") ++ nvshmem_lib = os.path.join(NVSHMEM_HOME, "lib") ++ nvlink_cmds = [ ++ NVLINK, ++ f"-arch={arch}", ++ f"-L{nvshmem_lib}", ++ "-lnvshmem_device", ++ fbin, ++ cuda_cubin, ++ "-o", ++ fbin_combined, ++ ] ++ try: ++ subprocess.run(nvlink_cmds, check=True, close_fds=False, stderr=flog) ++ except Exception as e: ++ import logging ++ logging.error(f"error runing nvlink: {nvlink_cmds}") ++ logging.exception(e) ++ ++ if use_nvcc: ++ with open(fbin_combined, 'rb') as f: ++ cubin = f.read() ++ else: ++ with open(fbin, 'rb') as f: ++ cubin = f.read() ++ if os.path.exists(fbin_combined): ++ os.remove(fbin_combined) + if os.path.exists(fbin): + os.remove(fbin) + return cubin +diff --git a/third_party/nvidia/language/cuda/__init__.py b/third_party/nvidia/language/cuda/__init__.py +index fbececf1d..39207bf60 100644 +--- a/third_party/nvidia/language/cuda/__init__.py ++++ b/third_party/nvidia/language/cuda/__init__.py +@@ -1,10 +1,12 @@ + from . import libdevice ++from . import libnvshmem_device + + from .utils import (globaltimer, num_threads, num_warps, smid, convert_custom_float8_sm70, convert_custom_float8_sm80) + from .gdc import (gdc_launch_dependents, gdc_wait) + + __all__ = [ + "libdevice", ++ 'libnvshmem_device', + "globaltimer", + "num_threads", + "num_warps", +diff --git a/third_party/nvidia/language/cuda/libnvshmem_device.py b/third_party/nvidia/language/cuda/libnvshmem_device.py +new file mode 100644 +index 000000000..3dec30b9a +--- /dev/null ++++ b/third_party/nvidia/language/cuda/libnvshmem_device.py +@@ -0,0 +1,31 @@ ++from triton.language import core ++import triton.language as tl ++ ++ ++def _pointer_type_hash(self): ++ return hash((self.name, self.element_ty, "tt_ptr")) ++ ++ ++def patch_hash_method_for_pointer_type(): ++ elem_dtype_list = tl.core.dtype.SINT_TYPES + tl.core.dtype.UINT_TYPES + tl.core.dtype.FP_TYPES + tl.core.dtype.OTHER_TYPES ++ for elem_dtype in elem_dtype_list: ++ ptr_ty = type(tl.core.pointer_type(tl.core.dtype(elem_dtype))) ++ ptr_ty.__hash__ = _pointer_type_hash ++ ++ ++patch_hash_method_for_pointer_type() ++ ++ ++@core.extern ++def simple_shift(dst, _semantic=None): ++ return core.extern_call( ++ "", # libname ++ "", # libpath ++ [dst], # args ++ {( ++ core.pointer_type(core.dtype("int32")), # arg_type_symbol_dict ++ ): ("simple_shift", ()), # void return type ++ }, ++ is_pure=False, ++ _semantic=_semantic, ++ ) +diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp +index 368a66629..0fc8242c5 100644 +--- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp ++++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp +@@ -222,6 +222,8 @@ struct ConvertTritonGPUToLLVM + targetInfo, benefit); + mlir::triton::populatePrintOpToLLVMPattern(typeConverter, patterns, + targetInfo, benefit); ++ mlir::triton::populateExternCallOpToLLVMPattern(typeConverter, patterns, ++ targetInfo, benefit); + mlir::triton::populateControlFlowOpToLLVMPattern(typeConverter, patterns, + targetInfo, benefit); + mlir::triton::NVIDIA::populateSPMDOpToLLVMPattern(typeConverter, patterns, diff --git a/third_party/nvidia/backend/compiler.py b/third_party/nvidia/backend/compiler.py index d6ea87518..c469c8a78 100644 --- a/third_party/nvidia/backend/compiler.py +++ b/third_party/nvidia/backend/compiler.py @@ -512,7 +512,7 @@ def make_ptx(self, src, metadata, opt, capability): assert len(names) == 1 metadata["name"] = names[0] # post-process - ptx_version = f'{ptx_version//10}.{ptx_version%10}' + ptx_version = '8.5' ret = re.sub(r'\.version \d+\.\d+', f'.version {ptx_version}', ret, flags=re.MULTILINE) ret = re.sub(r'\.target sm_\d+', f'.target sm_{capability}', ret, flags=re.MULTILINE) if not knobs.compilation.dump_ir_extract_di_local_variables: @@ -526,13 +526,15 @@ def make_ptx(self, src, metadata, opt, capability): return ret def make_cubin(self, src, metadata, opt, capability): - ptxas = get_ptxas(self.target.arch).path + ptxas = os.getenv("ptxas", "ptxas") with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.ptx') as fsrc, \ tempfile.NamedTemporaryFile(delete=False, mode='r', suffix='.log') as flog: fsrc.write(src) fsrc.flush() fbin = fsrc.name + '.o' + use_nvcc = os.getenv("USE_NVCC", '').lower() in ('1', 'true') + os.environ.pop("USE_NVCC", None) debug_info = [] if knobs.compilation.disable_line_info: # This option is ignored if used without -lineinfo @@ -553,9 +555,11 @@ def make_cubin(self, src, metadata, opt, capability): # Accept more ptxas options if provided ptx_extra_options = opt.ptx_options.split(" ") if opt.ptx_options else [] + # If use nvshmem, we need to compile the ptx file into a relocatable object file and then link it with nvshmem library + compile_only = ["-c"] if use_nvcc else [] ptxas_cmd = [ - ptxas, *debug_info, *fmad, '-v', *disable_opt, *ptx_extra_options, f'--gpu-name={arch}', fsrc.name, - '-o', fbin + ptxas, *compile_only, *debug_info, *fmad, '-v', *disable_opt, *ptx_extra_options, f'--gpu-name={arch}', + fsrc.name, '-o', fbin ] try: subprocess.run(ptxas_cmd, check=True, close_fds=False, stderr=flog) @@ -595,8 +599,37 @@ def make_cubin(self, src, metadata, opt, capability): """) raise PTXASError(error) - with open(fbin, 'rb') as f: - cubin = f.read() + if use_nvcc: + NVLINK = os.getenv("NVLINK", "nvlink") + NVSHMEM_HOME = os.getenv("NVSHMEM_HOME") + fbin_combined = fbin + ".combined.cubin" + cuda_cubin = os.getenv("CUDA_CUBIN") + nvshmem_lib = os.path.join(NVSHMEM_HOME, "lib") + nvlink_cmds = [ + NVLINK, + f"-arch={arch}", + f"-L{nvshmem_lib}", + "-lnvshmem_device", + fbin, + cuda_cubin, + "-o", + fbin_combined, + ] + try: + subprocess.run(nvlink_cmds, check=True, close_fds=False, stderr=flog) + except Exception as e: + import logging + logging.error(f"error runing nvlink: {nvlink_cmds}") + logging.exception(e) + + if use_nvcc: + with open(fbin_combined, 'rb') as f: + cubin = f.read() + else: + with open(fbin, 'rb') as f: + cubin = f.read() + if os.path.exists(fbin_combined): + os.remove(fbin_combined) if os.path.exists(fbin): os.remove(fbin) return cubin diff --git a/third_party/nvidia/language/cuda/__init__.py b/third_party/nvidia/language/cuda/__init__.py index fbececf1d..39207bf60 100644 --- a/third_party/nvidia/language/cuda/__init__.py +++ b/third_party/nvidia/language/cuda/__init__.py @@ -1,10 +1,12 @@ from . import libdevice +from . import libnvshmem_device from .utils import (globaltimer, num_threads, num_warps, smid, convert_custom_float8_sm70, convert_custom_float8_sm80) from .gdc import (gdc_launch_dependents, gdc_wait) __all__ = [ "libdevice", + 'libnvshmem_device', "globaltimer", "num_threads", "num_warps", diff --git a/third_party/nvidia/language/cuda/libnvshmem_device.py b/third_party/nvidia/language/cuda/libnvshmem_device.py new file mode 100644 index 000000000..3dec30b9a --- /dev/null +++ b/third_party/nvidia/language/cuda/libnvshmem_device.py @@ -0,0 +1,31 @@ +from triton.language import core +import triton.language as tl + + +def _pointer_type_hash(self): + return hash((self.name, self.element_ty, "tt_ptr")) + + +def patch_hash_method_for_pointer_type(): + elem_dtype_list = tl.core.dtype.SINT_TYPES + tl.core.dtype.UINT_TYPES + tl.core.dtype.FP_TYPES + tl.core.dtype.OTHER_TYPES + for elem_dtype in elem_dtype_list: + ptr_ty = type(tl.core.pointer_type(tl.core.dtype(elem_dtype))) + ptr_ty.__hash__ = _pointer_type_hash + + +patch_hash_method_for_pointer_type() + + +@core.extern +def simple_shift(dst, _semantic=None): + return core.extern_call( + "", # libname + "", # libpath + [dst], # args + {( + core.pointer_type(core.dtype("int32")), # arg_type_symbol_dict + ): ("simple_shift", ()), # void return type + }, + is_pure=False, + _semantic=_semantic, + ) diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp index 733f103c4..d65658e38 100644 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp @@ -228,6 +228,8 @@ struct ConvertTritonGPUToLLVM targetInfo, benefit); mlir::triton::populatePrintOpToLLVMPattern(typeConverter, patterns, targetInfo, benefit); + mlir::triton::populateExternCallOpToLLVMPattern(typeConverter, patterns, + targetInfo, benefit); mlir::triton::populateControlFlowOpToLLVMPattern(typeConverter, patterns, targetInfo, benefit); mlir::triton::NVIDIA::populateSPMDOpToLLVMPattern(typeConverter, patterns, From 9588e1332bc1105a93def937e37f25d419b941fa Mon Sep 17 00:00:00 2001 From: zyuli Date: Wed, 20 May 2026 09:17:47 +0800 Subject: [PATCH 2/8] add nvshmem/example --- env.sh | 29 -- .../experimental/tle/language/raw/__init__.py | 4 +- .../experimental/tle/language/raw/core.py | 8 +- .../experimental/tle/raw/cuda/runtime.py | 27 +- .../per-token-group-quant-fp8-clang.py | 281 +++++++++++++++++ .../per-token-group-quant-fp8-nvcc.py | 283 ++++++++++++++++++ .../per-token-group-quant-fp8.cu | 138 +++++++++ .../vectorization.cuh | 30 ++ .../vectorization_utils.cuh | 177 +++++++++++ .../raw/cuda/05-vector-add-template/clang.py | 98 ++++++ .../raw/cuda/05-vector-add-template/nvcc.py | 100 +++++++ .../vector-add-template.cu | 20 ++ .../vector-add-template.py | 50 ++++ .../00-gemm-allreduce/gemm-allreduce-host.cu | 113 +++++++ .../00-gemm-allreduce/gemm-allreduce.py | 139 +++++++++ .../nvshmem/00-gemm-allreduce/simple-gemm.cu | 51 ++++ .../02-ring-reduce/ring-reduce-device.cu | 71 +++++ .../02-ring-reduce/ring-reduce-host.cu | 74 +++++ .../raw/nvshmem/02-ring-reduce/ring-reduce.py | 147 +++++++++ .../03-ring-bcast/ring-bcast-device.cu | 23 ++ .../nvshmem/03-ring-bcast/ring-bcast-host.cu | 59 ++++ .../raw/nvshmem/03-ring-bcast/ring-bcast.py | 137 +++++++++ .../on-stream-device-accumulate.cu | 8 + .../on-stream-device-correct-accumulate.cu | 13 + .../nvshmem/04-on-stream/on-stream-host.cu | 69 +++++ .../tle/raw/nvshmem/04-on-stream/on-stream.py | 166 ++++++++++ .../nvshmem/05-put-block/put-block-device.cu | 23 ++ .../nvshmem/05-put-block/put-block-host.cu | 75 +++++ .../tle/raw/nvshmem/05-put-block/put-block.py | 128 ++++++++ third_party/nvidia/backend/compiler.py | 8 +- .../nvidia/language/cuda/libnvshmem_device.py | 190 +++++++++++- 31 files changed, 2692 insertions(+), 47 deletions(-) delete mode 100644 env.sh create mode 100644 python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8-clang.py create mode 100644 python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8-nvcc.py create mode 100644 python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8.cu create mode 100644 python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/vectorization.cuh create mode 100644 python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/vectorization_utils.cuh create mode 100644 python/tutorials/tle/raw/cuda/05-vector-add-template/clang.py create mode 100644 python/tutorials/tle/raw/cuda/05-vector-add-template/nvcc.py create mode 100644 python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.cu create mode 100644 python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.py create mode 100644 python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/gemm-allreduce-host.cu create mode 100644 python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/gemm-allreduce.py create mode 100644 python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/simple-gemm.cu create mode 100644 python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce-device.cu create mode 100644 python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce-host.cu create mode 100644 python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce.py create mode 100644 python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast-device.cu create mode 100644 python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast-host.cu create mode 100644 python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast.py create mode 100644 python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-device-accumulate.cu create mode 100644 python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-device-correct-accumulate.cu create mode 100644 python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-host.cu create mode 100644 python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream.py create mode 100644 python/tutorials/tle/raw/nvshmem/05-put-block/put-block-device.cu create mode 100644 python/tutorials/tle/raw/nvshmem/05-put-block/put-block-host.cu create mode 100644 python/tutorials/tle/raw/nvshmem/05-put-block/put-block.py diff --git a/env.sh b/env.sh deleted file mode 100644 index 1603032d6..000000000 --- a/env.sh +++ /dev/null @@ -1,29 +0,0 @@ -# libz.so -# ln -s /usr/lib/x86_64-linux-gnu/libz.so.1 /home/zyuli/local/lib/libz.so -export LIBRARY_PATH=/home/zyuli/local/lib:$LIBRARY_PATH -export LD_LIBRARY_PATH=/home/zyuli/local/lib:$LD_LIBRARY_PATH - -# highest priority -# set clang include path -export CPLUS_INCLUDE_PATH=/usr/include/c++/11:/usr/include/x86_64-linux-gnu/c++/11:$CPLUS_INCLUDE_PATH -export C_INCLUDE_PATH=/usr/include/x86_64-linux-gnu:$C_INCLUDE_PATH - -# cmake reads and stores into CMAKE_CXX_FLAGS -export CXXFLAGS="-Wno-gcc-install-dir-libstdcxx --gcc-install-dir=/usr/lib/gcc/x86_64-linux-gnu/11" - -# clang and lld -export PATH=/home/zyuli/local/clang-lld/bin:$PATH - -# mlir -export PYTHONPATH=/home/zyuli/local/flagtree-mlir/python_packages/mlir_core:$PYTHONPATH - -# llvm -export LLVM_INCLUDE_DIRS=/home/zyuli/llvm-project/build-mlir/include -export LLVM_LIBRARY_DIR=/home/zyuli/llvm-project/build-mlir/lib -export LLVM_SYSPATH=/home/zyuli/llvm-project/build-mlir - -# flagtree -export MAX_JOBS=32 -export TRITON_BUILD_WITH_CLANG_LLD=1 -export TRITON_BUILD_DIR=/home/zyuli/build/flagtree_extern_call - diff --git a/python/triton/experimental/tle/language/raw/__init__.py b/python/triton/experimental/tle/language/raw/__init__.py index 790c65cf3..d66f6cd38 100644 --- a/python/triton/experimental/tle/language/raw/__init__.py +++ b/python/triton/experimental/tle/language/raw/__init__.py @@ -1,3 +1,3 @@ -from .core import call, call_smem, call_nvshmem +from .core import call, call_smem -__all__ = ["call", "call_smem", "call_nvshmem"] +__all__ = ["call", "call_smem"] diff --git a/python/triton/experimental/tle/language/raw/core.py b/python/triton/experimental/tle/language/raw/core.py index 7e2538531..350fb210b 100644 --- a/python/triton/experimental/tle/language/raw/core.py +++ b/python/triton/experimental/tle/language/raw/core.py @@ -5,6 +5,9 @@ @builtin def call(func, args, _semantic=None): + if func.libs is not None: + func.make_cubin() + return context = _semantic.builder.get_context() llvm = func.make_llvm(context) handles = [arg.handle for arg in args] @@ -46,8 +49,3 @@ def call_smem(func, args, _semantic=None): return buffer_tensors[0] else: return tl.tuple(buffer_tensors) - - -@builtin -def call_nvshmem(func, outputs, inputs, _semantic=None): - func.make_cubin() diff --git a/python/triton/experimental/tle/raw/cuda/runtime.py b/python/triton/experimental/tle/raw/cuda/runtime.py index 955a366e6..1a3d2f0e1 100644 --- a/python/triton/experimental/tle/raw/cuda/runtime.py +++ b/python/triton/experimental/tle/raw/cuda/runtime.py @@ -11,6 +11,7 @@ # 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") NVCC = os.getenv("NVCC", "nvcc") +PTXAS = os.getenv("PTXAS", "ptxas") class CUDAJITFunction(object): @@ -20,17 +21,21 @@ def __init__(self, fn: Any, file: Path, *args, **kwargs) -> None: self.fn: Final[Any] = fn self.code: Final[str] = file.read_text() self.file: Final[Path] = file - self.libs = kwargs.get("library", {}) + self.libs = kwargs.get("library", None) self.__triton_builtin__: Final[bool] = True def make_llvm(self, mlir_context) -> str: + prop = torch.cuda.get_device_properties(torch.cuda.current_device()) + arch = f"--cuda-gpu-arch=sm_{prop.major}{prop.minor}" build = subprocess.run( [ CLANG, "-x", "cuda", "--cuda-device-only", + arch, "-emit-llvm", + "-I/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/torch/include", "-O2", "-S", "-", @@ -56,7 +61,25 @@ def make_cubin(self): include_flags = [f"-I{inc_dir}" for inc_dir in include_dirs] prop = torch.cuda.get_device_properties(torch.cuda.current_device()) arch = f"-arch=sm_{prop.major}{prop.minor}" - build = subprocess.run([NVCC, "-rdc=true", arch, *include_flags, "--extended-lambda", "-c", "-o", dst, src], + + # clang cubin + # build1 = subprocess.run([ + # CLANG, + # "-fgpu-rdc", + # "-c", + # "--cuda-device-only", + # "--cuda-gpu-arch=sm_90", + # "-O3", + # *include_flags, + # "-o", + # dst, + # src + # ], capture_output=True) + # print("clang cuda") + # assert build1.returncode == 0, (f"clang failed\nstderr:\n{build1.stderr.decode()}") + + # nvcc cubin + build = subprocess.run([NVCC, "-rdc=true", arch, "-O3", *include_flags, "--extended-lambda", "-c", "-o", dst, src], capture_output=True) assert build.returncode == 0, (f"nvcc failed\nstderr:\n{build.stderr.decode()}") # TODO: Remove the method of passing information by setting environment variables. diff --git a/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8-clang.py b/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8-clang.py new file mode 100644 index 000000000..3c56cf6a4 --- /dev/null +++ b/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8-clang.py @@ -0,0 +1,281 @@ +from pathlib import Path + +from typing import Optional +import logging +import torch +import triton +import triton.language as tl +from triton.experimental.tle.raw import dialect +import triton.experimental.tle.language.raw as tle_raw + +from triton.language.extra.cuda import libnvshmem_device + +torch.cuda.set_device(1) +DEVICE = triton.runtime.driver.active.get_active_torch_device() +logger = logging.getLogger(__name__) + + +@dialect(name="cuda", file=(Path(__file__).parent / "per-token-group-quant-fp8.cu").resolve()) + # library={"torch": "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/torch/"}) +def edsl(*args, **kwargs): + ... + + +@triton.jit +def test_kernel( + x_ptr, + x_q_ptr, + x_s_ptr, + group_size, + num_groups, + groups_per_block, + eps, + fp8_min, + fp8_max, +): + tle_raw.call(edsl, [x_q_ptr, x_s_ptr, x_ptr, group_size, num_groups, groups_per_block, eps, fp8_min, fp8_max]) + # libnvshmem_device.per_token_group_quant_8bit(x_ptr, x_q_ptr, x_s_ptr, group_size, num_groups, groups_per_block, eps, + # fp8_min, fp8_max) + + +def get_groups_per_block(num_groups: int) -> int: + # Removing this branch gives better performance. + # if (num_groups % 16 == 0): + # return 16 + if (num_groups % 8 == 0): + return 8 + elif (num_groups % 4 == 0): + return 4 + elif (num_groups % 2 == 0): + return 2 + else: + return 1 + + +def per_token_group_quant_fp8_tle( + x: torch.Tensor, + group_size: int, + eps: float = 1e-10, + dtype: Optional[torch.dtype] = None, + # column_major_scales: bool = False, + # scale_ue8m0: bool = False, +): + logger.debug("GEMS PER TOKEN GROUP QUANT FP8") + assert x.shape[-1] % group_size == 0, (f"the last dimension of `x` {x.shape[-1]} must be divisible " + f"by `group_size` {group_size}") + assert x.stride(-1) == 1, "`x` groups must be contiguous" + + fp8_dtype = torch.float8_e4m3fn if dtype is None else dtype + finfo = torch.finfo(fp8_dtype) + fp8_min = finfo.min + fp8_max = finfo.max + + x_q = torch.empty_like(x, device=x.device, dtype=fp8_dtype) + shape = x.shape[:-1] + (x.shape[-1] // group_size, ) + x_s = torch.empty(shape, device=x.device, dtype=torch.float32) + + # num_groups + num_groups = x.numel() // group_size + groups_per_block = get_groups_per_block(num_groups) + + # num_blocks + THREADS_PER_GROUP = 16 + num_blocks = num_groups // groups_per_block + num_warps = max(groups_per_block * THREADS_PER_GROUP // 32, 1) + + # The .cu device function uses `extern __shared__` for groups_per_block * group_size floats. + # Triton's compiler cannot infer this smem requirement from the extern_call, so we patch + # packed_metadata after warmup compilation to include the extra shared memory bytes. + smem_bytes = groups_per_block * group_size * x.element_size() # float32 = 4 bytes + + kernel = test_kernel.run( + x, + x_q, + x_s, + group_size, + num_groups, + groups_per_block, + eps, + fp8_min, + fp8_max, + grid=(num_blocks, ), + warmup=True, + num_warps=num_warps, + ) + + # Resolve async future if needed + if hasattr(kernel, "result"): + kernel = kernel.result() + + old_meta = kernel.packed_metadata + new_shared = max(old_meta[2], smem_bytes) + kernel.packed_metadata = old_meta[:2] + (new_shared, ) + old_meta[3:] + + test_kernel[(num_blocks, )]( + x, + x_q, + x_s, + group_size, + num_groups, + groups_per_block, + eps, + fp8_min, + fp8_max, + num_warps=num_warps, + ) + + return x_q, x_s + + +@triton.jit +def _per_token_group_quant_fp8( + y_ptr, + y_q_ptr, + y_s_ptr, + group_size, + y_num_columns, + y_row_stride, + eps, + fp8_min, + fp8_max, + scale_ue8m0, + BLOCK: tl.constexpr, +): + groups_per_row = y_num_columns // group_size + + g_id = tl.program_id(0) + row = g_id // groups_per_row + row_g_id = g_id % groups_per_row + + y_ptr += (row * y_row_stride) + (row_g_id * group_size) + y_q_ptr += g_id * group_size + y_s_ptr += g_id + + cols = tl.arange(0, BLOCK) + mask = cols < group_size + + y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32) + _absmax = tl.maximum(tl.max(tl.abs(y)), eps) + y_s = _absmax / fp8_max + + if scale_ue8m0: + y_s = tl.exp2(tl.ceil(tl.log2(tl.maximum(tl.abs(y_s), 1e-10)))) + + y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) + + tl.store(y_q_ptr + cols, y_q, mask=mask) + tl.store(y_s_ptr, y_s) + + +def per_token_group_quant_fp8_triton( + x: torch.Tensor, + group_size: int, + eps: float = 1e-10, + dtype: Optional[torch.dtype] = None, + column_major_scales: bool = False, + scale_ue8m0: bool = False, +): + logger.debug("GEMS PER TOKEN GROUP QUANT FP8") + # dtype: The dype of output tensor. Note that only `torch.float8_e4m3fn` + fp8_dtype = torch.float8_e4m3fn if dtype is None else dtype + assert x.shape[-1] % group_size == 0, (f"the last dimension of `x` {x.shape[-1]} must be divisible " + f"by `group_size` {group_size}") + assert x.stride(-1) == 1, "`x` groups must be contiguous" + + finfo = torch.finfo(fp8_dtype) + fp8_min = finfo.min + fp8_max = finfo.max + + x_q = torch.empty_like(x, device=x.device, dtype=fp8_dtype) + M = x.numel() // group_size + N = group_size + + if column_major_scales: + shape = (x.shape[-1] // group_size, ) + x.shape[:-1] + x_s = torch.empty(shape, device=x.device, dtype=torch.float32).permute(-1, -2) + else: + shape = x.shape[:-1] + (x.shape[-1] // group_size, ) + x_s = torch.empty(shape, device=x.device, dtype=torch.float32) + + BLOCK = triton.next_power_of_2(N) + num_warps = min(max(BLOCK // 256, 1), 8) + num_stages = 1 + _per_token_group_quant_fp8[(M, )]( + x, + x_q, + x_s, + group_size, + x.shape[1], + x.stride(0), + eps, + fp8_min=fp8_min, + fp8_max=fp8_max, + scale_ue8m0=scale_ue8m0, + BLOCK=BLOCK, + num_warps=num_warps, + num_stages=num_stages, + ) + + return x_q, x_s + + +if __name__ == "__main__": + x = torch.randn((16384, 32768), device=DEVICE, dtype=torch.float32) + group_size = 128 + + x_q_triton, x_s_triton = per_token_group_quant_fp8_tle(x, group_size) + # x_q_triton, x_s_triton = per_token_group_quant_fp8_triton(x, group_size) + # x_q_tle, x_s_tle = per_token_group_quant_fp8_tle(x, group_size) + + # q_tri = x_q_triton.to(torch.float32) + # q_tle = x_q_tle.to(torch.float32) + + # b_tri = x_q_triton.view(torch.int8).to(torch.int16) # promote to avoid int8 overflow + # b_tle = x_q_tle.view(torch.int8).to(torch.int16) + # bit_diff = (b_tri - b_tle).abs() # 0 = exact match, 1 = 1-ULP, >1 = real bug + # num_exact = bit_diff.eq(0).sum().item() + # num_1ulp = bit_diff.eq(1).sum().item() + # num_beyond = bit_diff.gt(1).sum().item() + # if num_beyond == 0: + # if num_1ulp == 0: + # print("✅ x_q Triton and TLE match (bit-exact)") + # else: + # print(f"✅ x_q Triton and TLE match (1-ULP diff={num_1ulp}, " + # f"expected from RTZ vs RTNE rounding)") + # else: + # q_tri = x_q_triton.to(torch.float32) + # q_tle = x_q_tle.to(torch.float32) + # float_diff = (q_tri - q_tle).abs() + # print(f"❌ x_q Triton and TLE differ: " + # f"exact={num_exact}, 1-ULP={num_1ulp}, >1-ULP={num_beyond}") + # beyond_idx = bit_diff.gt(1).nonzero() + # for idx in beyond_idx[:10]: + # r, c = idx[0].item(), idx[1].item() + # group_id = r * (q_tri.shape[1] // group_size) + c // group_size + # pos_in_group = c % group_size + # print(f" [{r},{c}] group={group_id} pos_in_group={pos_in_group}" + # f" triton_bits={b_tri[r,c].item():4d} tle_bits={b_tle[r,c].item():4d}" + # f" bit_diff={bit_diff[r,c].item()}" + # f" triton={q_tri[r,c].item():.1f} tle={q_tle[r,c].item():.1f}" + # f" float_diff={float_diff[r,c].item():.1f}" + # f" x={x[r,c].item():.6f} scale={x_s_triton.view(-1)[group_id].item():.8f}") + + # if torch.allclose(x_s_triton, x_s_tle, atol=0.125, rtol=0): + # print("✅ x_s Triton and TLE match") + # else: + # print("❌ x_s Triton and TLE differ") + + # # perf + # shapes = [(1024, 2048), (2048, 4096), (4096, 8192), (16384, 32768)] + # group_sizes = [64, 128] + # dtypes = [torch.float32] + # for shape in shapes: + # for dtype in dtypes: + # for group_size in group_sizes: + # x = torch.rand(shape, device=DEVICE, dtype=dtype) + # mean_ms_triton = triton.testing.do_bench(lambda: per_token_group_quant_fp8_triton(x, group_size)) + # mean_ms_tle = triton.testing.do_bench(lambda: per_token_group_quant_fp8_tle(x, group_size)) + + # print(f"\n========= Shape: {shape} Type: {dtype} Group size: {group_size} =========") + # print(f"Triton time: {mean_ms_triton:.4f} ms") + # print(f"TLE time: {mean_ms_tle:.4f} ms") diff --git a/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8-nvcc.py b/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8-nvcc.py new file mode 100644 index 000000000..da54552ce --- /dev/null +++ b/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8-nvcc.py @@ -0,0 +1,283 @@ +from pathlib import Path + +from typing import Optional +import logging +import torch +import triton +import triton.language as tl +from triton.experimental.tle.raw import dialect +import triton.experimental.tle.language.raw as tle_raw + +from triton.language.extra.cuda import libnvshmem_device + +torch.cuda.set_device(1) +DEVICE = triton.runtime.driver.active.get_active_torch_device() +logger = logging.getLogger(__name__) + + +@dialect(name="cuda", file=(Path(__file__).parent / "per-token-group-quant-fp8.cu").resolve(), + library={"torch": "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/torch/"}) +def edsl(*args, **kwargs): + ... + + +@triton.jit +def test_kernel( + x_ptr, + x_q_ptr, + x_s_ptr, + group_size, + num_groups, + groups_per_block, + eps, + fp8_min, + fp8_max, +): + tle_raw.call(edsl, []) + libnvshmem_device.per_token_group_quant_8bit(x_q_ptr, x_s_ptr, x_ptr, group_size, num_groups, groups_per_block, eps, + fp8_min, fp8_max) + + +def get_groups_per_block(num_groups: int) -> int: + # Removing this branch gives better performance. + # if (num_groups % 16 == 0): + # return 16 + if (num_groups % 8 == 0): + return 8 + elif (num_groups % 4 == 0): + return 4 + elif (num_groups % 2 == 0): + return 2 + else: + return 1 + + +def per_token_group_quant_fp8_tle( + x: torch.Tensor, + group_size: int, + eps: float = 1e-10, + dtype: Optional[torch.dtype] = None, + # column_major_scales: bool = False, + # scale_ue8m0: bool = False, +): + logger.debug("GEMS PER TOKEN GROUP QUANT FP8") + assert x.shape[-1] % group_size == 0, (f"the last dimension of `x` {x.shape[-1]} must be divisible " + f"by `group_size` {group_size}") + assert x.stride(-1) == 1, "`x` groups must be contiguous" + + fp8_dtype = torch.float8_e4m3fn if dtype is None else dtype + finfo = torch.finfo(fp8_dtype) + fp8_min = finfo.min + fp8_max = finfo.max + + x_q = torch.empty_like(x, device=x.device, dtype=fp8_dtype) + shape = x.shape[:-1] + (x.shape[-1] // group_size, ) + x_s = torch.empty(shape, device=x.device, dtype=torch.float32) + + # num_groups + num_groups = x.numel() // group_size + groups_per_block = get_groups_per_block(num_groups) + + # num_blocks + THREADS_PER_GROUP = 16 + num_blocks = num_groups // groups_per_block + num_warps = max(groups_per_block * THREADS_PER_GROUP // 32, 1) + + # The .cu device function uses `extern __shared__` for groups_per_block * group_size floats. + # Triton's compiler cannot infer this smem requirement from the extern_call, so we patch + # packed_metadata after warmup compilation to include the extra shared memory bytes. + smem_bytes = groups_per_block * group_size * x.element_size() # float32 = 4 bytes + + kernel = test_kernel.run( + x, + x_q, + x_s, + group_size, + num_groups, + groups_per_block, + eps, + fp8_min, + fp8_max, + grid=(num_blocks, ), + warmup=True, + num_warps=num_warps, + ) + + # Resolve async future if needed + if hasattr(kernel, "result"): + kernel = kernel.result() + + old_meta = kernel.packed_metadata + new_shared = max(old_meta[2], smem_bytes) + kernel.packed_metadata = old_meta[:2] + (new_shared, ) + old_meta[3:] + + test_kernel[(num_blocks, )]( + x, + x_q, + x_s, + group_size, + num_groups, + groups_per_block, + eps, + fp8_min, + fp8_max, + num_warps=num_warps, + ) + + return x_q, x_s + + +@triton.jit +def _per_token_group_quant_fp8( + y_ptr, + y_q_ptr, + y_s_ptr, + group_size, + y_num_columns, + y_row_stride, + eps, + fp8_min, + fp8_max, + scale_ue8m0, + BLOCK: tl.constexpr, +): + groups_per_row = y_num_columns // group_size + + g_id = tl.program_id(0) + row = g_id // groups_per_row + row_g_id = g_id % groups_per_row + + y_ptr += (row * y_row_stride) + (row_g_id * group_size) + y_q_ptr += g_id * group_size + y_s_ptr += g_id + + cols = tl.arange(0, BLOCK) + mask = cols < group_size + + y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32) + _absmax = tl.maximum(tl.max(tl.abs(y)), eps) + y_s = _absmax / fp8_max + + if scale_ue8m0: + y_s = tl.exp2(tl.ceil(tl.log2(tl.maximum(tl.abs(y_s), 1e-10)))) + + y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) + + tl.store(y_q_ptr + cols, y_q, mask=mask) + tl.store(y_s_ptr, y_s) + + +def per_token_group_quant_fp8_triton( + x: torch.Tensor, + group_size: int, + eps: float = 1e-10, + dtype: Optional[torch.dtype] = None, + column_major_scales: bool = False, + scale_ue8m0: bool = False, +): + logger.debug("GEMS PER TOKEN GROUP QUANT FP8") + # dtype: The dype of output tensor. Note that only `torch.float8_e4m3fn` + fp8_dtype = torch.float8_e4m3fn if dtype is None else dtype + assert x.shape[-1] % group_size == 0, (f"the last dimension of `x` {x.shape[-1]} must be divisible " + f"by `group_size` {group_size}") + assert x.stride(-1) == 1, "`x` groups must be contiguous" + + finfo = torch.finfo(fp8_dtype) + fp8_min = finfo.min + fp8_max = finfo.max + + x_q = torch.empty_like(x, device=x.device, dtype=fp8_dtype) + M = x.numel() // group_size + N = group_size + + if column_major_scales: + shape = (x.shape[-1] // group_size, ) + x.shape[:-1] + x_s = torch.empty(shape, device=x.device, dtype=torch.float32).permute(-1, -2) + else: + shape = x.shape[:-1] + (x.shape[-1] // group_size, ) + x_s = torch.empty(shape, device=x.device, dtype=torch.float32) + + BLOCK = triton.next_power_of_2(N) + num_warps = min(max(BLOCK // 256, 1), 8) + num_stages = 1 + _per_token_group_quant_fp8[(M, )]( + x, + x_q, + x_s, + group_size, + x.shape[1], + x.stride(0), + eps, + fp8_min=fp8_min, + fp8_max=fp8_max, + scale_ue8m0=scale_ue8m0, + BLOCK=BLOCK, + num_warps=num_warps, + num_stages=num_stages, + ) + + return x_q, x_s + + +if __name__ == "__main__": + x = torch.randn((16384, 32768), device=DEVICE, dtype=torch.float32) + group_size = 128 + + x_q_triton, x_s_triton = per_token_group_quant_fp8_tle(x, group_size) + # mean_ms_tle = triton.testing.do_bench(lambda: per_token_group_quant_fp8_tle(x, group_size)) + # print(f"TLE time: {mean_ms_tle:.4f} ms") + # x_q_triton, x_s_triton = per_token_group_quant_fp8_triton(x, group_size) + # x_q_tle, x_s_tle = per_token_group_quant_fp8_tle(x, group_size) + + # q_tri = x_q_triton.to(torch.float32) + # q_tle = x_q_tle.to(torch.float32) + + # b_tri = x_q_triton.view(torch.int8).to(torch.int16) # promote to avoid int8 overflow + # b_tle = x_q_tle.view(torch.int8).to(torch.int16) + # bit_diff = (b_tri - b_tle).abs() # 0 = exact match, 1 = 1-ULP, >1 = real bug + # num_exact = bit_diff.eq(0).sum().item() + # num_1ulp = bit_diff.eq(1).sum().item() + # num_beyond = bit_diff.gt(1).sum().item() + # if num_beyond == 0: + # if num_1ulp == 0: + # print("✅ x_q Triton and TLE match (bit-exact)") + # else: + # print(f"✅ x_q Triton and TLE match (1-ULP diff={num_1ulp}, " + # f"expected from RTZ vs RTNE rounding)") + # else: + # q_tri = x_q_triton.to(torch.float32) + # q_tle = x_q_tle.to(torch.float32) + # float_diff = (q_tri - q_tle).abs() + # print(f"❌ x_q Triton and TLE differ: " + # f"exact={num_exact}, 1-ULP={num_1ulp}, >1-ULP={num_beyond}") + # beyond_idx = bit_diff.gt(1).nonzero() + # for idx in beyond_idx[:10]: + # r, c = idx[0].item(), idx[1].item() + # group_id = r * (q_tri.shape[1] // group_size) + c // group_size + # pos_in_group = c % group_size + # print(f" [{r},{c}] group={group_id} pos_in_group={pos_in_group}" + # f" triton_bits={b_tri[r,c].item():4d} tle_bits={b_tle[r,c].item():4d}" + # f" bit_diff={bit_diff[r,c].item()}" + # f" triton={q_tri[r,c].item():.1f} tle={q_tle[r,c].item():.1f}" + # f" float_diff={float_diff[r,c].item():.1f}" + # f" x={x[r,c].item():.6f} scale={x_s_triton.view(-1)[group_id].item():.8f}") + + # if torch.allclose(x_s_triton, x_s_tle, atol=0.125, rtol=0): + # print("✅ x_s Triton and TLE match") + # else: + # print("❌ x_s Triton and TLE differ") + + # # perf + # shapes = [(1024, 2048), (2048, 4096), (4096, 8192), (16384, 32768)] + # group_sizes = [64, 128] + # dtypes = [torch.float32] + # for shape in shapes: + # for dtype in dtypes: + # for group_size in group_sizes: + # x = torch.rand(shape, device=DEVICE, dtype=dtype) + # mean_ms_triton = triton.testing.do_bench(lambda: per_token_group_quant_fp8_triton(x, group_size)) + # mean_ms_tle = triton.testing.do_bench(lambda: per_token_group_quant_fp8_tle(x, group_size)) + + # print(f"\n========= Shape: {shape} Type: {dtype} Group size: {group_size} =========") + # print(f"Triton time: {mean_ms_triton:.4f} ms") + # print(f"TLE time: {mean_ms_tle:.4f} ms") diff --git a/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8.cu b/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8.cu new file mode 100644 index 000000000..2add918ae --- /dev/null +++ b/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8.cu @@ -0,0 +1,138 @@ +#include "vectorization.cuh" +#include "vectorization_utils.cuh" +#include +#include + +__device__ __forceinline__ float GroupReduceMax(float val) { + unsigned mask = threadIdx.x % 32 >= 16 ? 0xffff0000 : 0x0000ffff; + + val = fmaxf(val, __shfl_xor_sync(mask, val, 8)); + val = fmaxf(val, __shfl_xor_sync(mask, val, 4)); + val = fmaxf(val, __shfl_xor_sync(mask, val, 2)); + val = fmaxf(val, __shfl_xor_sync(mask, val, 1)); + return val; +} + +// template +__device__ __forceinline__ float +ComputeGroupScale(const float *__restrict__ group_input, + float *__restrict__ smem_group, const int group_size, + const int lane_id, const int threads_per_group, + const float eps, const float max_8bit) { + float local_absmax = eps; + + constexpr int vec_size = 16 / sizeof(float); + + // copy global -> shared & compute absmax + auto scalar_op_cache = [&] __device__(float &dst, const float &src) { + float abs_v = fabsf(static_cast(src)); + local_absmax = fmaxf(local_absmax, abs_v); + dst = src; + }; + + vllm::vectorize_with_alignment(group_input, // in + smem_group, // out (shared) + group_size, // elements per group + lane_id, // thread id + threads_per_group, // stride in group + scalar_op_cache); // scalar handler + + local_absmax = GroupReduceMax(local_absmax); + + float y_s = local_absmax / max_8bit; + // if constexpr (SCALE_UE8M0) { + // y_s = exp2f(ceilf(log2f(fmaxf(fabsf(y_s), 1e-10f)))); + // } + + return y_s; +} + +// template +__device__ __forceinline__ void +QuantizeGroup(const float *__restrict__ smem_group, + __nv_fp8_e4m3 *__restrict__ group_output, const int group_size, + const int lane_id, const int threads_per_group, const float y_s, + const float min_8bit, const float max_8bit) { + constexpr int vec_size = 16 / sizeof(float); + + // quantize shared -> global 8-bit + auto scalar_op_quant = [&] __device__(__nv_fp8_e4m3 & dst, const float &src) { + float q = fminf(fmaxf(static_cast(src) / y_s, min_8bit), max_8bit); + dst = __nv_fp8_e4m3(q); + }; + + vllm::vectorize_with_alignment( + smem_group, // in (shared) + group_output, // out (global quant tensor) + group_size, // elements + lane_id, // tid + threads_per_group, // stride + scalar_op_quant); // scalar handler +} + +// T: float; DST_DTYPE: __nv_fp8_e4m3 +// template +// __global__ void per_token_group_quant_8bit_kernel( +extern "C" __device__ void per_token_group_quant_8bit( + __attribute__((address_space(1))) void *__restrict__ output_q, + __attribute__((address_space(1))) float *__restrict__ output_s, + __attribute__((address_space(1))) const float *__restrict__ input, + const int group_size, + const int num_groups, + const int groups_per_block, + const float eps, + const float min_8bit, + const float max_8bit) { + const int threads_per_group = 16; + const int64_t local_group_id = threadIdx.x / threads_per_group; + const int lane_id = threadIdx.x % threads_per_group; + + const int64_t block_group_id = blockIdx.x * groups_per_block; + const int64_t global_group_id = block_group_id + local_group_id; + const int64_t block_group_offset = global_group_id * group_size; + + static_assert(sizeof(float) % sizeof(float) == 0); + + const float *group_input = input + block_group_offset; + __nv_fp8_e4m3 *group_output = + static_cast<__nv_fp8_e4m3 *>(output_q) + block_group_offset; + float *scale_output; + + // bool IS_COLUMN_MAJOR = false; + // if (IS_COLUMN_MAJOR) { + // const int num_elems_per_pack = + // static_cast(sizeof(float) / sizeof(float)); + // const int scale_num_rows_element = scale_num_rows * num_elems_per_pack; + // const int row_idx = global_group_id / scale_num_rows_element; + // const int col_idx_raw = global_group_id % scale_num_rows_element; + // const int col_idx = col_idx_raw / num_elems_per_pack; + // const int pack_idx = col_idx_raw % num_elems_per_pack; + // scale_output = reinterpret_cast(output_s) + + // (col_idx * scale_stride * num_elems_per_pack + + // row_idx * num_elems_per_pack + pack_idx); + // } else { + scale_output = output_s + global_group_id; + // } + + // shared memory to cache each group's data to avoid double DRAM reads. + extern __shared__ __align__(16) char smem_raw[]; + float *smem = reinterpret_cast(smem_raw); + float *smem_group = smem + local_group_id * group_size; + + const float y_s = + ComputeGroupScale(group_input, smem_group, group_size, lane_id, + threads_per_group, eps, max_8bit); + + float y_s_quant = y_s; + + if (lane_id == 0) { + *scale_output = y_s_quant; + } + + __syncthreads(); + + QuantizeGroup(smem_group, group_output, group_size, lane_id, + threads_per_group, y_s, min_8bit, max_8bit); +} + diff --git a/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/vectorization.cuh b/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/vectorization.cuh new file mode 100644 index 000000000..624bb9fe1 --- /dev/null +++ b/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/vectorization.cuh @@ -0,0 +1,30 @@ +#pragma once +/** + * __device__ datatypes vectorized by 4 + */ + +// Include both AMD and NVIDIA fp8 types to avoid circular import +#include +#include + +namespace vllm { + +// Vectorization containers +template +struct __align__(vec_size * sizeof(scalar_t)) vec_n_t { + scalar_t val[vec_size]; +}; + +template +struct __align__(vec_size * sizeof(quant_type_t)) q8_n_t { + static_assert(std::is_same_v || + std::is_same_v || + std::is_same_v); + quant_type_t val[vec_size]; +}; + +template using vec4_t = vec_n_t; +template using q8x4_t = q8_n_t; + +} // namespace vllm + diff --git a/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/vectorization_utils.cuh b/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/vectorization_utils.cuh new file mode 100644 index 000000000..3323454db --- /dev/null +++ b/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/vectorization_utils.cuh @@ -0,0 +1,177 @@ +#pragma once +#include "vectorization.cuh" + +namespace vllm { + +template +struct DefaultVecOp { + ScaOp scalar_op; + + __device__ __forceinline__ void + operator()(vec_n_t &dst, + const vec_n_t &src) const { +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + scalar_op(dst.val[i], src.val[i]); + } + } +}; + +template +__device__ inline void +vectorize_with_alignment(const InT *in, OutT *out, int len, int tid, int stride, + VecOp &&vec_op, // vec_n_t -> vec_n_t + ScaOp &&scalar_op) { // InT -> OutT + static_assert(VEC_SIZE > 0 && (VEC_SIZE & (VEC_SIZE - 1)) == 0, + "VEC_SIZE must be a positive power-of-two"); + constexpr int WIDTH = VEC_SIZE * sizeof(InT); // eg: 64 B + uintptr_t addr = reinterpret_cast(in); + + // fast path when the whole region is already aligned + // Note: currently the output is guaranteed to be same as the input, so we + // don't check it here, comments here just for future reference. + bool can_vec = ((addr & (WIDTH - 1)) == 0) && ((len & (VEC_SIZE - 1)) == 0); + if (can_vec) { + int num_vec = len / VEC_SIZE; + + using vin_t = vec_n_t; + using vout_t = vec_n_t; + auto *v_in = reinterpret_cast(in); + auto *v_out = reinterpret_cast(out); + + for (int i = tid; i < num_vec; i += stride) { + vout_t tmp; + // Make a local copy of the entire pack + vin_t src = v_in[i]; // <- encourages a single vector ld + vec_op(tmp, src); + v_out[i] = tmp; // <- encourages a single vector st + } + return; + } + + int misalignment_offset = addr & (WIDTH - 1); // addr % 64 + int alignment_bytes = WIDTH - misalignment_offset; // 64 - (addr % 64) + int prefix_elems = alignment_bytes & (WIDTH - 1); // handle 64 + prefix_elems /= sizeof(InT); + prefix_elems = min(prefix_elems, len); // 0 ≤ prefix < 16 + + // 1. prefill the when it is unsafe to vectorize + for (int i = tid; i < prefix_elems; i += stride) { + scalar_op(out[i], in[i]); + } + + in += prefix_elems; + out += prefix_elems; + len -= prefix_elems; + + int num_vec = len / VEC_SIZE; + using vin_t = vec_n_t; + using vout_t = vec_n_t; + auto *v_in = reinterpret_cast(in); + auto *v_out = reinterpret_cast(out); + + // 2. vectorize the main part + for (int i = tid; i < num_vec; i += stride) { + vout_t tmp; + // Make a local copy of the entire pack + vin_t src = v_in[i]; // <- encourages a single vector ld + vec_op(tmp, src); + v_out[i] = tmp; // <- encourages a single vector st + } + + // 3. handle the tail + int tail_start = num_vec * VEC_SIZE; + for (int i = tid + tail_start; i < len; i += stride) { + scalar_op(out[i], in[i]); + } +} + +template +__device__ __forceinline__ void +vectorize_with_alignment(const InT *in, OutT *out, int len, int tid, int stride, + ScaOp &&scalar_op) { + using Vec = DefaultVecOp>; + vectorize_with_alignment(in, out, len, tid, stride, Vec{scalar_op}, + std::forward(scalar_op)); +} + +template struct DefaultReadVecOp { + ScaOp scalar_op; + + __device__ __forceinline__ void + operator()(const vec_n_t &src) const { +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + scalar_op(src.val[i]); + } + } +}; + +// read-only version: iterate over the input with alignment guarantees +template +__device__ inline void +vectorize_read_with_alignment(const InT *in, int len, int tid, int stride, + VecOp &&vec_op, ScaOp &&scalar_op) { + static_assert(VEC_SIZE > 0 && (VEC_SIZE & (VEC_SIZE - 1)) == 0, + "VEC_SIZE must be a positive power-of-two"); + constexpr int WIDTH = VEC_SIZE * sizeof(InT); + uintptr_t addr = reinterpret_cast(in); + + // fast path when the whole region is already aligned + bool can_vec = ((addr & (WIDTH - 1)) == 0) && ((len & (VEC_SIZE - 1)) == 0); + if (can_vec) { + int num_vec = len / VEC_SIZE; + + using vin_t = vec_n_t; + auto *v_in = reinterpret_cast(in); + + for (int i = tid; i < num_vec; i += stride) { + vin_t tmp = v_in[i]; + vec_op(tmp); + } + return; + } + + int misalignment_offset = addr & (WIDTH - 1); + int alignment_bytes = WIDTH - misalignment_offset; + int prefix_elems = alignment_bytes & (WIDTH - 1); + prefix_elems /= sizeof(InT); + prefix_elems = min(prefix_elems, len); + + // 1. handle the possibly unaligned prefix with scalar access. + for (int i = tid; i < prefix_elems; i += stride) { + scalar_op(in[i]); + } + + in += prefix_elems; + len -= prefix_elems; + + int num_vec = len / VEC_SIZE; + using vin_t = vec_n_t; + auto *v_in = reinterpret_cast(in); + + // 2. vectorized traversal of the main aligned region. + for (int i = tid; i < num_vec; i += stride) { + vec_op(v_in[i]); + } + + // 3. handle remaining tail elements. + int tail_start = num_vec * VEC_SIZE; + for (int i = tid + tail_start; i < len; i += stride) { + scalar_op(in[i]); + } +} + +// overload that requires only a scalar_op +template +__device__ __forceinline__ void +vectorize_read_with_alignment(const InT *in, int len, int tid, int stride, + ScaOp &&scalar_op) { + using Vec = DefaultReadVecOp>; + vectorize_read_with_alignment(in, len, tid, stride, Vec{scalar_op}, + std::forward(scalar_op)); +} + +} // namespace vllm + diff --git a/python/tutorials/tle/raw/cuda/05-vector-add-template/clang.py b/python/tutorials/tle/raw/cuda/05-vector-add-template/clang.py new file mode 100644 index 000000000..e09ff0ceb --- /dev/null +++ b/python/tutorials/tle/raw/cuda/05-vector-add-template/clang.py @@ -0,0 +1,98 @@ +from pathlib import Path + +import torch +import triton +import triton.language as tl +from triton.experimental.tle.raw import dialect +import triton.experimental.tle.language.raw as tle_raw + +DEVICE = triton.runtime.driver.active.get_active_torch_device() + + +@dialect(name="cuda", file=Path(__file__).parent / "01-vector-add.cu") +def edsl(*args, **kwargs): + ... + + +@triton.jit +def add_kernel_tle( + x_ptr, + y_ptr, + output_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, +): + tle_raw.call(edsl, [output_ptr, x_ptr, y_ptr, n_elements]) + + +def add_tle(x: torch.Tensor, y: torch.Tensor): + output = torch.empty_like(x) + assert x.device == DEVICE and y.device == DEVICE and output.device == DEVICE + n_elements = output.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), ) + add_kernel_tle[grid](x, y, output, n_elements, BLOCK_SIZE=1024) + return output + +@triton.jit +def add_kernel_triton(x_ptr, # *Pointer* to first input vector. + y_ptr, # *Pointer* to second input vector. + output_ptr, # *Pointer* to output vector. + n_elements, # Size of the vector. + BLOCK_SIZE: tl.constexpr, # Number of elements each program should process. + ): + pid = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0. + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = x + y + tl.store(output_ptr + offsets, output, mask=mask) + +def add_triton(x: torch.Tensor, y: torch.Tensor): + output = torch.empty_like(x) + assert x.device == DEVICE and y.device == DEVICE and output.device == DEVICE + n_elements = output.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) + add_kernel_triton[grid](x, y, output, n_elements, BLOCK_SIZE=1024) + return output + + +if __name__ == "__main__": + x = torch.randn((16384, 32768), device=DEVICE, dtype=torch.float32) + y = torch.randn((16384, 32768), device=DEVICE, dtype=torch.float32) + output_tle = add_tle(x, y) + # mean_ms_tle = triton.testing.do_bench(lambda: add_tle(x, y)) + # print(f"TLE time: {mean_ms_tle:.4f} ms") + output_triton = add_triton(x, y) + + if torch.allclose(output_triton, output_tle, atol=0.125, rtol=0): + print("✅ output Triton and TLE match") + else: + print("❌ output Triton and TLE differ") + + + # perf + dtype = torch.float32 + for i in range(27, 31, 1): + x = torch.randn(2 ** i, device=DEVICE, dtype=dtype) + y = torch.randn(2 ** i, device=DEVICE, dtype=dtype) + mean_ms_tle = triton.testing.do_bench(lambda: add_tle(x, y)) + mean_ms_triton = triton.testing.do_bench(lambda: add_triton(x, y)) + + print(f"\n======== x: 2 **{i} type: {dtype} =========") + print(f"Triton Time: {mean_ms_triton:.3f} ms") + print(f"TLE Time: {mean_ms_tle:.3f} ms") + + shapes = [(1024, 2048), (2048, 4096), (4096, 8192), (16384, 32768)] + dtypes = [torch.float32] + for shape in shapes: + for dtype in dtypes: + x = torch.rand(shape, device=DEVICE, dtype=dtype) + y = torch.rand(shape, device=DEVICE, dtype=dtype) + mean_ms_triton = triton.testing.do_bench(lambda: add_triton(x, y)) + mean_ms_tle = triton.testing.do_bench(lambda: add_tle(x, y)) + + print(f"\n========= Shape: {shape} Type: {dtype} =========") + print(f"Triton time: {mean_ms_triton:.4f} ms") + print(f"TLE time: {mean_ms_tle:.4f} ms") diff --git a/python/tutorials/tle/raw/cuda/05-vector-add-template/nvcc.py b/python/tutorials/tle/raw/cuda/05-vector-add-template/nvcc.py new file mode 100644 index 000000000..866c0f2b9 --- /dev/null +++ b/python/tutorials/tle/raw/cuda/05-vector-add-template/nvcc.py @@ -0,0 +1,100 @@ +from pathlib import Path + +import torch +import triton +import triton.language as tl +from triton.experimental.tle.raw import dialect +import triton.experimental.tle.language.raw as tle_raw +from triton.language.extra.cuda import libnvshmem_device + +DEVICE = triton.runtime.driver.active.get_active_torch_device() + + +@dialect(name="cuda", file=Path(__file__).parent / "01-vector-add.cu", library={}) +def edsl(*args, **kwargs): + ... + + +@triton.jit +def add_kernel_tle( + x_ptr, + y_ptr, + output_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, +): + tle_raw.call(edsl, []) + libnvshmem_device.vector_add(output_ptr, x_ptr, y_ptr, n_elements) + + +def add_tle(x: torch.Tensor, y: torch.Tensor): + output = torch.empty_like(x) + assert x.device == DEVICE and y.device == DEVICE and output.device == DEVICE + n_elements = output.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), ) + add_kernel_tle[grid](x, y, output, n_elements, BLOCK_SIZE=1024) + return output + +@triton.jit +def add_kernel_triton(x_ptr, # *Pointer* to first input vector. + y_ptr, # *Pointer* to second input vector. + output_ptr, # *Pointer* to output vector. + n_elements, # Size of the vector. + BLOCK_SIZE: tl.constexpr, # Number of elements each program should process. + ): + pid = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0. + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = x + y + tl.store(output_ptr + offsets, output, mask=mask) + +def add_triton(x: torch.Tensor, y: torch.Tensor): + output = torch.empty_like(x) + assert x.device == DEVICE and y.device == DEVICE and output.device == DEVICE + n_elements = output.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) + add_kernel_triton[grid](x, y, output, n_elements, BLOCK_SIZE=1024) + return output + + +if __name__ == "__main__": + x = torch.randn((16384, 32768), device=DEVICE, dtype=torch.float32) + y = torch.randn((16384, 32768), device=DEVICE, dtype=torch.float32) + output_tle = add_tle(x, y) + # mean_ms_tle = triton.testing.do_bench(lambda: add_tle(x, y)) + # print(f"TLE time: {mean_ms_tle:.4f} ms") + output_triton = add_triton(x, y) + + if torch.allclose(output_triton, output_tle, atol=0.125, rtol=0): + print("✅ output Triton and TLE match") + else: + print("❌ output Triton and TLE differ") + + + # perf + dtype = torch.float32 + for i in range(27, 31, 1): + x = torch.randn(2 ** i, device=DEVICE, dtype=dtype) + y = torch.randn(2 ** i, device=DEVICE, dtype=dtype) + mean_ms_tle = triton.testing.do_bench(lambda: add_tle(x, y)) + mean_ms_triton = triton.testing.do_bench(lambda: add_triton(x, y)) + + print(f"\n======== x: 2 **{i} type: {dtype} =========") + print(f"Triton Time: {mean_ms_triton:.3f} ms") + print(f"TLE Time: {mean_ms_tle:.3f} ms") + + shapes = [(1024, 2048), (2048, 4096), (4096, 8192), (16384, 32768)] + dtypes = [torch.float32] + for shape in shapes: + for dtype in dtypes: + x = torch.rand(shape, device=DEVICE, dtype=dtype) + y = torch.rand(shape, device=DEVICE, dtype=dtype) + mean_ms_triton = triton.testing.do_bench(lambda: add_triton(x, y)) + mean_ms_tle = triton.testing.do_bench(lambda: add_tle(x, y)) + + print(f"\n========= Shape: {shape} Type: {dtype} =========") + print(f"Triton time: {mean_ms_triton:.4f} ms") + print(f"TLE time: {mean_ms_tle:.4f} ms") diff --git a/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.cu b/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.cu new file mode 100644 index 000000000..1197db9d5 --- /dev/null +++ b/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.cu @@ -0,0 +1,20 @@ + +#define DEFINE_VECTOR_ADD(Type, FuncName) \ +extern "C" __device__ void FuncName( \ + __attribute__((address_space(1))) Type *C, \ + __attribute__((address_space(1))) const Type *A, \ + __attribute__((address_space(1))) const Type *B, \ + const int N) { \ + const int idx = blockIdx.x * blockDim.x + threadIdx.x; \ + for (int i = idx; i < N; i += blockDim.x * gridDim.x) { \ + C[i] = A[i] + B[i]; \ + } \ +} + +#ifdef ENABLE_FLOAT +DEFINE_VECTOR_ADD(float, vector_add_float) +#endif + +#ifdef ENABLE_INT +DEFINE_VECTOR_ADD(int, vector_add_int) +#endif \ No newline at end of file diff --git a/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.py b/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.py new file mode 100644 index 000000000..b14163640 --- /dev/null +++ b/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.py @@ -0,0 +1,50 @@ +from pathlib import Path + +import torch +import triton +import triton.language as tl +from triton.experimental.tle.raw import dialect +import triton.experimental.tle.language.raw as tle_raw +from triton.language.extra.cuda import libnvshmem_device + +DEVICE = triton.runtime.driver.active.get_active_torch_device() + + +@dialect(name="cuda", file=Path(__file__).parent / "01-vector-add.cu", + library={}) +def edsl(*args, **kwargs): + ... + + +@triton.jit +def add_kernel( + x_ptr, + y_ptr, + output_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, +): + tle_raw.call(edsl, [output_ptr, x_ptr, y_ptr, n_elements]) + libnvshmem_device.vector_add(output_ptr, x_ptr, y_ptr, n_elements) + + +def add(x: torch.Tensor, y: torch.Tensor): + output = torch.empty_like(x) + assert x.device == DEVICE and y.device == DEVICE and output.device == DEVICE + n_elements = output.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), ) + add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024) + return output + + +if __name__ == "__main__": + # dtype = torch.float32 + # x = torch.randn(2048, dtype=dtype, device=DEVICE) + # y = torch.randn(2048, dtype=dtype, device=DEVICE) + + dtype = torch.int32 + x = torch.randint(low=0, high=10, size=(2048, ), dtype=dtype, device=DEVICE) + y = torch.randint(low=0, high=10, size=(2048, ), dtype=dtype, device=DEVICE) + + z = add(x, y) + assert torch.allclose(x + y, z), (x + y, z) diff --git a/python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/gemm-allreduce-host.cu b/python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/gemm-allreduce-host.cu new file mode 100644 index 000000000..e9560d2ce --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/gemm-allreduce-host.cu @@ -0,0 +1,113 @@ +#include +#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 nvshmemx_cumodule_init_wrapper(CUmodule module) { + return nvshmemx_cumodule_init(module); +} + +extern "C" void gemm_allreduce_before_launch +( + int *mype, int *npes, int *mype_in_node, int *n_pes_in_node, + float **d_A, float **d_B, float **d_C, + float **h_A, float **h_B, + int M, int N, int K +) { + nvshmem_init(); + *mype = nvshmem_my_pe(); + *npes = nvshmem_n_pes(); + *mype_in_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE); + *n_pes_in_node = nvshmem_team_n_pes(NVSHMEMX_TEAM_NODE); + + CUDA_CHECK(cudaSetDevice(*mype_in_node)); + *d_A = (float *)nvshmem_malloc(sizeof(float) * M * K); + *d_B = (float *)nvshmem_malloc(sizeof(float) * K * N); + *d_C = (float *)nvshmem_malloc(sizeof(float) * M * N); + + *h_A = (float *)malloc(sizeof(float) * M * K); + *h_B = (float *)malloc(sizeof(float) * K * N); + + float val = static_cast(*mype + 1); + for (int i = 0; i < M * K; i++) { + (*h_A)[i] = val; + } + for (int i = 0; i < K * N; i++) { + (*h_B)[i] = val; + } + + CUDA_CHECK(cudaMemcpy(*d_A, *h_A, sizeof(float) * M * K, cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemcpy(*d_B, *h_B, sizeof(float) * K * N, cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemset(*d_C, 0, sizeof(float) * M * N)); +} + +extern "C" void gemm_allreduce_after_launch +( + int mype, int npes, + float *d_A, float *d_B, float *d_C, + float *h_A, float *h_B, + int M, int N, int K +) { + cudaDeviceSynchronize(); + std::vector h_C_local(M * N); + cudaMemcpy(h_C_local.data(), d_C, M * N * sizeof(float), cudaMemcpyDeviceToHost); + float val = static_cast(mype + 1); + float local_expected = K * val * val; + std::cout << "PE " << mype << " local GEMM (first 4): "; + for (int i = 0; i < 4; ++i) printf("%.1f ", h_C_local[i]); + std::cout << "[expected: " << local_expected << "]\n"; + + nvshmem_barrier_all(); + nvshmem_float_sum_reduce(NVSHMEM_TEAM_WORLD, d_C, d_C, M * N); + nvshmem_barrier_all(); + cudaDeviceSynchronize(); + std::vector h_C_final(M * N); + cudaMemcpy(h_C_final.data(), d_C, M * N * sizeof(float), cudaMemcpyDeviceToHost); + + if (mype == 0) { + float expected = 0.0f; + for (int p = 0; p < npes; ++p) { + float v = static_cast(p + 1); + expected += K * v * v; + } + + std::cout << "\n========================================\n"; + std::cout << " After AllReduce (PE 0, first 4): "; + for (int i = 0; i < 4; ++i) printf("%.1f ", h_C_final[i]); + std::cout << "\n Expected: " << expected << "\n"; + std::cout << " (PE0: " << K*1*1 << " + PE1: " << K*2*2 << " = " << expected << ")\n"; + + bool correct = true; + for (int i = 0; i < M * N && correct; ++i) { + if (std::fabs(h_C_final[i] - expected) > 0.1f) { + std::cout << " ERROR at index " << i << ": got " << h_C_final[i] + << ", expected " << expected << "\n"; + correct = false; + } + } + if (correct) std::cout << " ✓ All values correct!\n"; + std::cout << "========================================\n"; + } + + nvshmem_free(d_A); + nvshmem_free(d_B); + nvshmem_free(d_C); + free(h_A); + free(h_B); + + nvshmem_finalize(); +} \ No newline at end of file diff --git a/python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/gemm-allreduce.py b/python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/gemm-allreduce.py new file mode 100644 index 000000000..db6a427c0 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/gemm-allreduce.py @@ -0,0 +1,139 @@ +import os +import subprocess +import ctypes +import torch +import triton +import triton.knobs as knobs +import triton.experimental.tle.language.raw as tle_raw + +from pathlib import Path +from triton.experimental.tle.raw import dialect +from triton.language.extra.cuda import libnvshmem_device + + +@dialect( + name="cuda", + file=(Path(__file__).parent / "simple-gemm.cu").resolve(), + library={"nvshmem": "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem"}, +) +def edsl(*args, **kwargs): + ... + + +@triton.jit +def gemm_allreduce_kernel( + C, A, B, + m, n, k +): + tle_raw.call(edsl, []) + libnvshmem_device.tiled_gemm(C, A, B, m, n, k) + + +def cuda_host_compile(cuda_host_path, cuda_host_lib): + NVCC = os.getenv("NVCC", "nvcc") + NVSHMEM_HOME = "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem" + include_path = f"-I{os.path.join(NVSHMEM_HOME, 'include')}" + lib_path = f"-L{os.path.join(NVSHMEM_HOME, 'lib')}" + + prop = torch.cuda.get_device_properties(torch.cuda.current_device()) + arch = f"-arch=sm_{prop.major}{prop.minor}" + tmp_file = Path(cuda_host_lib).with_suffix('.so.tmp') + build = [ + NVCC, "-shared", "-Xcompiler", "-fPIC", "-rdc=true", arch, include_path, lib_path, "-lnvshmem_host", + "-lnvshmem_device", "-o", tmp_file, cuda_host_path + ] + build = subprocess.run(build, capture_output=True) + assert build.returncode == 0, (f"NVCC host failed\nstderr:\n{build.stderr.decode()}") + tmp_file.rename(cuda_host_lib) + + +def gemm_allreduce(): + cu_file = (Path(__file__).parent / "gemm-allreduce-host.cu").resolve() + lib_file = Path(cu_file).with_suffix('.so') + + rank = int(os.getenv("OMPI_COMM_WORLD_RANK", "0")) + if rank == 0: + cuda_host_compile(cu_file, lib_file) + + import time + timeout = 60 + start = time.time() + while True: + if lib_file.exists(): + try: + ctypes.CDLL(str(lib_file)) + break + except OSError: + pass + if time.time() - start > timeout: + raise RuntimeError(f"Timeout waiting for {lib_file}") + time.sleep(0.1) + + + lib = ctypes.CDLL(lib_file) + mype = ctypes.c_int() + npes = ctypes.c_int() + mype_in_node = ctypes.c_int() + npes_in_node = ctypes.c_int() + # stream = ctypes.c_void_p() + + d_A = ctypes.c_void_p() + d_B = ctypes.c_void_p() + d_C = ctypes.c_void_p() + h_A = ctypes.c_void_p() + h_B = ctypes.c_void_p() + + M, N, K = 64, 64, 64 + lib.gemm_allreduce_before_launch( + ctypes.byref(mype), ctypes.byref(npes), ctypes.byref(mype_in_node), ctypes.byref(npes_in_node), + ctypes.byref(d_A), ctypes.byref(d_B), ctypes.byref(d_C), + ctypes.byref(h_A), ctypes.byref(h_B), + M, N, K + ) + + # define tile_size, dtype, len(dtype), device + dtype = torch.float32 + TILE_SIZE = 16 + grid = (((M + TILE_SIZE - 1) // TILE_SIZE), ((N + TILE_SIZE - 1) // TILE_SIZE)) + num_warps = TILE_SIZE * TILE_SIZE // 32 + device = triton.runtime.driver.active.get_active_torch_device() + + d_A_storage = torch._C._construct_storage_from_data_pointer(d_A.value, device, dtype.itemsize * M * K) + d_A_tensor = torch.empty(0, dtype=dtype, device=device).set_(d_A_storage).view(M, K) + + d_B_storage = torch._C._construct_storage_from_data_pointer(d_B.value, device, dtype.itemsize * K * N) + d_B_tensor = torch.empty(0, dtype=dtype, device=device).set_(d_B_storage).view(K, N) + + d_C_storage = torch._C._construct_storage_from_data_pointer(d_C.value, device, dtype.itemsize * M * N) + d_C_tensor = torch.empty(0, dtype=dtype, device=device).set_(d_C_storage).view(M, N) + + # def cumodule_init_hook(*args, **kwargs): + # key = kwargs["key"] + # jit_function = kwargs["fn"].jit_function + # device = kwargs["compile"]["device"] + # kernel_cache = jit_function.device_caches[device][0] + # kernel = kernel_cache.get(key, None) + # assert kernel is not None + # kernel._init_handles() + # ret = lib.nvshmemx_cumodule_init_wrapper(ctypes.c_void_p(kernel.module)) + # assert ret == 0, f"nvshmemx_cumodule_init_wrapper failed: {ret}" + # knobs.runtime.jit_post_compile_hook = cumodule_init_hook + + gemm_allreduce_kernel[grid]( + d_C_tensor, + d_A_tensor, + d_B_tensor, + M, N, K, + num_warps = num_warps + ) + + lib.gemm_allreduce_after_launch( + mype.value, npes.value, + d_A, d_B, d_C, + h_A, h_B, + M, N, K + ) + + +if __name__ == "__main__": + gemm_allreduce() diff --git a/python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/simple-gemm.cu b/python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/simple-gemm.cu new file mode 100644 index 000000000..ad47c7918 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/simple-gemm.cu @@ -0,0 +1,51 @@ + + +#define TILE_SIZE 16 +extern "C" __device__ void tiled_gemm( + float* C, + const float* A, + const float* B, + int m, int n, int k +) { + __shared__ float As[TILE_SIZE][TILE_SIZE]; + __shared__ float Bs[TILE_SIZE][TILE_SIZE]; + + int block_row = blockIdx.y; + int block_col = blockIdx.x; + + int thread_row = threadIdx.x / TILE_SIZE; + int thread_col = threadIdx.x % TILE_SIZE; + + int C_row = block_row * TILE_SIZE + thread_row; + int C_col = block_col * TILE_SIZE + thread_col; + + float sum = 0.0f; + int num_tiles = (k + TILE_SIZE - 1) / TILE_SIZE; + for (int t = 0; t < num_tiles; ++t) { + int A_col = t * TILE_SIZE + thread_col; + if (C_row < m && A_col < k) { + As[thread_row][thread_col] = A[C_row * k + A_col]; + } else { + As[thread_row][thread_col] = 0.0f; + } + + int B_row = t * TILE_SIZE + thread_row; + if (B_row < k && C_col < n) { + Bs[thread_row][thread_col] = B[B_row * n + C_col]; + } else { + Bs[thread_row][thread_col] = 0.0f; + } + + __syncthreads(); + + for (int i = 0; i < TILE_SIZE; ++i) { + sum += As[thread_row][i] * Bs[i][thread_col]; + } + + __syncthreads(); + } + + if (C_row < m && C_col < n) { + C[C_row * n + C_col] = sum; + } +} \ No newline at end of file diff --git a/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce-device.cu b/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce-device.cu new file mode 100644 index 000000000..6dbb1ed51 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce-device.cu @@ -0,0 +1,71 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// perform Allreduce using ring +extern "C" __device__ void ring_reduce(int *dst, const int *src, int nreduce, uint64_t *signal, + int chunk_size) { + int mype = nvshmem_my_pe(); + int npes = nvshmem_n_pes(); + int peer = (mype + 1) % npes; + + // printf("mype: %d, npes: %d, peer: %d\n", mype, npes, peer); + + int thread_id = threadIdx.x; + int num_threads = blockDim.x; + int num_blocks = gridDim.x; + int block_idx = blockIdx.x; + int elems_per_block = nreduce / num_blocks; + + // Change src, dst, nreduce, signal to what this block is going to process + // Each CTA will work independently + if (elems_per_block * (blockIdx.x + 1) > nreduce) return; + src = src + block_idx * elems_per_block; + dst = dst + block_idx * elems_per_block; + nreduce = elems_per_block; + signal = signal + block_idx; + + int chunk_elems = chunk_size / sizeof(int); + int num_chunks = nreduce / chunk_elems; + + // reduce phase + for (int chunk = 0; chunk < num_chunks; chunk++) { + if (mype != 0) { + if (thread_id == 0) nvshmem_signal_wait_until(signal, NVSHMEM_CMP_GE, chunk + 1); + + __syncthreads(); + for (int i = thread_id; i < chunk_elems; i += num_threads) { + dst[i] = dst[i] + src[i]; + } + __syncthreads(); + } + if (thread_id == 0) + nvshmem_int_put_signal_nbi(dst, (mype == 0) ? src : dst, chunk_elems, signal, 1, + NVSHMEM_SIGNAL_ADD, peer); + src = src + chunk_elems; + dst = dst + chunk_elems; + } + + // Broadcast phase + dst = dst - num_chunks * chunk_elems; + if (thread_id == 0) { + for (int chunk = 0; chunk < num_chunks; chunk++) { + if (mype < npes - 1) { // Last pe already has the final result + nvshmem_signal_wait_until(signal, NVSHMEM_CMP_GE, + (mype == 0) ? chunk + 1 : num_chunks + chunk + 1); + } + if (mype < npes - 2) + nvshmem_int_put_signal_nbi(dst, dst, chunk_elems, signal, 1, NVSHMEM_SIGNAL_ADD, + peer); + dst = dst + chunk_elems; + } + *signal = 0; // reset for next iteration + } +} \ No newline at end of file diff --git a/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce-host.cu b/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce-host.cu new file mode 100644 index 000000000..e4392d0c5 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce-host.cu @@ -0,0 +1,74 @@ +#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 nvshmemx_cumodule_init_wrapper(CUmodule module) { + return nvshmemx_cumodule_init(module); +} + +extern "C" void ring_reduce_before_launch +( + int *mype, int *npes, int *mype_in_node, int *n_pes_in_node, + cudaStream_t *stream, + int **src, int **dst, int **data_h, uint64_t **signal, + int size +) { + + nvshmem_init(); + *mype = nvshmem_my_pe(); + *npes = nvshmem_n_pes(); + *mype_in_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE); + *n_pes_in_node = nvshmem_team_n_pes(NVSHMEMX_TEAM_NODE); + + CUDA_CHECK(cudaSetDevice(*mype_in_node)); + CUDA_CHECK(cudaStreamCreate(stream)); + + *src = (int *)nvshmem_malloc(sizeof(int) * size); + *dst = (int *)nvshmem_malloc(sizeof(int) * size); + *data_h = (int *)malloc(sizeof(int) * size); + *signal = (uint64_t *)nvshmem_calloc(*n_pes_in_node, sizeof(uint64_t)); + + for (size_t i = 0; i < size; i++) (*data_h)[i] = i; + CUDA_CHECK(cudaMemcpyAsync(*src, *data_h, sizeof(int) * size, cudaMemcpyHostToDevice, *stream)); + nvshmemx_barrier_all_on_stream(*stream); +} + +extern "C" void ring_reduce_after_launch +( + cudaStream_t stream, + void *src, void *dst, void *data_h, void *signal, + int mype, int npes, + int size +) { + nvshmemx_barrier_all_on_stream(stream); + CUDA_CHECK(cudaStreamSynchronize(stream)); + CUDA_CHECK(cudaMemcpy(data_h, dst, sizeof(int) * size, cudaMemcpyDeviceToHost)); + int *data = (int *)data_h; + for (size_t i = 0; i < size; i++) { + if (data[i] != (int)i * npes) { + // if (mype == 0) { + // printf("PE %d, data[%zu] = %d expected data[%zu] = %d\n", mype, i, data[i], i, (int)i * npes); + printf("PE %d error, data[%zu] = %d expected data[%zu] = %d\n", mype, i, data[i], i, (int)i * npes); + } + } + + nvshmem_free(dst); + nvshmem_free(src); + nvshmem_free(signal); + free(data_h); + + nvshmem_finalize(); +} \ No newline at end of file diff --git a/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce.py b/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce.py new file mode 100644 index 000000000..d7ac1da9c --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce.py @@ -0,0 +1,147 @@ +import os +import subprocess +import ctypes +import torch +import triton +import triton.knobs as knobs +import triton.experimental.tle.language.raw as tle_raw + +from pathlib import Path +from triton.experimental.tle.raw import dialect +from triton.language.extra.cuda import libnvshmem_device + + +@dialect( + name="cuda", + file=(Path(__file__).parent / "ring-reduce-device.cu").resolve(), + library={"nvshmem": "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem"}, +) +def edsl(*args, **kwargs): + ... + + +@triton.jit +def ring_reduce_kernel( + dst, + src, + nreduce, + signal, + chunk_size, +): + tle_raw.call(edsl, []) + libnvshmem_device.ring_reduce(dst, src, nreduce, signal, chunk_size) + + +def cuda_host_compile(cuda_host_path, cuda_host_lib): + NVCC = os.getenv("NVCC", "nvcc") + NVSHMEM_HOME = "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem" + include_path = f"-I{os.path.join(NVSHMEM_HOME, 'include')}" + lib_path = f"-L{os.path.join(NVSHMEM_HOME, 'lib')}" + + prop = torch.cuda.get_device_properties(torch.cuda.current_device()) + arch = f"-arch=sm_{prop.major}{prop.minor}" + tmp_file = Path(cuda_host_lib).with_suffix('.so.tmp') + build = [ + NVCC, "-shared", "-Xcompiler", "-fPIC", "-rdc=true", arch, include_path, lib_path, "-lnvshmem_host", + "-lnvshmem_device", "-o", tmp_file, cuda_host_path + ] + build = subprocess.run(build, capture_output=True) + assert build.returncode == 0, (f"NVCC host failed\nstderr:\n{build.stderr.decode()}") + tmp_file.rename(cuda_host_lib) + + +def ring_reduce(): + cu_file = (Path(__file__).parent / "ring-reduce-host.cu").resolve() + lib_file = Path(cu_file).with_suffix('.so') + + rank = int(os.getenv("OMPI_COMM_WORLD_RANK", "0")) + if rank == 0: + cuda_host_compile(cu_file, lib_file) + + import time + timeout = 60 + start = time.time() + while True: + if lib_file.exists(): + try: + ctypes.CDLL(str(lib_file)) + break + except OSError: + pass + if time.time() - start > timeout: + raise RuntimeError(f"Timeout waiting for {lib_file}") + time.sleep(0.1) + + M, N = 64, 8 + lib = ctypes.CDLL(lib_file) + mype = ctypes.c_int() + npes = ctypes.c_int() + mype_in_node = ctypes.c_int() + npes_in_node = ctypes.c_int() + + stream = ctypes.c_void_p() + + src = ctypes.c_void_p() + dst = ctypes.c_void_p() + data_h = ctypes.c_void_p() + signal = ctypes.c_void_p() + + size = M * N + lib.ring_reduce_before_launch( + ctypes.byref(mype), ctypes.byref(npes), ctypes.byref(mype_in_node), ctypes.byref(npes_in_node), + ctypes.byref(stream), + ctypes.byref(src), ctypes.byref(dst), ctypes.byref(data_h), ctypes.byref(signal), + size + ) + + # print("PE:", mype_in_node.value) + dtype = torch.int32 + num_blocks = npes_in_node.value + device = triton.runtime.driver.active.get_active_torch_device() + src_storage = torch._C._construct_storage_from_data_pointer(src.value, device, 4 * M * N) + src_tensor = torch.empty(0, dtype=torch.int32, device=device).set_(src_storage).view(M, N) + + dst_storage = torch._C._construct_storage_from_data_pointer(dst.value, device, 4 * M * N) + dst_tensor = torch.empty(0, dtype=torch.int32, device=device).set_(dst_storage).view(M, N) + + # data_h_storage = torch._C._construct_storage_from_data_pointer(data_h.value, device, 4 * M * N) + # data_h_tensor = torch.empty(0, dtype=torch.int32, device=device).set_(data_h_storage).view(M, N) + + signal_storage = torch._C._construct_storage_from_data_pointer(signal.value, device, 8 * num_blocks) + signal_tensor = torch.empty(0, dtype=torch.int64, device=device).set_(signal_storage).view(num_blocks, ) + + def cumodule_init_hook(*args, **kwargs): + key = kwargs["key"] + jit_function = kwargs["fn"].jit_function + device = kwargs["compile"]["device"] + kernel_cache = jit_function.device_caches[device][0] + kernel = kernel_cache.get(key, None) + assert kernel is not None + kernel._init_handles() + ret = lib.nvshmemx_cumodule_init_wrapper(ctypes.c_void_p(kernel.module)) + assert ret == 0, f"nvshmemx_cumodule_init_wrapper failed: {ret}" + knobs.runtime.jit_post_compile_hook = cumodule_init_hook + + curr_stream = torch.cuda.ExternalStream(stream.value, device=device) + # print(f"[PE {mype.value}] Using stream pointer: {hex(stream.value)}") + # print(f"[PE {mype.value}] Default stream before: {torch.cuda.default_stream().cuda_stream}") + with torch.cuda.stream(curr_stream): + chunk_size = int(M // num_blocks * dtype.itemsize) + ring_reduce_kernel[(num_blocks, )]( + dst_tensor, + src_tensor, + M * N, + signal_tensor, + chunk_size + ) + + lib.ring_reduce_after_launch( + stream, + src, dst, data_h, signal, + mype_in_node.value, npes_in_node.value, + size + ) + + +if __name__ == "__main__": + ring_reduce() diff --git a/python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast-device.cu b/python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast-device.cu new file mode 100644 index 000000000..7ca620d09 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast-device.cu @@ -0,0 +1,23 @@ +#include +#include +#include +#include +#include + +extern "C" __device__ void ring_bcast(int *data, int nelem, int root, uint64_t *psync) { + int mype = nvshmem_my_pe(); + int npes = nvshmem_n_pes(); + int peer = (mype + 1) % npes; + + if (mype == root) *psync = 1; + + nvshmem_signal_wait_until(psync, NVSHMEM_CMP_NE, 0); + + if (mype == npes - 1) return; + + nvshmem_int_put(data, data, nelem, peer); + nvshmem_fence(); + nvshmemx_signal_op(psync, 1, NVSHMEM_SIGNAL_SET, peer); + + *psync = 0; +} \ No newline at end of file diff --git a/python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast-host.cu b/python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast-host.cu new file mode 100644 index 000000000..9578bb945 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast-host.cu @@ -0,0 +1,59 @@ +#include +#include +#include +#include +#include + +extern "C" int nvshmemx_cumodule_init_wrapper(CUmodule module) { + return nvshmemx_cumodule_init(module); +} + +extern "C" void ring_bcast_before_launch +( + int *mype, int *npes, int *mype_in_node, int *n_pes_in_node, + cudaStream_t *stream, + int **data, int **data_h, uint64_t **psync, + int data_len +) { + + nvshmem_init(); + *mype = nvshmem_my_pe(); + *npes = nvshmem_n_pes(); + *mype_in_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE); + *n_pes_in_node = nvshmem_team_n_pes(NVSHMEMX_TEAM_NODE); + + cudaSetDevice(*mype_in_node); + cudaStreamCreate(stream); + + *data = (int *)nvshmem_malloc(sizeof(int) * data_len); + *data_h = (int *)malloc(sizeof(int) * data_len); + *psync = (uint64_t *)nvshmem_calloc(1, sizeof(uint64_t)); + + for (size_t i = 0; i < data_len; i++) (*data_h)[i] = *mype + i; + cudaMemcpyAsync(*data, *data_h, sizeof(int) * data_len, cudaMemcpyHostToDevice, *stream); + nvshmemx_barrier_all_on_stream(*stream); +} + +extern "C" void ring_bcast_after_launch +( + cudaStream_t stream, + void *data, void *data_h, void *psync, + int mype, int npes, + int data_len +) { + nvshmemx_barrier_all_on_stream(stream); + cudaMemcpyAsync(data_h, data, sizeof(int) * data_len, cudaMemcpyDeviceToHost, stream); + cudaStreamSynchronize(stream); + + int *h_ptr = (int *)data_h; + for (size_t i = 0; i < data_len; i++) { + if (h_ptr[i] != (int) i) + printf("PE %d error, data[%zu] = %d expected data[%zu] = %d\n", mype, i, h_ptr[i], i, (int) i); + } + + nvshmem_free(data); + nvshmem_free(psync); + free(data_h); + + nvshmem_finalize(); +} \ No newline at end of file diff --git a/python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast.py b/python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast.py new file mode 100644 index 000000000..80e325ea9 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast.py @@ -0,0 +1,137 @@ +import os +import subprocess +import ctypes +import torch +import triton +import triton.knobs as knobs +import triton.experimental.tle.language.raw as tle_raw + +from pathlib import Path +from triton.experimental.tle.raw import dialect +from triton.language.extra.cuda import libnvshmem_device + + +@dialect( + name="cuda", + file=(Path(__file__).parent / "ring-bcast-device.cu").resolve(), + library={"nvshmem": "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem"}, +) +def edsl(*args, **kwargs): + ... + + +@triton.jit +def ring_bcast_kernel( + data, + nelem, + root, + psync, +): + tle_raw.call(edsl, []) + libnvshmem_device.ring_bcast(data, nelem, root, psync) + + +def cuda_host_compile(cuda_host_path, cuda_host_lib): + NVCC = os.getenv("NVCC", "nvcc") + NVSHMEM_HOME = "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem" + include_path = f"-I{os.path.join(NVSHMEM_HOME, 'include')}" + lib_path = f"-L{os.path.join(NVSHMEM_HOME, 'lib')}" + + prop = torch.cuda.get_device_properties(torch.cuda.current_device()) + arch = f"-arch=sm_{prop.major}{prop.minor}" + tmp_file = Path(cuda_host_lib).with_suffix('.so.tmp') + build = [ + NVCC, "-shared", "-Xcompiler", "-fPIC", "-rdc=true", arch, include_path, lib_path, "-lnvshmem_host", + "-lnvshmem_device", "-o", tmp_file, cuda_host_path + ] + build = subprocess.run(build, capture_output=True) + assert build.returncode == 0, (f"NVCC host failed\nstderr:\n{build.stderr.decode()}") + tmp_file.rename(cuda_host_lib) + + +def ring_bcast(): + cu_file = (Path(__file__).parent / "ring-bcast-host.cu").resolve() + lib_file = Path(cu_file).with_suffix('.so') + + rank = int(os.getenv("OMPI_COMM_WORLD_RANK", "0")) + if rank == 0: + cuda_host_compile(cu_file, lib_file) + + import time + timeout = 60 + start = time.time() + while True: + if lib_file.exists(): + try: + ctypes.CDLL(str(lib_file)) + break + except OSError: + pass + if time.time() - start > timeout: + raise RuntimeError(f"Timeout waiting for {lib_file}") + time.sleep(0.1) + + lib = ctypes.CDLL(lib_file) + mype = ctypes.c_int() + npes = ctypes.c_int() + mype_in_node = ctypes.c_int() + npes_in_node = ctypes.c_int() + + stream = ctypes.c_void_p() + + data = ctypes.c_void_p() + data_h = ctypes.c_void_p() + psync = ctypes.c_void_p() + + data_len = 32 + root = 0 + lib.ring_bcast_before_launch( + ctypes.byref(mype), ctypes.byref(npes), ctypes.byref(mype_in_node), ctypes.byref(npes_in_node), + ctypes.byref(stream), + ctypes.byref(data), ctypes.byref(data_h), ctypes.byref(psync), + data_len + ) + + # print("PE:", mype_in_node.value) + dtype = torch.int32 + num_blocks = 1 + num_warps = 1 + device = triton.runtime.driver.active.get_active_torch_device() + + data_storage = torch._C._construct_storage_from_data_pointer(data.value, device, dtype.itemsize * data_len) + data_tensor = torch.empty(0, dtype=dtype, device=device).set_(data_storage).view(data_len, ) + + psync_storage = torch._C._construct_storage_from_data_pointer(psync.value, device, 8 * 1) + psync_tensor = torch.empty(0, dtype=torch.uint64, device=device).set_(psync_storage).view(1, ) + + def cumodule_init_hook(*args, **kwargs): + key = kwargs["key"] + jit_function = kwargs["fn"].jit_function + device = kwargs["compile"]["device"] + kernel_cache = jit_function.device_caches[device][0] + kernel = kernel_cache.get(key, None) + assert kernel is not None + kernel._init_handles() + ret = lib.nvshmemx_cumodule_init_wrapper(ctypes.c_void_p(kernel.module)) + assert ret == 0, f"nvshmemx_cumodule_init_wrapper failed: {ret}" + knobs.runtime.jit_post_compile_hook = cumodule_init_hook + + curr_stream = torch.cuda.ExternalStream(stream.value, device=device) + with torch.cuda.stream(curr_stream): + ring_bcast_kernel[(num_blocks, )]( + data_tensor, + data_len, + root, + psync_tensor + ) + + lib.ring_bcast_after_launch( + stream, + data, data_h, psync, + mype.value, npes.value, + data_len + ) + + +if __name__ == "__main__": + ring_bcast() diff --git a/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-device-accumulate.cu b/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-device-accumulate.cu new file mode 100644 index 000000000..316497b6e --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-device-accumulate.cu @@ -0,0 +1,8 @@ +#include + +extern "C" __device__ void accumulate(int *input, int *partial_sum) { + int index = threadIdx.x; + if (0 == index) *partial_sum = 0; + __syncthreads(); + atomicAdd(partial_sum, input[index]); +} diff --git a/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-device-correct-accumulate.cu b/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-device-correct-accumulate.cu new file mode 100644 index 000000000..00b90e339 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-device-correct-accumulate.cu @@ -0,0 +1,13 @@ +#include + +#define THRESHOLD 42 +#define CORRECTION 7 +extern "C" __device__ void correct_accumulate(int *input, int *partial_sum, int *full_sum) { + int index = threadIdx.x; + if (*full_sum > THRESHOLD) { + input[index] = input[index] - CORRECTION; + } + if (0 == index) *partial_sum = 0; + __syncthreads(); + atomicAdd(partial_sum, input[index]); +} \ No newline at end of file diff --git a/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-host.cu b/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-host.cu new file mode 100644 index 000000000..1fa8f34fb --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-host.cu @@ -0,0 +1,69 @@ +#include +#include "nvshmem.h" +#include "nvshmemx.h" + +#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 nvshmemx_cumodule_init_wrapper(CUmodule module) { + return nvshmemx_cumodule_init(module); +} + + +extern "C" void on_stream_before_launch_accumulate +( + int *mype, int *npes, int *mype_in_node, int *n_pes_in_node, + cudaStream_t *stream, + int **input, int **partial_sum, int **full_sum, + int input_nelems +) { + + nvshmem_init(); + *mype = nvshmem_my_pe(); + *npes = nvshmem_n_pes(); + *mype_in_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE); + *n_pes_in_node = nvshmem_team_n_pes(NVSHMEMX_TEAM_NODE); + + CUDA_CHECK(cudaSetDevice(*mype_in_node)); + CUDA_CHECK(cudaStreamCreate(stream)); + + *input = (int *)nvshmem_malloc(sizeof(int) * input_nelems); + *partial_sum = (int *)nvshmem_malloc(sizeof(int)); + *full_sum = (int *)nvshmem_malloc(sizeof(int)); +} + +extern "C" void on_stream_before_launch_correct_accumulate +( + cudaStream_t *stream, + int **partial_sum, int **full_sum, + int to_all_nelems +) { + nvshmemx_int_sum_reduce_on_stream(NVSHMEM_TEAM_WORLD, *full_sum, *partial_sum, to_all_nelems, *stream); +} + + +extern "C" void on_stream_after_launch_correct_accumulate +( + int mype, int npes, + cudaStream_t stream, + void *input, void *partial_sum, void *full_sum +) { + CUDA_CHECK(cudaStreamSynchronize(stream)); + printf("[%d of %d] run complete \n", mype, npes); + CUDA_CHECK(cudaStreamDestroy(stream)); + + nvshmem_free(input); + nvshmem_free(partial_sum); + nvshmem_free(full_sum); + + nvshmem_finalize(); +} \ No newline at end of file diff --git a/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream.py b/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream.py new file mode 100644 index 000000000..56a8928e0 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream.py @@ -0,0 +1,166 @@ +import os +import subprocess +import ctypes +import torch +import triton +import triton.knobs as knobs +import triton.experimental.tle.language.raw as tle_raw + +from pathlib import Path +from triton.experimental.tle.raw import dialect +from triton.language.extra.cuda import libnvshmem_device + + +@dialect( + name="cuda", + file=(Path(__file__).parent / "on-stream-device-accumulate.cu").resolve(), + library={"nvshmem": "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem"}, +) +def edsl_accumulate(*args, **kwargs): + ... + + +@dialect( + name="cuda", + file=(Path(__file__).parent / "on-stream-device-correct-accumulate.cu").resolve(), + library={"nvshmem": "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem"}, +) +def edsl_correct_accumulate(*args, **kwargs): + ... + + +@triton.jit +def accumulate_kernel( + input, + partial_sum +): + tle_raw.call(edsl_accumulate, []) + libnvshmem_device.accumulate(input, partial_sum) + + +@triton.jit +def correct_accumulate_kernel( + input, + partial_sum, + full_sum +): + tle_raw.call(edsl_correct_accumulate, []) + libnvshmem_device.correct_accumulate(input, partial_sum, full_sum) + + +def cuda_host_compile(cuda_host_path, cuda_host_lib): + NVCC = os.getenv("NVCC", "nvcc") + NVSHMEM_HOME = "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem" + include_path = f"-I{os.path.join(NVSHMEM_HOME, 'include')}" + lib_path = f"-L{os.path.join(NVSHMEM_HOME, 'lib')}" + + prop = torch.cuda.get_device_properties(torch.cuda.current_device()) + arch = f"-arch=sm_{prop.major}{prop.minor}" + tmp_file = Path(cuda_host_lib).with_suffix('.so.tmp') + build = [ + NVCC, "-shared", "-Xcompiler", "-fPIC", "-rdc=true", arch, include_path, lib_path, "-lnvshmem_host", + "-lnvshmem_device", "-o", tmp_file, cuda_host_path + ] + build = subprocess.run(build, capture_output=True) + assert build.returncode == 0, (f"NVCC host failed\nstderr:\n{build.stderr.decode()}") + tmp_file.rename(cuda_host_lib) + + +def on_stream(): + cu_file = (Path(__file__).parent / "on-stream-host.cu").resolve() + lib_file = Path(cu_file).with_suffix('.so') + + rank = int(os.getenv("OMPI_COMM_WORLD_RANK", "0")) + if rank == 0: + cuda_host_compile(cu_file, lib_file) + + import time + timeout = 60 + start = time.time() + while True: + if lib_file.exists(): + try: + ctypes.CDLL(str(lib_file)) + break + except OSError: + pass + if time.time() - start > timeout: + raise RuntimeError(f"Timeout waiting for {lib_file}") + time.sleep(0.1) + + lib = ctypes.CDLL(lib_file) + mype = ctypes.c_int() + npes = ctypes.c_int() + mype_in_node = ctypes.c_int() + npes_in_node = ctypes.c_int() + + stream = ctypes.c_void_p() + + input = ctypes.c_void_p() + partial_sum = ctypes.c_void_p() + full_sum = ctypes.c_void_p() + + input_nelems = 512 + to_all_nelems = 1 + lib.on_stream_before_launch_accumulate( + ctypes.byref(mype), ctypes.byref(npes), ctypes.byref(mype_in_node), ctypes.byref(npes_in_node), + ctypes.byref(stream), + ctypes.byref(input), ctypes.byref(partial_sum), ctypes.byref(full_sum), + input_nelems + ) + + dtype = torch.int32 + num_blocks = 1 + num_warps = (input_nelems + 31) // 32 + device = triton.runtime.driver.active.get_active_torch_device() + + input_storage = torch._C._construct_storage_from_data_pointer(input.value, device, dtype.itemsize * input_nelems) + input_tensor = torch.empty(0, dtype=torch.int32, device=device).set_(input_storage).view(input_nelems, ) + partial_sum_storage = torch._C._construct_storage_from_data_pointer(partial_sum.value, device, dtype.itemsize) + partial_sum_tensor = torch.empty(0, dtype=torch.int32, device=device).set_(partial_sum_storage).view(1, ) + full_sum_storage = torch._C._construct_storage_from_data_pointer(full_sum.value, device, dtype.itemsize) + full_sum_tensor = torch.empty(0, dtype=torch.int32, device=device).set_(full_sum_storage).view(1, ) + + def cumodule_init_hook(*args, **kwargs): + key = kwargs["key"] + jit_function = kwargs["fn"].jit_function + device = kwargs["compile"]["device"] + kernel_cache = jit_function.device_caches[device][0] + kernel = kernel_cache.get(key, None) + assert kernel is not None + kernel._init_handles() + ret = lib.nvshmemx_cumodule_init_wrapper(ctypes.c_void_p(kernel.module)) + assert ret == 0, f"nvshmemx_cumodule_init_wrapper failed: {ret}" + knobs.runtime.jit_post_compile_hook = cumodule_init_hook + + curr_stream = torch.cuda.ExternalStream(stream.value, device=device) + with torch.cuda.stream(curr_stream): + accumulate_kernel[(num_blocks, )]( + input_tensor, + partial_sum_tensor, + num_warps=num_warps + ) + + lib.on_stream_before_launch_correct_accumulate( + ctypes.byref(stream), + ctypes.byref(partial_sum), ctypes.byref(full_sum), + to_all_nelems + ) + + with torch.cuda.stream(curr_stream): + correct_accumulate_kernel[(num_blocks, )]( + input_tensor, + partial_sum_tensor, + full_sum_tensor, + num_warps=num_warps + ) + + lib.on_stream_after_launch_correct_accumulate( + mype_in_node.value, npes_in_node.value, + stream, + input, partial_sum, full_sum + ) + + +if __name__ == "__main__": + on_stream() \ No newline at end of file diff --git a/python/tutorials/tle/raw/nvshmem/05-put-block/put-block-device.cu b/python/tutorials/tle/raw/nvshmem/05-put-block/put-block-device.cu new file mode 100644 index 000000000..6ce11ff82 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/05-put-block/put-block-device.cu @@ -0,0 +1,23 @@ +#include +#include "nvshmem.h" +#include "nvshmemx.h" + +extern "C" __device__ void set_and_shift(float *send_data, float *recv_data, int num_elems, int mype, int npes) { + int thread_idx = blockIdx.x * blockDim.x + threadIdx.x; + /* set the corresponding element of send_data */ + if (thread_idx < num_elems) send_data[thread_idx] = mype; + + int peer = (mype + 1) % npes; + /* Every thread in block 0 calls nvshmemx_float_put_block. Alternatively, + every thread can call shmem_float_p, but shmem_float_p has a disadvantage + that when the destination GPU is connected via IB, there will be one rma + message for every single element which can be detrimental to performance. + And the disadvantage with shmem_float_put is that when the destination GPU is p2p + connected, it cannot leverage multiple threads to copy the data to the destination + GPU. */ + int block_offset = blockIdx.x * blockDim.x; + nvshmemx_float_put_block(recv_data + block_offset, send_data + block_offset, + min(blockDim.x, num_elems - block_offset), + peer); /* All threads in a block call the API + with the same arguments */ +} \ No newline at end of file diff --git a/python/tutorials/tle/raw/nvshmem/05-put-block/put-block-host.cu b/python/tutorials/tle/raw/nvshmem/05-put-block/put-block-host.cu new file mode 100644 index 000000000..2ecf04f7c --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/05-put-block/put-block-host.cu @@ -0,0 +1,75 @@ +#include +#include "nvshmem.h" +#include "nvshmemx.h" + +#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 nvshmemx_cumodule_init_wrapper(CUmodule module) { + return nvshmemx_cumodule_init(module); +} + + +extern "C" void put_block_before_launch +( + int *mype, int *npes, int *mype_in_node, int *n_pes_in_node, + float **send_data, float **recv_data, + int num_elems +) { + + nvshmem_init(); + *mype = nvshmem_my_pe(); + *npes = nvshmem_n_pes(); + *mype_in_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE); + *n_pes_in_node = nvshmem_team_n_pes(NVSHMEMX_TEAM_NODE); + + CUDA_CHECK(cudaSetDevice(*mype_in_node)); + + *send_data = (float *)nvshmem_malloc(sizeof(float) * num_elems); + *recv_data = (float *)nvshmem_malloc(sizeof(float) * num_elems); +} + + +extern "C" void put_block_after_launch +( + int mype, int npes, + void *send_data, void *recv_data, + int num_elems +) { + CUDA_CHECK(cudaGetLastError()); + CUDA_CHECK(cudaDeviceSynchronize()); + nvshmem_barrier_all(); + + /* Do data validation */ + float *host = new float[num_elems]; + CUDA_CHECK(cudaMemcpy(host, recv_data, num_elems * sizeof(float), cudaMemcpyDefault)); + int ref = (mype - 1 + npes) % npes; + bool success = true; + for (int i = 0; i < num_elems; ++i) { + if (host[i] != ref) { + printf("Error at %d of rank %d: %f\n", i, mype, host[i]); + success = false; + break; + } + } + + if (success) { + printf("[%d of %d] run complete \n", mype, npes); + } else { + printf("[%d of %d] run failure \n", mype, npes); + } + + nvshmem_free(send_data); + nvshmem_free(recv_data); + + nvshmem_finalize(); +} \ No newline at end of file diff --git a/python/tutorials/tle/raw/nvshmem/05-put-block/put-block.py b/python/tutorials/tle/raw/nvshmem/05-put-block/put-block.py new file mode 100644 index 000000000..48162524f --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/05-put-block/put-block.py @@ -0,0 +1,128 @@ +import os +import subprocess +import ctypes +import torch +import triton +import triton.knobs as knobs +import triton.experimental.tle.language.raw as tle_raw + +from pathlib import Path +from triton.experimental.tle.raw import dialect +from triton.language.extra.cuda import libnvshmem_device + + +@dialect( + name="cuda", + file=(Path(__file__).parent / "put-block-device.cu").resolve(), + library={"nvshmem": "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem"}, +) +def edsl(*args, **kwargs): + ... + + + +@triton.jit +def set_and_shift_kernel( + send_data, recv_data, num_elems, mype, npes +): + tle_raw.call(edsl, []) + libnvshmem_device.set_and_shift(send_data, recv_data, num_elems, mype, npes) + + +def cuda_host_compile(cuda_host_path, cuda_host_lib): + NVCC = os.getenv("NVCC", "nvcc") + NVSHMEM_HOME = "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem" + include_path = f"-I{os.path.join(NVSHMEM_HOME, 'include')}" + lib_path = f"-L{os.path.join(NVSHMEM_HOME, 'lib')}" + + prop = torch.cuda.get_device_properties(torch.cuda.current_device()) + arch = f"-arch=sm_{prop.major}{prop.minor}" + tmp_file = Path(cuda_host_lib).with_suffix('.so.tmp') + build = [ + NVCC, "-shared", "-Xcompiler", "-fPIC", "-rdc=true", arch, include_path, lib_path, "-lnvshmem_host", + "-lnvshmem_device", "-o", tmp_file, cuda_host_path + ] + build = subprocess.run(build, capture_output=True) + assert build.returncode == 0, (f"NVCC host failed\nstderr:\n{build.stderr.decode()}") + tmp_file.rename(cuda_host_lib) + + +def put_block(): + cu_file = (Path(__file__).parent / "put-block-host.cu").resolve() + lib_file = Path(cu_file).with_suffix('.so') + + rank = int(os.getenv("OMPI_COMM_WORLD_RANK", "0")) + if rank == 0: + cuda_host_compile(cu_file, lib_file) + + import time + timeout = 60 + start = time.time() + while True: + if lib_file.exists(): + try: + ctypes.CDLL(str(lib_file)) + break + except OSError: + pass + if time.time() - start > timeout: + raise RuntimeError(f"Timeout waiting for {lib_file}") + time.sleep(0.1) + + lib = ctypes.CDLL(lib_file) + mype = ctypes.c_int() + npes = ctypes.c_int() + mype_in_node = ctypes.c_int() + npes_in_node = ctypes.c_int() + + send_data = ctypes.c_void_p() + recv_data = ctypes.c_void_p() + + num_elems = 8192 + lib.put_block_before_launch( + ctypes.byref(mype), ctypes.byref(npes), ctypes.byref(mype_in_node), ctypes.byref(npes_in_node), + ctypes.byref(send_data), ctypes.byref(recv_data), + num_elems + ) + + dtype = torch.float32 + THREADS_PER_BLOCK = 1024 + num_blocks = num_elems // THREADS_PER_BLOCK + num_warps = (THREADS_PER_BLOCK + 31) // 32 + device = triton.runtime.driver.active.get_active_torch_device() + + send_data_storage = torch._C._construct_storage_from_data_pointer(send_data.value, device, dtype.itemsize * num_elems) + send_data_tensor = torch.empty(0, dtype=dtype, device=device).set_(send_data_storage).view(num_elems, ) + recv_data_storage = torch._C._construct_storage_from_data_pointer(recv_data.value, device, dtype.itemsize * num_elems) + recv_data_tensor = torch.empty(0, dtype=dtype, device=device).set_(recv_data_storage).view(num_elems, ) + + def cumodule_init_hook(*args, **kwargs): + key = kwargs["key"] + jit_function = kwargs["fn"].jit_function + device = kwargs["compile"]["device"] + kernel_cache = jit_function.device_caches[device][0] + kernel = kernel_cache.get(key, None) + assert kernel is not None + kernel._init_handles() + ret = lib.nvshmemx_cumodule_init_wrapper(ctypes.c_void_p(kernel.module)) + assert ret == 0, f"nvshmemx_cumodule_init_wrapper failed: {ret}" + knobs.runtime.jit_post_compile_hook = cumodule_init_hook + + set_and_shift_kernel[(num_blocks, )]( + send_data_tensor, + recv_data_tensor, + num_elems, + mype.value, + npes.value, + num_warps=num_warps + ) + + lib.put_block_after_launch( + mype.value, npes.value, + send_data, recv_data, + num_elems + ) + + +if __name__ == "__main__": + put_block() \ No newline at end of file diff --git a/third_party/nvidia/backend/compiler.py b/third_party/nvidia/backend/compiler.py index c469c8a78..d83616896 100644 --- a/third_party/nvidia/backend/compiler.py +++ b/third_party/nvidia/backend/compiler.py @@ -512,7 +512,7 @@ def make_ptx(self, src, metadata, opt, capability): assert len(names) == 1 metadata["name"] = names[0] # post-process - ptx_version = '8.5' + ptx_version = f'{ptx_version//10}.{ptx_version%10}' ret = re.sub(r'\.version \d+\.\d+', f'.version {ptx_version}', ret, flags=re.MULTILINE) ret = re.sub(r'\.target sm_\d+', f'.target sm_{capability}', ret, flags=re.MULTILINE) if not knobs.compilation.dump_ir_extract_di_local_variables: @@ -526,7 +526,7 @@ def make_ptx(self, src, metadata, opt, capability): return ret def make_cubin(self, src, metadata, opt, capability): - ptxas = os.getenv("ptxas", "ptxas") + ptxas = get_ptxas(self.target.arch).path with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.ptx') as fsrc, \ tempfile.NamedTemporaryFile(delete=False, mode='r', suffix='.log') as flog: fsrc.write(src) @@ -625,11 +625,11 @@ def make_cubin(self, src, metadata, opt, capability): if use_nvcc: with open(fbin_combined, 'rb') as f: cubin = f.read() + if os.path.exists(fbin_combined): + os.remove(fbin_combined) else: with open(fbin, 'rb') as f: cubin = f.read() - if os.path.exists(fbin_combined): - os.remove(fbin_combined) if os.path.exists(fbin): os.remove(fbin) return cubin diff --git a/third_party/nvidia/language/cuda/libnvshmem_device.py b/third_party/nvidia/language/cuda/libnvshmem_device.py index 3dec30b9a..799a298d9 100644 --- a/third_party/nvidia/language/cuda/libnvshmem_device.py +++ b/third_party/nvidia/language/cuda/libnvshmem_device.py @@ -15,17 +15,197 @@ def patch_hash_method_for_pointer_type(): patch_hash_method_for_pointer_type() - +# 01-simple-shift @core.extern def simple_shift(dst, _semantic=None): return core.extern_call( "", # libname "", # libpath [dst], # args - {( - core.pointer_type(core.dtype("int32")), # arg_type_symbol_dict - ): ("simple_shift", ()), # void return type - }, + { + ( + core.pointer_type(core.dtype("int32")), # arg_type_symbol_dict + ): ("simple_shift", ()), # void return type + }, + is_pure=False, + _semantic=_semantic, + ) + +# 02-ring-reduce +@core.extern +def ring_reduce(dst, src, nreduce, signal, chunk_size, _semantic=None): + return core.extern_call( + "", + "", + [dst, + src, + tl.cast(nreduce, tl.int32, _semantic=_semantic), + signal, + tl.cast(chunk_size, tl.int32, _semantic=_semantic)], + { + ( + core.pointer_type(core.dtype("int32")), + core.pointer_type(core.dtype("int32")), + core.dtype("int32"), + core.pointer_type(core.dtype("int64")), + core.dtype("int32"), + ) : ("ring_reduce", ()), + }, + is_pure=False, + _semantic=_semantic, + ) + +# 03-ring-bcast +@core.extern +def ring_bcast(data, nelem, root, psync, _semantic=None): + return core.extern_call( + "", + "", + [data, nelem, root, psync], + { + ( + core.pointer_type(core.dtype("int32")), + core.dtype("int32"), + core.dtype("int32"), + core.pointer_type(core.dtype("uint64")), + ) : ("ring_bcast", ()), + }, + is_pure=False, + _semantic=_semantic, + ) + +# 04-on-stream +@core.extern +def accumulate(input, partial_sum, _semantic=None): + return core.extern_call( + "", # libname + "", # libpath + [input, partial_sum], + { + ( + core.pointer_type(core.dtype("int32")), + core.pointer_type(core.dtype("int32")), + ): ("accumulate", ()), + }, + is_pure=False, + _semantic=_semantic, + ) + +@core.extern +def correct_accumulate(input, partial_sum, full_sum, _semantic=None): + return core.extern_call( + "", # libname + "", # libpath + [input, partial_sum, full_sum], + { + ( + core.pointer_type(core.dtype("int32")), + core.pointer_type(core.dtype("int32")), + core.pointer_type(core.dtype("int32")), + ): ("correct_accumulate", ()), + }, + is_pure=False, + _semantic=_semantic, + ) + +# 05-put-block +@core.extern +def set_and_shift(send_data, recv_data, num_elems, mype, npes, _semantic=None): + return core.extern_call( + "", # libname + "", # libpath + [send_data, recv_data, num_elems, mype, npes], + { + ( + core.pointer_type(core.dtype("fp32")), + core.pointer_type(core.dtype("fp32")), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + ): ("set_and_shift", ()), + }, + is_pure=False, + _semantic=_semantic, + ) + +# 00-gemm-allreduce +@core.extern +def tiled_gemm(C, A, B, m, n, k, _semantic=None): + return core.extern_call( + "", + "", + [ + C, A, B, + m, n, k + ], + { + ( + core.pointer_type(core.dtype("fp32")), + core.pointer_type(core.dtype("fp32")), + core.pointer_type(core.dtype("fp32")), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + ) : ("tiled_gemm", ()), + }, + is_pure=False, + _semantic=_semantic, + ) + +@core.extern +def per_token_group_quant_8bit(x_q_ptr, x_s_ptr, x_ptr, group_size, num_groups, groups_per_block, eps, fp8_min, fp8_max, + _semantic=None): + return core.extern_call( + "", + "", + [ + tl.cast(x_q_ptr, tl.pointer_type(core.dtype("void")), _semantic=_semantic), + x_s_ptr, + x_ptr, + group_size, + num_groups, + groups_per_block, + eps, + fp8_min, + fp8_max + ], + { + ( + core.pointer_type(core.dtype("void")), + core.pointer_type(core.dtype("fp32")), + core.pointer_type(core.dtype("fp32")), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("fp32"), + core.dtype("fp32"), + core.dtype("fp32"), + ): ("per_token_group_quant_8bit", ()), + }, + is_pure=False, + _semantic=_semantic, + ) + +@core.extern +def vector_add(C, A, B, N, _semantic=None): + return core.extern_call( + "", + "", + [C, A, B, N], + { + ( + core.pointer_type(core.dtype("int32")), + core.pointer_type(core.dtype("int32")), + core.pointer_type(core.dtype("int32")), + core.dtype("int32"), + ) : ("vector_add_int", ()), + ( + core.pointer_type(core.dtype("fp32")), + core.pointer_type(core.dtype("fp32")), + core.pointer_type(core.dtype("fp32")), + core.dtype("int32"), + ) : ("vector_add_float", ()), + }, is_pure=False, _semantic=_semantic, ) From fd24dcad080a93a156b510c8bae8ea7edc660cf3 Mon Sep 17 00:00:00 2001 From: zyuli Date: Sat, 23 May 2026 09:15:13 +0800 Subject: [PATCH 3/8] add macro define parameters --- .../experimental/tle/raw/cuda/runtime.py | 21 +++++++++++++++++-- .../vector-add-template.cu | 16 ++++++++------ .../vector-add-template.py | 6 ++++-- 3 files changed, 33 insertions(+), 10 deletions(-) diff --git a/python/triton/experimental/tle/raw/cuda/runtime.py b/python/triton/experimental/tle/raw/cuda/runtime.py index 1a3d2f0e1..2d0259002 100644 --- a/python/triton/experimental/tle/raw/cuda/runtime.py +++ b/python/triton/experimental/tle/raw/cuda/runtime.py @@ -22,6 +22,7 @@ def __init__(self, fn: Any, file: Path, *args, **kwargs) -> None: self.code: Final[str] = file.read_text() self.file: Final[Path] = file self.libs = kwargs.get("library", None) + self.macros = kwargs.get("macro", None) self.__triton_builtin__: Final[bool] = True def make_llvm(self, mlir_context) -> str: @@ -59,6 +60,11 @@ def make_cubin(self): os.environ[(lib_name + "_home").upper()] = lib_path include_dirs.append(os.path.join(lib_path, "include")) include_flags = [f"-I{inc_dir}" for inc_dir in include_dirs] + + macro_flags = [] + for macro_name, macro_value in self.macros.items(): + macro_flags.append(f"-D{macro_name}={macro_value}") + prop = torch.cuda.get_device_properties(torch.cuda.current_device()) arch = f"-arch=sm_{prop.major}{prop.minor}" @@ -78,10 +84,21 @@ def make_cubin(self): # print("clang cuda") # assert build1.returncode == 0, (f"clang failed\nstderr:\n{build1.stderr.decode()}") - # nvcc cubin - build = subprocess.run([NVCC, "-rdc=true", arch, "-O3", *include_flags, "--extended-lambda", "-c", "-o", dst, src], + # nvcc -> cubin + build = subprocess.run([NVCC, "-rdc=true", arch, "-O3", *include_flags, *macro_flags, "--extended-lambda", "-c", "-o", dst, src], capture_output=True) assert build.returncode == 0, (f"nvcc failed\nstderr:\n{build.stderr.decode()}") + + # nvcc -> ptx -> ptxas -> cubin + # max_reg_per_block = 65536 + # maxnreg = max_reg_per_block // (4 * 32) + # NVCC_GENCODE = f"-gencode=arch=compute_{capability}{suffix},code={arch}" + # nvcc_cmd = [ + # NVCC, "-rdc=true", f"-maxrregcount={maxnreg}", "-ccbin", "g++", NVCC_GENCODE, + # "-I/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/torch/include", + # fsrc.name, "-ptx", "-c", "-o", fptx.name + # ] + # TODO: Remove the method of passing information by setting environment variables. os.environ["USE_NVCC"] = 'True' os.environ["CUDA_CUBIN"] = str(dst) diff --git a/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.cu b/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.cu index 1197db9d5..518258452 100644 --- a/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.cu +++ b/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.cu @@ -11,10 +11,14 @@ extern "C" __device__ void FuncName( \ } \ } -#ifdef ENABLE_FLOAT -DEFINE_VECTOR_ADD(float, vector_add_float) -#endif +#define CONCAT(a, b) a ## b +#define VECTOR_ADD_FUNC_NAME(type) CONCAT(vector_add_, type) +DEFINE_VECTOR_ADD(VECTOR_ELEM_TYPE, VECTOR_ADD_FUNC_NAME(VECTOR_ELEM_TYPE)) -#ifdef ENABLE_INT -DEFINE_VECTOR_ADD(int, vector_add_int) -#endif \ No newline at end of file +// #ifdef ENABLE_FLOAT +// DEFINE_VECTOR_ADD(float, vector_add_float) +// #endif + +// #ifdef ENABLE_INT +// DEFINE_VECTOR_ADD(int, vector_add_int) +// #endif \ No newline at end of file diff --git a/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.py b/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.py index b14163640..61af29193 100644 --- a/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.py +++ b/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.py @@ -10,8 +10,10 @@ DEVICE = triton.runtime.driver.active.get_active_torch_device() -@dialect(name="cuda", file=Path(__file__).parent / "01-vector-add.cu", - library={}) +@dialect(name="cuda", + file=Path(__file__).parent / "vector-add-template.cu", + library={"nvshmem": "/home/zyl/zyuli/envs/nvshmem/lib/python3.12/site-packages/nvidia/nvshmem"}, + macro={"VECTOR_ELEM_TYPE": "int"}) def edsl(*args, **kwargs): ... From 396925ce320556586b03388d80a049266018960b Mon Sep 17 00:00:00 2001 From: zyuli Date: Sat, 23 May 2026 14:15:47 +0800 Subject: [PATCH 4/8] merge tle_raw.call and libdevice.call --- .../experimental/tle/language/raw/core.py | 24 ++++++- .../experimental/tle/raw/cuda/runtime.py | 63 +++++++++++++------ .../vector-add-return-extern-call.py | 20 ++++++ .../06-vector-add-return/vector-add-return.cu | 12 ++++ .../06-vector-add-return/vector-add-return.py | 60 ++++++++++++++++++ 5 files changed, 159 insertions(+), 20 deletions(-) create mode 100644 python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return-extern-call.py create mode 100644 python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.cu create mode 100644 python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.py diff --git a/python/triton/experimental/tle/language/raw/core.py b/python/triton/experimental/tle/language/raw/core.py index 350fb210b..c1a2d5e4d 100644 --- a/python/triton/experimental/tle/language/raw/core.py +++ b/python/triton/experimental/tle/language/raw/core.py @@ -2,12 +2,34 @@ from triton.language.core import builtin, tensor from triton.experimental.tle.language.gpu import buffered_tensor +import importlib.util + +def _pointer_type_hash(self): + return hash((self.name, self.element_ty, "tt_ptr")) + +def patch_hash_method_for_pointer_type(): + elem_dtype_list = tl.core.dtype.SINT_TYPES + tl.core.dtype.UINT_TYPES + tl.core.dtype.FP_TYPES + tl.core.dtype.OTHER_TYPES + for elem_dtype in elem_dtype_list: + ptr_ty = type(tl.core.pointer_type(tl.core.dtype(elem_dtype))) + ptr_ty.__hash__ = _pointer_type_hash + +def import_from_path(file_path): + module_name = f"_imported_{abs(hash(file_path))}" + spec = importlib.util.spec_from_file_location(module_name, file_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module @builtin def call(func, args, _semantic=None): if func.libs is not None: func.make_cubin() - return + patch_hash_method_for_pointer_type() + module = import_from_path(func.extern) + target_fn = getattr(module, func.extern_func_name) + ret = target_fn(*args, _semantic=_semantic) + return ret + context = _semantic.builder.get_context() llvm = func.make_llvm(context) handles = [arg.handle for arg in args] diff --git a/python/triton/experimental/tle/raw/cuda/runtime.py b/python/triton/experimental/tle/raw/cuda/runtime.py index 2d0259002..c9074fb06 100644 --- a/python/triton/experimental/tle/raw/cuda/runtime.py +++ b/python/triton/experimental/tle/raw/cuda/runtime.py @@ -21,6 +21,8 @@ def __init__(self, fn: Any, file: Path, *args, **kwargs) -> None: self.fn: Final[Any] = fn self.code: Final[str] = file.read_text() self.file: Final[Path] = file + self.extern: Final[Path] = kwargs.get("extern", None) + self.extern_func_name = kwargs.get("extern_func_name", None) self.libs = kwargs.get("library", None) self.macros = kwargs.get("macro", None) self.__triton_builtin__: Final[bool] = True @@ -52,21 +54,20 @@ def make_llvm(self, mlir_context) -> str: return f"{module}" def make_cubin(self): - src = self.file - dst = Path(src).with_suffix('.o') + fsrc = self.file + fdst = Path(fsrc).with_suffix('.cubin') include_dirs = [] for lib_name, lib_path in self.libs.items(): # TODO: Remove the method of passing information by setting environment variables. os.environ[(lib_name + "_home").upper()] = lib_path include_dirs.append(os.path.join(lib_path, "include")) include_flags = [f"-I{inc_dir}" for inc_dir in include_dirs] - - macro_flags = [] - for macro_name, macro_value in self.macros.items(): - macro_flags.append(f"-D{macro_name}={macro_value}") + macro_flags = [f"-D{macro_name}={macro_value}" for macro_name, macro_value in self.macros.items()] if self.macros is not None else [] prop = torch.cuda.get_device_properties(torch.cuda.current_device()) - arch = f"-arch=sm_{prop.major}{prop.minor}" + capability = prop.major * 10 + prop.minor + suffix = "a" if capability >= 90 else "" + arch = f"sm_{capability}{suffix}" # clang cubin # build1 = subprocess.run([ @@ -85,21 +86,45 @@ def make_cubin(self): # assert build1.returncode == 0, (f"clang failed\nstderr:\n{build1.stderr.decode()}") # nvcc -> cubin - build = subprocess.run([NVCC, "-rdc=true", arch, "-O3", *include_flags, *macro_flags, "--extended-lambda", "-c", "-o", dst, src], - capture_output=True) - assert build.returncode == 0, (f"nvcc failed\nstderr:\n{build.stderr.decode()}") + # build = subprocess.run([NVCC, "-rdc=true", "-arch={arch}", "-O3", *include_flags, *macro_flags, "--extended-lambda", "-c", "-o", dst, src], + # capture_output=True) + # assert build.returncode == 0, (f"nvcc failed\nstderr:\n{build.stderr.decode()}") # nvcc -> ptx -> ptxas -> cubin - # max_reg_per_block = 65536 - # maxnreg = max_reg_per_block // (4 * 32) - # NVCC_GENCODE = f"-gencode=arch=compute_{capability}{suffix},code={arch}" - # nvcc_cmd = [ - # NVCC, "-rdc=true", f"-maxrregcount={maxnreg}", "-ccbin", "g++", NVCC_GENCODE, - # "-I/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/torch/include", - # fsrc.name, "-ptx", "-c", "-o", fptx.name - # ] + fptx = Path(fsrc).with_suffix('.ptx') + max_reg_per_block = 65536 + num_warps = 4 + + maxnreg = max_reg_per_block // (num_warps * 32) + NVCC_GENCODE = f"-gencode=arch=compute_{capability}{suffix},code={arch}" + nvcc_cmd = [ + NVCC, + "-rdc=true", + "--extended-lambda", + f"-maxrregcount={maxnreg}", + "-ccbin", + "g++", + NVCC_GENCODE, + *include_flags, + *macro_flags, + "-ptx", + "-c", + "-o", + fptx, + fsrc + ] + try: + subprocess.run(nvcc_cmd, check=True, close_fds=False) + except subprocess.CalledProcessError as e: + raise RuntimeError(f"PTX generation failed: {e}") + + ptxas_cmd = [PTXAS, "-c", fptx, f"--gpu-name={arch}", f"-maxrregcount={maxnreg}", "-o", fdst] + try: + subprocess.run(ptxas_cmd, check=True, close_fds=False) + except subprocess.CalledProcessError as e: + raise RuntimeError(f"PTX assembly failed for {arch}: {e}") # TODO: Remove the method of passing information by setting environment variables. os.environ["USE_NVCC"] = 'True' - os.environ["CUDA_CUBIN"] = str(dst) + os.environ["CUDA_CUBIN"] = str(fdst) return diff --git a/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return-extern-call.py b/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return-extern-call.py new file mode 100644 index 000000000..f6bbeef08 --- /dev/null +++ b/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return-extern-call.py @@ -0,0 +1,20 @@ +import triton.language as tl +import triton.language.core as core + +@core.extern +def vector_add_return(C, A, B, N, _semantic=None): + return core.extern_call( + "", + "", + [C, A, B, N], + { + ( + core.pointer_type(core.dtype("fp32")), + core.pointer_type(core.dtype("fp32")), + core.pointer_type(core.dtype("fp32")), + core.dtype("int32"), + ) : ("vector_add_float_return", (core.dtype("uint64"))), + }, + is_pure=False, + _semantic=_semantic, + ) \ No newline at end of file diff --git a/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.cu b/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.cu new file mode 100644 index 000000000..6d6599592 --- /dev/null +++ b/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.cu @@ -0,0 +1,12 @@ +extern "C" __device__ float* vector_add_float_return( + __attribute__((address_space(1))) float *C, + __attribute__((address_space(1))) const float *A, + __attribute__((address_space(1))) const float *B, + const int N) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + for (int i = idx; i < N; i += blockDim.x * gridDim.x) { + C[i] = A[i] + B[i]; + } + + return C; +} diff --git a/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.py b/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.py new file mode 100644 index 000000000..4f2783b8f --- /dev/null +++ b/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.py @@ -0,0 +1,60 @@ +from pathlib import Path + +import torch +import triton +import triton.language as tl +from triton.experimental.tle.raw import dialect +import triton.experimental.tle.language.raw as tle_raw +from triton.language.extra.cuda import libnvshmem_device + +DEVICE = triton.runtime.driver.active.get_active_torch_device() + + +@dialect( + name="cuda", + file=Path(__file__).parent / "vector-add-return.cu", + extern=Path(__file__).parent / "vector-add-return-extern-call.py", + extern_func_name="vector_add_return", + library={"nvshmem": "/home/zyl/zyuli/envs/nvshmem/lib/python3.12/site-packages/nvidia/nvshmem"} +) +def edsl(*args, **kwargs): + ... + + +@triton.jit +def add_kernel( + x_ptr, + y_ptr, + output_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + + ret = tle_raw.call(edsl, [output_ptr, x_ptr, y_ptr, n_elements]) + base = ret.to(tl.pointer_type(tl.float32), bitcast=True) + x = tl.load(base + offsets) + y = tl.full((BLOCK_SIZE, ), 1.0, tl.float32) + output = x + y + tl.store(output_ptr + offsets, output) + + + +def add(x: torch.Tensor, y: torch.Tensor): + output = torch.empty_like(x) + assert x.device == DEVICE and y.device == DEVICE and output.device == DEVICE + n_elements = output.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), ) + add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024) + return output + + +if __name__ == "__main__": + x = torch.randn(2048, device=DEVICE) + y = torch.randn(2048, device=DEVICE) + extra = torch.ones(2048, device=DEVICE) + z = add(x, y) + print(z) + assert torch.allclose(x + y + extra, z), (x + y + extra, z) From c56b174ca67e846c4aa16f7e4b332ab9a62d6659 Mon Sep 17 00:00:00 2001 From: zyuli Date: Wed, 27 May 2026 16:16:02 +0800 Subject: [PATCH 5/8] refactor cuda jit nvcc compile --- .../experimental/tle/language/raw/core.py | 2 +- .../experimental/tle/raw/cuda/runtime.py | 206 +++++++++++------- .../vector-add-template-extern-call.py | 26 +++ .../vector-add-template.cu | 10 +- .../vector-add-template.py | 22 +- .../vector-add-return-extern-call.py | 2 +- .../06-vector-add-return/vector-add-return.cu | 4 +- .../06-vector-add-return/vector-add-return.py | 3 +- .../ring-reduce-device-extern-call.py | 25 +++ .../raw/nvshmem/02-ring-reduce/ring-reduce.py | 18 +- .../tle/raw/nvshmem/04-on-stream/on-stream.py | 4 +- .../nvidia/language/cuda/libnvshmem_device.py | 18 ++ 12 files changed, 226 insertions(+), 114 deletions(-) create mode 100644 python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template-extern-call.py create mode 100644 python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce-device-extern-call.py diff --git a/python/triton/experimental/tle/language/raw/core.py b/python/triton/experimental/tle/language/raw/core.py index c1a2d5e4d..51836be76 100644 --- a/python/triton/experimental/tle/language/raw/core.py +++ b/python/triton/experimental/tle/language/raw/core.py @@ -22,7 +22,7 @@ def import_from_path(file_path): @builtin def call(func, args, _semantic=None): - if func.libs is not None: + if (func.compiler).lower() == "nvcc": func.make_cubin() patch_hash_method_for_pointer_type() module = import_from_path(func.extern) diff --git a/python/triton/experimental/tle/raw/cuda/runtime.py b/python/triton/experimental/tle/raw/cuda/runtime.py index c9074fb06..ad26f9432 100644 --- a/python/triton/experimental/tle/raw/cuda/runtime.py +++ b/python/triton/experimental/tle/raw/cuda/runtime.py @@ -5,6 +5,9 @@ from typing import Any, Final import torch +from functools import partial +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] @@ -12,8 +15,120 @@ CLANG = os.getenv("CLANG", "clang") NVCC = os.getenv("NVCC", "nvcc") PTXAS = os.getenv("PTXAS", "ptxas") +NVLINK = os.getenv("NVLINK", "nvlink") +def make_cubin_inspection_hook(cuda_self, triton_self, stages, options, language, capability): + from triton.backends.nvidia.compiler import sm_arch_from_capability, get_ptxas + def make_cubin(self, src, metadata, opt, capability): + # cuda compile + fsrc_cuda = cuda_self.file + fbin_cuda = Path(fsrc_cuda).with_suffix('.o') + + include_flags = [] + link_flags = [] + if cuda_self.libs is not None: + for _, lib_path in cuda_self.libs.items(): + include_flags.append(f"-I{os.path.join(lib_path, "include")}") + link_flags.append(f"-L{os.path.join(lib_path, "lib")}") + + link_libs = [f"-l{lib}" for lib in cuda_self.links] if cuda_self.links is not None else [] + macro_flags = [f"-D{macro_name}={macro_value}" for macro_name, macro_value in cuda_self.macros.items()] if cuda_self.macros is not None else [] + + prop = torch.cuda.get_device_properties(torch.cuda.current_device()) + capability = prop.major * 10 + prop.minor + suffix = "a" if capability >= 90 else "" + arch = f"sm_{capability}{suffix}" + + build = subprocess.run([ + NVCC, + "-rdc=true", + f"-arch={arch}", + *include_flags, + *macro_flags, + "--extended-lambda", + "-c", + "-o", + fbin_cuda, + fsrc_cuda], + capture_output=True) + assert build.returncode == 0, (f"nvcc failed\nstderr:\n{build.stderr.decode()}") + + # compile triton + import tempfile + import signal + from triton.runtime.errors import PTXASError + with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.ptx') as fsrc_triton, \ + tempfile.NamedTemporaryFile(delete=False, mode='r', suffix='.log') as flog: + fsrc_triton.write(src) + fsrc_triton.flush() + fbin_triton = fsrc_triton.name + '.o' + fbin_combined = fbin_triton + '.combined.cubin' + + compile_only_cmds = ['-c'] + line_info = ["-lineinfo", "-suppress-debug-info"] if knobs.compilation.disable_line_info else ["-lineinfo"] + fmad = [] if opt.enable_fp_fusion else ['--fmad=false'] + disable_opt = ['--opt-level', '0'] if knobs.nvidia.disable_ptxas_opt else [] + ptx_extra_options = opt.ptx_options.split(" ") if opt.ptx_options else [] + + ptxas_cmd = [ + PTXAS, *compile_only_cmds, *line_info, *fmad, '-v', *disable_opt, *ptx_extra_options, + f'--gpu-name={arch}', fsrc_triton.name, '-o', fbin_triton + ] + + try: + subprocess.run(ptxas_cmd, check=True, close_fds=False, stderr=flog) + if os.path.exists(fsrc_triton.name): + os.remove(fsrc_triton.name) + if os.path.exists(flog.name): + os.remove(flog.name) + except subprocess.CalledProcessError as e: + with open(flog.name) as log_file: + log = log_file.read() + if os.path.exists(flog.name): + os.remove(flog.name) + + if e.returncode == 255: + error = 'Internal Triton PTX codegen error' + elif e.returncode == 128 + signal.SIGSEGV: + error = '`ptxas` raised SIGSEGV' + else: + error = f'`ptxas` failed with error code {e.returncode}' + raise PTXASError(f"{error}\n" + f"`ptxas` stderr:\n{log}\n" + f'Repro command: {" ".join(ptxas_cmd)}\n') + + # nvlink + nvlink_cmds = [ + NVLINK, + f"-arch={arch}", + *link_flags, + *link_libs, + fbin_triton, + fbin_cuda, + "-o", + fbin_combined, + ] + + try: + subprocess.run(nvlink_cmds, check=True, close_fds=False, stderr=flog) + except Exception as e: + import logging + logging.error(f"error runing nvlink: {nvlink_cmds}") + logging.exception(e) + + with open(fbin_combined, 'rb') as f: + cubin = f.read() + if os.path.exists(fbin_combined): + os.remove(fbin_combined) + if os.path.exists(fbin_triton): + os.remove(fbin_triton) + + return cubin + + stages["cubin"] = lambda src, metadata: make_cubin(triton_self, src, metadata, options, triton_self.target.arch) + + class CUDAJITFunction(object): def __init__(self, fn: Any, file: Path, *args, **kwargs) -> None: @@ -21,22 +136,30 @@ def __init__(self, fn: Any, file: Path, *args, **kwargs) -> None: self.fn: Final[Any] = fn self.code: Final[str] = file.read_text() self.file: Final[Path] = file + self.compiler = kwargs.get("compiler", "clang") self.extern: Final[Path] = kwargs.get("extern", None) self.extern_func_name = kwargs.get("extern_func_name", None) - self.libs = kwargs.get("library", None) - self.macros = kwargs.get("macro", None) + self.libs = kwargs.get("libs", None) + self.links = kwargs.get("links", None) + self.macros = kwargs.get("macros", None) self.__triton_builtin__: Final[bool] = True + + if (self.compiler).lower() == "nvcc": + bound_hook = partial(make_cubin_inspection_hook, self) + knobs.runtime.add_stages_inspection_hook = bound_hook def make_llvm(self, mlir_context) -> str: prop = torch.cuda.get_device_properties(torch.cuda.current_device()) - arch = f"--cuda-gpu-arch=sm_{prop.major}{prop.minor}" + capability = prop.major * 10 + prop.minor + suffix = "a" if capability >= 90 else "" + arch = f"sm_{capability}{suffix}" build = subprocess.run( [ CLANG, "-x", "cuda", "--cuda-device-only", - arch, + f"--cuda-gpu-arch={arch}", "-emit-llvm", "-I/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/torch/include", "-O2", @@ -54,77 +177,4 @@ def make_llvm(self, mlir_context) -> str: return f"{module}" def make_cubin(self): - fsrc = self.file - fdst = Path(fsrc).with_suffix('.cubin') - include_dirs = [] - for lib_name, lib_path in self.libs.items(): - # TODO: Remove the method of passing information by setting environment variables. - os.environ[(lib_name + "_home").upper()] = lib_path - include_dirs.append(os.path.join(lib_path, "include")) - include_flags = [f"-I{inc_dir}" for inc_dir in include_dirs] - macro_flags = [f"-D{macro_name}={macro_value}" for macro_name, macro_value in self.macros.items()] if self.macros is not None else [] - - prop = torch.cuda.get_device_properties(torch.cuda.current_device()) - capability = prop.major * 10 + prop.minor - suffix = "a" if capability >= 90 else "" - arch = f"sm_{capability}{suffix}" - - # clang cubin - # build1 = subprocess.run([ - # CLANG, - # "-fgpu-rdc", - # "-c", - # "--cuda-device-only", - # "--cuda-gpu-arch=sm_90", - # "-O3", - # *include_flags, - # "-o", - # dst, - # src - # ], capture_output=True) - # print("clang cuda") - # assert build1.returncode == 0, (f"clang failed\nstderr:\n{build1.stderr.decode()}") - - # nvcc -> cubin - # build = subprocess.run([NVCC, "-rdc=true", "-arch={arch}", "-O3", *include_flags, *macro_flags, "--extended-lambda", "-c", "-o", dst, src], - # capture_output=True) - # assert build.returncode == 0, (f"nvcc failed\nstderr:\n{build.stderr.decode()}") - - # nvcc -> ptx -> ptxas -> cubin - fptx = Path(fsrc).with_suffix('.ptx') - max_reg_per_block = 65536 - num_warps = 4 - - maxnreg = max_reg_per_block // (num_warps * 32) - NVCC_GENCODE = f"-gencode=arch=compute_{capability}{suffix},code={arch}" - nvcc_cmd = [ - NVCC, - "-rdc=true", - "--extended-lambda", - f"-maxrregcount={maxnreg}", - "-ccbin", - "g++", - NVCC_GENCODE, - *include_flags, - *macro_flags, - "-ptx", - "-c", - "-o", - fptx, - fsrc - ] - try: - subprocess.run(nvcc_cmd, check=True, close_fds=False) - except subprocess.CalledProcessError as e: - raise RuntimeError(f"PTX generation failed: {e}") - - ptxas_cmd = [PTXAS, "-c", fptx, f"--gpu-name={arch}", f"-maxrregcount={maxnreg}", "-o", fdst] - try: - subprocess.run(ptxas_cmd, check=True, close_fds=False) - except subprocess.CalledProcessError as e: - raise RuntimeError(f"PTX assembly failed for {arch}: {e}") - - # TODO: Remove the method of passing information by setting environment variables. - os.environ["USE_NVCC"] = 'True' - os.environ["CUDA_CUBIN"] = str(fdst) - return + ... \ No newline at end of file diff --git a/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template-extern-call.py b/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template-extern-call.py new file mode 100644 index 000000000..f495a8d59 --- /dev/null +++ b/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template-extern-call.py @@ -0,0 +1,26 @@ +import triton.language as tl +import triton.language.core as core + +@core.extern +def vector_add(C, A, B, N, _semantic=None): + return core.extern_call( + "", + "", + [C, A, B, N], + { + ( + core.pointer_type(core.dtype("int32")), + core.pointer_type(core.dtype("int32")), + core.pointer_type(core.dtype("int32")), + core.dtype("int32"), + ) : ("vector_add_int", ()), + ( + core.pointer_type(core.dtype("fp32")), + core.pointer_type(core.dtype("fp32")), + core.pointer_type(core.dtype("fp32")), + core.dtype("int32"), + ) : ("vector_add_float", ()), + }, + is_pure=False, + _semantic=_semantic, + ) \ No newline at end of file diff --git a/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.cu b/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.cu index 518258452..ade9be51b 100644 --- a/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.cu +++ b/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.cu @@ -13,12 +13,4 @@ extern "C" __device__ void FuncName( \ #define CONCAT(a, b) a ## b #define VECTOR_ADD_FUNC_NAME(type) CONCAT(vector_add_, type) -DEFINE_VECTOR_ADD(VECTOR_ELEM_TYPE, VECTOR_ADD_FUNC_NAME(VECTOR_ELEM_TYPE)) - -// #ifdef ENABLE_FLOAT -// DEFINE_VECTOR_ADD(float, vector_add_float) -// #endif - -// #ifdef ENABLE_INT -// DEFINE_VECTOR_ADD(int, vector_add_int) -// #endif \ No newline at end of file +DEFINE_VECTOR_ADD(VECTOR_ELEM_TYPE, VECTOR_ADD_FUNC_NAME(VECTOR_ELEM_TYPE)) \ No newline at end of file diff --git a/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.py b/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.py index 61af29193..297423695 100644 --- a/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.py +++ b/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.py @@ -10,10 +10,12 @@ DEVICE = triton.runtime.driver.active.get_active_torch_device() -@dialect(name="cuda", +@dialect(name="cuda", + compiler="nvcc", file=Path(__file__).parent / "vector-add-template.cu", - library={"nvshmem": "/home/zyl/zyuli/envs/nvshmem/lib/python3.12/site-packages/nvidia/nvshmem"}, - macro={"VECTOR_ELEM_TYPE": "int"}) + extern=Path(__file__).parent / "vector-add-template-extern-call.py", + extern_func_name="vector_add", + macros={"VECTOR_ELEM_TYPE": "float"}) def edsl(*args, **kwargs): ... @@ -27,7 +29,6 @@ def add_kernel( BLOCK_SIZE: tl.constexpr, ): tle_raw.call(edsl, [output_ptr, x_ptr, y_ptr, n_elements]) - libnvshmem_device.vector_add(output_ptr, x_ptr, y_ptr, n_elements) def add(x: torch.Tensor, y: torch.Tensor): @@ -40,13 +41,14 @@ def add(x: torch.Tensor, y: torch.Tensor): if __name__ == "__main__": - # dtype = torch.float32 - # x = torch.randn(2048, dtype=dtype, device=DEVICE) - # y = torch.randn(2048, dtype=dtype, device=DEVICE) + dtype = torch.float32 + x = torch.randn(2048, dtype=dtype, device=DEVICE) + y = torch.randn(2048, dtype=dtype, device=DEVICE) - dtype = torch.int32 - x = torch.randint(low=0, high=10, size=(2048, ), dtype=dtype, device=DEVICE) - y = torch.randint(low=0, high=10, size=(2048, ), dtype=dtype, device=DEVICE) + # dtype = torch.int32 + # x = torch.randint(low=0, high=10, size=(2048, ), dtype=dtype, device=DEVICE) + # y = torch.randint(low=0, high=10, size=(2048, ), dtype=dtype, device=DEVICE) z = add(x, y) + print(z) assert torch.allclose(x + y, z), (x + y, z) diff --git a/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return-extern-call.py b/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return-extern-call.py index f6bbeef08..b786c6fa3 100644 --- a/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return-extern-call.py +++ b/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return-extern-call.py @@ -13,7 +13,7 @@ def vector_add_return(C, A, B, N, _semantic=None): core.pointer_type(core.dtype("fp32")), core.pointer_type(core.dtype("fp32")), core.dtype("int32"), - ) : ("vector_add_float_return", (core.dtype("uint64"))), + ) : ("vector_add_float_return", (core.dtype("int32"))), }, is_pure=False, _semantic=_semantic, diff --git a/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.cu b/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.cu index 6d6599592..f5386777b 100644 --- a/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.cu +++ b/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.cu @@ -1,4 +1,4 @@ -extern "C" __device__ float* vector_add_float_return( +extern "C" __device__ int vector_add_float_return( __attribute__((address_space(1))) float *C, __attribute__((address_space(1))) const float *A, __attribute__((address_space(1))) const float *B, @@ -8,5 +8,5 @@ extern "C" __device__ float* vector_add_float_return( C[i] = A[i] + B[i]; } - return C; + return idx; } diff --git a/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.py b/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.py index 4f2783b8f..d2d9719b0 100644 --- a/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.py +++ b/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.py @@ -12,10 +12,10 @@ @dialect( name="cuda", + compiler="nvcc", file=Path(__file__).parent / "vector-add-return.cu", extern=Path(__file__).parent / "vector-add-return-extern-call.py", extern_func_name="vector_add_return", - library={"nvshmem": "/home/zyl/zyuli/envs/nvshmem/lib/python3.12/site-packages/nvidia/nvshmem"} ) def edsl(*args, **kwargs): ... @@ -34,6 +34,7 @@ def add_kernel( offsets = block_start + tl.arange(0, BLOCK_SIZE) ret = tle_raw.call(edsl, [output_ptr, x_ptr, y_ptr, n_elements]) + tl.device_print("ret", ret) base = ret.to(tl.pointer_type(tl.float32), bitcast=True) x = tl.load(base + offsets) y = tl.full((BLOCK_SIZE, ), 1.0, tl.float32) diff --git a/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce-device-extern-call.py b/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce-device-extern-call.py new file mode 100644 index 000000000..00d3268b0 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce-device-extern-call.py @@ -0,0 +1,25 @@ +import triton.language as tl +import triton.language.core as core + +@core.extern +def ring_reduce(dst, src, nreduce, signal, chunk_size, _semantic=None): + return core.extern_call( + "", + "", + [dst, + src, + tl.cast(nreduce, tl.int32, _semantic=_semantic), + signal, + tl.cast(chunk_size, tl.int32, _semantic=_semantic)], + { + ( + core.pointer_type(core.dtype("int32")), + core.pointer_type(core.dtype("int32")), + core.dtype("int32"), + core.pointer_type(core.dtype("int64")), + core.dtype("int32"), + ) : ("ring_reduce", ()), + }, + is_pure=False, + _semantic=_semantic, + ) diff --git a/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce.py b/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce.py index d7ac1da9c..67d6c8309 100644 --- a/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce.py +++ b/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce.py @@ -13,8 +13,12 @@ @dialect( name="cuda", - file=(Path(__file__).parent / "ring-reduce-device.cu").resolve(), - library={"nvshmem": "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem"}, + compiler="nvcc", + file=(Path(__file__).parent / "ring-reduce-device.cu"), + extern=(Path(__file__).parent / "ring-reduce-device-extern-call.py"), + extern_func_name="ring_reduce", + libs={"nvshmem": "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem"}, + links=["nvshmem_device"] ) def edsl(*args, **kwargs): ... @@ -28,8 +32,8 @@ def ring_reduce_kernel( signal, chunk_size, ): - tle_raw.call(edsl, []) - libnvshmem_device.ring_reduce(dst, src, nreduce, signal, chunk_size) + tle_raw.call(edsl, [dst, src, nreduce, signal, chunk_size]) + # libnvshmem_device.ring_reduce(dst, src, nreduce, signal, chunk_size) def cuda_host_compile(cuda_host_path, cuda_host_lib): @@ -94,7 +98,6 @@ def ring_reduce(): size ) - # print("PE:", mype_in_node.value) dtype = torch.int32 num_blocks = npes_in_node.value device = triton.runtime.driver.active.get_active_torch_device() @@ -104,9 +107,6 @@ def ring_reduce(): dst_storage = torch._C._construct_storage_from_data_pointer(dst.value, device, 4 * M * N) dst_tensor = torch.empty(0, dtype=torch.int32, device=device).set_(dst_storage).view(M, N) - # data_h_storage = torch._C._construct_storage_from_data_pointer(data_h.value, device, 4 * M * N) - # data_h_tensor = torch.empty(0, dtype=torch.int32, device=device).set_(data_h_storage).view(M, N) - signal_storage = torch._C._construct_storage_from_data_pointer(signal.value, device, 8 * num_blocks) signal_tensor = torch.empty(0, dtype=torch.int64, device=device).set_(signal_storage).view(num_blocks, ) @@ -123,8 +123,6 @@ def cumodule_init_hook(*args, **kwargs): knobs.runtime.jit_post_compile_hook = cumodule_init_hook curr_stream = torch.cuda.ExternalStream(stream.value, device=device) - # print(f"[PE {mype.value}] Using stream pointer: {hex(stream.value)}") - # print(f"[PE {mype.value}] Default stream before: {torch.cuda.default_stream().cuda_stream}") with torch.cuda.stream(curr_stream): chunk_size = int(M // num_blocks * dtype.itemsize) ring_reduce_kernel[(num_blocks, )]( diff --git a/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream.py b/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream.py index 56a8928e0..03899665c 100644 --- a/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream.py +++ b/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream.py @@ -14,7 +14,7 @@ @dialect( name="cuda", file=(Path(__file__).parent / "on-stream-device-accumulate.cu").resolve(), - library={"nvshmem": "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem"}, + libs={"nvshmem": "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem"}, ) def edsl_accumulate(*args, **kwargs): ... @@ -23,7 +23,7 @@ def edsl_accumulate(*args, **kwargs): @dialect( name="cuda", file=(Path(__file__).parent / "on-stream-device-correct-accumulate.cu").resolve(), - library={"nvshmem": "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem"}, + libs={"nvshmem": "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem"}, ) def edsl_correct_accumulate(*args, **kwargs): ... diff --git a/third_party/nvidia/language/cuda/libnvshmem_device.py b/third_party/nvidia/language/cuda/libnvshmem_device.py index 799a298d9..c184c0dd7 100644 --- a/third_party/nvidia/language/cuda/libnvshmem_device.py +++ b/third_party/nvidia/language/cuda/libnvshmem_device.py @@ -209,3 +209,21 @@ def vector_add(C, A, B, N, _semantic=None): is_pure=False, _semantic=_semantic, ) + +@core.extern +def vector_add_return(C, A, B, N, _semantic=None): + return core.extern_call( + "", + "", + [C, A, B, N], + { + ( + core.pointer_type(core.dtype("fp32")), + core.pointer_type(core.dtype("fp32")), + core.pointer_type(core.dtype("fp32")), + core.dtype("int32"), + ) : ("vector_add_float_return", (core.dtype("uint64"))), + }, + is_pure=False, + _semantic=_semantic, + ) \ No newline at end of file From 2f7ef638ccfdae97e02af1373af012b788ebe8c7 Mon Sep 17 00:00:00 2001 From: zyuli Date: Wed, 10 Jun 2026 00:31:58 +0000 Subject: [PATCH 6/8] register make_cubin during the first initialization --- clean_cache.sh | 1 - .../experimental/tle/raw/cuda/runtime.py | 2 +- .../nvshmem/01-simple-shift/simple-shift.py | 16 ++++--- .../ring-reduce-device-extern-call.py | 27 +++++++----- .../raw/nvshmem/02-ring-reduce/ring-reduce.py | 44 +++++++++---------- 5 files changed, 49 insertions(+), 41 deletions(-) delete mode 100755 clean_cache.sh diff --git a/clean_cache.sh b/clean_cache.sh deleted file mode 100755 index 8cdaf2438..000000000 --- a/clean_cache.sh +++ /dev/null @@ -1 +0,0 @@ -rm -rf /home/zyuli/.triton/cache/* diff --git a/python/triton/experimental/tle/raw/cuda/runtime.py b/python/triton/experimental/tle/raw/cuda/runtime.py index ad26f9432..9910a3638 100644 --- a/python/triton/experimental/tle/raw/cuda/runtime.py +++ b/python/triton/experimental/tle/raw/cuda/runtime.py @@ -144,7 +144,7 @@ def __init__(self, fn: Any, file: Path, *args, **kwargs) -> None: self.macros = kwargs.get("macros", None) self.__triton_builtin__: Final[bool] = True - if (self.compiler).lower() == "nvcc": + if (self.compiler).lower() == "nvcc" and (knobs.runtime.add_stages_inspection_hook is None): bound_hook = partial(make_cubin_inspection_hook, self) knobs.runtime.add_stages_inspection_hook = bound_hook 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 2f87d70f6..04bdc8f70 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 @@ -10,11 +10,16 @@ from triton.experimental.tle.raw import dialect from triton.language.extra.cuda import libnvshmem_device +NVSHMEM_HOME = "/data/zyuli/miniconda3/envs/flagtree_triton_v3.6.x/lib/python3.12/site-packages/nvidia/nvshmem" @dialect( name="cuda", - file=(Path(__file__).parent / "simple-shift-device.cu").resolve(), - library={"nvshmem": "/home/zyl/zyuli/envs/nvshmem/lib/python3.12/site-packages/nvidia/nvshmem"}, + compiler="nvcc", + file=(Path(__file__).parent / "simple-shift-device.cu"), + extern=(Path(__file__).parent / "simple-shift-device-extern-call.py"), + extern_func_name="simple_shift", + libs={"nvshmem": NVSHMEM_HOME}, + links=["nvshmem_device"] ) def edsl(*args, **kwargs): ... @@ -23,13 +28,12 @@ def edsl(*args, **kwargs): @triton.jit def simple_shift_kernel(destination_ptr, ): # TODO: Combine with tle_raw.call, then dispatch - tle_raw.call_nvshmem(edsl, [], [destination_ptr]) - libnvshmem_device.simple_shift(destination_ptr) + tle_raw.call(edsl, [destination_ptr]) + # libnvshmem_device.simple_shift(destination_ptr) def cuda_host_compile(cuda_host_path, cuda_host_lib): NVCC = os.getenv("NVCC", "nvcc") - NVSHMEM_HOME = "/home/zyl/zyuli/envs/nvshmem/lib/python3.12/site-packages/nvidia/nvshmem" include_path = f"-I{os.path.join(NVSHMEM_HOME, 'include')}" lib_path = f"-L{os.path.join(NVSHMEM_HOME, 'lib')}" @@ -103,9 +107,9 @@ def cumodule_init_hook(*args, **kwargs): kernel._init_handles() ret = lib.nvshmemx_cumodule_init_wrapper(ctypes.c_void_p(kernel.module)) assert ret == 0, f"nvshmemx_cumodule_init_wrapper failed: {ret}" - knobs.runtime.jit_post_compile_hook = cumodule_init_hook + # extern_libs = {'simple-shift': 'simple-shift-device-opti.bc'} simple_shift_kernel[(1, )](dest_tensor) stream_ptr = stream.cuda_stream diff --git a/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce-device-extern-call.py b/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce-device-extern-call.py index 00d3268b0..65271655a 100644 --- a/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce-device-extern-call.py +++ b/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce-device-extern-call.py @@ -1,25 +1,30 @@ +# Generated from ring-reduce-device.cu; do not edit manually. import triton.language as tl import triton.language.core as core + @core.extern def ring_reduce(dst, src, nreduce, signal, chunk_size, _semantic=None): return core.extern_call( "", "", - [dst, - src, - tl.cast(nreduce, tl.int32, _semantic=_semantic), - signal, - tl.cast(chunk_size, tl.int32, _semantic=_semantic)], + [ + dst, + src, + tl.cast(nreduce, tl.int32, _semantic=_semantic), + signal, + tl.cast(chunk_size, tl.int32, _semantic=_semantic), + ], { ( - core.pointer_type(core.dtype("int32")), - core.pointer_type(core.dtype("int32")), - core.dtype("int32"), - core.pointer_type(core.dtype("int64")), - core.dtype("int32"), - ) : ("ring_reduce", ()), + core.pointer_type(core.dtype("int32")), + core.pointer_type(core.dtype("int32")), + core.dtype("int32"), + core.pointer_type(core.dtype("uint64")), + core.dtype("int32"), + ): ("ring_reduce", ()), }, is_pure=False, _semantic=_semantic, ) + diff --git a/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce.py b/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce.py index 67d6c8309..a3b01caba 100644 --- a/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce.py +++ b/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce.py @@ -10,14 +10,14 @@ from triton.experimental.tle.raw import dialect from triton.language.extra.cuda import libnvshmem_device - +NVSHMEM_HOME = "/data/zyuli/miniconda3/envs/flagtree_triton_v3.6.x/lib/python3.12/site-packages/nvidia/nvshmem" @dialect( name="cuda", compiler="nvcc", file=(Path(__file__).parent / "ring-reduce-device.cu"), extern=(Path(__file__).parent / "ring-reduce-device-extern-call.py"), extern_func_name="ring_reduce", - libs={"nvshmem": "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem"}, + libs={"nvshmem": NVSHMEM_HOME}, links=["nvshmem_device"] ) def edsl(*args, **kwargs): @@ -33,12 +33,10 @@ def ring_reduce_kernel( chunk_size, ): tle_raw.call(edsl, [dst, src, nreduce, signal, chunk_size]) - # libnvshmem_device.ring_reduce(dst, src, nreduce, signal, chunk_size) def cuda_host_compile(cuda_host_path, cuda_host_lib): NVCC = os.getenv("NVCC", "nvcc") - NVSHMEM_HOME = "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem" include_path = f"-I{os.path.join(NVSHMEM_HOME, 'include')}" lib_path = f"-L{os.path.join(NVSHMEM_HOME, 'lib')}" @@ -58,23 +56,24 @@ def ring_reduce(): cu_file = (Path(__file__).parent / "ring-reduce-host.cu").resolve() lib_file = Path(cu_file).with_suffix('.so') - rank = int(os.getenv("OMPI_COMM_WORLD_RANK", "0")) - if rank == 0: - cuda_host_compile(cu_file, lib_file) - - import time - timeout = 60 - start = time.time() - while True: - if lib_file.exists(): - try: - ctypes.CDLL(str(lib_file)) - break - except OSError: - pass - if time.time() - start > timeout: - raise RuntimeError(f"Timeout waiting for {lib_file}") - time.sleep(0.1) + # rank = int(os.getenv("OMPI_COMM_WORLD_RANK", "0")) + rank = int(os.environ["PMI_RANK"]) + # if rank == 0: + # cuda_host_compile(cu_file, lib_file) + + # import time + # timeout = 60 + # start = time.time() + # while True: + # if lib_file.exists(): + # try: + # ctypes.CDLL(str(lib_file)) + # break + # except OSError: + # pass + # if time.time() - start > timeout: + # raise RuntimeError(f"Timeout waiting for {lib_file}") + # time.sleep(0.1) M, N = 64, 8 lib = ctypes.CDLL(lib_file) @@ -108,7 +107,7 @@ def ring_reduce(): dst_tensor = torch.empty(0, dtype=torch.int32, device=device).set_(dst_storage).view(M, N) signal_storage = torch._C._construct_storage_from_data_pointer(signal.value, device, 8 * num_blocks) - signal_tensor = torch.empty(0, dtype=torch.int64, device=device).set_(signal_storage).view(num_blocks, ) + signal_tensor = torch.empty(0, dtype=torch.uint64, device=device).set_(signal_storage).view(num_blocks, ) def cumodule_init_hook(*args, **kwargs): key = kwargs["key"] @@ -133,6 +132,7 @@ def cumodule_init_hook(*args, **kwargs): chunk_size ) + print(f"PE {mype.value}: {npes.value}") lib.ring_reduce_after_launch( stream, src, dst, data_h, signal, From 8cea58e799ec15a8e202da4634f8444a2944333b Mon Sep 17 00:00:00 2001 From: zyuli Date: Wed, 10 Jun 2026 03:47:52 +0000 Subject: [PATCH 7/8] add test case --- .../experimental/tle/raw/cuda/runtime.py | 1 - .../per-token-group-quant-fp8-clang.py | 281 ------- .../per-token-group-quant-fp8-nvcc.py | 283 ------- .../per-token-group-quant-fp8.cu | 138 ---- .../vectorization.cuh | 30 - .../vectorization_utils.cuh | 177 ----- .../raw/cuda/05-vector-add-template/clang.py | 98 --- .../raw/cuda/05-vector-add-template/nvcc.py | 100 --- .../vector-add-template-extern-call.py | 26 - .../vector-add-template.cu | 16 - .../vector-add-template.py | 54 -- .../vector-add-return-extern-call.py | 20 - .../06-vector-add-return/vector-add-return.cu | 12 - .../06-vector-add-return/vector-add-return.py | 61 -- .../00-gemm-allreduce/gemm-allreduce-host.cu | 113 --- .../00-gemm-allreduce/gemm-allreduce.py | 139 ---- .../nvshmem/00-gemm-allreduce/simple-gemm.cu | 51 -- .../nvshmem/01-simple-shift/simple-shift.py | 41 +- .../ring-reduce-device-extern-call.py | 30 - .../raw/nvshmem/02-ring-reduce/ring-reduce.py | 35 - .../03-overlap-in-one-kernel/no-overlap.py | 288 ++++++++ .../03-overlap-in-one-kernel/overlap.py | 279 +++++++ .../ring-reduce-host.cu | 79 ++ .../ring-reduce-put.cu | 165 +++++ .../03-ring-bcast/ring-bcast-device.cu | 23 - .../nvshmem/03-ring-bcast/ring-bcast-host.cu | 59 -- .../raw/nvshmem/03-ring-bcast/ring-bcast.py | 137 ---- .../on-stream-device-accumulate.cu | 8 - .../on-stream-device-correct-accumulate.cu | 13 - .../nvshmem/04-on-stream/on-stream-host.cu | 69 -- .../tle/raw/nvshmem/04-on-stream/on-stream.py | 166 ----- .../04-overlap-in-mulit-kernel/no-overlap.py | 5 + .../04-overlap-in-mulit-kernel/overlap.py | 5 + .../ring-reduce-host.cu | 78 ++ .../ring-reduce-put.cu | 84 +++ .../separate_stream_common.py | 507 +++++++++++++ .../nvshmem/05-put-block/put-block-device.cu | 23 - .../nvshmem/05-put-block/put-block-host.cu | 75 -- .../tle/raw/nvshmem/05-put-block/put-block.py | 128 ---- .../nvshmem/common/generate_extern_call.py | 207 ++++++ .../tle/raw/nvshmem/common/prepare_nvshmem.py | 163 ++++ .../tutorials/tle/raw/nvshmem/run_nvshmem.sh | 138 ++++ support_nvshmem.patch | 697 ------------------ third_party/nvidia/backend/compiler.py | 41 +- third_party/nvidia/language/cuda/__init__.py | 2 - .../nvidia/language/cuda/libnvshmem_device.py | 229 ------ 46 files changed, 2004 insertions(+), 3370 deletions(-) delete mode 100644 python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8-clang.py delete mode 100644 python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8-nvcc.py delete mode 100644 python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8.cu delete mode 100644 python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/vectorization.cuh delete mode 100644 python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/vectorization_utils.cuh delete mode 100644 python/tutorials/tle/raw/cuda/05-vector-add-template/clang.py delete mode 100644 python/tutorials/tle/raw/cuda/05-vector-add-template/nvcc.py delete mode 100644 python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template-extern-call.py delete mode 100644 python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.cu delete mode 100644 python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.py delete mode 100644 python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return-extern-call.py delete mode 100644 python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.cu delete mode 100644 python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.py delete mode 100644 python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/gemm-allreduce-host.cu delete mode 100644 python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/gemm-allreduce.py delete mode 100644 python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/simple-gemm.cu delete mode 100644 python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce-device-extern-call.py create mode 100644 python/tutorials/tle/raw/nvshmem/03-overlap-in-one-kernel/no-overlap.py create mode 100644 python/tutorials/tle/raw/nvshmem/03-overlap-in-one-kernel/overlap.py create mode 100644 python/tutorials/tle/raw/nvshmem/03-overlap-in-one-kernel/ring-reduce-host.cu create mode 100644 python/tutorials/tle/raw/nvshmem/03-overlap-in-one-kernel/ring-reduce-put.cu delete mode 100644 python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast-device.cu delete mode 100644 python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast-host.cu delete mode 100644 python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast.py delete mode 100644 python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-device-accumulate.cu delete mode 100644 python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-device-correct-accumulate.cu delete mode 100644 python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-host.cu delete mode 100644 python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream.py create mode 100644 python/tutorials/tle/raw/nvshmem/04-overlap-in-mulit-kernel/no-overlap.py create mode 100644 python/tutorials/tle/raw/nvshmem/04-overlap-in-mulit-kernel/overlap.py create mode 100644 python/tutorials/tle/raw/nvshmem/04-overlap-in-mulit-kernel/ring-reduce-host.cu create mode 100644 python/tutorials/tle/raw/nvshmem/04-overlap-in-mulit-kernel/ring-reduce-put.cu create mode 100644 python/tutorials/tle/raw/nvshmem/04-overlap-in-mulit-kernel/separate_stream_common.py delete mode 100644 python/tutorials/tle/raw/nvshmem/05-put-block/put-block-device.cu delete mode 100644 python/tutorials/tle/raw/nvshmem/05-put-block/put-block-host.cu delete mode 100644 python/tutorials/tle/raw/nvshmem/05-put-block/put-block.py create mode 100644 python/tutorials/tle/raw/nvshmem/common/generate_extern_call.py create mode 100644 python/tutorials/tle/raw/nvshmem/common/prepare_nvshmem.py create mode 100755 python/tutorials/tle/raw/nvshmem/run_nvshmem.sh delete mode 100644 support_nvshmem.patch delete mode 100644 third_party/nvidia/language/cuda/libnvshmem_device.py diff --git a/python/triton/experimental/tle/raw/cuda/runtime.py b/python/triton/experimental/tle/raw/cuda/runtime.py index 9910a3638..c3f21bc89 100644 --- a/python/triton/experimental/tle/raw/cuda/runtime.py +++ b/python/triton/experimental/tle/raw/cuda/runtime.py @@ -161,7 +161,6 @@ def make_llvm(self, mlir_context) -> str: "--cuda-device-only", f"--cuda-gpu-arch={arch}", "-emit-llvm", - "-I/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/torch/include", "-O2", "-S", "-", diff --git a/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8-clang.py b/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8-clang.py deleted file mode 100644 index 3c56cf6a4..000000000 --- a/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8-clang.py +++ /dev/null @@ -1,281 +0,0 @@ -from pathlib import Path - -from typing import Optional -import logging -import torch -import triton -import triton.language as tl -from triton.experimental.tle.raw import dialect -import triton.experimental.tle.language.raw as tle_raw - -from triton.language.extra.cuda import libnvshmem_device - -torch.cuda.set_device(1) -DEVICE = triton.runtime.driver.active.get_active_torch_device() -logger = logging.getLogger(__name__) - - -@dialect(name="cuda", file=(Path(__file__).parent / "per-token-group-quant-fp8.cu").resolve()) - # library={"torch": "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/torch/"}) -def edsl(*args, **kwargs): - ... - - -@triton.jit -def test_kernel( - x_ptr, - x_q_ptr, - x_s_ptr, - group_size, - num_groups, - groups_per_block, - eps, - fp8_min, - fp8_max, -): - tle_raw.call(edsl, [x_q_ptr, x_s_ptr, x_ptr, group_size, num_groups, groups_per_block, eps, fp8_min, fp8_max]) - # libnvshmem_device.per_token_group_quant_8bit(x_ptr, x_q_ptr, x_s_ptr, group_size, num_groups, groups_per_block, eps, - # fp8_min, fp8_max) - - -def get_groups_per_block(num_groups: int) -> int: - # Removing this branch gives better performance. - # if (num_groups % 16 == 0): - # return 16 - if (num_groups % 8 == 0): - return 8 - elif (num_groups % 4 == 0): - return 4 - elif (num_groups % 2 == 0): - return 2 - else: - return 1 - - -def per_token_group_quant_fp8_tle( - x: torch.Tensor, - group_size: int, - eps: float = 1e-10, - dtype: Optional[torch.dtype] = None, - # column_major_scales: bool = False, - # scale_ue8m0: bool = False, -): - logger.debug("GEMS PER TOKEN GROUP QUANT FP8") - assert x.shape[-1] % group_size == 0, (f"the last dimension of `x` {x.shape[-1]} must be divisible " - f"by `group_size` {group_size}") - assert x.stride(-1) == 1, "`x` groups must be contiguous" - - fp8_dtype = torch.float8_e4m3fn if dtype is None else dtype - finfo = torch.finfo(fp8_dtype) - fp8_min = finfo.min - fp8_max = finfo.max - - x_q = torch.empty_like(x, device=x.device, dtype=fp8_dtype) - shape = x.shape[:-1] + (x.shape[-1] // group_size, ) - x_s = torch.empty(shape, device=x.device, dtype=torch.float32) - - # num_groups - num_groups = x.numel() // group_size - groups_per_block = get_groups_per_block(num_groups) - - # num_blocks - THREADS_PER_GROUP = 16 - num_blocks = num_groups // groups_per_block - num_warps = max(groups_per_block * THREADS_PER_GROUP // 32, 1) - - # The .cu device function uses `extern __shared__` for groups_per_block * group_size floats. - # Triton's compiler cannot infer this smem requirement from the extern_call, so we patch - # packed_metadata after warmup compilation to include the extra shared memory bytes. - smem_bytes = groups_per_block * group_size * x.element_size() # float32 = 4 bytes - - kernel = test_kernel.run( - x, - x_q, - x_s, - group_size, - num_groups, - groups_per_block, - eps, - fp8_min, - fp8_max, - grid=(num_blocks, ), - warmup=True, - num_warps=num_warps, - ) - - # Resolve async future if needed - if hasattr(kernel, "result"): - kernel = kernel.result() - - old_meta = kernel.packed_metadata - new_shared = max(old_meta[2], smem_bytes) - kernel.packed_metadata = old_meta[:2] + (new_shared, ) + old_meta[3:] - - test_kernel[(num_blocks, )]( - x, - x_q, - x_s, - group_size, - num_groups, - groups_per_block, - eps, - fp8_min, - fp8_max, - num_warps=num_warps, - ) - - return x_q, x_s - - -@triton.jit -def _per_token_group_quant_fp8( - y_ptr, - y_q_ptr, - y_s_ptr, - group_size, - y_num_columns, - y_row_stride, - eps, - fp8_min, - fp8_max, - scale_ue8m0, - BLOCK: tl.constexpr, -): - groups_per_row = y_num_columns // group_size - - g_id = tl.program_id(0) - row = g_id // groups_per_row - row_g_id = g_id % groups_per_row - - y_ptr += (row * y_row_stride) + (row_g_id * group_size) - y_q_ptr += g_id * group_size - y_s_ptr += g_id - - cols = tl.arange(0, BLOCK) - mask = cols < group_size - - y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32) - _absmax = tl.maximum(tl.max(tl.abs(y)), eps) - y_s = _absmax / fp8_max - - if scale_ue8m0: - y_s = tl.exp2(tl.ceil(tl.log2(tl.maximum(tl.abs(y_s), 1e-10)))) - - y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) - - tl.store(y_q_ptr + cols, y_q, mask=mask) - tl.store(y_s_ptr, y_s) - - -def per_token_group_quant_fp8_triton( - x: torch.Tensor, - group_size: int, - eps: float = 1e-10, - dtype: Optional[torch.dtype] = None, - column_major_scales: bool = False, - scale_ue8m0: bool = False, -): - logger.debug("GEMS PER TOKEN GROUP QUANT FP8") - # dtype: The dype of output tensor. Note that only `torch.float8_e4m3fn` - fp8_dtype = torch.float8_e4m3fn if dtype is None else dtype - assert x.shape[-1] % group_size == 0, (f"the last dimension of `x` {x.shape[-1]} must be divisible " - f"by `group_size` {group_size}") - assert x.stride(-1) == 1, "`x` groups must be contiguous" - - finfo = torch.finfo(fp8_dtype) - fp8_min = finfo.min - fp8_max = finfo.max - - x_q = torch.empty_like(x, device=x.device, dtype=fp8_dtype) - M = x.numel() // group_size - N = group_size - - if column_major_scales: - shape = (x.shape[-1] // group_size, ) + x.shape[:-1] - x_s = torch.empty(shape, device=x.device, dtype=torch.float32).permute(-1, -2) - else: - shape = x.shape[:-1] + (x.shape[-1] // group_size, ) - x_s = torch.empty(shape, device=x.device, dtype=torch.float32) - - BLOCK = triton.next_power_of_2(N) - num_warps = min(max(BLOCK // 256, 1), 8) - num_stages = 1 - _per_token_group_quant_fp8[(M, )]( - x, - x_q, - x_s, - group_size, - x.shape[1], - x.stride(0), - eps, - fp8_min=fp8_min, - fp8_max=fp8_max, - scale_ue8m0=scale_ue8m0, - BLOCK=BLOCK, - num_warps=num_warps, - num_stages=num_stages, - ) - - return x_q, x_s - - -if __name__ == "__main__": - x = torch.randn((16384, 32768), device=DEVICE, dtype=torch.float32) - group_size = 128 - - x_q_triton, x_s_triton = per_token_group_quant_fp8_tle(x, group_size) - # x_q_triton, x_s_triton = per_token_group_quant_fp8_triton(x, group_size) - # x_q_tle, x_s_tle = per_token_group_quant_fp8_tle(x, group_size) - - # q_tri = x_q_triton.to(torch.float32) - # q_tle = x_q_tle.to(torch.float32) - - # b_tri = x_q_triton.view(torch.int8).to(torch.int16) # promote to avoid int8 overflow - # b_tle = x_q_tle.view(torch.int8).to(torch.int16) - # bit_diff = (b_tri - b_tle).abs() # 0 = exact match, 1 = 1-ULP, >1 = real bug - # num_exact = bit_diff.eq(0).sum().item() - # num_1ulp = bit_diff.eq(1).sum().item() - # num_beyond = bit_diff.gt(1).sum().item() - # if num_beyond == 0: - # if num_1ulp == 0: - # print("✅ x_q Triton and TLE match (bit-exact)") - # else: - # print(f"✅ x_q Triton and TLE match (1-ULP diff={num_1ulp}, " - # f"expected from RTZ vs RTNE rounding)") - # else: - # q_tri = x_q_triton.to(torch.float32) - # q_tle = x_q_tle.to(torch.float32) - # float_diff = (q_tri - q_tle).abs() - # print(f"❌ x_q Triton and TLE differ: " - # f"exact={num_exact}, 1-ULP={num_1ulp}, >1-ULP={num_beyond}") - # beyond_idx = bit_diff.gt(1).nonzero() - # for idx in beyond_idx[:10]: - # r, c = idx[0].item(), idx[1].item() - # group_id = r * (q_tri.shape[1] // group_size) + c // group_size - # pos_in_group = c % group_size - # print(f" [{r},{c}] group={group_id} pos_in_group={pos_in_group}" - # f" triton_bits={b_tri[r,c].item():4d} tle_bits={b_tle[r,c].item():4d}" - # f" bit_diff={bit_diff[r,c].item()}" - # f" triton={q_tri[r,c].item():.1f} tle={q_tle[r,c].item():.1f}" - # f" float_diff={float_diff[r,c].item():.1f}" - # f" x={x[r,c].item():.6f} scale={x_s_triton.view(-1)[group_id].item():.8f}") - - # if torch.allclose(x_s_triton, x_s_tle, atol=0.125, rtol=0): - # print("✅ x_s Triton and TLE match") - # else: - # print("❌ x_s Triton and TLE differ") - - # # perf - # shapes = [(1024, 2048), (2048, 4096), (4096, 8192), (16384, 32768)] - # group_sizes = [64, 128] - # dtypes = [torch.float32] - # for shape in shapes: - # for dtype in dtypes: - # for group_size in group_sizes: - # x = torch.rand(shape, device=DEVICE, dtype=dtype) - # mean_ms_triton = triton.testing.do_bench(lambda: per_token_group_quant_fp8_triton(x, group_size)) - # mean_ms_tle = triton.testing.do_bench(lambda: per_token_group_quant_fp8_tle(x, group_size)) - - # print(f"\n========= Shape: {shape} Type: {dtype} Group size: {group_size} =========") - # print(f"Triton time: {mean_ms_triton:.4f} ms") - # print(f"TLE time: {mean_ms_tle:.4f} ms") diff --git a/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8-nvcc.py b/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8-nvcc.py deleted file mode 100644 index da54552ce..000000000 --- a/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8-nvcc.py +++ /dev/null @@ -1,283 +0,0 @@ -from pathlib import Path - -from typing import Optional -import logging -import torch -import triton -import triton.language as tl -from triton.experimental.tle.raw import dialect -import triton.experimental.tle.language.raw as tle_raw - -from triton.language.extra.cuda import libnvshmem_device - -torch.cuda.set_device(1) -DEVICE = triton.runtime.driver.active.get_active_torch_device() -logger = logging.getLogger(__name__) - - -@dialect(name="cuda", file=(Path(__file__).parent / "per-token-group-quant-fp8.cu").resolve(), - library={"torch": "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/torch/"}) -def edsl(*args, **kwargs): - ... - - -@triton.jit -def test_kernel( - x_ptr, - x_q_ptr, - x_s_ptr, - group_size, - num_groups, - groups_per_block, - eps, - fp8_min, - fp8_max, -): - tle_raw.call(edsl, []) - libnvshmem_device.per_token_group_quant_8bit(x_q_ptr, x_s_ptr, x_ptr, group_size, num_groups, groups_per_block, eps, - fp8_min, fp8_max) - - -def get_groups_per_block(num_groups: int) -> int: - # Removing this branch gives better performance. - # if (num_groups % 16 == 0): - # return 16 - if (num_groups % 8 == 0): - return 8 - elif (num_groups % 4 == 0): - return 4 - elif (num_groups % 2 == 0): - return 2 - else: - return 1 - - -def per_token_group_quant_fp8_tle( - x: torch.Tensor, - group_size: int, - eps: float = 1e-10, - dtype: Optional[torch.dtype] = None, - # column_major_scales: bool = False, - # scale_ue8m0: bool = False, -): - logger.debug("GEMS PER TOKEN GROUP QUANT FP8") - assert x.shape[-1] % group_size == 0, (f"the last dimension of `x` {x.shape[-1]} must be divisible " - f"by `group_size` {group_size}") - assert x.stride(-1) == 1, "`x` groups must be contiguous" - - fp8_dtype = torch.float8_e4m3fn if dtype is None else dtype - finfo = torch.finfo(fp8_dtype) - fp8_min = finfo.min - fp8_max = finfo.max - - x_q = torch.empty_like(x, device=x.device, dtype=fp8_dtype) - shape = x.shape[:-1] + (x.shape[-1] // group_size, ) - x_s = torch.empty(shape, device=x.device, dtype=torch.float32) - - # num_groups - num_groups = x.numel() // group_size - groups_per_block = get_groups_per_block(num_groups) - - # num_blocks - THREADS_PER_GROUP = 16 - num_blocks = num_groups // groups_per_block - num_warps = max(groups_per_block * THREADS_PER_GROUP // 32, 1) - - # The .cu device function uses `extern __shared__` for groups_per_block * group_size floats. - # Triton's compiler cannot infer this smem requirement from the extern_call, so we patch - # packed_metadata after warmup compilation to include the extra shared memory bytes. - smem_bytes = groups_per_block * group_size * x.element_size() # float32 = 4 bytes - - kernel = test_kernel.run( - x, - x_q, - x_s, - group_size, - num_groups, - groups_per_block, - eps, - fp8_min, - fp8_max, - grid=(num_blocks, ), - warmup=True, - num_warps=num_warps, - ) - - # Resolve async future if needed - if hasattr(kernel, "result"): - kernel = kernel.result() - - old_meta = kernel.packed_metadata - new_shared = max(old_meta[2], smem_bytes) - kernel.packed_metadata = old_meta[:2] + (new_shared, ) + old_meta[3:] - - test_kernel[(num_blocks, )]( - x, - x_q, - x_s, - group_size, - num_groups, - groups_per_block, - eps, - fp8_min, - fp8_max, - num_warps=num_warps, - ) - - return x_q, x_s - - -@triton.jit -def _per_token_group_quant_fp8( - y_ptr, - y_q_ptr, - y_s_ptr, - group_size, - y_num_columns, - y_row_stride, - eps, - fp8_min, - fp8_max, - scale_ue8m0, - BLOCK: tl.constexpr, -): - groups_per_row = y_num_columns // group_size - - g_id = tl.program_id(0) - row = g_id // groups_per_row - row_g_id = g_id % groups_per_row - - y_ptr += (row * y_row_stride) + (row_g_id * group_size) - y_q_ptr += g_id * group_size - y_s_ptr += g_id - - cols = tl.arange(0, BLOCK) - mask = cols < group_size - - y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32) - _absmax = tl.maximum(tl.max(tl.abs(y)), eps) - y_s = _absmax / fp8_max - - if scale_ue8m0: - y_s = tl.exp2(tl.ceil(tl.log2(tl.maximum(tl.abs(y_s), 1e-10)))) - - y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) - - tl.store(y_q_ptr + cols, y_q, mask=mask) - tl.store(y_s_ptr, y_s) - - -def per_token_group_quant_fp8_triton( - x: torch.Tensor, - group_size: int, - eps: float = 1e-10, - dtype: Optional[torch.dtype] = None, - column_major_scales: bool = False, - scale_ue8m0: bool = False, -): - logger.debug("GEMS PER TOKEN GROUP QUANT FP8") - # dtype: The dype of output tensor. Note that only `torch.float8_e4m3fn` - fp8_dtype = torch.float8_e4m3fn if dtype is None else dtype - assert x.shape[-1] % group_size == 0, (f"the last dimension of `x` {x.shape[-1]} must be divisible " - f"by `group_size` {group_size}") - assert x.stride(-1) == 1, "`x` groups must be contiguous" - - finfo = torch.finfo(fp8_dtype) - fp8_min = finfo.min - fp8_max = finfo.max - - x_q = torch.empty_like(x, device=x.device, dtype=fp8_dtype) - M = x.numel() // group_size - N = group_size - - if column_major_scales: - shape = (x.shape[-1] // group_size, ) + x.shape[:-1] - x_s = torch.empty(shape, device=x.device, dtype=torch.float32).permute(-1, -2) - else: - shape = x.shape[:-1] + (x.shape[-1] // group_size, ) - x_s = torch.empty(shape, device=x.device, dtype=torch.float32) - - BLOCK = triton.next_power_of_2(N) - num_warps = min(max(BLOCK // 256, 1), 8) - num_stages = 1 - _per_token_group_quant_fp8[(M, )]( - x, - x_q, - x_s, - group_size, - x.shape[1], - x.stride(0), - eps, - fp8_min=fp8_min, - fp8_max=fp8_max, - scale_ue8m0=scale_ue8m0, - BLOCK=BLOCK, - num_warps=num_warps, - num_stages=num_stages, - ) - - return x_q, x_s - - -if __name__ == "__main__": - x = torch.randn((16384, 32768), device=DEVICE, dtype=torch.float32) - group_size = 128 - - x_q_triton, x_s_triton = per_token_group_quant_fp8_tle(x, group_size) - # mean_ms_tle = triton.testing.do_bench(lambda: per_token_group_quant_fp8_tle(x, group_size)) - # print(f"TLE time: {mean_ms_tle:.4f} ms") - # x_q_triton, x_s_triton = per_token_group_quant_fp8_triton(x, group_size) - # x_q_tle, x_s_tle = per_token_group_quant_fp8_tle(x, group_size) - - # q_tri = x_q_triton.to(torch.float32) - # q_tle = x_q_tle.to(torch.float32) - - # b_tri = x_q_triton.view(torch.int8).to(torch.int16) # promote to avoid int8 overflow - # b_tle = x_q_tle.view(torch.int8).to(torch.int16) - # bit_diff = (b_tri - b_tle).abs() # 0 = exact match, 1 = 1-ULP, >1 = real bug - # num_exact = bit_diff.eq(0).sum().item() - # num_1ulp = bit_diff.eq(1).sum().item() - # num_beyond = bit_diff.gt(1).sum().item() - # if num_beyond == 0: - # if num_1ulp == 0: - # print("✅ x_q Triton and TLE match (bit-exact)") - # else: - # print(f"✅ x_q Triton and TLE match (1-ULP diff={num_1ulp}, " - # f"expected from RTZ vs RTNE rounding)") - # else: - # q_tri = x_q_triton.to(torch.float32) - # q_tle = x_q_tle.to(torch.float32) - # float_diff = (q_tri - q_tle).abs() - # print(f"❌ x_q Triton and TLE differ: " - # f"exact={num_exact}, 1-ULP={num_1ulp}, >1-ULP={num_beyond}") - # beyond_idx = bit_diff.gt(1).nonzero() - # for idx in beyond_idx[:10]: - # r, c = idx[0].item(), idx[1].item() - # group_id = r * (q_tri.shape[1] // group_size) + c // group_size - # pos_in_group = c % group_size - # print(f" [{r},{c}] group={group_id} pos_in_group={pos_in_group}" - # f" triton_bits={b_tri[r,c].item():4d} tle_bits={b_tle[r,c].item():4d}" - # f" bit_diff={bit_diff[r,c].item()}" - # f" triton={q_tri[r,c].item():.1f} tle={q_tle[r,c].item():.1f}" - # f" float_diff={float_diff[r,c].item():.1f}" - # f" x={x[r,c].item():.6f} scale={x_s_triton.view(-1)[group_id].item():.8f}") - - # if torch.allclose(x_s_triton, x_s_tle, atol=0.125, rtol=0): - # print("✅ x_s Triton and TLE match") - # else: - # print("❌ x_s Triton and TLE differ") - - # # perf - # shapes = [(1024, 2048), (2048, 4096), (4096, 8192), (16384, 32768)] - # group_sizes = [64, 128] - # dtypes = [torch.float32] - # for shape in shapes: - # for dtype in dtypes: - # for group_size in group_sizes: - # x = torch.rand(shape, device=DEVICE, dtype=dtype) - # mean_ms_triton = triton.testing.do_bench(lambda: per_token_group_quant_fp8_triton(x, group_size)) - # mean_ms_tle = triton.testing.do_bench(lambda: per_token_group_quant_fp8_tle(x, group_size)) - - # print(f"\n========= Shape: {shape} Type: {dtype} Group size: {group_size} =========") - # print(f"Triton time: {mean_ms_triton:.4f} ms") - # print(f"TLE time: {mean_ms_tle:.4f} ms") diff --git a/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8.cu b/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8.cu deleted file mode 100644 index 2add918ae..000000000 --- a/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/per-token-group-quant-fp8.cu +++ /dev/null @@ -1,138 +0,0 @@ -#include "vectorization.cuh" -#include "vectorization_utils.cuh" -#include -#include - -__device__ __forceinline__ float GroupReduceMax(float val) { - unsigned mask = threadIdx.x % 32 >= 16 ? 0xffff0000 : 0x0000ffff; - - val = fmaxf(val, __shfl_xor_sync(mask, val, 8)); - val = fmaxf(val, __shfl_xor_sync(mask, val, 4)); - val = fmaxf(val, __shfl_xor_sync(mask, val, 2)); - val = fmaxf(val, __shfl_xor_sync(mask, val, 1)); - return val; -} - -// template -__device__ __forceinline__ float -ComputeGroupScale(const float *__restrict__ group_input, - float *__restrict__ smem_group, const int group_size, - const int lane_id, const int threads_per_group, - const float eps, const float max_8bit) { - float local_absmax = eps; - - constexpr int vec_size = 16 / sizeof(float); - - // copy global -> shared & compute absmax - auto scalar_op_cache = [&] __device__(float &dst, const float &src) { - float abs_v = fabsf(static_cast(src)); - local_absmax = fmaxf(local_absmax, abs_v); - dst = src; - }; - - vllm::vectorize_with_alignment(group_input, // in - smem_group, // out (shared) - group_size, // elements per group - lane_id, // thread id - threads_per_group, // stride in group - scalar_op_cache); // scalar handler - - local_absmax = GroupReduceMax(local_absmax); - - float y_s = local_absmax / max_8bit; - // if constexpr (SCALE_UE8M0) { - // y_s = exp2f(ceilf(log2f(fmaxf(fabsf(y_s), 1e-10f)))); - // } - - return y_s; -} - -// template -__device__ __forceinline__ void -QuantizeGroup(const float *__restrict__ smem_group, - __nv_fp8_e4m3 *__restrict__ group_output, const int group_size, - const int lane_id, const int threads_per_group, const float y_s, - const float min_8bit, const float max_8bit) { - constexpr int vec_size = 16 / sizeof(float); - - // quantize shared -> global 8-bit - auto scalar_op_quant = [&] __device__(__nv_fp8_e4m3 & dst, const float &src) { - float q = fminf(fmaxf(static_cast(src) / y_s, min_8bit), max_8bit); - dst = __nv_fp8_e4m3(q); - }; - - vllm::vectorize_with_alignment( - smem_group, // in (shared) - group_output, // out (global quant tensor) - group_size, // elements - lane_id, // tid - threads_per_group, // stride - scalar_op_quant); // scalar handler -} - -// T: float; DST_DTYPE: __nv_fp8_e4m3 -// template -// __global__ void per_token_group_quant_8bit_kernel( -extern "C" __device__ void per_token_group_quant_8bit( - __attribute__((address_space(1))) void *__restrict__ output_q, - __attribute__((address_space(1))) float *__restrict__ output_s, - __attribute__((address_space(1))) const float *__restrict__ input, - const int group_size, - const int num_groups, - const int groups_per_block, - const float eps, - const float min_8bit, - const float max_8bit) { - const int threads_per_group = 16; - const int64_t local_group_id = threadIdx.x / threads_per_group; - const int lane_id = threadIdx.x % threads_per_group; - - const int64_t block_group_id = blockIdx.x * groups_per_block; - const int64_t global_group_id = block_group_id + local_group_id; - const int64_t block_group_offset = global_group_id * group_size; - - static_assert(sizeof(float) % sizeof(float) == 0); - - const float *group_input = input + block_group_offset; - __nv_fp8_e4m3 *group_output = - static_cast<__nv_fp8_e4m3 *>(output_q) + block_group_offset; - float *scale_output; - - // bool IS_COLUMN_MAJOR = false; - // if (IS_COLUMN_MAJOR) { - // const int num_elems_per_pack = - // static_cast(sizeof(float) / sizeof(float)); - // const int scale_num_rows_element = scale_num_rows * num_elems_per_pack; - // const int row_idx = global_group_id / scale_num_rows_element; - // const int col_idx_raw = global_group_id % scale_num_rows_element; - // const int col_idx = col_idx_raw / num_elems_per_pack; - // const int pack_idx = col_idx_raw % num_elems_per_pack; - // scale_output = reinterpret_cast(output_s) + - // (col_idx * scale_stride * num_elems_per_pack + - // row_idx * num_elems_per_pack + pack_idx); - // } else { - scale_output = output_s + global_group_id; - // } - - // shared memory to cache each group's data to avoid double DRAM reads. - extern __shared__ __align__(16) char smem_raw[]; - float *smem = reinterpret_cast(smem_raw); - float *smem_group = smem + local_group_id * group_size; - - const float y_s = - ComputeGroupScale(group_input, smem_group, group_size, lane_id, - threads_per_group, eps, max_8bit); - - float y_s_quant = y_s; - - if (lane_id == 0) { - *scale_output = y_s_quant; - } - - __syncthreads(); - - QuantizeGroup(smem_group, group_output, group_size, lane_id, - threads_per_group, y_s, min_8bit, max_8bit); -} - diff --git a/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/vectorization.cuh b/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/vectorization.cuh deleted file mode 100644 index 624bb9fe1..000000000 --- a/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/vectorization.cuh +++ /dev/null @@ -1,30 +0,0 @@ -#pragma once -/** - * __device__ datatypes vectorized by 4 - */ - -// Include both AMD and NVIDIA fp8 types to avoid circular import -#include -#include - -namespace vllm { - -// Vectorization containers -template -struct __align__(vec_size * sizeof(scalar_t)) vec_n_t { - scalar_t val[vec_size]; -}; - -template -struct __align__(vec_size * sizeof(quant_type_t)) q8_n_t { - static_assert(std::is_same_v || - std::is_same_v || - std::is_same_v); - quant_type_t val[vec_size]; -}; - -template using vec4_t = vec_n_t; -template using q8x4_t = q8_n_t; - -} // namespace vllm - diff --git a/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/vectorization_utils.cuh b/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/vectorization_utils.cuh deleted file mode 100644 index 3323454db..000000000 --- a/python/tutorials/tle/raw/cuda/04-per-token-group-quant-fp8/vectorization_utils.cuh +++ /dev/null @@ -1,177 +0,0 @@ -#pragma once -#include "vectorization.cuh" - -namespace vllm { - -template -struct DefaultVecOp { - ScaOp scalar_op; - - __device__ __forceinline__ void - operator()(vec_n_t &dst, - const vec_n_t &src) const { -#pragma unroll - for (int i = 0; i < VEC_SIZE; ++i) { - scalar_op(dst.val[i], src.val[i]); - } - } -}; - -template -__device__ inline void -vectorize_with_alignment(const InT *in, OutT *out, int len, int tid, int stride, - VecOp &&vec_op, // vec_n_t -> vec_n_t - ScaOp &&scalar_op) { // InT -> OutT - static_assert(VEC_SIZE > 0 && (VEC_SIZE & (VEC_SIZE - 1)) == 0, - "VEC_SIZE must be a positive power-of-two"); - constexpr int WIDTH = VEC_SIZE * sizeof(InT); // eg: 64 B - uintptr_t addr = reinterpret_cast(in); - - // fast path when the whole region is already aligned - // Note: currently the output is guaranteed to be same as the input, so we - // don't check it here, comments here just for future reference. - bool can_vec = ((addr & (WIDTH - 1)) == 0) && ((len & (VEC_SIZE - 1)) == 0); - if (can_vec) { - int num_vec = len / VEC_SIZE; - - using vin_t = vec_n_t; - using vout_t = vec_n_t; - auto *v_in = reinterpret_cast(in); - auto *v_out = reinterpret_cast(out); - - for (int i = tid; i < num_vec; i += stride) { - vout_t tmp; - // Make a local copy of the entire pack - vin_t src = v_in[i]; // <- encourages a single vector ld - vec_op(tmp, src); - v_out[i] = tmp; // <- encourages a single vector st - } - return; - } - - int misalignment_offset = addr & (WIDTH - 1); // addr % 64 - int alignment_bytes = WIDTH - misalignment_offset; // 64 - (addr % 64) - int prefix_elems = alignment_bytes & (WIDTH - 1); // handle 64 - prefix_elems /= sizeof(InT); - prefix_elems = min(prefix_elems, len); // 0 ≤ prefix < 16 - - // 1. prefill the when it is unsafe to vectorize - for (int i = tid; i < prefix_elems; i += stride) { - scalar_op(out[i], in[i]); - } - - in += prefix_elems; - out += prefix_elems; - len -= prefix_elems; - - int num_vec = len / VEC_SIZE; - using vin_t = vec_n_t; - using vout_t = vec_n_t; - auto *v_in = reinterpret_cast(in); - auto *v_out = reinterpret_cast(out); - - // 2. vectorize the main part - for (int i = tid; i < num_vec; i += stride) { - vout_t tmp; - // Make a local copy of the entire pack - vin_t src = v_in[i]; // <- encourages a single vector ld - vec_op(tmp, src); - v_out[i] = tmp; // <- encourages a single vector st - } - - // 3. handle the tail - int tail_start = num_vec * VEC_SIZE; - for (int i = tid + tail_start; i < len; i += stride) { - scalar_op(out[i], in[i]); - } -} - -template -__device__ __forceinline__ void -vectorize_with_alignment(const InT *in, OutT *out, int len, int tid, int stride, - ScaOp &&scalar_op) { - using Vec = DefaultVecOp>; - vectorize_with_alignment(in, out, len, tid, stride, Vec{scalar_op}, - std::forward(scalar_op)); -} - -template struct DefaultReadVecOp { - ScaOp scalar_op; - - __device__ __forceinline__ void - operator()(const vec_n_t &src) const { -#pragma unroll - for (int i = 0; i < VEC_SIZE; ++i) { - scalar_op(src.val[i]); - } - } -}; - -// read-only version: iterate over the input with alignment guarantees -template -__device__ inline void -vectorize_read_with_alignment(const InT *in, int len, int tid, int stride, - VecOp &&vec_op, ScaOp &&scalar_op) { - static_assert(VEC_SIZE > 0 && (VEC_SIZE & (VEC_SIZE - 1)) == 0, - "VEC_SIZE must be a positive power-of-two"); - constexpr int WIDTH = VEC_SIZE * sizeof(InT); - uintptr_t addr = reinterpret_cast(in); - - // fast path when the whole region is already aligned - bool can_vec = ((addr & (WIDTH - 1)) == 0) && ((len & (VEC_SIZE - 1)) == 0); - if (can_vec) { - int num_vec = len / VEC_SIZE; - - using vin_t = vec_n_t; - auto *v_in = reinterpret_cast(in); - - for (int i = tid; i < num_vec; i += stride) { - vin_t tmp = v_in[i]; - vec_op(tmp); - } - return; - } - - int misalignment_offset = addr & (WIDTH - 1); - int alignment_bytes = WIDTH - misalignment_offset; - int prefix_elems = alignment_bytes & (WIDTH - 1); - prefix_elems /= sizeof(InT); - prefix_elems = min(prefix_elems, len); - - // 1. handle the possibly unaligned prefix with scalar access. - for (int i = tid; i < prefix_elems; i += stride) { - scalar_op(in[i]); - } - - in += prefix_elems; - len -= prefix_elems; - - int num_vec = len / VEC_SIZE; - using vin_t = vec_n_t; - auto *v_in = reinterpret_cast(in); - - // 2. vectorized traversal of the main aligned region. - for (int i = tid; i < num_vec; i += stride) { - vec_op(v_in[i]); - } - - // 3. handle remaining tail elements. - int tail_start = num_vec * VEC_SIZE; - for (int i = tid + tail_start; i < len; i += stride) { - scalar_op(in[i]); - } -} - -// overload that requires only a scalar_op -template -__device__ __forceinline__ void -vectorize_read_with_alignment(const InT *in, int len, int tid, int stride, - ScaOp &&scalar_op) { - using Vec = DefaultReadVecOp>; - vectorize_read_with_alignment(in, len, tid, stride, Vec{scalar_op}, - std::forward(scalar_op)); -} - -} // namespace vllm - diff --git a/python/tutorials/tle/raw/cuda/05-vector-add-template/clang.py b/python/tutorials/tle/raw/cuda/05-vector-add-template/clang.py deleted file mode 100644 index e09ff0ceb..000000000 --- a/python/tutorials/tle/raw/cuda/05-vector-add-template/clang.py +++ /dev/null @@ -1,98 +0,0 @@ -from pathlib import Path - -import torch -import triton -import triton.language as tl -from triton.experimental.tle.raw import dialect -import triton.experimental.tle.language.raw as tle_raw - -DEVICE = triton.runtime.driver.active.get_active_torch_device() - - -@dialect(name="cuda", file=Path(__file__).parent / "01-vector-add.cu") -def edsl(*args, **kwargs): - ... - - -@triton.jit -def add_kernel_tle( - x_ptr, - y_ptr, - output_ptr, - n_elements, - BLOCK_SIZE: tl.constexpr, -): - tle_raw.call(edsl, [output_ptr, x_ptr, y_ptr, n_elements]) - - -def add_tle(x: torch.Tensor, y: torch.Tensor): - output = torch.empty_like(x) - assert x.device == DEVICE and y.device == DEVICE and output.device == DEVICE - n_elements = output.numel() - grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), ) - add_kernel_tle[grid](x, y, output, n_elements, BLOCK_SIZE=1024) - return output - -@triton.jit -def add_kernel_triton(x_ptr, # *Pointer* to first input vector. - y_ptr, # *Pointer* to second input vector. - output_ptr, # *Pointer* to output vector. - n_elements, # Size of the vector. - BLOCK_SIZE: tl.constexpr, # Number of elements each program should process. - ): - pid = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0. - block_start = pid * BLOCK_SIZE - offsets = block_start + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - x = tl.load(x_ptr + offsets, mask=mask) - y = tl.load(y_ptr + offsets, mask=mask) - output = x + y - tl.store(output_ptr + offsets, output, mask=mask) - -def add_triton(x: torch.Tensor, y: torch.Tensor): - output = torch.empty_like(x) - assert x.device == DEVICE and y.device == DEVICE and output.device == DEVICE - n_elements = output.numel() - grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) - add_kernel_triton[grid](x, y, output, n_elements, BLOCK_SIZE=1024) - return output - - -if __name__ == "__main__": - x = torch.randn((16384, 32768), device=DEVICE, dtype=torch.float32) - y = torch.randn((16384, 32768), device=DEVICE, dtype=torch.float32) - output_tle = add_tle(x, y) - # mean_ms_tle = triton.testing.do_bench(lambda: add_tle(x, y)) - # print(f"TLE time: {mean_ms_tle:.4f} ms") - output_triton = add_triton(x, y) - - if torch.allclose(output_triton, output_tle, atol=0.125, rtol=0): - print("✅ output Triton and TLE match") - else: - print("❌ output Triton and TLE differ") - - - # perf - dtype = torch.float32 - for i in range(27, 31, 1): - x = torch.randn(2 ** i, device=DEVICE, dtype=dtype) - y = torch.randn(2 ** i, device=DEVICE, dtype=dtype) - mean_ms_tle = triton.testing.do_bench(lambda: add_tle(x, y)) - mean_ms_triton = triton.testing.do_bench(lambda: add_triton(x, y)) - - print(f"\n======== x: 2 **{i} type: {dtype} =========") - print(f"Triton Time: {mean_ms_triton:.3f} ms") - print(f"TLE Time: {mean_ms_tle:.3f} ms") - - shapes = [(1024, 2048), (2048, 4096), (4096, 8192), (16384, 32768)] - dtypes = [torch.float32] - for shape in shapes: - for dtype in dtypes: - x = torch.rand(shape, device=DEVICE, dtype=dtype) - y = torch.rand(shape, device=DEVICE, dtype=dtype) - mean_ms_triton = triton.testing.do_bench(lambda: add_triton(x, y)) - mean_ms_tle = triton.testing.do_bench(lambda: add_tle(x, y)) - - print(f"\n========= Shape: {shape} Type: {dtype} =========") - print(f"Triton time: {mean_ms_triton:.4f} ms") - print(f"TLE time: {mean_ms_tle:.4f} ms") diff --git a/python/tutorials/tle/raw/cuda/05-vector-add-template/nvcc.py b/python/tutorials/tle/raw/cuda/05-vector-add-template/nvcc.py deleted file mode 100644 index 866c0f2b9..000000000 --- a/python/tutorials/tle/raw/cuda/05-vector-add-template/nvcc.py +++ /dev/null @@ -1,100 +0,0 @@ -from pathlib import Path - -import torch -import triton -import triton.language as tl -from triton.experimental.tle.raw import dialect -import triton.experimental.tle.language.raw as tle_raw -from triton.language.extra.cuda import libnvshmem_device - -DEVICE = triton.runtime.driver.active.get_active_torch_device() - - -@dialect(name="cuda", file=Path(__file__).parent / "01-vector-add.cu", library={}) -def edsl(*args, **kwargs): - ... - - -@triton.jit -def add_kernel_tle( - x_ptr, - y_ptr, - output_ptr, - n_elements, - BLOCK_SIZE: tl.constexpr, -): - tle_raw.call(edsl, []) - libnvshmem_device.vector_add(output_ptr, x_ptr, y_ptr, n_elements) - - -def add_tle(x: torch.Tensor, y: torch.Tensor): - output = torch.empty_like(x) - assert x.device == DEVICE and y.device == DEVICE and output.device == DEVICE - n_elements = output.numel() - grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), ) - add_kernel_tle[grid](x, y, output, n_elements, BLOCK_SIZE=1024) - return output - -@triton.jit -def add_kernel_triton(x_ptr, # *Pointer* to first input vector. - y_ptr, # *Pointer* to second input vector. - output_ptr, # *Pointer* to output vector. - n_elements, # Size of the vector. - BLOCK_SIZE: tl.constexpr, # Number of elements each program should process. - ): - pid = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0. - block_start = pid * BLOCK_SIZE - offsets = block_start + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - x = tl.load(x_ptr + offsets, mask=mask) - y = tl.load(y_ptr + offsets, mask=mask) - output = x + y - tl.store(output_ptr + offsets, output, mask=mask) - -def add_triton(x: torch.Tensor, y: torch.Tensor): - output = torch.empty_like(x) - assert x.device == DEVICE and y.device == DEVICE and output.device == DEVICE - n_elements = output.numel() - grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) - add_kernel_triton[grid](x, y, output, n_elements, BLOCK_SIZE=1024) - return output - - -if __name__ == "__main__": - x = torch.randn((16384, 32768), device=DEVICE, dtype=torch.float32) - y = torch.randn((16384, 32768), device=DEVICE, dtype=torch.float32) - output_tle = add_tle(x, y) - # mean_ms_tle = triton.testing.do_bench(lambda: add_tle(x, y)) - # print(f"TLE time: {mean_ms_tle:.4f} ms") - output_triton = add_triton(x, y) - - if torch.allclose(output_triton, output_tle, atol=0.125, rtol=0): - print("✅ output Triton and TLE match") - else: - print("❌ output Triton and TLE differ") - - - # perf - dtype = torch.float32 - for i in range(27, 31, 1): - x = torch.randn(2 ** i, device=DEVICE, dtype=dtype) - y = torch.randn(2 ** i, device=DEVICE, dtype=dtype) - mean_ms_tle = triton.testing.do_bench(lambda: add_tle(x, y)) - mean_ms_triton = triton.testing.do_bench(lambda: add_triton(x, y)) - - print(f"\n======== x: 2 **{i} type: {dtype} =========") - print(f"Triton Time: {mean_ms_triton:.3f} ms") - print(f"TLE Time: {mean_ms_tle:.3f} ms") - - shapes = [(1024, 2048), (2048, 4096), (4096, 8192), (16384, 32768)] - dtypes = [torch.float32] - for shape in shapes: - for dtype in dtypes: - x = torch.rand(shape, device=DEVICE, dtype=dtype) - y = torch.rand(shape, device=DEVICE, dtype=dtype) - mean_ms_triton = triton.testing.do_bench(lambda: add_triton(x, y)) - mean_ms_tle = triton.testing.do_bench(lambda: add_tle(x, y)) - - print(f"\n========= Shape: {shape} Type: {dtype} =========") - print(f"Triton time: {mean_ms_triton:.4f} ms") - print(f"TLE time: {mean_ms_tle:.4f} ms") diff --git a/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template-extern-call.py b/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template-extern-call.py deleted file mode 100644 index f495a8d59..000000000 --- a/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template-extern-call.py +++ /dev/null @@ -1,26 +0,0 @@ -import triton.language as tl -import triton.language.core as core - -@core.extern -def vector_add(C, A, B, N, _semantic=None): - return core.extern_call( - "", - "", - [C, A, B, N], - { - ( - core.pointer_type(core.dtype("int32")), - core.pointer_type(core.dtype("int32")), - core.pointer_type(core.dtype("int32")), - core.dtype("int32"), - ) : ("vector_add_int", ()), - ( - core.pointer_type(core.dtype("fp32")), - core.pointer_type(core.dtype("fp32")), - core.pointer_type(core.dtype("fp32")), - core.dtype("int32"), - ) : ("vector_add_float", ()), - }, - is_pure=False, - _semantic=_semantic, - ) \ No newline at end of file diff --git a/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.cu b/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.cu deleted file mode 100644 index ade9be51b..000000000 --- a/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.cu +++ /dev/null @@ -1,16 +0,0 @@ - -#define DEFINE_VECTOR_ADD(Type, FuncName) \ -extern "C" __device__ void FuncName( \ - __attribute__((address_space(1))) Type *C, \ - __attribute__((address_space(1))) const Type *A, \ - __attribute__((address_space(1))) const Type *B, \ - const int N) { \ - const int idx = blockIdx.x * blockDim.x + threadIdx.x; \ - for (int i = idx; i < N; i += blockDim.x * gridDim.x) { \ - C[i] = A[i] + B[i]; \ - } \ -} - -#define CONCAT(a, b) a ## b -#define VECTOR_ADD_FUNC_NAME(type) CONCAT(vector_add_, type) -DEFINE_VECTOR_ADD(VECTOR_ELEM_TYPE, VECTOR_ADD_FUNC_NAME(VECTOR_ELEM_TYPE)) \ No newline at end of file diff --git a/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.py b/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.py deleted file mode 100644 index 297423695..000000000 --- a/python/tutorials/tle/raw/cuda/05-vector-add-template/vector-add-template.py +++ /dev/null @@ -1,54 +0,0 @@ -from pathlib import Path - -import torch -import triton -import triton.language as tl -from triton.experimental.tle.raw import dialect -import triton.experimental.tle.language.raw as tle_raw -from triton.language.extra.cuda import libnvshmem_device - -DEVICE = triton.runtime.driver.active.get_active_torch_device() - - -@dialect(name="cuda", - compiler="nvcc", - file=Path(__file__).parent / "vector-add-template.cu", - extern=Path(__file__).parent / "vector-add-template-extern-call.py", - extern_func_name="vector_add", - macros={"VECTOR_ELEM_TYPE": "float"}) -def edsl(*args, **kwargs): - ... - - -@triton.jit -def add_kernel( - x_ptr, - y_ptr, - output_ptr, - n_elements, - BLOCK_SIZE: tl.constexpr, -): - tle_raw.call(edsl, [output_ptr, x_ptr, y_ptr, n_elements]) - - -def add(x: torch.Tensor, y: torch.Tensor): - output = torch.empty_like(x) - assert x.device == DEVICE and y.device == DEVICE and output.device == DEVICE - n_elements = output.numel() - grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), ) - add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024) - return output - - -if __name__ == "__main__": - dtype = torch.float32 - x = torch.randn(2048, dtype=dtype, device=DEVICE) - y = torch.randn(2048, dtype=dtype, device=DEVICE) - - # dtype = torch.int32 - # x = torch.randint(low=0, high=10, size=(2048, ), dtype=dtype, device=DEVICE) - # y = torch.randint(low=0, high=10, size=(2048, ), dtype=dtype, device=DEVICE) - - z = add(x, y) - print(z) - assert torch.allclose(x + y, z), (x + y, z) diff --git a/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return-extern-call.py b/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return-extern-call.py deleted file mode 100644 index b786c6fa3..000000000 --- a/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return-extern-call.py +++ /dev/null @@ -1,20 +0,0 @@ -import triton.language as tl -import triton.language.core as core - -@core.extern -def vector_add_return(C, A, B, N, _semantic=None): - return core.extern_call( - "", - "", - [C, A, B, N], - { - ( - core.pointer_type(core.dtype("fp32")), - core.pointer_type(core.dtype("fp32")), - core.pointer_type(core.dtype("fp32")), - core.dtype("int32"), - ) : ("vector_add_float_return", (core.dtype("int32"))), - }, - is_pure=False, - _semantic=_semantic, - ) \ No newline at end of file diff --git a/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.cu b/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.cu deleted file mode 100644 index f5386777b..000000000 --- a/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.cu +++ /dev/null @@ -1,12 +0,0 @@ -extern "C" __device__ int vector_add_float_return( - __attribute__((address_space(1))) float *C, - __attribute__((address_space(1))) const float *A, - __attribute__((address_space(1))) const float *B, - const int N) { - const int idx = blockIdx.x * blockDim.x + threadIdx.x; - for (int i = idx; i < N; i += blockDim.x * gridDim.x) { - C[i] = A[i] + B[i]; - } - - return idx; -} diff --git a/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.py b/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.py deleted file mode 100644 index d2d9719b0..000000000 --- a/python/tutorials/tle/raw/cuda/06-vector-add-return/vector-add-return.py +++ /dev/null @@ -1,61 +0,0 @@ -from pathlib import Path - -import torch -import triton -import triton.language as tl -from triton.experimental.tle.raw import dialect -import triton.experimental.tle.language.raw as tle_raw -from triton.language.extra.cuda import libnvshmem_device - -DEVICE = triton.runtime.driver.active.get_active_torch_device() - - -@dialect( - name="cuda", - compiler="nvcc", - file=Path(__file__).parent / "vector-add-return.cu", - extern=Path(__file__).parent / "vector-add-return-extern-call.py", - extern_func_name="vector_add_return", -) -def edsl(*args, **kwargs): - ... - - -@triton.jit -def add_kernel( - x_ptr, - y_ptr, - output_ptr, - n_elements, - BLOCK_SIZE: tl.constexpr, -): - pid = tl.program_id(axis=0) - block_start = pid * BLOCK_SIZE - offsets = block_start + tl.arange(0, BLOCK_SIZE) - - ret = tle_raw.call(edsl, [output_ptr, x_ptr, y_ptr, n_elements]) - tl.device_print("ret", ret) - base = ret.to(tl.pointer_type(tl.float32), bitcast=True) - x = tl.load(base + offsets) - y = tl.full((BLOCK_SIZE, ), 1.0, tl.float32) - output = x + y - tl.store(output_ptr + offsets, output) - - - -def add(x: torch.Tensor, y: torch.Tensor): - output = torch.empty_like(x) - assert x.device == DEVICE and y.device == DEVICE and output.device == DEVICE - n_elements = output.numel() - grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), ) - add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024) - return output - - -if __name__ == "__main__": - x = torch.randn(2048, device=DEVICE) - y = torch.randn(2048, device=DEVICE) - extra = torch.ones(2048, device=DEVICE) - z = add(x, y) - print(z) - assert torch.allclose(x + y + extra, z), (x + y + extra, z) diff --git a/python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/gemm-allreduce-host.cu b/python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/gemm-allreduce-host.cu deleted file mode 100644 index e9560d2ce..000000000 --- a/python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/gemm-allreduce-host.cu +++ /dev/null @@ -1,113 +0,0 @@ -#include -#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 nvshmemx_cumodule_init_wrapper(CUmodule module) { - return nvshmemx_cumodule_init(module); -} - -extern "C" void gemm_allreduce_before_launch -( - int *mype, int *npes, int *mype_in_node, int *n_pes_in_node, - float **d_A, float **d_B, float **d_C, - float **h_A, float **h_B, - int M, int N, int K -) { - nvshmem_init(); - *mype = nvshmem_my_pe(); - *npes = nvshmem_n_pes(); - *mype_in_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE); - *n_pes_in_node = nvshmem_team_n_pes(NVSHMEMX_TEAM_NODE); - - CUDA_CHECK(cudaSetDevice(*mype_in_node)); - *d_A = (float *)nvshmem_malloc(sizeof(float) * M * K); - *d_B = (float *)nvshmem_malloc(sizeof(float) * K * N); - *d_C = (float *)nvshmem_malloc(sizeof(float) * M * N); - - *h_A = (float *)malloc(sizeof(float) * M * K); - *h_B = (float *)malloc(sizeof(float) * K * N); - - float val = static_cast(*mype + 1); - for (int i = 0; i < M * K; i++) { - (*h_A)[i] = val; - } - for (int i = 0; i < K * N; i++) { - (*h_B)[i] = val; - } - - CUDA_CHECK(cudaMemcpy(*d_A, *h_A, sizeof(float) * M * K, cudaMemcpyHostToDevice)); - CUDA_CHECK(cudaMemcpy(*d_B, *h_B, sizeof(float) * K * N, cudaMemcpyHostToDevice)); - CUDA_CHECK(cudaMemset(*d_C, 0, sizeof(float) * M * N)); -} - -extern "C" void gemm_allreduce_after_launch -( - int mype, int npes, - float *d_A, float *d_B, float *d_C, - float *h_A, float *h_B, - int M, int N, int K -) { - cudaDeviceSynchronize(); - std::vector h_C_local(M * N); - cudaMemcpy(h_C_local.data(), d_C, M * N * sizeof(float), cudaMemcpyDeviceToHost); - float val = static_cast(mype + 1); - float local_expected = K * val * val; - std::cout << "PE " << mype << " local GEMM (first 4): "; - for (int i = 0; i < 4; ++i) printf("%.1f ", h_C_local[i]); - std::cout << "[expected: " << local_expected << "]\n"; - - nvshmem_barrier_all(); - nvshmem_float_sum_reduce(NVSHMEM_TEAM_WORLD, d_C, d_C, M * N); - nvshmem_barrier_all(); - cudaDeviceSynchronize(); - std::vector h_C_final(M * N); - cudaMemcpy(h_C_final.data(), d_C, M * N * sizeof(float), cudaMemcpyDeviceToHost); - - if (mype == 0) { - float expected = 0.0f; - for (int p = 0; p < npes; ++p) { - float v = static_cast(p + 1); - expected += K * v * v; - } - - std::cout << "\n========================================\n"; - std::cout << " After AllReduce (PE 0, first 4): "; - for (int i = 0; i < 4; ++i) printf("%.1f ", h_C_final[i]); - std::cout << "\n Expected: " << expected << "\n"; - std::cout << " (PE0: " << K*1*1 << " + PE1: " << K*2*2 << " = " << expected << ")\n"; - - bool correct = true; - for (int i = 0; i < M * N && correct; ++i) { - if (std::fabs(h_C_final[i] - expected) > 0.1f) { - std::cout << " ERROR at index " << i << ": got " << h_C_final[i] - << ", expected " << expected << "\n"; - correct = false; - } - } - if (correct) std::cout << " ✓ All values correct!\n"; - std::cout << "========================================\n"; - } - - nvshmem_free(d_A); - nvshmem_free(d_B); - nvshmem_free(d_C); - free(h_A); - free(h_B); - - nvshmem_finalize(); -} \ No newline at end of file diff --git a/python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/gemm-allreduce.py b/python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/gemm-allreduce.py deleted file mode 100644 index db6a427c0..000000000 --- a/python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/gemm-allreduce.py +++ /dev/null @@ -1,139 +0,0 @@ -import os -import subprocess -import ctypes -import torch -import triton -import triton.knobs as knobs -import triton.experimental.tle.language.raw as tle_raw - -from pathlib import Path -from triton.experimental.tle.raw import dialect -from triton.language.extra.cuda import libnvshmem_device - - -@dialect( - name="cuda", - file=(Path(__file__).parent / "simple-gemm.cu").resolve(), - library={"nvshmem": "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem"}, -) -def edsl(*args, **kwargs): - ... - - -@triton.jit -def gemm_allreduce_kernel( - C, A, B, - m, n, k -): - tle_raw.call(edsl, []) - libnvshmem_device.tiled_gemm(C, A, B, m, n, k) - - -def cuda_host_compile(cuda_host_path, cuda_host_lib): - NVCC = os.getenv("NVCC", "nvcc") - NVSHMEM_HOME = "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem" - include_path = f"-I{os.path.join(NVSHMEM_HOME, 'include')}" - lib_path = f"-L{os.path.join(NVSHMEM_HOME, 'lib')}" - - prop = torch.cuda.get_device_properties(torch.cuda.current_device()) - arch = f"-arch=sm_{prop.major}{prop.minor}" - tmp_file = Path(cuda_host_lib).with_suffix('.so.tmp') - build = [ - NVCC, "-shared", "-Xcompiler", "-fPIC", "-rdc=true", arch, include_path, lib_path, "-lnvshmem_host", - "-lnvshmem_device", "-o", tmp_file, cuda_host_path - ] - build = subprocess.run(build, capture_output=True) - assert build.returncode == 0, (f"NVCC host failed\nstderr:\n{build.stderr.decode()}") - tmp_file.rename(cuda_host_lib) - - -def gemm_allreduce(): - cu_file = (Path(__file__).parent / "gemm-allreduce-host.cu").resolve() - lib_file = Path(cu_file).with_suffix('.so') - - rank = int(os.getenv("OMPI_COMM_WORLD_RANK", "0")) - if rank == 0: - cuda_host_compile(cu_file, lib_file) - - import time - timeout = 60 - start = time.time() - while True: - if lib_file.exists(): - try: - ctypes.CDLL(str(lib_file)) - break - except OSError: - pass - if time.time() - start > timeout: - raise RuntimeError(f"Timeout waiting for {lib_file}") - time.sleep(0.1) - - - lib = ctypes.CDLL(lib_file) - mype = ctypes.c_int() - npes = ctypes.c_int() - mype_in_node = ctypes.c_int() - npes_in_node = ctypes.c_int() - # stream = ctypes.c_void_p() - - d_A = ctypes.c_void_p() - d_B = ctypes.c_void_p() - d_C = ctypes.c_void_p() - h_A = ctypes.c_void_p() - h_B = ctypes.c_void_p() - - M, N, K = 64, 64, 64 - lib.gemm_allreduce_before_launch( - ctypes.byref(mype), ctypes.byref(npes), ctypes.byref(mype_in_node), ctypes.byref(npes_in_node), - ctypes.byref(d_A), ctypes.byref(d_B), ctypes.byref(d_C), - ctypes.byref(h_A), ctypes.byref(h_B), - M, N, K - ) - - # define tile_size, dtype, len(dtype), device - dtype = torch.float32 - TILE_SIZE = 16 - grid = (((M + TILE_SIZE - 1) // TILE_SIZE), ((N + TILE_SIZE - 1) // TILE_SIZE)) - num_warps = TILE_SIZE * TILE_SIZE // 32 - device = triton.runtime.driver.active.get_active_torch_device() - - d_A_storage = torch._C._construct_storage_from_data_pointer(d_A.value, device, dtype.itemsize * M * K) - d_A_tensor = torch.empty(0, dtype=dtype, device=device).set_(d_A_storage).view(M, K) - - d_B_storage = torch._C._construct_storage_from_data_pointer(d_B.value, device, dtype.itemsize * K * N) - d_B_tensor = torch.empty(0, dtype=dtype, device=device).set_(d_B_storage).view(K, N) - - d_C_storage = torch._C._construct_storage_from_data_pointer(d_C.value, device, dtype.itemsize * M * N) - d_C_tensor = torch.empty(0, dtype=dtype, device=device).set_(d_C_storage).view(M, N) - - # def cumodule_init_hook(*args, **kwargs): - # key = kwargs["key"] - # jit_function = kwargs["fn"].jit_function - # device = kwargs["compile"]["device"] - # kernel_cache = jit_function.device_caches[device][0] - # kernel = kernel_cache.get(key, None) - # assert kernel is not None - # kernel._init_handles() - # ret = lib.nvshmemx_cumodule_init_wrapper(ctypes.c_void_p(kernel.module)) - # assert ret == 0, f"nvshmemx_cumodule_init_wrapper failed: {ret}" - # knobs.runtime.jit_post_compile_hook = cumodule_init_hook - - gemm_allreduce_kernel[grid]( - d_C_tensor, - d_A_tensor, - d_B_tensor, - M, N, K, - num_warps = num_warps - ) - - lib.gemm_allreduce_after_launch( - mype.value, npes.value, - d_A, d_B, d_C, - h_A, h_B, - M, N, K - ) - - -if __name__ == "__main__": - gemm_allreduce() diff --git a/python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/simple-gemm.cu b/python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/simple-gemm.cu deleted file mode 100644 index ad47c7918..000000000 --- a/python/tutorials/tle/raw/nvshmem/00-gemm-allreduce/simple-gemm.cu +++ /dev/null @@ -1,51 +0,0 @@ - - -#define TILE_SIZE 16 -extern "C" __device__ void tiled_gemm( - float* C, - const float* A, - const float* B, - int m, int n, int k -) { - __shared__ float As[TILE_SIZE][TILE_SIZE]; - __shared__ float Bs[TILE_SIZE][TILE_SIZE]; - - int block_row = blockIdx.y; - int block_col = blockIdx.x; - - int thread_row = threadIdx.x / TILE_SIZE; - int thread_col = threadIdx.x % TILE_SIZE; - - int C_row = block_row * TILE_SIZE + thread_row; - int C_col = block_col * TILE_SIZE + thread_col; - - float sum = 0.0f; - int num_tiles = (k + TILE_SIZE - 1) / TILE_SIZE; - for (int t = 0; t < num_tiles; ++t) { - int A_col = t * TILE_SIZE + thread_col; - if (C_row < m && A_col < k) { - As[thread_row][thread_col] = A[C_row * k + A_col]; - } else { - As[thread_row][thread_col] = 0.0f; - } - - int B_row = t * TILE_SIZE + thread_row; - if (B_row < k && C_col < n) { - Bs[thread_row][thread_col] = B[B_row * n + C_col]; - } else { - Bs[thread_row][thread_col] = 0.0f; - } - - __syncthreads(); - - for (int i = 0; i < TILE_SIZE; ++i) { - sum += As[thread_row][i] * Bs[i][thread_col]; - } - - __syncthreads(); - } - - if (C_row < m && C_col < n) { - C[C_row * n + C_col] = sum; - } -} \ No newline at end of file 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 04bdc8f70..5e4448e65 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 @@ -8,7 +8,6 @@ from pathlib import Path from triton.experimental.tle.raw import dialect -from triton.language.extra.cuda import libnvshmem_device NVSHMEM_HOME = "/data/zyuli/miniconda3/envs/flagtree_triton_v3.6.x/lib/python3.12/site-packages/nvidia/nvshmem" @@ -27,49 +26,14 @@ def edsl(*args, **kwargs): @triton.jit def simple_shift_kernel(destination_ptr, ): - # TODO: Combine with tle_raw.call, then dispatch tle_raw.call(edsl, [destination_ptr]) - # libnvshmem_device.simple_shift(destination_ptr) - - -def cuda_host_compile(cuda_host_path, cuda_host_lib): - NVCC = os.getenv("NVCC", "nvcc") - include_path = f"-I{os.path.join(NVSHMEM_HOME, 'include')}" - lib_path = f"-L{os.path.join(NVSHMEM_HOME, 'lib')}" - - prop = torch.cuda.get_device_properties(torch.cuda.current_device()) - arch = f"-arch=sm_{prop.major}{prop.minor}" - tmp_file = Path(cuda_host_lib).with_suffix('.so.tmp') - build = [ - NVCC, "-shared", "-Xcompiler", "-fPIC", "-rdc=true", arch, include_path, lib_path, "-lnvshmem_host", - "-lnvshmem_device", "-o", tmp_file, cuda_host_path - ] - build = subprocess.run(build, capture_output=True) - assert build.returncode == 0, (f"NVCC host failed\nstderr:\n{build.stderr.decode()}") - tmp_file.rename(cuda_host_lib) def simpe_shift(): cu_file = (Path(__file__).parent / "simple-shift-host.cu").resolve() lib_file = Path(cu_file).with_suffix('.so') - - rank = int(os.getenv("OMPI_COMM_WORLD_RANK", "0")) - if rank == 0: - cuda_host_compile(cu_file, lib_file) - - import time - timeout = 60 - start = time.time() - while True: - if lib_file.exists(): - try: - ctypes.CDLL(str(lib_file)) - break - except OSError: - pass - if time.time() - start > timeout: - raise RuntimeError(f"Timeout waiting for {lib_file}") - time.sleep(0.1) + + rank = int(os.getenv("PMI_RANK", "0")) lib = ctypes.CDLL(lib_file) lib.nvshmem_init_wrapper.argtypes = [] @@ -109,7 +73,6 @@ def cumodule_init_hook(*args, **kwargs): assert ret == 0, f"nvshmemx_cumodule_init_wrapper failed: {ret}" knobs.runtime.jit_post_compile_hook = cumodule_init_hook - # extern_libs = {'simple-shift': 'simple-shift-device-opti.bc'} simple_shift_kernel[(1, )](dest_tensor) stream_ptr = stream.cuda_stream diff --git a/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce-device-extern-call.py b/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce-device-extern-call.py deleted file mode 100644 index 65271655a..000000000 --- a/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce-device-extern-call.py +++ /dev/null @@ -1,30 +0,0 @@ -# Generated from ring-reduce-device.cu; do not edit manually. -import triton.language as tl -import triton.language.core as core - - -@core.extern -def ring_reduce(dst, src, nreduce, signal, chunk_size, _semantic=None): - return core.extern_call( - "", - "", - [ - dst, - src, - tl.cast(nreduce, tl.int32, _semantic=_semantic), - signal, - tl.cast(chunk_size, tl.int32, _semantic=_semantic), - ], - { - ( - core.pointer_type(core.dtype("int32")), - core.pointer_type(core.dtype("int32")), - core.dtype("int32"), - core.pointer_type(core.dtype("uint64")), - core.dtype("int32"), - ): ("ring_reduce", ()), - }, - is_pure=False, - _semantic=_semantic, - ) - diff --git a/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce.py b/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce.py index a3b01caba..041b44ca6 100644 --- a/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce.py +++ b/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce.py @@ -8,7 +8,6 @@ from pathlib import Path from triton.experimental.tle.raw import dialect -from triton.language.extra.cuda import libnvshmem_device NVSHMEM_HOME = "/data/zyuli/miniconda3/envs/flagtree_triton_v3.6.x/lib/python3.12/site-packages/nvidia/nvshmem" @dialect( @@ -34,46 +33,12 @@ def ring_reduce_kernel( ): tle_raw.call(edsl, [dst, src, nreduce, signal, chunk_size]) - -def cuda_host_compile(cuda_host_path, cuda_host_lib): - NVCC = os.getenv("NVCC", "nvcc") - include_path = f"-I{os.path.join(NVSHMEM_HOME, 'include')}" - lib_path = f"-L{os.path.join(NVSHMEM_HOME, 'lib')}" - - prop = torch.cuda.get_device_properties(torch.cuda.current_device()) - arch = f"-arch=sm_{prop.major}{prop.minor}" - tmp_file = Path(cuda_host_lib).with_suffix('.so.tmp') - build = [ - NVCC, "-shared", "-Xcompiler", "-fPIC", "-rdc=true", arch, include_path, lib_path, "-lnvshmem_host", - "-lnvshmem_device", "-o", tmp_file, cuda_host_path - ] - build = subprocess.run(build, capture_output=True) - assert build.returncode == 0, (f"NVCC host failed\nstderr:\n{build.stderr.decode()}") - tmp_file.rename(cuda_host_lib) - - def ring_reduce(): cu_file = (Path(__file__).parent / "ring-reduce-host.cu").resolve() lib_file = Path(cu_file).with_suffix('.so') # rank = int(os.getenv("OMPI_COMM_WORLD_RANK", "0")) rank = int(os.environ["PMI_RANK"]) - # if rank == 0: - # cuda_host_compile(cu_file, lib_file) - - # import time - # timeout = 60 - # start = time.time() - # while True: - # if lib_file.exists(): - # try: - # ctypes.CDLL(str(lib_file)) - # break - # except OSError: - # pass - # if time.time() - start > timeout: - # raise RuntimeError(f"Timeout waiting for {lib_file}") - # time.sleep(0.1) M, N = 64, 8 lib = ctypes.CDLL(lib_file) diff --git a/python/tutorials/tle/raw/nvshmem/03-overlap-in-one-kernel/no-overlap.py b/python/tutorials/tle/raw/nvshmem/03-overlap-in-one-kernel/no-overlap.py new file mode 100644 index 000000000..c5a0d4129 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/03-overlap-in-one-kernel/no-overlap.py @@ -0,0 +1,288 @@ +import os +import subprocess +import ctypes +import torch +import triton +import triton.knobs as knobs +import triton.experimental.tle.language.raw as tle_raw + +from pathlib import Path +from triton.experimental.tle.raw import dialect + +NVSHMEM_HOME = "/data/zyuli/miniconda3/envs/flagtree_triton_v3.6.x/lib/python3.12/site-packages/nvidia/nvshmem" + +# MAT_M = 16 +# MAT_N = 32 +# MAT_K = 16 + +MAT_M = 128 +MAT_N = 256 +MAT_K = 128 + +# ------------------------------------------------------------------ +# 通信数据规模 +# SIZE 为每个 PE 的 int32 元素总数 +# NUM_CHUNKS 为每个 block 内的 chunk 数量 +# ------------------------------------------------------------------ +SIZE = 1024 * 1024 * 32 +NUM_CHUNKS = 1 + +# ================================================================== +# dialect 注册 +# 每个 device 函数对应一个 dialect，供 triton kernel 通过 tle_raw.call 调用 +# ================================================================== + +def _make_dialect(extern_func_name): + return dialect( + name="cuda", + compiler="nvcc", + file=(Path(__file__).parent / "ring-reduce-put.cu"), + extern=(Path(__file__).parent / "ring-reduce-put-extern-call.py"), + extern_func_name=extern_func_name, + libs={"nvshmem": NVSHMEM_HOME}, + links=["nvshmem_device"] + ) + +@_make_dialect("ring_reduce_put_one_chunk_timed") +def edsl_put_one_chunk_timed(*args, **kwargs): ... + +@_make_dialect("ring_reduce_wait_timed") +def edsl_wait_timed(*args, **kwargs): ... + +@_make_dialect("local_matmul") +def edsl_local_matmul(*args, **kwargs): ... + + + +@triton.jit +def no_overlap_kernel( + dst, + src, + nreduce: triton.language.constexpr, + signal, + put_cycles, + wait_cycles, + chunk_size: triton.language.constexpr, + num_chunks: triton.language.constexpr, + expected: triton.language.constexpr, + mat_A, + mat_B, + mat_C, + mat_M: triton.language.constexpr, + mat_N: triton.language.constexpr, + mat_K: triton.language.constexpr, +): + for chunk in range(num_chunks): + tle_raw.call(edsl_put_one_chunk_timed, + [dst, src, nreduce, signal, chunk_size, chunk, + put_cycles, num_chunks]) + tle_raw.call(edsl_wait_timed, [signal, expected, wait_cycles, num_chunks, 0]) + + for chunk in range(num_chunks): + tle_raw.call(edsl_local_matmul, + [mat_A, mat_B, mat_C, mat_M, mat_N, mat_K]) + + +def setup_nvshmem(lib, size): + mype = ctypes.c_int() + npes = ctypes.c_int() + mype_in_node = ctypes.c_int() + npes_in_node = ctypes.c_int() + stream = ctypes.c_void_p() + src = ctypes.c_void_p() + dst = ctypes.c_void_p() + data_h = ctypes.c_void_p() + signal = ctypes.c_void_p() + + lib.ring_reduce_before_launch( + ctypes.byref(mype), ctypes.byref(npes), + ctypes.byref(mype_in_node), ctypes.byref(npes_in_node), + ctypes.byref(stream), + ctypes.byref(src), ctypes.byref(dst), ctypes.byref(data_h), ctypes.byref(signal), + size + ) + return mype, npes, mype_in_node, npes_in_node, stream, src, dst, data_h, signal + + +def make_tensors(src, dst, signal, num_blocks, size, device): + src_storage = torch._C._construct_storage_from_data_pointer( + src.value, device, 4 * size) + src_tensor = torch.empty(0, dtype=torch.int32, device=device).set_( + src_storage).view(size) + + dst_storage = torch._C._construct_storage_from_data_pointer( + dst.value, device, 4 * size) + dst_tensor = torch.empty(0, dtype=torch.int32, device=device).set_( + dst_storage).view(size) + + signal_storage = torch._C._construct_storage_from_data_pointer( + signal.value, device, 8 * num_blocks) + signal_tensor = torch.empty(0, dtype=torch.uint64, device=device).set_( + signal_storage).view(num_blocks) + + return src_tensor, dst_tensor, signal_tensor + + +def install_cumodule_hook(lib): + def cumodule_init_hook(*args, **kwargs): + key = kwargs["key"] + jit_function = kwargs["fn"].jit_function + dev = kwargs["compile"]["device"] + kernel_cache = jit_function.device_caches[dev][0] + kernel = kernel_cache.get(key, None) + assert kernel is not None + kernel._init_handles() + ret = lib.nvshmemx_cumodule_init_wrapper(ctypes.c_void_p(kernel.module)) + assert ret == 0, f"nvshmemx_cumodule_init_wrapper failed: {ret}" + knobs.runtime.jit_post_compile_hook = cumodule_init_hook + + +def reset_signal_on_stream(signal_tensor, curr_stream): + with torch.cuda.stream(curr_stream): + signal_tensor.zero_() + curr_stream.synchronize() + + +def reset_signal_collective(lib, signal_tensor, curr_stream): + lib.ring_reduce_barrier_all() + reset_signal_on_stream(signal_tensor, curr_stream) + lib.ring_reduce_barrier_all() + + +def reset_tensors_on_stream(curr_stream, *tensors): + with torch.cuda.stream(curr_stream): + for tensor in tensors: + tensor.zero_() + curr_stream.synchronize() + + +def format_cycle_stats(tensor, clock_rate_khz): + tensor = tensor.to(torch.int64) + values = tensor[tensor > 0].detach().cpu().to(torch.float64) + if values.numel() == 0: + return "n/a" + us = values / clock_rate_khz * 1000.0 + return ( + f"avg {us.mean().item():.3f} us, " + f"min {us.min().item():.3f} us, " + f"max {us.max().item():.3f} us " + f"({int(values.mean().item())} cycles avg)" + ) + + +def run_experiment(): + # ---- host 库编译 ---- + cu_file = (Path(__file__).parent / "ring-reduce-host.cu").resolve() + lib_file = cu_file.with_suffix(".so") + + rank = int(os.getenv("PMI_RANK", "0")) + + lib = ctypes.CDLL(str(lib_file)) + lib.ring_reduce_barrier_all.argtypes = [] + lib.ring_reduce_barrier_all.restype = None + + # ---- nvshmem 初始化 ---- + size = SIZE + (mype, npes, mype_in_node, npes_in_node, + stream, src, dst, data_h, signal) = setup_nvshmem(lib, size) + + num_blocks = npes_in_node.value + device = triton.runtime.driver.active.get_active_torch_device() + + src_tensor, dst_tensor, signal_tensor = make_tensors( + src, dst, signal, num_blocks, size, device) + + # ---- 矩阵乘法 buffer（每个 PE 分配，所有 block 共享同一组指针） + # 让不同 block 的 matmul 独立，按 block 偏移寻址 + # mat_A, mat_B, mat_C 各 num_blocks 份，每份 MAT_M*MAT_N / MAT_N*MAT_K 大小 + mat_A = torch.randn(num_blocks, MAT_M, MAT_N, dtype=torch.float32, device=device) + mat_B = torch.randn(num_blocks, MAT_N, MAT_K, dtype=torch.float32, device=device) + mat_C = torch.zeros(num_blocks, MAT_M, MAT_K, dtype=torch.float32, device=device) + no_overlap_put_cycles = torch.zeros(num_blocks, NUM_CHUNKS, dtype=torch.uint64, device=device) + no_overlap_wait_cycles = torch.zeros(num_blocks, NUM_CHUNKS, dtype=torch.uint64, device=device) + + # chunk 参数 + elems_per_block = size // num_blocks + chunk_elems = elems_per_block // NUM_CHUNKS + chunk_size = chunk_elems * 4 # bytes（int32） + + assert chunk_elems > 0, "SIZE 太小或 NUM_CHUNKS 太大, chunk_elems = 0" + assert elems_per_block % NUM_CHUNKS == 0, \ + f"elems_per_block={elems_per_block} 不能被 NUM_CHUNKS={NUM_CHUNKS} 整除" + + # ---- install cumodule hook ---- + install_cumodule_hook(lib) + + # ---- CUDA event 计时 ---- + curr_stream = torch.cuda.ExternalStream(stream.value, device=device) + num_warps = 4 + + # ---- 预热（warmup） ---- + with torch.cuda.stream(curr_stream): + no_overlap_kernel[(num_blocks,)]( + dst_tensor, src_tensor, size, signal_tensor, + no_overlap_put_cycles, no_overlap_wait_cycles, + chunk_size, NUM_CHUNKS, + NUM_CHUNKS, + mat_A, mat_B, mat_C, MAT_M, MAT_N, MAT_K, + num_warps=num_warps + ) + curr_stream.synchronize() + reset_signal_collective(lib, signal_tensor, curr_stream) + reset_tensors_on_stream(curr_stream, no_overlap_put_cycles, no_overlap_wait_cycles) + + evt_start_no = torch.cuda.Event(enable_timing=True) + evt_end_no = torch.cuda.Event(enable_timing=True) + with torch.cuda.stream(curr_stream): + evt_start_no.record(curr_stream) + no_overlap_kernel[(num_blocks,)]( + dst_tensor, src_tensor, size, signal_tensor, + no_overlap_put_cycles, no_overlap_wait_cycles, + chunk_size, NUM_CHUNKS, + NUM_CHUNKS, + mat_A, mat_B, mat_C, MAT_M, MAT_N, MAT_K, + num_warps=num_warps + ) + evt_end_no.record(curr_stream) + + curr_stream.synchronize() + time_no_overlap = evt_start_no.elapsed_time(evt_end_no) # ms + lib.ring_reduce_barrier_all() + + clock_rate_khz = torch.cuda.get_device_properties(device).clock_rate + no_overlap_put_stats = format_cycle_stats(no_overlap_put_cycles, clock_rate_khz) + no_overlap_wait_stats = format_cycle_stats(no_overlap_wait_cycles, clock_rate_khz) + + # ---- 结果打印 ---- + pe = mype_in_node.value + if pe == 0 or pe == 1: + print("\n" + "=" * 60) + print(" 通信-计算重叠 vs 无重叠 耗时对比") + print("\n" + "=" * 30 + f" PE {pe}: " + "=" * 30) + print(f" 配置:") + print(f" PEs (GPUs) : {npes_in_node.value}") + print(f" 数据规模 : {size} int32 elements ({size * 4 / 1024:.1f} KB per PE)") + print(f" chunk 数 : {NUM_CHUNKS} per block") + print(f" chunk 大小 : {chunk_size} bytes ({chunk_elems} int32 each)") + print(f" matmul 规模 : {MAT_M}×{MAT_N} × {MAT_N}×{MAT_K} (float32)") + print(f" 计时 launch : 1 per scheme") + print("-" * 60) + print(f" 方案二 no-overlap: {time_no_overlap:.3f} ms") + print("-" * 60) + print(" device clock64 统计:") + print(f" no-overlap put_nbi 调用 : {no_overlap_put_stats}") + print(f" no-overlap quiet + wait/chunk : {no_overlap_wait_stats}") + print() + + # ---- 清理 ---- + lib.ring_reduce_after_launch( + stream, + src, dst, data_h, signal, + mype.value, npes.value, + size + ) + + +if __name__ == "__main__": + run_experiment() + diff --git a/python/tutorials/tle/raw/nvshmem/03-overlap-in-one-kernel/overlap.py b/python/tutorials/tle/raw/nvshmem/03-overlap-in-one-kernel/overlap.py new file mode 100644 index 000000000..ac9999f0a --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/03-overlap-in-one-kernel/overlap.py @@ -0,0 +1,279 @@ +import os +import subprocess +import ctypes +import torch +import triton +import triton.knobs as knobs +import triton.experimental.tle.language.raw as tle_raw + +from pathlib import Path +from triton.experimental.tle.raw import dialect + +NVSHMEM_HOME = "/data/zyuli/miniconda3/envs/flagtree_triton_v3.6.x/lib/python3.12/site-packages/nvidia/nvshmem" + + +# MAT_M = 16 +# MAT_N = 32 +# MAT_K = 16 + +MAT_M = 128 +MAT_N = 256 +MAT_K = 128 + +# ------------------------------------------------------------------ +# 通信数据规模 +# SIZE 为每个 PE 的 int32 元素总数 +# NUM_CHUNKS 为每个 block 内的 chunk 数量 +# ------------------------------------------------------------------ +SIZE = 1024 * 1024 * 32 +NUM_CHUNKS = 1 + + +def _make_dialect(extern_func_name): + return dialect( + name="cuda", + compiler="nvcc", + file=(Path(__file__).parent / "ring-reduce-put.cu"), + extern=(Path(__file__).parent / "ring-reduce-put-extern-call.py"), + extern_func_name=extern_func_name, + libs={"nvshmem": NVSHMEM_HOME}, + links=["nvshmem_device"] + ) + +@_make_dialect("ring_reduce_put_one_chunk_timed") +def edsl_put_one_chunk_timed(*args, **kwargs): ... + +@_make_dialect("ring_reduce_wait_timed") +def edsl_wait_timed(*args, **kwargs): ... + +@_make_dialect("local_matmul") +def edsl_local_matmul(*args, **kwargs): ... + + +@triton.jit +def overlap_kernel( + dst, + src, + nreduce: triton.language.constexpr, + signal, + put_cycles, + wait_cycles, + chunk_size: triton.language.constexpr, + num_chunks: triton.language.constexpr, + expected, + mat_A, + mat_B, + mat_C, + mat_M: triton.language.constexpr, + mat_N: triton.language.constexpr, + mat_K: triton.language.constexpr, +): + for chunk in range(num_chunks): + tle_raw.call(edsl_put_one_chunk_timed, + [dst, src, nreduce, signal, chunk_size, chunk, + put_cycles, num_chunks]) + tle_raw.call(edsl_local_matmul, + [mat_A, mat_B, mat_C, mat_M, mat_N, mat_K]) + + tle_raw.call(edsl_wait_timed, [signal, expected, wait_cycles, num_chunks, 0]) + + +def setup_nvshmem(lib, size): + mype = ctypes.c_int() + npes = ctypes.c_int() + mype_in_node = ctypes.c_int() + npes_in_node = ctypes.c_int() + stream = ctypes.c_void_p() + src = ctypes.c_void_p() + dst = ctypes.c_void_p() + data_h = ctypes.c_void_p() + signal = ctypes.c_void_p() + + lib.ring_reduce_before_launch( + ctypes.byref(mype), ctypes.byref(npes), + ctypes.byref(mype_in_node), ctypes.byref(npes_in_node), + ctypes.byref(stream), + ctypes.byref(src), ctypes.byref(dst), ctypes.byref(data_h), ctypes.byref(signal), + size + ) + return mype, npes, mype_in_node, npes_in_node, stream, src, dst, data_h, signal + + +def make_tensors(src, dst, signal, num_blocks, size, device): + src_storage = torch._C._construct_storage_from_data_pointer( + src.value, device, 4 * size) + src_tensor = torch.empty(0, dtype=torch.int32, device=device).set_( + src_storage).view(size) + + dst_storage = torch._C._construct_storage_from_data_pointer( + dst.value, device, 4 * size) + dst_tensor = torch.empty(0, dtype=torch.int32, device=device).set_( + dst_storage).view(size) + + signal_storage = torch._C._construct_storage_from_data_pointer( + signal.value, device, 8 * num_blocks) + signal_tensor = torch.empty(0, dtype=torch.uint64, device=device).set_( + signal_storage).view(num_blocks) + + return src_tensor, dst_tensor, signal_tensor + + +def install_cumodule_hook(lib): + def cumodule_init_hook(*args, **kwargs): + key = kwargs["key"] + jit_function = kwargs["fn"].jit_function + dev = kwargs["compile"]["device"] + kernel_cache = jit_function.device_caches[dev][0] + kernel = kernel_cache.get(key, None) + assert kernel is not None + kernel._init_handles() + ret = lib.nvshmemx_cumodule_init_wrapper(ctypes.c_void_p(kernel.module)) + assert ret == 0, f"nvshmemx_cumodule_init_wrapper failed: {ret}" + knobs.runtime.jit_post_compile_hook = cumodule_init_hook + + +def reset_signal_on_stream(signal_tensor, curr_stream): + with torch.cuda.stream(curr_stream): + signal_tensor.zero_() + curr_stream.synchronize() + + +def reset_signal_collective(lib, signal_tensor, curr_stream): + lib.ring_reduce_barrier_all() + reset_signal_on_stream(signal_tensor, curr_stream) + lib.ring_reduce_barrier_all() + + +def reset_tensors_on_stream(curr_stream, *tensors): + with torch.cuda.stream(curr_stream): + for tensor in tensors: + tensor.zero_() + curr_stream.synchronize() + + +def format_cycle_stats(tensor, clock_rate_khz): + tensor = tensor.to(torch.int64) + values = tensor[tensor > 0].detach().cpu().to(torch.float64) + if values.numel() == 0: + return "n/a" + us = values / clock_rate_khz * 1000.0 + return ( + f"avg {us.mean().item():.3f} us, " + f"min {us.min().item():.3f} us, " + f"max {us.max().item():.3f} us " + f"({int(values.mean().item())} cycles avg)" + ) + + +def run_experiment(): + cu_file = (Path(__file__).parent / "ring-reduce-host.cu").resolve() + lib_file = cu_file.with_suffix(".so") + + rank = int(os.getenv("PMI_RANK", "0")) + lib = ctypes.CDLL(str(lib_file)) + lib.ring_reduce_barrier_all.argtypes = [] + lib.ring_reduce_barrier_all.restype = None + + # ---- nvshmem 初始化 ---- + size = SIZE + (mype, npes, mype_in_node, npes_in_node, + stream, src, dst, data_h, signal) = setup_nvshmem(lib, size) + + num_blocks = npes_in_node.value + device = triton.runtime.driver.active.get_active_torch_device() + + src_tensor, dst_tensor, signal_tensor = make_tensors( + src, dst, signal, num_blocks, size, device) + + # ---- 矩阵乘法 buffer（每个 PE 分配，所有 block 共享同一组指针） + # 为了让不同 block 的 matmul 真正，按 block 偏移寻址 + # mat_A, mat_B, mat_C 各 num_blocks 份，每份 MAT_M*MAT_N / MAT_N*MAT_K 大小 + mat_A = torch.randn(num_blocks, MAT_M, MAT_N, dtype=torch.float32, device=device) + mat_B = torch.randn(num_blocks, MAT_N, MAT_K, dtype=torch.float32, device=device) + mat_C = torch.zeros(num_blocks, MAT_M, MAT_K, dtype=torch.float32, device=device) + overlap_put_cycles = torch.zeros(num_blocks, NUM_CHUNKS, dtype=torch.uint64, device=device) + overlap_wait_cycles = torch.zeros(num_blocks, NUM_CHUNKS, dtype=torch.uint64, device=device) + + # chunk 参数 + elems_per_block = size // num_blocks + chunk_elems = elems_per_block // NUM_CHUNKS + chunk_size = chunk_elems * 4 # bytes（int32） + + assert chunk_elems > 0, "SIZE 太小或 NUM_CHUNKS 太大，chunk_elems = 0" + assert elems_per_block % NUM_CHUNKS == 0, \ + f"elems_per_block={elems_per_block} 不能被 NUM_CHUNKS={NUM_CHUNKS} 整除" + + # ---- install cumodule hook ---- + install_cumodule_hook(lib) + + # ---- CUDA event 计时 ---- + curr_stream = torch.cuda.ExternalStream(stream.value, device=device) + num_warps = 4 + + # ---- 预热（warmup） ---- + with torch.cuda.stream(curr_stream): + overlap_kernel[(num_blocks,)]( + dst_tensor, src_tensor, size, signal_tensor, + overlap_put_cycles, overlap_wait_cycles, + chunk_size, NUM_CHUNKS, + NUM_CHUNKS, + mat_A, mat_B, mat_C, MAT_M, MAT_N, MAT_K, + num_warps=num_warps + ) + curr_stream.synchronize() + reset_signal_collective(lib, signal_tensor, curr_stream) + reset_tensors_on_stream(curr_stream, overlap_put_cycles, overlap_wait_cycles) + + evt_start_ov = torch.cuda.Event(enable_timing=True) + evt_end_ov = torch.cuda.Event(enable_timing=True) + with torch.cuda.stream(curr_stream): + evt_start_ov.record(curr_stream) + overlap_kernel[(num_blocks,)]( + dst_tensor, src_tensor, size, signal_tensor, + overlap_put_cycles, overlap_wait_cycles, + chunk_size, NUM_CHUNKS, + NUM_CHUNKS, + mat_A, mat_B, mat_C, MAT_M, MAT_N, MAT_K, + num_warps=num_warps + ) + evt_end_ov.record(curr_stream) + + curr_stream.synchronize() + time_overlap = evt_start_ov.elapsed_time(evt_end_ov) # ms + lib.ring_reduce_barrier_all() + + clock_rate_khz = torch.cuda.get_device_properties(device).clock_rate + overlap_put_stats = format_cycle_stats(overlap_put_cycles, clock_rate_khz) + overlap_wait_stats = format_cycle_stats(overlap_wait_cycles, clock_rate_khz) + + pe = mype_in_node.value + if pe == 0 or pe == 1: + print("\n" + "=" * 30 + f" PE {pe}: " + "=" * 30) + print(" 通信-计算重叠 vs 无重叠 耗时对比") + print("=" * 60) + print(f" 配置:") + print(f" PEs (GPUs) : {npes_in_node.value}") + print(f" 数据规模 : {size} int32 elements ({size * 4 / 1024:.1f} KB per PE)") + print(f" chunk 数 : {NUM_CHUNKS} per block") + print(f" chunk 大小 : {chunk_size} bytes ({chunk_elems} int32 each)") + print(f" matmul 规模 : {MAT_M}×{MAT_N} × {MAT_N}×{MAT_K} (float32)") + print(f" 计时 launch : 1 per scheme") + print("-" * 60) + print(f" 方案一 overlap : {time_overlap:.3f} ms") + print("-" * 60) + print(" device clock64 统计:") + print(f" overlap put_nbi 调用 : {overlap_put_stats}") + print(f" overlap quiet + wait 收口 : {overlap_wait_stats}") + print() + + lib.ring_reduce_after_launch( + stream, + src, dst, data_h, signal, + mype.value, npes.value, + size + ) + + +if __name__ == "__main__": + run_experiment() + diff --git a/python/tutorials/tle/raw/nvshmem/03-overlap-in-one-kernel/ring-reduce-host.cu b/python/tutorials/tle/raw/nvshmem/03-overlap-in-one-kernel/ring-reduce-host.cu new file mode 100644 index 000000000..ad9a39d6d --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/03-overlap-in-one-kernel/ring-reduce-host.cu @@ -0,0 +1,79 @@ +#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 nvshmemx_cumodule_init_wrapper(CUmodule module) { + return nvshmemx_cumodule_init(module); +} + +extern "C" void ring_reduce_barrier_all() { + nvshmem_barrier_all(); +} + +extern "C" void ring_reduce_before_launch +( + int *mype, int *npes, int *mype_in_node, int *n_pes_in_node, + cudaStream_t *stream, + int **src, int **dst, int **data_h, uint64_t **signal, + int size +) { + + nvshmem_init(); + *mype = nvshmem_my_pe(); + *npes = nvshmem_n_pes(); + *mype_in_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE); + *n_pes_in_node = nvshmem_team_n_pes(NVSHMEMX_TEAM_NODE); + + CUDA_CHECK(cudaSetDevice(*mype_in_node)); + CUDA_CHECK(cudaStreamCreate(stream)); + + *src = (int *)nvshmem_malloc(sizeof(int) * size); + *dst = (int *)nvshmem_malloc(sizeof(int) * size); + *data_h = (int *)malloc(sizeof(int) * size); + *signal = (uint64_t *)nvshmem_calloc(*n_pes_in_node, sizeof(uint64_t)); + + for (size_t i = 0; i < size; i++) (*data_h)[i] = i; + CUDA_CHECK(cudaMemcpyAsync(*src, *data_h, sizeof(int) * size, cudaMemcpyHostToDevice, *stream)); + nvshmemx_barrier_all_on_stream(*stream); +} + +extern "C" void ring_reduce_after_launch +( + cudaStream_t stream, + void *src, void *dst, void *data_h, void *signal, + int mype, int npes, + int size +) { + CUDA_CHECK(cudaStreamSynchronize(stream)); + nvshmem_barrier_all(); + CUDA_CHECK(cudaMemcpy(data_h, dst, sizeof(int) * size, cudaMemcpyDeviceToHost)); + // int *data = (int *)data_h; + // for (size_t i = 0; i < size; i++) { + // if (data[i] != (int)i * npes) { + // // if (mype == 0) { + // // printf("PE %d, data[%zu] = %d expected data[%zu] = %d\n", mype, i, data[i], i, (int)i * npes); + // printf("PE %d error, data[%zu] = %d expected data[%zu] = %d\n", mype, i, data[i], i, (int)i * npes); + // } + // } + + nvshmem_free(dst); + nvshmem_free(src); + nvshmem_free(signal); + free(data_h); + + nvshmem_finalize(); +} + diff --git a/python/tutorials/tle/raw/nvshmem/03-overlap-in-one-kernel/ring-reduce-put.cu b/python/tutorials/tle/raw/nvshmem/03-overlap-in-one-kernel/ring-reduce-put.cu new file mode 100644 index 000000000..6c1608c2d --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/03-overlap-in-one-kernel/ring-reduce-put.cu @@ -0,0 +1,165 @@ +#include +#include +#include +#include +#include + +// ============================================================ +// 通信函数（分 chunk 逐步发起）：用于方案一 overlap 版本 +// 每次只发起第 chunk_idx 个 chunk 的 NBI put+signal，立即返回 +// ============================================================ +extern "C" __device__ void ring_reduce_put_one_chunk( + int *dst, const int *src, int nreduce, + uint64_t *signal, int chunk_size, int chunk_idx) +{ + int mype = nvshmem_my_pe(); + int npes = nvshmem_n_pes(); + int peer = (mype + 1) % npes; + + int block_idx = blockIdx.x; + int num_blocks = gridDim.x; + int elems_per_block = nreduce / num_blocks; + + if (elems_per_block * (block_idx + 1) > nreduce) return; + + src = src + block_idx * elems_per_block; + dst = dst + block_idx * elems_per_block; + signal = signal + block_idx; + + int chunk_elems = chunk_size / sizeof(int); + + if (threadIdx.x == 0) { + nvshmem_int_put_signal_nbi( + dst + chunk_idx * chunk_elems, + src + chunk_idx * chunk_elems, + chunk_elems, + signal, + 1, + NVSHMEM_SIGNAL_ADD, + peer + ); + } +} + +// ============================================================ +// 等待函数：阻塞直到所有 chunk 均已到达（signal >= num_chunks） +// ============================================================ +extern "C" __device__ void ring_reduce_put_one_chunk_timed( + int *dst, const int *src, int nreduce, + uint64_t *signal, int chunk_size, int chunk_idx, + uint64_t *put_cycles, int num_chunks) +{ + int mype = nvshmem_my_pe(); + int npes = nvshmem_n_pes(); + int peer = (mype + 1) % npes; + + int block_idx = blockIdx.x; + int num_blocks = gridDim.x; + int elems_per_block = nreduce / num_blocks; + + if (elems_per_block * (block_idx + 1) > nreduce) return; + + src = src + block_idx * elems_per_block; + dst = dst + block_idx * elems_per_block; + signal = signal + block_idx; + + int chunk_elems = chunk_size / sizeof(int); + + __syncthreads(); + if (threadIdx.x == 0) { + uint64_t t0 = clock64(); + nvshmem_int_put_signal_nbi( + dst + chunk_idx * chunk_elems, + src + chunk_idx * chunk_elems, + chunk_elems, + signal, + 1, + NVSHMEM_SIGNAL_ADD, + peer + ); + uint64_t t1 = clock64(); + put_cycles[block_idx * num_chunks + chunk_idx] = t1 - t0; + } + __syncthreads(); +} + +extern "C" __device__ void ring_reduce_wait(uint64_t *signal, uint64_t expected) { + int block_idx = blockIdx.x; + signal = signal + block_idx; + + if (threadIdx.x == 0) { + nvshmem_quiet(); + nvshmem_signal_wait_until(signal, NVSHMEM_CMP_GE, expected); + // *signal = 0; + } + __syncthreads(); +} + +// ============================================================ +// 等待函数（等待所有 chunk 到达并 reset signal）：overlap 版本收口 +// ============================================================ +extern "C" __device__ void ring_reduce_wait_timed( + uint64_t *signal, uint64_t expected, + uint64_t *wait_cycles, int num_chunks, int wait_idx) +{ + int block_idx = blockIdx.x; + signal = signal + block_idx; + + if (threadIdx.x == 0) { + uint64_t t0 = clock64(); + nvshmem_quiet(); + nvshmem_signal_wait_until(signal, NVSHMEM_CMP_GE, expected); + uint64_t t1 = clock64(); + wait_cycles[block_idx * num_chunks + wait_idx] = t1 - t0; + } + __syncthreads(); +} + +// ============================================================ +// 本地矩阵乘法计算：C[M×K] = A[M×N] × B[N×K]（模拟重计算） +// +// 参数说明： +// A, B, C : float* 指针，指向整个 [num_blocks, M, N/K] 数组的首地址 +// 本函数自动按 blockIdx.x 偏移，取属于本 block 的 slice +// M/N/K : 矩阵维度 +// +// 目的：模拟真实 transformer 中的计算密集型操作， +// 使得通信-计算重叠能显著减少端到端延迟 +// ============================================================ +extern "C" __device__ void local_matmul( + const float *A, const float *B, float *C, + int M, int N, int K) +{ + int block_idx = blockIdx.x; + + // 每个 block 只操作自己的 slice：偏移到 block_idx 对应的子矩阵 + A = A + block_idx * M * N; + B = B + block_idx * N * K; + C = C + block_idx * M * K; + + int tid = threadIdx.x; + int nthreads = blockDim.x; + + // 线程 tid 负责 C 的第 tid 行（stride 为 nthreads） + // 内层循环按 col 方向遍历： + // - A[row*N + k] 在 k 固定时按 row 变化，线程间 coalesced + // - B[k*K + col] 在 col 方向连续，内层按 col 遍历为 coalesced + // - C[row*K + col] 内层按 col 连续写，coalesced + // 先将本线程负责的 C 行清零，确保多次调用时结果正确 + for (int row = tid; row < M; row += nthreads) { + for (int col = 0; col < K; col++) { + C[row * K + col] = 0.0f; + } + } + // C[row][col] = sum_k A[row][k] * B[k][col] + // 外层 k 循环让 a_val 可被寄存器缓存；内层 col 遍历 B/C 的连续内存（coalesced） + for (int row = tid; row < M; row += nthreads) { + for (int k = 0; k < N; k++) { + float a_val = A[row * N + k]; + for (int col = 0; col < K; col++) { + C[row * K + col] += a_val * B[k * K + col]; + } + } + } + __syncthreads(); +} \ No newline at end of file diff --git a/python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast-device.cu b/python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast-device.cu deleted file mode 100644 index 7ca620d09..000000000 --- a/python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast-device.cu +++ /dev/null @@ -1,23 +0,0 @@ -#include -#include -#include -#include -#include - -extern "C" __device__ void ring_bcast(int *data, int nelem, int root, uint64_t *psync) { - int mype = nvshmem_my_pe(); - int npes = nvshmem_n_pes(); - int peer = (mype + 1) % npes; - - if (mype == root) *psync = 1; - - nvshmem_signal_wait_until(psync, NVSHMEM_CMP_NE, 0); - - if (mype == npes - 1) return; - - nvshmem_int_put(data, data, nelem, peer); - nvshmem_fence(); - nvshmemx_signal_op(psync, 1, NVSHMEM_SIGNAL_SET, peer); - - *psync = 0; -} \ No newline at end of file diff --git a/python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast-host.cu b/python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast-host.cu deleted file mode 100644 index 9578bb945..000000000 --- a/python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast-host.cu +++ /dev/null @@ -1,59 +0,0 @@ -#include -#include -#include -#include -#include - -extern "C" int nvshmemx_cumodule_init_wrapper(CUmodule module) { - return nvshmemx_cumodule_init(module); -} - -extern "C" void ring_bcast_before_launch -( - int *mype, int *npes, int *mype_in_node, int *n_pes_in_node, - cudaStream_t *stream, - int **data, int **data_h, uint64_t **psync, - int data_len -) { - - nvshmem_init(); - *mype = nvshmem_my_pe(); - *npes = nvshmem_n_pes(); - *mype_in_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE); - *n_pes_in_node = nvshmem_team_n_pes(NVSHMEMX_TEAM_NODE); - - cudaSetDevice(*mype_in_node); - cudaStreamCreate(stream); - - *data = (int *)nvshmem_malloc(sizeof(int) * data_len); - *data_h = (int *)malloc(sizeof(int) * data_len); - *psync = (uint64_t *)nvshmem_calloc(1, sizeof(uint64_t)); - - for (size_t i = 0; i < data_len; i++) (*data_h)[i] = *mype + i; - cudaMemcpyAsync(*data, *data_h, sizeof(int) * data_len, cudaMemcpyHostToDevice, *stream); - nvshmemx_barrier_all_on_stream(*stream); -} - -extern "C" void ring_bcast_after_launch -( - cudaStream_t stream, - void *data, void *data_h, void *psync, - int mype, int npes, - int data_len -) { - nvshmemx_barrier_all_on_stream(stream); - cudaMemcpyAsync(data_h, data, sizeof(int) * data_len, cudaMemcpyDeviceToHost, stream); - cudaStreamSynchronize(stream); - - int *h_ptr = (int *)data_h; - for (size_t i = 0; i < data_len; i++) { - if (h_ptr[i] != (int) i) - printf("PE %d error, data[%zu] = %d expected data[%zu] = %d\n", mype, i, h_ptr[i], i, (int) i); - } - - nvshmem_free(data); - nvshmem_free(psync); - free(data_h); - - nvshmem_finalize(); -} \ No newline at end of file diff --git a/python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast.py b/python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast.py deleted file mode 100644 index 80e325ea9..000000000 --- a/python/tutorials/tle/raw/nvshmem/03-ring-bcast/ring-bcast.py +++ /dev/null @@ -1,137 +0,0 @@ -import os -import subprocess -import ctypes -import torch -import triton -import triton.knobs as knobs -import triton.experimental.tle.language.raw as tle_raw - -from pathlib import Path -from triton.experimental.tle.raw import dialect -from triton.language.extra.cuda import libnvshmem_device - - -@dialect( - name="cuda", - file=(Path(__file__).parent / "ring-bcast-device.cu").resolve(), - library={"nvshmem": "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem"}, -) -def edsl(*args, **kwargs): - ... - - -@triton.jit -def ring_bcast_kernel( - data, - nelem, - root, - psync, -): - tle_raw.call(edsl, []) - libnvshmem_device.ring_bcast(data, nelem, root, psync) - - -def cuda_host_compile(cuda_host_path, cuda_host_lib): - NVCC = os.getenv("NVCC", "nvcc") - NVSHMEM_HOME = "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem" - include_path = f"-I{os.path.join(NVSHMEM_HOME, 'include')}" - lib_path = f"-L{os.path.join(NVSHMEM_HOME, 'lib')}" - - prop = torch.cuda.get_device_properties(torch.cuda.current_device()) - arch = f"-arch=sm_{prop.major}{prop.minor}" - tmp_file = Path(cuda_host_lib).with_suffix('.so.tmp') - build = [ - NVCC, "-shared", "-Xcompiler", "-fPIC", "-rdc=true", arch, include_path, lib_path, "-lnvshmem_host", - "-lnvshmem_device", "-o", tmp_file, cuda_host_path - ] - build = subprocess.run(build, capture_output=True) - assert build.returncode == 0, (f"NVCC host failed\nstderr:\n{build.stderr.decode()}") - tmp_file.rename(cuda_host_lib) - - -def ring_bcast(): - cu_file = (Path(__file__).parent / "ring-bcast-host.cu").resolve() - lib_file = Path(cu_file).with_suffix('.so') - - rank = int(os.getenv("OMPI_COMM_WORLD_RANK", "0")) - if rank == 0: - cuda_host_compile(cu_file, lib_file) - - import time - timeout = 60 - start = time.time() - while True: - if lib_file.exists(): - try: - ctypes.CDLL(str(lib_file)) - break - except OSError: - pass - if time.time() - start > timeout: - raise RuntimeError(f"Timeout waiting for {lib_file}") - time.sleep(0.1) - - lib = ctypes.CDLL(lib_file) - mype = ctypes.c_int() - npes = ctypes.c_int() - mype_in_node = ctypes.c_int() - npes_in_node = ctypes.c_int() - - stream = ctypes.c_void_p() - - data = ctypes.c_void_p() - data_h = ctypes.c_void_p() - psync = ctypes.c_void_p() - - data_len = 32 - root = 0 - lib.ring_bcast_before_launch( - ctypes.byref(mype), ctypes.byref(npes), ctypes.byref(mype_in_node), ctypes.byref(npes_in_node), - ctypes.byref(stream), - ctypes.byref(data), ctypes.byref(data_h), ctypes.byref(psync), - data_len - ) - - # print("PE:", mype_in_node.value) - dtype = torch.int32 - num_blocks = 1 - num_warps = 1 - device = triton.runtime.driver.active.get_active_torch_device() - - data_storage = torch._C._construct_storage_from_data_pointer(data.value, device, dtype.itemsize * data_len) - data_tensor = torch.empty(0, dtype=dtype, device=device).set_(data_storage).view(data_len, ) - - psync_storage = torch._C._construct_storage_from_data_pointer(psync.value, device, 8 * 1) - psync_tensor = torch.empty(0, dtype=torch.uint64, device=device).set_(psync_storage).view(1, ) - - def cumodule_init_hook(*args, **kwargs): - key = kwargs["key"] - jit_function = kwargs["fn"].jit_function - device = kwargs["compile"]["device"] - kernel_cache = jit_function.device_caches[device][0] - kernel = kernel_cache.get(key, None) - assert kernel is not None - kernel._init_handles() - ret = lib.nvshmemx_cumodule_init_wrapper(ctypes.c_void_p(kernel.module)) - assert ret == 0, f"nvshmemx_cumodule_init_wrapper failed: {ret}" - knobs.runtime.jit_post_compile_hook = cumodule_init_hook - - curr_stream = torch.cuda.ExternalStream(stream.value, device=device) - with torch.cuda.stream(curr_stream): - ring_bcast_kernel[(num_blocks, )]( - data_tensor, - data_len, - root, - psync_tensor - ) - - lib.ring_bcast_after_launch( - stream, - data, data_h, psync, - mype.value, npes.value, - data_len - ) - - -if __name__ == "__main__": - ring_bcast() diff --git a/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-device-accumulate.cu b/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-device-accumulate.cu deleted file mode 100644 index 316497b6e..000000000 --- a/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-device-accumulate.cu +++ /dev/null @@ -1,8 +0,0 @@ -#include - -extern "C" __device__ void accumulate(int *input, int *partial_sum) { - int index = threadIdx.x; - if (0 == index) *partial_sum = 0; - __syncthreads(); - atomicAdd(partial_sum, input[index]); -} diff --git a/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-device-correct-accumulate.cu b/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-device-correct-accumulate.cu deleted file mode 100644 index 00b90e339..000000000 --- a/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-device-correct-accumulate.cu +++ /dev/null @@ -1,13 +0,0 @@ -#include - -#define THRESHOLD 42 -#define CORRECTION 7 -extern "C" __device__ void correct_accumulate(int *input, int *partial_sum, int *full_sum) { - int index = threadIdx.x; - if (*full_sum > THRESHOLD) { - input[index] = input[index] - CORRECTION; - } - if (0 == index) *partial_sum = 0; - __syncthreads(); - atomicAdd(partial_sum, input[index]); -} \ No newline at end of file diff --git a/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-host.cu b/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-host.cu deleted file mode 100644 index 1fa8f34fb..000000000 --- a/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream-host.cu +++ /dev/null @@ -1,69 +0,0 @@ -#include -#include "nvshmem.h" -#include "nvshmemx.h" - -#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 nvshmemx_cumodule_init_wrapper(CUmodule module) { - return nvshmemx_cumodule_init(module); -} - - -extern "C" void on_stream_before_launch_accumulate -( - int *mype, int *npes, int *mype_in_node, int *n_pes_in_node, - cudaStream_t *stream, - int **input, int **partial_sum, int **full_sum, - int input_nelems -) { - - nvshmem_init(); - *mype = nvshmem_my_pe(); - *npes = nvshmem_n_pes(); - *mype_in_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE); - *n_pes_in_node = nvshmem_team_n_pes(NVSHMEMX_TEAM_NODE); - - CUDA_CHECK(cudaSetDevice(*mype_in_node)); - CUDA_CHECK(cudaStreamCreate(stream)); - - *input = (int *)nvshmem_malloc(sizeof(int) * input_nelems); - *partial_sum = (int *)nvshmem_malloc(sizeof(int)); - *full_sum = (int *)nvshmem_malloc(sizeof(int)); -} - -extern "C" void on_stream_before_launch_correct_accumulate -( - cudaStream_t *stream, - int **partial_sum, int **full_sum, - int to_all_nelems -) { - nvshmemx_int_sum_reduce_on_stream(NVSHMEM_TEAM_WORLD, *full_sum, *partial_sum, to_all_nelems, *stream); -} - - -extern "C" void on_stream_after_launch_correct_accumulate -( - int mype, int npes, - cudaStream_t stream, - void *input, void *partial_sum, void *full_sum -) { - CUDA_CHECK(cudaStreamSynchronize(stream)); - printf("[%d of %d] run complete \n", mype, npes); - CUDA_CHECK(cudaStreamDestroy(stream)); - - nvshmem_free(input); - nvshmem_free(partial_sum); - nvshmem_free(full_sum); - - nvshmem_finalize(); -} \ No newline at end of file diff --git a/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream.py b/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream.py deleted file mode 100644 index 03899665c..000000000 --- a/python/tutorials/tle/raw/nvshmem/04-on-stream/on-stream.py +++ /dev/null @@ -1,166 +0,0 @@ -import os -import subprocess -import ctypes -import torch -import triton -import triton.knobs as knobs -import triton.experimental.tle.language.raw as tle_raw - -from pathlib import Path -from triton.experimental.tle.raw import dialect -from triton.language.extra.cuda import libnvshmem_device - - -@dialect( - name="cuda", - file=(Path(__file__).parent / "on-stream-device-accumulate.cu").resolve(), - libs={"nvshmem": "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem"}, -) -def edsl_accumulate(*args, **kwargs): - ... - - -@dialect( - name="cuda", - file=(Path(__file__).parent / "on-stream-device-correct-accumulate.cu").resolve(), - libs={"nvshmem": "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem"}, -) -def edsl_correct_accumulate(*args, **kwargs): - ... - - -@triton.jit -def accumulate_kernel( - input, - partial_sum -): - tle_raw.call(edsl_accumulate, []) - libnvshmem_device.accumulate(input, partial_sum) - - -@triton.jit -def correct_accumulate_kernel( - input, - partial_sum, - full_sum -): - tle_raw.call(edsl_correct_accumulate, []) - libnvshmem_device.correct_accumulate(input, partial_sum, full_sum) - - -def cuda_host_compile(cuda_host_path, cuda_host_lib): - NVCC = os.getenv("NVCC", "nvcc") - NVSHMEM_HOME = "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem" - include_path = f"-I{os.path.join(NVSHMEM_HOME, 'include')}" - lib_path = f"-L{os.path.join(NVSHMEM_HOME, 'lib')}" - - prop = torch.cuda.get_device_properties(torch.cuda.current_device()) - arch = f"-arch=sm_{prop.major}{prop.minor}" - tmp_file = Path(cuda_host_lib).with_suffix('.so.tmp') - build = [ - NVCC, "-shared", "-Xcompiler", "-fPIC", "-rdc=true", arch, include_path, lib_path, "-lnvshmem_host", - "-lnvshmem_device", "-o", tmp_file, cuda_host_path - ] - build = subprocess.run(build, capture_output=True) - assert build.returncode == 0, (f"NVCC host failed\nstderr:\n{build.stderr.decode()}") - tmp_file.rename(cuda_host_lib) - - -def on_stream(): - cu_file = (Path(__file__).parent / "on-stream-host.cu").resolve() - lib_file = Path(cu_file).with_suffix('.so') - - rank = int(os.getenv("OMPI_COMM_WORLD_RANK", "0")) - if rank == 0: - cuda_host_compile(cu_file, lib_file) - - import time - timeout = 60 - start = time.time() - while True: - if lib_file.exists(): - try: - ctypes.CDLL(str(lib_file)) - break - except OSError: - pass - if time.time() - start > timeout: - raise RuntimeError(f"Timeout waiting for {lib_file}") - time.sleep(0.1) - - lib = ctypes.CDLL(lib_file) - mype = ctypes.c_int() - npes = ctypes.c_int() - mype_in_node = ctypes.c_int() - npes_in_node = ctypes.c_int() - - stream = ctypes.c_void_p() - - input = ctypes.c_void_p() - partial_sum = ctypes.c_void_p() - full_sum = ctypes.c_void_p() - - input_nelems = 512 - to_all_nelems = 1 - lib.on_stream_before_launch_accumulate( - ctypes.byref(mype), ctypes.byref(npes), ctypes.byref(mype_in_node), ctypes.byref(npes_in_node), - ctypes.byref(stream), - ctypes.byref(input), ctypes.byref(partial_sum), ctypes.byref(full_sum), - input_nelems - ) - - dtype = torch.int32 - num_blocks = 1 - num_warps = (input_nelems + 31) // 32 - device = triton.runtime.driver.active.get_active_torch_device() - - input_storage = torch._C._construct_storage_from_data_pointer(input.value, device, dtype.itemsize * input_nelems) - input_tensor = torch.empty(0, dtype=torch.int32, device=device).set_(input_storage).view(input_nelems, ) - partial_sum_storage = torch._C._construct_storage_from_data_pointer(partial_sum.value, device, dtype.itemsize) - partial_sum_tensor = torch.empty(0, dtype=torch.int32, device=device).set_(partial_sum_storage).view(1, ) - full_sum_storage = torch._C._construct_storage_from_data_pointer(full_sum.value, device, dtype.itemsize) - full_sum_tensor = torch.empty(0, dtype=torch.int32, device=device).set_(full_sum_storage).view(1, ) - - def cumodule_init_hook(*args, **kwargs): - key = kwargs["key"] - jit_function = kwargs["fn"].jit_function - device = kwargs["compile"]["device"] - kernel_cache = jit_function.device_caches[device][0] - kernel = kernel_cache.get(key, None) - assert kernel is not None - kernel._init_handles() - ret = lib.nvshmemx_cumodule_init_wrapper(ctypes.c_void_p(kernel.module)) - assert ret == 0, f"nvshmemx_cumodule_init_wrapper failed: {ret}" - knobs.runtime.jit_post_compile_hook = cumodule_init_hook - - curr_stream = torch.cuda.ExternalStream(stream.value, device=device) - with torch.cuda.stream(curr_stream): - accumulate_kernel[(num_blocks, )]( - input_tensor, - partial_sum_tensor, - num_warps=num_warps - ) - - lib.on_stream_before_launch_correct_accumulate( - ctypes.byref(stream), - ctypes.byref(partial_sum), ctypes.byref(full_sum), - to_all_nelems - ) - - with torch.cuda.stream(curr_stream): - correct_accumulate_kernel[(num_blocks, )]( - input_tensor, - partial_sum_tensor, - full_sum_tensor, - num_warps=num_warps - ) - - lib.on_stream_after_launch_correct_accumulate( - mype_in_node.value, npes_in_node.value, - stream, - input, partial_sum, full_sum - ) - - -if __name__ == "__main__": - on_stream() \ No newline at end of file diff --git a/python/tutorials/tle/raw/nvshmem/04-overlap-in-mulit-kernel/no-overlap.py b/python/tutorials/tle/raw/nvshmem/04-overlap-in-mulit-kernel/no-overlap.py new file mode 100644 index 000000000..d0d61a288 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/04-overlap-in-mulit-kernel/no-overlap.py @@ -0,0 +1,5 @@ +from separate_stream_common import run_experiment + + +if __name__ == "__main__": + run_experiment("no-overlap") diff --git a/python/tutorials/tle/raw/nvshmem/04-overlap-in-mulit-kernel/overlap.py b/python/tutorials/tle/raw/nvshmem/04-overlap-in-mulit-kernel/overlap.py new file mode 100644 index 000000000..769888afa --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/04-overlap-in-mulit-kernel/overlap.py @@ -0,0 +1,5 @@ +from separate_stream_common import run_experiment + + +if __name__ == "__main__": + run_experiment("overlap") diff --git a/python/tutorials/tle/raw/nvshmem/04-overlap-in-mulit-kernel/ring-reduce-host.cu b/python/tutorials/tle/raw/nvshmem/04-overlap-in-mulit-kernel/ring-reduce-host.cu new file mode 100644 index 000000000..93105360c --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/04-overlap-in-mulit-kernel/ring-reduce-host.cu @@ -0,0 +1,78 @@ +#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 nvshmemx_cumodule_init_wrapper(CUmodule module) { + return nvshmemx_cumodule_init(module); +} + +extern "C" void ring_reduce_barrier_all() { + nvshmem_barrier_all(); +} + +extern "C" void ring_reduce_before_launch +( + int *mype, int *npes, int *mype_in_node, int *n_pes_in_node, + cudaStream_t *stream, + int **src, int **dst, int **data_h, uint64_t **signal, + int size +) { + + nvshmem_init(); + *mype = nvshmem_my_pe(); + *npes = nvshmem_n_pes(); + *mype_in_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE); + *n_pes_in_node = nvshmem_team_n_pes(NVSHMEMX_TEAM_NODE); + + CUDA_CHECK(cudaSetDevice(*mype_in_node)); + CUDA_CHECK(cudaStreamCreate(stream)); + + *src = (int *)nvshmem_malloc(sizeof(int) * size); + *dst = (int *)nvshmem_malloc(sizeof(int) * size); + *data_h = (int *)malloc(sizeof(int) * size); + *signal = (uint64_t *)nvshmem_calloc(*n_pes_in_node, sizeof(uint64_t)); + + for (size_t i = 0; i < size; i++) (*data_h)[i] = i; + CUDA_CHECK(cudaMemcpyAsync(*src, *data_h, sizeof(int) * size, cudaMemcpyHostToDevice, *stream)); + nvshmemx_barrier_all_on_stream(*stream); +} + +extern "C" void ring_reduce_after_launch +( + cudaStream_t stream, + void *src, void *dst, void *data_h, void *signal, + int mype, int npes, + int size +) { + CUDA_CHECK(cudaStreamSynchronize(stream)); + nvshmem_barrier_all(); + CUDA_CHECK(cudaMemcpy(data_h, dst, sizeof(int) * size, cudaMemcpyDeviceToHost)); + // int *data = (int *)data_h; + // for (size_t i = 0; i < size; i++) { + // if (data[i] != (int)i * npes) { + // // if (mype == 0) { + // // printf("PE %d, data[%zu] = %d expected data[%zu] = %d\n", mype, i, data[i], i, (int)i * npes); + // printf("PE %d error, data[%zu] = %d expected data[%zu] = %d\n", mype, i, data[i], i, (int)i * npes); + // } + // } + + nvshmem_free(dst); + nvshmem_free(src); + nvshmem_free(signal); + free(data_h); + + nvshmem_finalize(); +} diff --git a/python/tutorials/tle/raw/nvshmem/04-overlap-in-mulit-kernel/ring-reduce-put.cu b/python/tutorials/tle/raw/nvshmem/04-overlap-in-mulit-kernel/ring-reduce-put.cu new file mode 100644 index 000000000..8f45314fd --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/04-overlap-in-mulit-kernel/ring-reduce-put.cu @@ -0,0 +1,84 @@ +#include +#include +#include +#include + +extern "C" __device__ void ring_reduce_put_timed( + int *dst, const int *src, int nreduce, + uint64_t *signal, uint64_t *put_cycles) +{ + int mype = nvshmem_my_pe(); + int npes = nvshmem_n_pes(); + int peer = (mype + 1) % npes; + + int block_idx = blockIdx.x; + int num_blocks = gridDim.x; + int elems_per_block = nreduce / num_blocks; + + if (elems_per_block * (block_idx + 1) > nreduce) return; + + src += block_idx * elems_per_block; + dst += block_idx * elems_per_block; + signal += block_idx; + + if (threadIdx.x == 0) { + uint64_t t0 = clock64(); + nvshmem_int_put_signal_nbi( + dst, + src, + elems_per_block, + signal, + 1, + NVSHMEM_SIGNAL_ADD, + peer + ); + uint64_t t1 = clock64(); + put_cycles[block_idx] = t1 - t0; + } +} + +extern "C" __device__ void ring_reduce_wait_timed( + uint64_t *signal, uint64_t *wait_cycles) +{ + int block_idx = blockIdx.x; + signal += block_idx; + + if (threadIdx.x == 0) { + uint64_t t0 = clock64(); + nvshmem_quiet(); + nvshmem_signal_wait_until(signal, NVSHMEM_CMP_GE, 1); + uint64_t t1 = clock64(); + wait_cycles[block_idx] = t1 - t0; + } + __syncthreads(); +} + +extern "C" __device__ void local_matmul( + const float *A, const float *B, float *C, + int M, int N, int K) +{ + int block_idx = blockIdx.x; + + A += block_idx * M * N; + B += block_idx * N * K; + C += block_idx * M * K; + + int tid = threadIdx.x; + int nthreads = blockDim.x; + + for (int row = tid; row < M; row += nthreads) { + for (int col = 0; col < K; col++) { + C[row * K + col] = 0.0f; + } + } + + for (int row = tid; row < M; row += nthreads) { + for (int k = 0; k < N; k++) { + float a_val = A[row * N + k]; + for (int col = 0; col < K; col++) { + C[row * K + col] += a_val * B[k * K + col]; + } + } + } + __syncthreads(); +} diff --git a/python/tutorials/tle/raw/nvshmem/04-overlap-in-mulit-kernel/separate_stream_common.py b/python/tutorials/tle/raw/nvshmem/04-overlap-in-mulit-kernel/separate_stream_common.py new file mode 100644 index 000000000..8b2e84d78 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/04-overlap-in-mulit-kernel/separate_stream_common.py @@ -0,0 +1,507 @@ +import ctypes +import os +import subprocess +from pathlib import Path + +import torch +import triton +import triton.experimental.tle.language.raw as tle_raw +import triton.knobs as knobs +from triton.experimental.tle.raw import dialect + + +NVSHMEM_HOME = "/data/zyuli/miniconda3/envs/flagtree_triton_v3.6.x/lib/python3.12/site-packages/nvidia/nvshmem" + +SIZE = 1024 * 1024 * 32 +MAT_M = 128 +MAT_N = 256 +MAT_K = 128 + + +def _make_dialect(extern_func_name): + return dialect( + name="cuda", + compiler="nvcc", + file=(Path(__file__).parent / "ring-reduce-put.cu"), + extern=(Path(__file__).parent / "ring-reduce-put-extern-call.py"), + extern_func_name=extern_func_name, + libs={"nvshmem": NVSHMEM_HOME}, + links=["nvshmem_device"], + ) + + +@_make_dialect("ring_reduce_put_timed") +def edsl_put_timed(*args, **kwargs): ... + + +@_make_dialect("ring_reduce_wait_timed") +def edsl_wait_timed(*args, **kwargs): ... + + +@_make_dialect("local_matmul") +def edsl_local_matmul(*args, **kwargs): ... + + +@triton.jit +def communication_kernel( + dst, + src, + nreduce: triton.language.constexpr, + signal, + put_cycles, + wait_cycles, +): + tle_raw.call( + edsl_put_timed, + [dst, src, nreduce, signal, put_cycles], + ) + tle_raw.call( + edsl_wait_timed, + [signal, wait_cycles], + ) + + +@triton.jit +def matmul_kernel( + mat_A, + mat_B, + mat_C, + mat_M: triton.language.constexpr, + mat_N: triton.language.constexpr, + mat_K: triton.language.constexpr, +): + tle_raw.call( + edsl_local_matmul, + [mat_A, mat_B, mat_C, mat_M, mat_N, mat_K], + ) + + +def cuda_host_compile(cuda_host_path, cuda_host_lib): + nvcc = os.getenv("NVCC", "nvcc") + include_path = f"-I{os.path.join(NVSHMEM_HOME, 'include')}" + lib_path = f"-L{os.path.join(NVSHMEM_HOME, 'lib')}" + + prop = torch.cuda.get_device_properties(torch.cuda.current_device()) + arch = f"-arch=sm_{prop.major}{prop.minor}" + tmp_file = Path(cuda_host_lib).with_suffix(".so.tmp") + build = [ + nvcc, + "-shared", + "-Xcompiler", + "-fPIC", + "-rdc=true", + arch, + include_path, + lib_path, + "-lnvshmem_host", + "-lnvshmem_device", + "-o", + tmp_file, + cuda_host_path, + ] + result = subprocess.run(build, capture_output=True) + assert result.returncode == 0, ( + f"NVCC host failed\nstderr:\n{result.stderr.decode()}" + ) + tmp_file.rename(cuda_host_lib) + + +def wait_for_lib(lib_file, timeout=60): + import time + + start = time.time() + while True: + if lib_file.exists(): + try: + ctypes.CDLL(str(lib_file)) + return + except OSError: + pass + if time.time() - start > timeout: + raise RuntimeError(f"Timeout waiting for {lib_file}") + time.sleep(0.1) + + +def setup_nvshmem(lib, size): + mype = ctypes.c_int() + npes = ctypes.c_int() + mype_in_node = ctypes.c_int() + npes_in_node = ctypes.c_int() + stream = ctypes.c_void_p() + src = ctypes.c_void_p() + dst = ctypes.c_void_p() + data_h = ctypes.c_void_p() + signal = ctypes.c_void_p() + + lib.ring_reduce_before_launch( + ctypes.byref(mype), + ctypes.byref(npes), + ctypes.byref(mype_in_node), + ctypes.byref(npes_in_node), + ctypes.byref(stream), + ctypes.byref(src), + ctypes.byref(dst), + ctypes.byref(data_h), + ctypes.byref(signal), + size, + ) + return ( + mype, + npes, + mype_in_node, + npes_in_node, + stream, + src, + dst, + data_h, + signal, + ) + + +def make_tensors(src, dst, signal, size, device): + src_storage = torch._C._construct_storage_from_data_pointer( + src.value, device, 4 * size + ) + src_tensor = torch.empty(0, dtype=torch.int32, device=device).set_( + src_storage + ).view(size) + + dst_storage = torch._C._construct_storage_from_data_pointer( + dst.value, device, 4 * size + ) + dst_tensor = torch.empty(0, dtype=torch.int32, device=device).set_( + dst_storage + ).view(size) + + signal_storage = torch._C._construct_storage_from_data_pointer( + signal.value, device, 8 + ) + signal_tensor = torch.empty(0, dtype=torch.uint64, device=device).set_( + signal_storage + ).view(1) + + return src_tensor, dst_tensor, signal_tensor + + +def install_cumodule_hook(lib): + def cumodule_init_hook(*args, **kwargs): + key = kwargs["key"] + jit_function = kwargs["fn"].jit_function + dev = kwargs["compile"]["device"] + kernel_cache = jit_function.device_caches[dev][0] + kernel = kernel_cache.get(key) + assert kernel is not None + kernel._init_handles() + ret = lib.nvshmemx_cumodule_init_wrapper(ctypes.c_void_p(kernel.module)) + assert ret == 0, f"nvshmemx_cumodule_init_wrapper failed: {ret}" + + knobs.runtime.jit_post_compile_hook = cumodule_init_hook + + +def reset_signal_collective(lib, signal_tensor, comm_stream): + lib.ring_reduce_barrier_all() + with torch.cuda.stream(comm_stream): + signal_tensor.zero_() + comm_stream.synchronize() + lib.ring_reduce_barrier_all() + + +def reset_tensors_on_stream(stream, *tensors): + with torch.cuda.stream(stream): + for tensor in tensors: + tensor.zero_() + stream.synchronize() + + +def format_cycle_stats(tensor, clock_rate_khz): + tensor = tensor.to(torch.int64) + values = tensor[tensor > 0].detach().cpu().to(torch.float64) + if values.numel() == 0: + return "n/a" + us = values / clock_rate_khz * 1000.0 + return ( + f"{us.mean().item():.3f} us " + f"({int(values.mean().item())} cycles)" + ) + + +def launch_communication( + stream, + args, + wait_event=None, + done_event=None, +): + with torch.cuda.stream(stream): + if wait_event is not None: + stream.wait_event(wait_event) + communication_kernel[(1,)]( + *args, + num_warps=4, + ) + if done_event is not None: + done_event.record(stream) + + +def launch_matmul( + stream, + grid, + args, + wait_event=None, + done_event=None, +): + with torch.cuda.stream(stream): + if wait_event is not None: + stream.wait_event(wait_event) + matmul_kernel[grid]( + *args, + num_warps=4, + ) + if done_event is not None: + done_event.record(stream) + + +def run_overlap( + timing_stream, + comm_stream, + compute_stream, + comm_args, + matmul_grid, + matmul_args, +): + gate = torch.cuda.Event(enable_timing=True) + comm_done = torch.cuda.Event() + compute_done = torch.cuda.Event() + all_done = torch.cuda.Event(enable_timing=True) + + launch_communication( + comm_stream, + comm_args, + wait_event=gate, + done_event=comm_done, + ) + launch_matmul( + compute_stream, + matmul_grid, + matmul_args, + wait_event=gate, + done_event=compute_done, + ) + + with torch.cuda.stream(timing_stream): + gate.record(timing_stream) + timing_stream.wait_event(comm_done) + timing_stream.wait_event(compute_done) + all_done.record(timing_stream) + + all_done.synchronize() + return gate.elapsed_time(all_done) + + +def run_no_overlap( + timing_stream, + comm_stream, + compute_stream, + comm_args, + matmul_grid, + matmul_args, +): + gate = torch.cuda.Event(enable_timing=True) + comm_done = torch.cuda.Event() + compute_done = torch.cuda.Event() + all_done = torch.cuda.Event(enable_timing=True) + + launch_communication( + comm_stream, + comm_args, + wait_event=gate, + done_event=comm_done, + ) + launch_matmul( + compute_stream, + matmul_grid, + matmul_args, + wait_event=comm_done, + done_event=compute_done, + ) + + with torch.cuda.stream(timing_stream): + gate.record(timing_stream) + timing_stream.wait_event(compute_done) + all_done.record(timing_stream) + + all_done.synchronize() + return gate.elapsed_time(all_done) + + +def run_experiment(mode): + if mode not in {"overlap", "no-overlap"}: + raise ValueError(f"Unsupported mode: {mode}") + cu_file = (Path(__file__).parent / "ring-reduce-host.cu").resolve() + lib_file = cu_file.with_suffix(".so") + + rank = int(os.getenv("MPI_RANK", "0")) + lib = ctypes.CDLL(str(lib_file)) + lib.ring_reduce_barrier_all.argtypes = [] + lib.ring_reduce_barrier_all.restype = None + + ( + mype, + npes, + mype_in_node, + npes_in_node, + stream, + src, + dst, + data_h, + signal, + ) = setup_nvshmem(lib, SIZE) + + device = triton.runtime.driver.active.get_active_torch_device() + props = torch.cuda.get_device_properties(device) + compute_blocks = max(1, props.multi_processor_count - 1) + + src_tensor, dst_tensor, signal_tensor = make_tensors( + src, dst, signal, SIZE, device + ) + + mat_A = torch.randn( + compute_blocks, MAT_M, MAT_N, dtype=torch.float32, device=device + ) + mat_B = torch.randn( + compute_blocks, MAT_N, MAT_K, dtype=torch.float32, device=device + ) + mat_C = torch.zeros( + compute_blocks, MAT_M, MAT_K, dtype=torch.float32, device=device + ) + + overlap_put_cycles = torch.zeros(1, dtype=torch.uint64, device=device) + overlap_wait_cycles = torch.zeros(1, dtype=torch.uint64, device=device) + no_overlap_put_cycles = torch.zeros(1, dtype=torch.uint64, device=device) + no_overlap_wait_cycles = torch.zeros(1, dtype=torch.uint64, device=device) + + install_cumodule_hook(lib) + + comm_stream = torch.cuda.ExternalStream(stream.value, device=device) + compute_stream = torch.cuda.Stream(device=device, priority=0) + timing_stream = torch.cuda.Stream(device=device, priority=0) + matmul_grid = (compute_blocks,) + + overlap_comm_args = ( + dst_tensor, + src_tensor, + SIZE, + signal_tensor, + overlap_put_cycles, + overlap_wait_cycles, + ) + no_overlap_comm_args = ( + dst_tensor, + src_tensor, + SIZE, + signal_tensor, + no_overlap_put_cycles, + no_overlap_wait_cycles, + ) + matmul_args = ( + mat_A, + mat_B, + mat_C, + MAT_M, + MAT_N, + MAT_K, + ) + + torch.cuda.synchronize(device) + + selected_comm_args = ( + overlap_comm_args if mode == "overlap" else no_overlap_comm_args + ) + if mode == "overlap": + run_overlap( + timing_stream, + comm_stream, + compute_stream, + selected_comm_args, + matmul_grid, + matmul_args, + ) + else: + run_no_overlap( + timing_stream, + comm_stream, + compute_stream, + selected_comm_args, + matmul_grid, + matmul_args, + ) + reset_signal_collective(lib, signal_tensor, comm_stream) + + selected_put_cycles = ( + overlap_put_cycles if mode == "overlap" else no_overlap_put_cycles + ) + selected_wait_cycles = ( + overlap_wait_cycles if mode == "overlap" else no_overlap_wait_cycles + ) + reset_tensors_on_stream( + comm_stream, selected_put_cycles, selected_wait_cycles + ) + lib.ring_reduce_barrier_all() + + if mode == "overlap": + elapsed_ms = run_overlap( + timing_stream, + comm_stream, + compute_stream, + selected_comm_args, + matmul_grid, + matmul_args, + ) + else: + elapsed_ms = run_no_overlap( + timing_stream, + comm_stream, + compute_stream, + selected_comm_args, + matmul_grid, + matmul_args, + ) + + lib.ring_reduce_barrier_all() + + clock_rate_khz = props.clock_rate + put_stats = format_cycle_stats(selected_put_cycles, clock_rate_khz) + wait_stats = format_cycle_stats(selected_wait_cycles, clock_rate_khz) + + if mype.value == 0 or mype.value == 1: + print("\n" + "=" * 64) + print(f" Separate-stream {mode}") + print("=" * 64) + print(f" PEs (GPUs) : {npes.value}") + print(f" GPU SMs : {props.multi_processor_count}") + print(" Communication : 1 block on communication stream") + print(f" Matmul : {compute_blocks} blocks on compute stream") + print( + f" Message size : {SIZE * 4 / 1024 / 1024:.1f} MB per PE" + ) + print( + f" Per-block matmul : {MAT_M}x{MAT_N} x {MAT_N}x{MAT_K}" + ) + print("-" * 64) + print(f" elapsed : {elapsed_ms:.3f} ms") + print("-" * 64) + print(f" put_nbi : {put_stats}") + print(f" quiet/wait : {wait_stats}") + print("=" * 64) + + compute_stream.synchronize() + timing_stream.synchronize() + lib.ring_reduce_after_launch( + stream, + src, + dst, + data_h, + signal, + mype.value, + npes.value, + SIZE, + ) diff --git a/python/tutorials/tle/raw/nvshmem/05-put-block/put-block-device.cu b/python/tutorials/tle/raw/nvshmem/05-put-block/put-block-device.cu deleted file mode 100644 index 6ce11ff82..000000000 --- a/python/tutorials/tle/raw/nvshmem/05-put-block/put-block-device.cu +++ /dev/null @@ -1,23 +0,0 @@ -#include -#include "nvshmem.h" -#include "nvshmemx.h" - -extern "C" __device__ void set_and_shift(float *send_data, float *recv_data, int num_elems, int mype, int npes) { - int thread_idx = blockIdx.x * blockDim.x + threadIdx.x; - /* set the corresponding element of send_data */ - if (thread_idx < num_elems) send_data[thread_idx] = mype; - - int peer = (mype + 1) % npes; - /* Every thread in block 0 calls nvshmemx_float_put_block. Alternatively, - every thread can call shmem_float_p, but shmem_float_p has a disadvantage - that when the destination GPU is connected via IB, there will be one rma - message for every single element which can be detrimental to performance. - And the disadvantage with shmem_float_put is that when the destination GPU is p2p - connected, it cannot leverage multiple threads to copy the data to the destination - GPU. */ - int block_offset = blockIdx.x * blockDim.x; - nvshmemx_float_put_block(recv_data + block_offset, send_data + block_offset, - min(blockDim.x, num_elems - block_offset), - peer); /* All threads in a block call the API - with the same arguments */ -} \ No newline at end of file diff --git a/python/tutorials/tle/raw/nvshmem/05-put-block/put-block-host.cu b/python/tutorials/tle/raw/nvshmem/05-put-block/put-block-host.cu deleted file mode 100644 index 2ecf04f7c..000000000 --- a/python/tutorials/tle/raw/nvshmem/05-put-block/put-block-host.cu +++ /dev/null @@ -1,75 +0,0 @@ -#include -#include "nvshmem.h" -#include "nvshmemx.h" - -#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 nvshmemx_cumodule_init_wrapper(CUmodule module) { - return nvshmemx_cumodule_init(module); -} - - -extern "C" void put_block_before_launch -( - int *mype, int *npes, int *mype_in_node, int *n_pes_in_node, - float **send_data, float **recv_data, - int num_elems -) { - - nvshmem_init(); - *mype = nvshmem_my_pe(); - *npes = nvshmem_n_pes(); - *mype_in_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE); - *n_pes_in_node = nvshmem_team_n_pes(NVSHMEMX_TEAM_NODE); - - CUDA_CHECK(cudaSetDevice(*mype_in_node)); - - *send_data = (float *)nvshmem_malloc(sizeof(float) * num_elems); - *recv_data = (float *)nvshmem_malloc(sizeof(float) * num_elems); -} - - -extern "C" void put_block_after_launch -( - int mype, int npes, - void *send_data, void *recv_data, - int num_elems -) { - CUDA_CHECK(cudaGetLastError()); - CUDA_CHECK(cudaDeviceSynchronize()); - nvshmem_barrier_all(); - - /* Do data validation */ - float *host = new float[num_elems]; - CUDA_CHECK(cudaMemcpy(host, recv_data, num_elems * sizeof(float), cudaMemcpyDefault)); - int ref = (mype - 1 + npes) % npes; - bool success = true; - for (int i = 0; i < num_elems; ++i) { - if (host[i] != ref) { - printf("Error at %d of rank %d: %f\n", i, mype, host[i]); - success = false; - break; - } - } - - if (success) { - printf("[%d of %d] run complete \n", mype, npes); - } else { - printf("[%d of %d] run failure \n", mype, npes); - } - - nvshmem_free(send_data); - nvshmem_free(recv_data); - - nvshmem_finalize(); -} \ No newline at end of file diff --git a/python/tutorials/tle/raw/nvshmem/05-put-block/put-block.py b/python/tutorials/tle/raw/nvshmem/05-put-block/put-block.py deleted file mode 100644 index 48162524f..000000000 --- a/python/tutorials/tle/raw/nvshmem/05-put-block/put-block.py +++ /dev/null @@ -1,128 +0,0 @@ -import os -import subprocess -import ctypes -import torch -import triton -import triton.knobs as knobs -import triton.experimental.tle.language.raw as tle_raw - -from pathlib import Path -from triton.experimental.tle.raw import dialect -from triton.language.extra.cuda import libnvshmem_device - - -@dialect( - name="cuda", - file=(Path(__file__).parent / "put-block-device.cu").resolve(), - library={"nvshmem": "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem"}, -) -def edsl(*args, **kwargs): - ... - - - -@triton.jit -def set_and_shift_kernel( - send_data, recv_data, num_elems, mype, npes -): - tle_raw.call(edsl, []) - libnvshmem_device.set_and_shift(send_data, recv_data, num_elems, mype, npes) - - -def cuda_host_compile(cuda_host_path, cuda_host_lib): - NVCC = os.getenv("NVCC", "nvcc") - NVSHMEM_HOME = "/home/zyuli/miniconda3/envs/flagtree/lib/python3.12/site-packages/nvidia/nvshmem" - include_path = f"-I{os.path.join(NVSHMEM_HOME, 'include')}" - lib_path = f"-L{os.path.join(NVSHMEM_HOME, 'lib')}" - - prop = torch.cuda.get_device_properties(torch.cuda.current_device()) - arch = f"-arch=sm_{prop.major}{prop.minor}" - tmp_file = Path(cuda_host_lib).with_suffix('.so.tmp') - build = [ - NVCC, "-shared", "-Xcompiler", "-fPIC", "-rdc=true", arch, include_path, lib_path, "-lnvshmem_host", - "-lnvshmem_device", "-o", tmp_file, cuda_host_path - ] - build = subprocess.run(build, capture_output=True) - assert build.returncode == 0, (f"NVCC host failed\nstderr:\n{build.stderr.decode()}") - tmp_file.rename(cuda_host_lib) - - -def put_block(): - cu_file = (Path(__file__).parent / "put-block-host.cu").resolve() - lib_file = Path(cu_file).with_suffix('.so') - - rank = int(os.getenv("OMPI_COMM_WORLD_RANK", "0")) - if rank == 0: - cuda_host_compile(cu_file, lib_file) - - import time - timeout = 60 - start = time.time() - while True: - if lib_file.exists(): - try: - ctypes.CDLL(str(lib_file)) - break - except OSError: - pass - if time.time() - start > timeout: - raise RuntimeError(f"Timeout waiting for {lib_file}") - time.sleep(0.1) - - lib = ctypes.CDLL(lib_file) - mype = ctypes.c_int() - npes = ctypes.c_int() - mype_in_node = ctypes.c_int() - npes_in_node = ctypes.c_int() - - send_data = ctypes.c_void_p() - recv_data = ctypes.c_void_p() - - num_elems = 8192 - lib.put_block_before_launch( - ctypes.byref(mype), ctypes.byref(npes), ctypes.byref(mype_in_node), ctypes.byref(npes_in_node), - ctypes.byref(send_data), ctypes.byref(recv_data), - num_elems - ) - - dtype = torch.float32 - THREADS_PER_BLOCK = 1024 - num_blocks = num_elems // THREADS_PER_BLOCK - num_warps = (THREADS_PER_BLOCK + 31) // 32 - device = triton.runtime.driver.active.get_active_torch_device() - - send_data_storage = torch._C._construct_storage_from_data_pointer(send_data.value, device, dtype.itemsize * num_elems) - send_data_tensor = torch.empty(0, dtype=dtype, device=device).set_(send_data_storage).view(num_elems, ) - recv_data_storage = torch._C._construct_storage_from_data_pointer(recv_data.value, device, dtype.itemsize * num_elems) - recv_data_tensor = torch.empty(0, dtype=dtype, device=device).set_(recv_data_storage).view(num_elems, ) - - def cumodule_init_hook(*args, **kwargs): - key = kwargs["key"] - jit_function = kwargs["fn"].jit_function - device = kwargs["compile"]["device"] - kernel_cache = jit_function.device_caches[device][0] - kernel = kernel_cache.get(key, None) - assert kernel is not None - kernel._init_handles() - ret = lib.nvshmemx_cumodule_init_wrapper(ctypes.c_void_p(kernel.module)) - assert ret == 0, f"nvshmemx_cumodule_init_wrapper failed: {ret}" - knobs.runtime.jit_post_compile_hook = cumodule_init_hook - - set_and_shift_kernel[(num_blocks, )]( - send_data_tensor, - recv_data_tensor, - num_elems, - mype.value, - npes.value, - num_warps=num_warps - ) - - lib.put_block_after_launch( - mype.value, npes.value, - send_data, recv_data, - num_elems - ) - - -if __name__ == "__main__": - put_block() \ No newline at end of file diff --git a/python/tutorials/tle/raw/nvshmem/common/generate_extern_call.py b/python/tutorials/tle/raw/nvshmem/common/generate_extern_call.py new file mode 100644 index 000000000..f2cbb5598 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/common/generate_extern_call.py @@ -0,0 +1,207 @@ +"""Generate Triton extern-call wrappers from CUDA device functions.""" + +import re +from dataclasses import dataclass +from pathlib import Path + + +_DTYPES = { + "bool": "int1", + "char": "int8", + "signed char": "int8", + "unsigned char": "uint8", + "int8_t": "int8", + "uint8_t": "uint8", + "short": "int16", + "short int": "int16", + "signed short": "int16", + "unsigned short": "uint16", + "int16_t": "int16", + "uint16_t": "uint16", + "int": "int32", + "signed": "int32", + "signed int": "int32", + "unsigned": "uint32", + "unsigned int": "uint32", + "int32_t": "int32", + "uint32_t": "uint32", + "long": "int64", + "long int": "int64", + "long long": "int64", + "long long int": "int64", + "unsigned long": "uint64", + "unsigned long int": "uint64", + "unsigned long long": "uint64", + "unsigned long long int": "uint64", + "int64_t": "int64", + "uint64_t": "uint64", + "size_t": "uint64", + "float": "fp32", + "double": "fp64", +} + +_QUALIFIERS = re.compile( + r"\b(?:const|volatile|restrict|__restrict__|__restrict|__const__|" + r"__device__|__forceinline__|inline|static)\b" +) +_ATTRIBUTE = re.compile( + r"__attribute__\s*\(\((?:[^()]|\([^()]*\))*\)\)", + re.DOTALL, +) +_COMMENT = re.compile(r"//[^\n]*|/\*.*?\*/", re.DOTALL) +_FUNCTION = re.compile( + r'extern\s+"C"\s+(?P.*?)\b(?P[A-Za-z_]\w*)\s*' + r"\((?P.*?)\)\s*(?:noexcept\s*)?\{", + re.DOTALL, +) + + +@dataclass(frozen=True) +class CudaType: + dtype: str + pointer: bool = False + + @property + def triton_type(self) -> str: + dtype = f'core.dtype("{self.dtype}")' + return f"core.pointer_type({dtype})" if self.pointer else dtype + + +@dataclass(frozen=True) +class Parameter: + name: str + type: CudaType + + @property + def argument(self) -> str: + if self.type.pointer: + return self.name + return f"tl.cast({self.name}, tl.{self.type.dtype}, _semantic=_semantic)" + + +@dataclass(frozen=True) +class Function: + name: str + parameters: tuple[Parameter, ...] + return_type: CudaType | None + + +def _split_parameters(parameters: str) -> list[str]: + result = [] + start = 0 + depth = 0 + for index, char in enumerate(parameters): + if char in "([": + depth += 1 + elif char in ")]": + depth -= 1 + elif char == "," and depth == 0: + result.append(parameters[start:index].strip()) + start = index + 1 + tail = parameters[start:].strip() + if tail and tail != "void": + result.append(tail) + return result + + +def _parse_type(spelling: str, *, context: str) -> CudaType | None: + spelling = _ATTRIBUTE.sub(" ", spelling) + pointer = "*" in spelling or "[" in spelling + spelling = re.sub(r"\[[^\]]*\]", " ", spelling) + spelling = spelling.replace("*", " ") + spelling = _QUALIFIERS.sub(" ", spelling) + spelling = " ".join(spelling.split()) + if spelling == "void" and not pointer: + return None + if spelling == "void" and pointer: + return CudaType("uint8", pointer=True) + dtype = _DTYPES.get(spelling) + if dtype is None: + raise ValueError(f"unsupported CUDA type {spelling!r} in {context}") + return CudaType(dtype, pointer=pointer) + + +def _parse_parameter(parameter: str, function_name: str) -> Parameter: + parameter = _ATTRIBUTE.sub(" ", parameter).strip() + match = re.search(r"([A-Za-z_]\w*)\s*(?:\[[^\]]*\])?\s*$", parameter) + if match is None: + raise ValueError(f"cannot parse parameter {parameter!r} in {function_name}") + name = match.group(1) + type_spelling = parameter[:match.start(1)] + parameter[match.end(1):] + parsed_type = _parse_type(type_spelling, context=f"{function_name}.{name}") + if parsed_type is None: + raise ValueError(f"parameter {name!r} in {function_name} cannot have type void") + return Parameter(name, parsed_type) + + +def parse_cuda_functions(source: str) -> tuple[Function, ...]: + source = _COMMENT.sub("", source) + functions = [] + for match in _FUNCTION.finditer(source): + if "__device__" not in match.group("prefix"): + continue + prefix = _QUALIFIERS.sub(" ", _ATTRIBUTE.sub(" ", match.group("prefix"))) + return_type = _parse_type(prefix, context=f"{match.group('name')} return type") + parameters = tuple( + _parse_parameter(parameter, match.group("name")) + for parameter in _split_parameters(match.group("params")) + ) + functions.append(Function(match.group("name"), parameters, return_type)) + return tuple(functions) + + +def _render_function(function: Function) -> str: + names = ", ".join(parameter.name for parameter in function.parameters) + signature = f"{names}, _semantic=None" if names else "_semantic=None" + arguments = ",\n ".join(parameter.argument for parameter in function.parameters) + if arguments: + arguments = f"[\n {arguments},\n ]" + else: + arguments = "[]" + types = "\n".join(f" {parameter.type.triton_type}," for parameter in function.parameters) + returns = "()" if function.return_type is None else f"({function.return_type.triton_type},)" + return ( + f"@core.extern\n" + f"def {function.name}({signature}):\n" + f" return core.extern_call(\n" + f' "",\n' + f' "",\n' + f" {arguments},\n" + f" {{\n" + f" (\n{types}\n" + f' ): ("{function.name}", {returns}),\n' + f" }},\n" + f" is_pure=False,\n" + f" _semantic=_semantic,\n" + f" )\n" + ) + + +def generate( + cuda_file: str | Path, + output_file: str | Path, + required_function: str | None = None, +) -> Path: + cuda_file = Path(cuda_file) + output_file = Path(output_file) + functions = parse_cuda_functions(cuda_file.read_text(encoding="utf-8")) + if not functions: + raise ValueError(f"no extern \"C\" CUDA functions found in {cuda_file}") + names = {function.name for function in functions} + if required_function is not None and required_function not in names: + raise ValueError( + f"extern function {required_function!r} was not found in {cuda_file}; " + f"found: {', '.join(sorted(names))}" + ) + content = ( + "# Generated from " + f"{cuda_file.name}; do not edit manually.\n" + "import triton.language as tl\n" + "import triton.language.core as core\n\n\n" + + "\n\n".join(_render_function(function) for function in functions) + + "\n" + ) + output_file.parent.mkdir(parents=True, exist_ok=True) + if not output_file.exists() or output_file.read_text(encoding="utf-8") != content: + output_file.write_text(content, encoding="utf-8") + return output_file diff --git a/python/tutorials/tle/raw/nvshmem/common/prepare_nvshmem.py b/python/tutorials/tle/raw/nvshmem/common/prepare_nvshmem.py new file mode 100644 index 000000000..4aad0e5f9 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/common/prepare_nvshmem.py @@ -0,0 +1,163 @@ +#!/usr/bin/env python3 + +import argparse +import ast +import importlib.util +import os +import subprocess +import sys +import tempfile +from pathlib import Path + + +NVSHMEM_ROOT = Path(__file__).resolve().parents[1] +GENERATOR_FILE = Path(__file__).with_name("generate_extern_call.py") +_SPEC = importlib.util.spec_from_file_location("_tle_generate_extern_call", GENERATOR_FILE) +if _SPEC is None or _SPEC.loader is None: + raise ImportError(f"cannot load extern call generator from {GENERATOR_FILE}") +_GENERATOR = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _GENERATOR +_DONT_WRITE_BYTECODE = sys.dont_write_bytecode +try: + sys.dont_write_bytecode = True + _SPEC.loader.exec_module(_GENERATOR) +finally: + sys.dont_write_bytecode = _DONT_WRITE_BYTECODE +generate = _GENERATOR.generate + + +def _last_string(node): + if isinstance(node, ast.Constant) and isinstance(node.value, str): + return node.value + for child in reversed(list(ast.iter_child_nodes(node))): + value = _last_string(child) + if value is not None: + return value + return None + + +def _dialect_file_pairs(example_dir): + pairs = set() + for python_file in example_dir.glob("*.py"): + tree = ast.parse(python_file.read_text(encoding="utf-8"), filename=str(python_file)) + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + name = getattr(node.func, "id", None) or getattr(node.func, "attr", None) + if name != "dialect": + continue + keywords = {keyword.arg: keyword.value for keyword in node.keywords if keyword.arg} + cuda_name = _last_string(keywords.get("file")) if "file" in keywords else None + extern_name = _last_string(keywords.get("extern")) if "extern" in keywords else None + if cuda_name and extern_name: + pairs.add((example_dir / cuda_name, example_dir / extern_name)) + return pairs + + +def generate_extern_files(example_dir): + pairs = _dialect_file_pairs(example_dir) + generated = [] + for cuda_file, output_file in sorted(pairs, key=lambda pair: str(pair[1])): + if not cuda_file.is_file(): + raise FileNotFoundError(f"CUDA device file does not exist: {cuda_file}") + generate(cuda_file, output_file) + generated.append(output_file) + return generated + + +def detect_arch(explicit_arch): + arch = explicit_arch or "sm_80" + arch = arch.removeprefix("-arch=").removeprefix("sm_") + return f"sm_{arch}" + + +def resolve_nvshmem_home(explicit_home): + if not explicit_home: + raise ValueError( + "NVSHMEM_HOME is required; set the environment variable or pass --nvshmem-home" + ) + home = Path(explicit_home).expanduser().resolve() + + if not (home / "include").is_dir() or not (home / "lib").is_dir(): + raise FileNotFoundError( + f"invalid NVSHMEM_HOME {home}: expected include/ and lib/ directories" + ) + return home + + +def compile_host_files(example_dir, nvshmem_home, arch, force): + nvcc = os.getenv("NVCC", "nvcc") + outputs = [] + for cuda_file in sorted(example_dir.glob("*-host.cu")): + output_file = cuda_file.with_suffix(".so") + if ( + not force + and output_file.exists() + and output_file.stat().st_mtime_ns >= cuda_file.stat().st_mtime_ns + ): + outputs.append(output_file) + continue + + temporary = tempfile.NamedTemporaryFile( + prefix=f".{output_file.name}.", + suffix=".tmp", + dir=output_file.parent, + delete=False, + ) + temporary_path = Path(temporary.name) + temporary.close() + 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(cuda_file), + ] + try: + subprocess.run(command, check=True) + os.replace(temporary_path, output_file) + finally: + temporary_path.unlink(missing_ok=True) + outputs.append(output_file) + return outputs + + +def main(): + parser = argparse.ArgumentParser( + description="Generate Triton extern calls and compile NVSHMEM host libraries." + ) + parser.add_argument("target", type=Path, help="NVSHMEM example Python file") + parser.add_argument("--nvshmem-home", default=os.getenv("NVSHMEM_HOME")) + parser.add_argument("--arch", default="sm_80", help="CUDA architecture (default: sm_80)") + parser.add_argument("--force", action="store_true", help="always rebuild host libraries") + args = parser.parse_args() + + target = args.target.expanduser().resolve() + if not target.is_file(): + parser.error(f"target does not exist: {target}") + if NVSHMEM_ROOT not in target.parents: + parser.error(f"target must be below {NVSHMEM_ROOT}") + + nvshmem_home = resolve_nvshmem_home(args.nvshmem_home) + arch = detect_arch(args.arch) + generated = generate_extern_files(target.parent) + libraries = compile_host_files(target.parent, nvshmem_home, arch, args.force) + + for path in generated: + print(f"[prepare] extern: {path}") + for path in libraries: + print(f"[prepare] host library: {path}") + print(f"[prepare] NVSHMEM_HOME={nvshmem_home}") + print(f"[prepare] CUDA architecture={arch}") + + +if __name__ == "__main__": + main() diff --git a/python/tutorials/tle/raw/nvshmem/run_nvshmem.sh b/python/tutorials/tle/raw/nvshmem/run_nvshmem.sh new file mode 100755 index 000000000..61464ce2e --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/run_nvshmem.sh @@ -0,0 +1,138 @@ +#!/usr/bin/env bash + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +PYTHON_BIN="${PYTHON:-python}" +NVSHMRUN_BIN="${NVSHMRUN:-nvshmrun}" +NP=2 +ARCH="" +FORCE=0 +NVSHMEM_HOME_ARG="${NVSHMEM_HOME:-}" + +usage() { + cat <<'EOF' +Usage: + run_nvshmem.sh [options] [-- ] + +Options: + --np N Number of NVSHMEM processes (default: 2) + --python PATH Python executable (default: $PYTHON or python) + --nvshmrun PATH nvshmrun executable (default: $NVSHMRUN or nvshmrun) + --nvshmem-home PATH NVSHMEM installation root (or set NVSHMEM_HOME) + --arch ARCH CUDA architecture (default: sm_80) + --force Rebuild the host shared library + -h, --help Show this help + +Examples: + ./run_nvshmem.sh --np 2 02-ring-reduce/ring-reduce.py + ./run_nvshmem.sh --np 2 06-ring-reduce-overlap/overlap.py + ./run_nvshmem.sh --np 4 ring-reduce.py +EOF +} + +while (($#)); do + case "$1" in + --np) + NP="$2" + shift 2 + ;; + --python) + PYTHON_BIN="$2" + shift 2 + ;; + --nvshmrun) + NVSHMRUN_BIN="$2" + shift 2 + ;; + --nvshmem-home) + NVSHMEM_HOME_ARG="$2" + shift 2 + ;; + --arch) + ARCH="$2" + shift 2 + ;; + --force) + FORCE=1 + shift + ;; + -h|--help) + usage + exit 0 + ;; + --) + shift + break + ;; + -*) + echo "Unknown option: $1" >&2 + usage >&2 + exit 2 + ;; + *) + TARGET_INPUT="$1" + shift + break + ;; + esac +done + +if [[ -z "${TARGET_INPUT:-}" ]]; then + usage >&2 + exit 2 +fi + +if [[ ! "$NP" =~ ^[1-9][0-9]*$ ]]; then + echo "Process count must be a positive integer: $NP" >&2 + exit 2 +fi + +if [[ -z "$NVSHMEM_HOME_ARG" ]]; then + echo "NVSHMEM_HOME is required; set it or pass --nvshmem-home PATH" >&2 + exit 2 +fi + +if [[ "${1:-}" == "--" ]]; then + shift +fi + +if [[ -f "$TARGET_INPUT" ]]; then + TARGET="$(cd -- "$(dirname -- "$TARGET_INPUT")" && pwd)/$(basename -- "$TARGET_INPUT")" +elif [[ -f "$SCRIPT_DIR/$TARGET_INPUT" ]]; then + TARGET="$SCRIPT_DIR/$TARGET_INPUT" +else + mapfile -t MATCHES < <(find "$SCRIPT_DIR" -mindepth 2 -maxdepth 2 -type f -name "$TARGET_INPUT") + if ((${#MATCHES[@]} != 1)); then + echo "Could not uniquely resolve example: $TARGET_INPUT" >&2 + exit 2 + fi + TARGET="${MATCHES[0]}" +fi + +if ! NVSHMRUN_PATH="$(command -v "$NVSHMRUN_BIN")"; then + echo "Could not find nvshmrun executable: $NVSHMRUN_BIN" >&2 + exit 2 +fi +NVSHMRUN_BIN="$(cd -- "$(dirname -- "$NVSHMRUN_PATH")" && pwd)/$(basename -- "$NVSHMRUN_PATH")" + +if [[ -n "$NVSHMEM_HOME_ARG" ]]; then + NVSHMEM_HOME_ARG="$(cd -- "$NVSHMEM_HOME_ARG" && pwd)" +fi + +PREPARE_ARGS=( + "$SCRIPT_DIR/common/prepare_nvshmem.py" + "$TARGET" + "--nvshmem-home" + "$NVSHMEM_HOME_ARG" +) +[[ -n "$ARCH" ]] && PREPARE_ARGS+=("--arch" "$ARCH") +((FORCE)) && PREPARE_ARGS+=("--force") + +"$PYTHON_BIN" "${PREPARE_ARGS[@]}" + +export NVSHMEM_HOME="$NVSHMEM_HOME_ARG" +export LD_LIBRARY_PATH="$NVSHMEM_HOME/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" + +cd -- "$(dirname -- "$TARGET")" +exec "$NVSHMRUN_BIN" -np "$NP" "$PYTHON_BIN" "$(basename -- "$TARGET")" "$@" diff --git a/support_nvshmem.patch b/support_nvshmem.patch deleted file mode 100644 index f6d472e0d..000000000 --- a/support_nvshmem.patch +++ /dev/null @@ -1,697 +0,0 @@ -diff --git a/include/triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h b/include/triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h -index 680bf0e04..69102ed61 100644 ---- a/include/triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h -+++ b/include/triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h -@@ -102,6 +102,11 @@ void populatePrintOpToLLVMPattern(LLVMTypeConverter &typeConverter, - const TargetInfoBase &targetInfo, - PatternBenefit benefit); - -+void populateExternCallOpToLLVMPattern(LLVMTypeConverter &typeConverter, -+ RewritePatternSet &patterns, -+ const TargetInfoBase &targetInfo, -+ PatternBenefit benefit); -+ - void populateInstrumentationToLLVMPatterns(LLVMTypeConverter &typeConverter, - const TargetInfoBase &targetInfo, - RewritePatternSet &patterns, -diff --git a/include/triton/Dialect/Triton/IR/TritonOps.td b/include/triton/Dialect/Triton/IR/TritonOps.td -index 7fd215f9d..95239b3cb 100644 ---- a/include/triton/Dialect/Triton/IR/TritonOps.td -+++ b/include/triton/Dialect/Triton/IR/TritonOps.td -@@ -845,6 +845,32 @@ def TT_MapElementwiseReturnOp: TT_Op<"map_elementwise.return", - let assemblyFormat = "attr-dict ($result^ `:` type($result))?"; - } - -+// -+// External Call op -+// -+def TT_ExternCallOp : TT_Op<"extern_call", [ -+ DeclareOpInterfaceMethods, -+ ConditionallySpeculatable, -+]> { -+ -+ let description = [{ -+ call an external function $symbol implemented in $libpath/$libname with $args -+ return $libpath/$libname:$symbol($args...) -+ }]; -+ -+ let arguments = (ins Variadic:$srcs, StrAttr:$libname, StrAttr:$libpath, StrAttr:$symbol, BoolAttr:$pure); -+ -+ let results = (outs Variadic:$result); -+ -+ let assemblyFormat = "operands attr-dict `:` functional-type(operands, $result)"; -+ -+ let extraClassDeclaration = [{ -+ // Interface method for ConditionallySpeculatable. -+ Speculation::Speculatability getSpeculatability(); -+ }]; -+ -+} -+ - // - // External Elementwise op - // -diff --git a/lib/Conversion/TritonGPUToLLVM/CMakeLists.txt b/lib/Conversion/TritonGPUToLLVM/CMakeLists.txt -index d4f49c8d1..44e52cf37 100644 ---- a/lib/Conversion/TritonGPUToLLVM/CMakeLists.txt -+++ b/lib/Conversion/TritonGPUToLLVM/CMakeLists.txt -@@ -7,6 +7,7 @@ add_triton_library(TritonGPUToLLVM - AssertOpToLLVM.cpp - ControlFlowOpToLLVM.cpp - ConvertLayoutOpToLLVM.cpp -+ ExternCallOpToLLVM.cpp - ElementwiseOpToLLVM.cpp - FuncOpToLLVM.cpp - GatherOpToLLVM.cpp -diff --git a/lib/Conversion/TritonGPUToLLVM/ExternCallOpToLLVM.cpp b/lib/Conversion/TritonGPUToLLVM/ExternCallOpToLLVM.cpp -new file mode 100644 -index 000000000..81b6bf256 ---- /dev/null -+++ b/lib/Conversion/TritonGPUToLLVM/ExternCallOpToLLVM.cpp -@@ -0,0 +1,61 @@ -+#include "mlir/Conversion/LLVMCommon/Pattern.h" -+#include "mlir/Dialect/GPU/IR/GPUDialect.h" -+#include "mlir/Support/LLVM.h" -+#include "triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h" -+#include "triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h" -+#include "triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h" -+#include "triton/Conversion/TritonGPUToLLVM/Utility.h" -+#include "triton/Dialect/TritonGPU/IR/Dialect.h" -+ -+namespace { -+ -+class ExternCallOpConversion -+ : public ConvertOpToLLVMPattern { -+public: -+ ExternCallOpConversion(const LLVMTypeConverter &converter, -+ const PatternBenefit &benefit) -+ : ConvertOpToLLVMPattern(converter, benefit) {} -+ -+ LogicalResult -+ matchAndRewrite(triton::ExternCallOp op, OpAdaptor adaptor, -+ ConversionPatternRewriter &rewriter) const override { -+ auto loc = op->getLoc(); -+ -+ if (op->getNumResults() > 1) { -+ llvm::errs() << "ExternCallConversion does not support multi outs."; -+ return failure(); -+ } -+ -+ LLVM::LLVMVoidType voidTy = void_ty(op->getContext()); -+ auto newOperands = adaptor.getOperands(); -+ Type retType = -+ op->getNumResults() == 0 -+ ? voidTy -+ : this->getTypeConverter()->convertType(op->getResult(0).getType()); -+ std::string funcName = op.getSymbol().str(); -+ StringRef libname = op.getLibname(); -+ StringRef libpath = op.getLibpath(); -+ -+ Operation *externCallOp; -+ Type funcType = mlir::triton::gpu::getFunctionType(retType, newOperands); -+ LLVM::LLVMFuncOp funcOp = mlir::triton::gpu::appendOrGetExternFuncOp( -+ rewriter, op, funcName, funcType, libname, libpath); -+ externCallOp = LLVM::createLLVMCallOp(rewriter, loc, funcOp, newOperands); -+ -+ if (op->getNumResults() == 0) { -+ rewriter.eraseOp(op); -+ } else { -+ rewriter.replaceOp(op, externCallOp->getResult(0)); -+ } -+ -+ return success(); -+ } -+}; -+ -+} // namespace -+ -+void mlir::triton::populateExternCallOpToLLVMPattern( -+ LLVMTypeConverter &typeConverter, RewritePatternSet &patterns, -+ const TargetInfoBase &targetInfo, PatternBenefit benefit) { -+ patterns.add(typeConverter, benefit); -+} -diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp -index b49871486..45c613465 100644 ---- a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp -+++ b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp -@@ -600,6 +600,7 @@ void populateTritonPatterns(TritonGPUTypeConverter &typeConverter, - GenericOpPattern, - GenericOpPattern, - GenericOpPattern, -+ GenericOpPattern, - GenericOpPattern, - GenericOpPattern, - GenericOpPattern, -diff --git a/lib/Dialect/Triton/IR/Ops.cpp b/lib/Dialect/Triton/IR/Ops.cpp -index 27fa26554..854ae1f2b 100644 ---- a/lib/Dialect/Triton/IR/Ops.cpp -+++ b/lib/Dialect/Triton/IR/Ops.cpp -@@ -1313,6 +1313,24 @@ Speculation::Speculatability ExternElementwiseOp::getSpeculatability() { - return Speculation::NotSpeculatable; - } - -+// -- ExternCallOp -- -+void ExternCallOp::getEffects( -+ SmallVectorImpl> -+ &effects) { -+ if (getPure()) -+ return; -+ effects.emplace_back(MemoryEffects::Write::get(), -+ SideEffects::DefaultResource::get()); -+ effects.emplace_back(MemoryEffects::Read::get(), -+ SideEffects::DefaultResource::get()); -+} -+ -+Speculation::Speculatability ExternCallOp::getSpeculatability() { -+ if (getPure()) -+ return Speculation::Speculatable; -+ return Speculation::NotSpeculatable; -+} -+ - // -- GatherOp -- - LogicalResult GatherOp::verify() { - RankedTensorType indicesTy = getIndices().getType(); -diff --git a/python/src/ir.cc b/python/src/ir.cc -index e5591dcb0..9abfe8472 100644 ---- a/python/src/ir.cc -+++ b/python/src/ir.cc -@@ -1673,6 +1673,14 @@ void init_triton_ir(py::module &&m) { - return self.create(retType, argList, libName, - libPath, symbol, isPure); - }) -+ .def("create_extern_call", -+ [](TritonOpBuilder &self, const std::string &libName, -+ const std::string &libPath, const std::string &symbol, -+ std::vector &argList, const std::vector &retTypes, -+ bool isPure) -> OpState { -+ return self.create(retTypes, argList, libName, -+ libPath, symbol, isPure); -+ }) - // Built-in instruction - .def("create_get_program_id", - [](TritonOpBuilder &self, int axis) -> Value { -diff --git a/python/triton/experimental/tle/language/raw/__init__.py b/python/triton/experimental/tle/language/raw/__init__.py -index d66f6cd38..790c65cf3 100644 ---- a/python/triton/experimental/tle/language/raw/__init__.py -+++ b/python/triton/experimental/tle/language/raw/__init__.py -@@ -1,3 +1,3 @@ --from .core import call, call_smem -+from .core import call, call_smem, call_nvshmem - --__all__ = ["call", "call_smem"] -+__all__ = ["call", "call_smem", "call_nvshmem"] -diff --git a/python/triton/experimental/tle/language/raw/core.py b/python/triton/experimental/tle/language/raw/core.py -index d7029380e..7e2538531 100644 ---- a/python/triton/experimental/tle/language/raw/core.py -+++ b/python/triton/experimental/tle/language/raw/core.py -@@ -46,3 +46,8 @@ def call_smem(func, args, _semantic=None): - return buffer_tensors[0] - else: - return tl.tuple(buffer_tensors) -+ -+ -+@builtin -+def call_nvshmem(func, outputs, inputs, _semantic=None): -+ func.make_cubin() -diff --git a/python/triton/experimental/tle/raw/cuda/runtime.py b/python/triton/experimental/tle/raw/cuda/runtime.py -index 9015f2bf7..955a366e6 100644 ---- a/python/triton/experimental/tle/raw/cuda/runtime.py -+++ b/python/triton/experimental/tle/raw/cuda/runtime.py -@@ -3,20 +3,24 @@ import os - from pathlib import Path - import subprocess - from typing import Any, Final -+import torch - - from triton._C.libtriton import llvm # pyright: ignore[reportMissingImports] - from triton._C.libtriton.tle.llvm import parse_llvm_ir # pyright: ignore[reportMissingImports] - - # 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") -+NVCC = os.getenv("NVCC", "nvcc") - - - class CUDAJITFunction(object): - - def __init__(self, fn: Any, file: Path, *args, **kwargs) -> None: -- super().__init__(*args, **kwargs) -+ super().__init__() - self.fn: Final[Any] = fn - self.code: Final[str] = file.read_text() -+ self.file: Final[Path] = file -+ self.libs = kwargs.get("library", {}) - self.__triton_builtin__: Final[bool] = True - - def make_llvm(self, mlir_context) -> str: -@@ -40,3 +44,22 @@ class CUDAJITFunction(object): - llvm_context = llvm.context() - module = parse_llvm_ir(build.stdout.decode(), llvm_context, mlir_context) - return f"{module}" -+ -+ def make_cubin(self): -+ src = self.file -+ dst = Path(src).with_suffix('.o') -+ include_dirs = [] -+ for lib_name, lib_path in self.libs.items(): -+ # TODO: Remove the method of passing information by setting environment variables. -+ os.environ[(lib_name + "_home").upper()] = lib_path -+ include_dirs.append(os.path.join(lib_path, "include")) -+ include_flags = [f"-I{inc_dir}" for inc_dir in include_dirs] -+ prop = torch.cuda.get_device_properties(torch.cuda.current_device()) -+ arch = f"-arch=sm_{prop.major}{prop.minor}" -+ build = subprocess.run([NVCC, "-rdc=true", arch, *include_flags, "--extended-lambda", "-c", "-o", dst, src], -+ capture_output=True) -+ assert build.returncode == 0, (f"nvcc failed\nstderr:\n{build.stderr.decode()}") -+ # TODO: Remove the method of passing information by setting environment variables. -+ os.environ["USE_NVCC"] = 'True' -+ os.environ["CUDA_CUBIN"] = str(dst) -+ return -diff --git a/python/triton/language/core.py b/python/triton/language/core.py -index 439671c3a..5a9c56b12 100644 ---- a/python/triton/language/core.py -+++ b/python/triton/language/core.py -@@ -3433,6 +3433,94 @@ def binary_op_type_legalization(lhs, rhs, semantic): - return semantic.binary_op_type_checking_impl(lhs, rhs) - - -+def dispatch_ec(func, lib_name: str, lib_path: str, args: list, arg_type_symbol_dict: dict, is_pure: bool, -+ _semantic=None): -+ ''' -+ Dispatch a function to a library -+ :param func: the function to dispatch -+ :param lib_name: the name of the library -+ :param lib_path: the path of the library -+ :param args: the arguments of the function -+ :param arg_type_symbol_dict: the type of the arguments -+ :param ret_shape: the shape of the return value -+ :param _semantic: the builder -+ :return: the return value of the function -+ ''' -+ if len(arg_type_symbol_dict) == 0: -+ raise ValueError("arg_type_symbol_dict is empty") -+ -+ num_args = len(list(arg_type_symbol_dict.keys())[0]) -+ if len(args) != num_args: -+ raise ValueError(f"length of input args does not match." -+ f"Expect {len(args)}, got {num_args}") -+ -+ arg_types = [] -+ arg_list = [] -+ for arg in args: -+ if isinstance(arg, tensor): -+ arg_types.append(arg.dtype) -+ arg_list.append(arg.handle) -+ else: -+ arg_types.append(type(arg)) -+ arg_list.append(arg) -+ arg_types = tuple(arg_types) -+ -+ if arg_types not in arg_type_symbol_dict: -+ raise ValueError(f"input arg type does not match." -+ f"Expect one of {arg_type_symbol_dict.keys()}, got {arg_types}") -+ else: -+ symbol = arg_type_symbol_dict[arg_types][0] -+ ret_types = arg_type_symbol_dict[arg_types][1] -+ if not isinstance(ret_types, (builtins.list, builtins.tuple)): -+ ret_types = [ret_types] -+ -+ if symbol == "": -+ raise ValueError("Symbol can not be empty") -+ call = func(lib_name, lib_path, symbol, arg_list, [ret_type.to_ir(_semantic.builder) for ret_type in ret_types], -+ is_pure) -+ -+ if len(ret_types) == 0: -+ return tensor(call, void) -+ if len(ret_types) == 1: -+ return tensor(call.get_result(0), ret_types[0]) -+ return tuple(tensor(call.get_result(i), ty) for i, ty in enumerate(ret_types)) -+ -+ -+@builtin -+def extern_call(lib_name: str, lib_path: str, args: list, arg_type_symbol_dict: dict, is_pure: bool, _semantic=None): -+ ''' -+ Dispatch an function to a library -+ :param lib_name: the name of the library -+ :param lib_path: the path of the library -+ :param args: the arguments of the function -+ :param arg_type_symbol_dict: the type of the arguments -+ :param is_pure: whether the function is pure -+ :param _semantic: the semantic -+ :return: the return value of the function -+ ''' -+ dispatch_args = args.copy() -+ all_scalar = True -+ arg_types = [] -+ for i in builtins.range(len(dispatch_args)): -+ dispatch_args[i] = _semantic.to_tensor(dispatch_args[i]) -+ arg_types.append(dispatch_args[i].dtype) -+ if dispatch_args[i].type.is_block(): -+ all_scalar = False -+ if not all_scalar: -+ raise ValueError("extern call only support inputs with scalr type") -+ -+ if len(arg_type_symbol_dict) == 0: -+ raise ValueError("arg_type_symbol_dict is empty") -+ -+ num_args = len(list(arg_type_symbol_dict.keys())[0]) -+ if len(args) != num_args: -+ raise ValueError(f"length of input args does not match." -+ f"Expect {len(args)}, got {num_args}") -+ -+ func = _semantic.builder.create_extern_call -+ return dispatch_ec(func, lib_name, lib_path, dispatch_args, arg_type_symbol_dict, is_pure, _semantic) -+ -+ - def extern(fn): - """A decorator for external functions.""" - return builtin(fn) -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..f2910883b ---- /dev/null -+++ b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-device.cu -@@ -0,0 +1,8 @@ -+#include -+ -+extern "C" __device__ void simple_shift(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..be9ea236a ---- /dev/null -+++ b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift-host.cu -@@ -0,0 +1,31 @@ -+#include -+#include -+#include -+#include -+#include -+ -+extern "C" void nvshmem_init_wrapper() { nvshmem_init(); } -+ -+extern "C" int nvshmemx_cumodule_init_wrapper(CUmodule module) { -+ return nvshmemx_cumodule_init(module); -+} -+ -+extern "C" int nvshmem_team_mype_wrapper() { -+ int mype_node = nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE); -+ return mype_node; -+} -+ -+// TODO: Adapt to different data types -+extern "C" int *nvshmem_alloc_wrapper(int size) { -+ int *destination = (int *)nvshmem_malloc(sizeof(int) * size); -+ return destination; -+} -+ -+extern "C" void nvshmemx_barrier_warpper(cudaStream_t stream) { -+ nvshmemx_barrier_all_on_stream(stream); -+} -+ -+extern "C" void nvshmem_finalize_wrapper(int *dest) { -+ nvshmem_free(dest); -+ 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..2f87d70f6 ---- /dev/null -+++ b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py -@@ -0,0 +1,122 @@ -+import os -+import subprocess -+import ctypes -+import torch -+import triton -+import triton.knobs as knobs -+import triton.experimental.tle.language.raw as tle_raw -+ -+from pathlib import Path -+from triton.experimental.tle.raw import dialect -+from triton.language.extra.cuda import libnvshmem_device -+ -+ -+@dialect( -+ name="cuda", -+ file=(Path(__file__).parent / "simple-shift-device.cu").resolve(), -+ library={"nvshmem": "/home/zyl/zyuli/envs/nvshmem/lib/python3.12/site-packages/nvidia/nvshmem"}, -+) -+def edsl(*args, **kwargs): -+ ... -+ -+ -+@triton.jit -+def simple_shift_kernel(destination_ptr, ): -+ # TODO: Combine with tle_raw.call, then dispatch -+ tle_raw.call_nvshmem(edsl, [], [destination_ptr]) -+ libnvshmem_device.simple_shift(destination_ptr) -+ -+ -+def cuda_host_compile(cuda_host_path, cuda_host_lib): -+ NVCC = os.getenv("NVCC", "nvcc") -+ NVSHMEM_HOME = "/home/zyl/zyuli/envs/nvshmem/lib/python3.12/site-packages/nvidia/nvshmem" -+ include_path = f"-I{os.path.join(NVSHMEM_HOME, 'include')}" -+ lib_path = f"-L{os.path.join(NVSHMEM_HOME, 'lib')}" -+ -+ prop = torch.cuda.get_device_properties(torch.cuda.current_device()) -+ arch = f"-arch=sm_{prop.major}{prop.minor}" -+ tmp_file = Path(cuda_host_lib).with_suffix('.so.tmp') -+ build = [ -+ NVCC, "-shared", "-Xcompiler", "-fPIC", "-rdc=true", arch, include_path, lib_path, "-lnvshmem_host", -+ "-lnvshmem_device", "-o", tmp_file, cuda_host_path -+ ] -+ build = subprocess.run(build, capture_output=True) -+ assert build.returncode == 0, (f"NVCC host failed\nstderr:\n{build.stderr.decode()}") -+ tmp_file.rename(cuda_host_lib) -+ -+ -+def simpe_shift(): -+ cu_file = (Path(__file__).parent / "simple-shift-host.cu").resolve() -+ lib_file = Path(cu_file).with_suffix('.so') -+ -+ rank = int(os.getenv("OMPI_COMM_WORLD_RANK", "0")) -+ if rank == 0: -+ cuda_host_compile(cu_file, lib_file) -+ -+ import time -+ timeout = 60 -+ start = time.time() -+ while True: -+ if lib_file.exists(): -+ try: -+ ctypes.CDLL(str(lib_file)) -+ break -+ except OSError: -+ pass -+ if time.time() - start > timeout: -+ raise RuntimeError(f"Timeout waiting for {lib_file}") -+ time.sleep(0.1) -+ -+ lib = ctypes.CDLL(lib_file) -+ lib.nvshmem_init_wrapper.argtypes = [] -+ lib.nvshmem_init_wrapper.restype = None -+ lib.nvshmemx_cumodule_init_wrapper.argtypes = [ctypes.c_void_p] -+ lib.nvshmemx_cumodule_init_wrapper.restype = ctypes.c_int -+ lib.nvshmem_team_mype_wrapper.argtypes = [] -+ lib.nvshmem_team_mype_wrapper.restype = ctypes.c_int -+ lib.nvshmem_alloc_wrapper.argtypes = [ctypes.c_int] -+ lib.nvshmem_alloc_wrapper.restype = ctypes.POINTER(ctypes.c_int) -+ lib.nvshmemx_barrier_warpper.argtypes = [ctypes.c_void_p] -+ lib.nvshmemx_barrier_warpper.restype = None -+ lib.nvshmem_finalize_wrapper.argtypes = [ctypes.POINTER(ctypes.c_int)] -+ lib.nvshmem_finalize_wrapper.restype = None -+ -+ lib.nvshmem_init_wrapper() -+ mype_node = lib.nvshmem_team_mype_wrapper() -+ torch.cuda.set_device(mype_node) -+ device = triton.runtime.driver.active.get_active_torch_device() -+ stream = torch.cuda.Stream() -+ -+ dest = lib.nvshmem_alloc_wrapper(1) -+ dest_addr = ctypes.cast(dest, ctypes.c_void_p).value -+ storage = torch._C._construct_storage_from_data_pointer(dest_addr, device, 4) -+ dest_tensor = torch.empty(0, dtype=torch.int32, device=device).set_(storage).view(1) -+ msg = torch.empty((1, ), dtype=torch.int32, pin_memory=True) -+ -+ def cumodule_init_hook(*args, **kwargs): -+ key = kwargs["key"] -+ jit_function = kwargs["fn"].jit_function -+ device = kwargs["compile"]["device"] -+ kernel_cache = jit_function.device_caches[device][0] -+ kernel = kernel_cache.get(key, None) -+ assert kernel is not None -+ kernel._init_handles() -+ ret = lib.nvshmemx_cumodule_init_wrapper(ctypes.c_void_p(kernel.module)) -+ assert ret == 0, f"nvshmemx_cumodule_init_wrapper failed: {ret}" -+ -+ knobs.runtime.jit_post_compile_hook = cumodule_init_hook -+ -+ simple_shift_kernel[(1, )](dest_tensor) -+ -+ stream_ptr = stream.cuda_stream -+ lib.nvshmemx_barrier_warpper(ctypes.c_void_p(stream_ptr)) -+ with torch.cuda.stream(stream): -+ msg.copy_(dest_tensor, non_blocking=True) -+ stream.synchronize() -+ -+ lib.nvshmem_finalize_wrapper(dest) -+ print(f"Rank {mype_node}: {msg}") -+ -+ -+if __name__ == "__main__": -+ simpe_shift() -diff --git a/third_party/nvidia/backend/compiler.py b/third_party/nvidia/backend/compiler.py -index 632b79075..482af7b32 100644 ---- a/third_party/nvidia/backend/compiler.py -+++ b/third_party/nvidia/backend/compiler.py -@@ -498,7 +498,7 @@ class CUDABackend(BaseBackend): - assert len(names) == 1 - metadata["name"] = names[0] - # post-process -- ptx_version = f'{ptx_version//10}.{ptx_version%10}' -+ ptx_version = '8.5' - ret = re.sub(r'\.version \d+\.\d+', f'.version {ptx_version}', ret, flags=re.MULTILINE) - ret = re.sub(r'\.target sm_\d+', f'.target sm_{capability}', ret, flags=re.MULTILINE) - if not knobs.compilation.dump_ir_extract_di_local_variables: -@@ -512,13 +512,15 @@ class CUDABackend(BaseBackend): - return ret - - def make_cubin(self, src, metadata, opt, capability): -- ptxas = get_ptxas(self.target.arch).path -+ ptxas = os.getenv("ptxas", "ptxas") - with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.ptx') as fsrc, \ - tempfile.NamedTemporaryFile(delete=False, mode='r', suffix='.log') as flog: - fsrc.write(src) - fsrc.flush() - fbin = fsrc.name + '.o' - -+ use_nvcc = os.getenv("USE_NVCC", '').lower() in ('1', 'true') -+ os.environ.pop("USE_NVCC", None) - debug_info = [] - if knobs.compilation.disable_line_info: - # This option is ignored if used without -lineinfo -@@ -539,9 +541,11 @@ class CUDABackend(BaseBackend): - # Accept more ptxas options if provided - ptx_extra_options = opt.ptx_options.split(" ") if opt.ptx_options else [] - -+ # If use nvshmem, we need to compile the ptx file into a relocatable object file and then link it with nvshmem library -+ compile_only = ["-c"] if use_nvcc else [] - ptxas_cmd = [ -- ptxas, *debug_info, *fmad, '-v', *disable_opt, *ptx_extra_options, f'--gpu-name={arch}', fsrc.name, -- '-o', fbin -+ ptxas, *compile_only, *debug_info, *fmad, '-v', *disable_opt, *ptx_extra_options, f'--gpu-name={arch}', -+ fsrc.name, '-o', fbin - ] - try: - subprocess.run(ptxas_cmd, check=True, close_fds=False, stderr=flog) -@@ -581,8 +585,37 @@ please share the reproducer above with Triton project. - """) - raise PTXASError(error) - -- with open(fbin, 'rb') as f: -- cubin = f.read() -+ if use_nvcc: -+ NVLINK = os.getenv("NVLINK", "nvlink") -+ NVSHMEM_HOME = os.getenv("NVSHMEM_HOME") -+ fbin_combined = fbin + ".combined.cubin" -+ cuda_cubin = os.getenv("CUDA_CUBIN") -+ nvshmem_lib = os.path.join(NVSHMEM_HOME, "lib") -+ nvlink_cmds = [ -+ NVLINK, -+ f"-arch={arch}", -+ f"-L{nvshmem_lib}", -+ "-lnvshmem_device", -+ fbin, -+ cuda_cubin, -+ "-o", -+ fbin_combined, -+ ] -+ try: -+ subprocess.run(nvlink_cmds, check=True, close_fds=False, stderr=flog) -+ except Exception as e: -+ import logging -+ logging.error(f"error runing nvlink: {nvlink_cmds}") -+ logging.exception(e) -+ -+ if use_nvcc: -+ with open(fbin_combined, 'rb') as f: -+ cubin = f.read() -+ else: -+ with open(fbin, 'rb') as f: -+ cubin = f.read() -+ if os.path.exists(fbin_combined): -+ os.remove(fbin_combined) - if os.path.exists(fbin): - os.remove(fbin) - return cubin -diff --git a/third_party/nvidia/language/cuda/__init__.py b/third_party/nvidia/language/cuda/__init__.py -index fbececf1d..39207bf60 100644 ---- a/third_party/nvidia/language/cuda/__init__.py -+++ b/third_party/nvidia/language/cuda/__init__.py -@@ -1,10 +1,12 @@ - from . import libdevice -+from . import libnvshmem_device - - from .utils import (globaltimer, num_threads, num_warps, smid, convert_custom_float8_sm70, convert_custom_float8_sm80) - from .gdc import (gdc_launch_dependents, gdc_wait) - - __all__ = [ - "libdevice", -+ 'libnvshmem_device', - "globaltimer", - "num_threads", - "num_warps", -diff --git a/third_party/nvidia/language/cuda/libnvshmem_device.py b/third_party/nvidia/language/cuda/libnvshmem_device.py -new file mode 100644 -index 000000000..3dec30b9a ---- /dev/null -+++ b/third_party/nvidia/language/cuda/libnvshmem_device.py -@@ -0,0 +1,31 @@ -+from triton.language import core -+import triton.language as tl -+ -+ -+def _pointer_type_hash(self): -+ return hash((self.name, self.element_ty, "tt_ptr")) -+ -+ -+def patch_hash_method_for_pointer_type(): -+ elem_dtype_list = tl.core.dtype.SINT_TYPES + tl.core.dtype.UINT_TYPES + tl.core.dtype.FP_TYPES + tl.core.dtype.OTHER_TYPES -+ for elem_dtype in elem_dtype_list: -+ ptr_ty = type(tl.core.pointer_type(tl.core.dtype(elem_dtype))) -+ ptr_ty.__hash__ = _pointer_type_hash -+ -+ -+patch_hash_method_for_pointer_type() -+ -+ -+@core.extern -+def simple_shift(dst, _semantic=None): -+ return core.extern_call( -+ "", # libname -+ "", # libpath -+ [dst], # args -+ {( -+ core.pointer_type(core.dtype("int32")), # arg_type_symbol_dict -+ ): ("simple_shift", ()), # void return type -+ }, -+ is_pure=False, -+ _semantic=_semantic, -+ ) -diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp -index 368a66629..0fc8242c5 100644 ---- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp -+++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp -@@ -222,6 +222,8 @@ struct ConvertTritonGPUToLLVM - targetInfo, benefit); - mlir::triton::populatePrintOpToLLVMPattern(typeConverter, patterns, - targetInfo, benefit); -+ mlir::triton::populateExternCallOpToLLVMPattern(typeConverter, patterns, -+ targetInfo, benefit); - mlir::triton::populateControlFlowOpToLLVMPattern(typeConverter, patterns, - targetInfo, benefit); - mlir::triton::NVIDIA::populateSPMDOpToLLVMPattern(typeConverter, patterns, diff --git a/third_party/nvidia/backend/compiler.py b/third_party/nvidia/backend/compiler.py index d83616896..d6ea87518 100644 --- a/third_party/nvidia/backend/compiler.py +++ b/third_party/nvidia/backend/compiler.py @@ -533,8 +533,6 @@ def make_cubin(self, src, metadata, opt, capability): fsrc.flush() fbin = fsrc.name + '.o' - use_nvcc = os.getenv("USE_NVCC", '').lower() in ('1', 'true') - os.environ.pop("USE_NVCC", None) debug_info = [] if knobs.compilation.disable_line_info: # This option is ignored if used without -lineinfo @@ -555,11 +553,9 @@ def make_cubin(self, src, metadata, opt, capability): # Accept more ptxas options if provided ptx_extra_options = opt.ptx_options.split(" ") if opt.ptx_options else [] - # If use nvshmem, we need to compile the ptx file into a relocatable object file and then link it with nvshmem library - compile_only = ["-c"] if use_nvcc else [] ptxas_cmd = [ - ptxas, *compile_only, *debug_info, *fmad, '-v', *disable_opt, *ptx_extra_options, f'--gpu-name={arch}', - fsrc.name, '-o', fbin + ptxas, *debug_info, *fmad, '-v', *disable_opt, *ptx_extra_options, f'--gpu-name={arch}', fsrc.name, + '-o', fbin ] try: subprocess.run(ptxas_cmd, check=True, close_fds=False, stderr=flog) @@ -599,37 +595,8 @@ def make_cubin(self, src, metadata, opt, capability): """) raise PTXASError(error) - if use_nvcc: - NVLINK = os.getenv("NVLINK", "nvlink") - NVSHMEM_HOME = os.getenv("NVSHMEM_HOME") - fbin_combined = fbin + ".combined.cubin" - cuda_cubin = os.getenv("CUDA_CUBIN") - nvshmem_lib = os.path.join(NVSHMEM_HOME, "lib") - nvlink_cmds = [ - NVLINK, - f"-arch={arch}", - f"-L{nvshmem_lib}", - "-lnvshmem_device", - fbin, - cuda_cubin, - "-o", - fbin_combined, - ] - try: - subprocess.run(nvlink_cmds, check=True, close_fds=False, stderr=flog) - except Exception as e: - import logging - logging.error(f"error runing nvlink: {nvlink_cmds}") - logging.exception(e) - - if use_nvcc: - with open(fbin_combined, 'rb') as f: - cubin = f.read() - if os.path.exists(fbin_combined): - os.remove(fbin_combined) - else: - with open(fbin, 'rb') as f: - cubin = f.read() + with open(fbin, 'rb') as f: + cubin = f.read() if os.path.exists(fbin): os.remove(fbin) return cubin diff --git a/third_party/nvidia/language/cuda/__init__.py b/third_party/nvidia/language/cuda/__init__.py index 39207bf60..fbececf1d 100644 --- a/third_party/nvidia/language/cuda/__init__.py +++ b/third_party/nvidia/language/cuda/__init__.py @@ -1,12 +1,10 @@ from . import libdevice -from . import libnvshmem_device from .utils import (globaltimer, num_threads, num_warps, smid, convert_custom_float8_sm70, convert_custom_float8_sm80) from .gdc import (gdc_launch_dependents, gdc_wait) __all__ = [ "libdevice", - 'libnvshmem_device', "globaltimer", "num_threads", "num_warps", diff --git a/third_party/nvidia/language/cuda/libnvshmem_device.py b/third_party/nvidia/language/cuda/libnvshmem_device.py deleted file mode 100644 index c184c0dd7..000000000 --- a/third_party/nvidia/language/cuda/libnvshmem_device.py +++ /dev/null @@ -1,229 +0,0 @@ -from triton.language import core -import triton.language as tl - - -def _pointer_type_hash(self): - return hash((self.name, self.element_ty, "tt_ptr")) - - -def patch_hash_method_for_pointer_type(): - elem_dtype_list = tl.core.dtype.SINT_TYPES + tl.core.dtype.UINT_TYPES + tl.core.dtype.FP_TYPES + tl.core.dtype.OTHER_TYPES - for elem_dtype in elem_dtype_list: - ptr_ty = type(tl.core.pointer_type(tl.core.dtype(elem_dtype))) - ptr_ty.__hash__ = _pointer_type_hash - - -patch_hash_method_for_pointer_type() - -# 01-simple-shift -@core.extern -def simple_shift(dst, _semantic=None): - return core.extern_call( - "", # libname - "", # libpath - [dst], # args - { - ( - core.pointer_type(core.dtype("int32")), # arg_type_symbol_dict - ): ("simple_shift", ()), # void return type - }, - is_pure=False, - _semantic=_semantic, - ) - -# 02-ring-reduce -@core.extern -def ring_reduce(dst, src, nreduce, signal, chunk_size, _semantic=None): - return core.extern_call( - "", - "", - [dst, - src, - tl.cast(nreduce, tl.int32, _semantic=_semantic), - signal, - tl.cast(chunk_size, tl.int32, _semantic=_semantic)], - { - ( - core.pointer_type(core.dtype("int32")), - core.pointer_type(core.dtype("int32")), - core.dtype("int32"), - core.pointer_type(core.dtype("int64")), - core.dtype("int32"), - ) : ("ring_reduce", ()), - }, - is_pure=False, - _semantic=_semantic, - ) - -# 03-ring-bcast -@core.extern -def ring_bcast(data, nelem, root, psync, _semantic=None): - return core.extern_call( - "", - "", - [data, nelem, root, psync], - { - ( - core.pointer_type(core.dtype("int32")), - core.dtype("int32"), - core.dtype("int32"), - core.pointer_type(core.dtype("uint64")), - ) : ("ring_bcast", ()), - }, - is_pure=False, - _semantic=_semantic, - ) - -# 04-on-stream -@core.extern -def accumulate(input, partial_sum, _semantic=None): - return core.extern_call( - "", # libname - "", # libpath - [input, partial_sum], - { - ( - core.pointer_type(core.dtype("int32")), - core.pointer_type(core.dtype("int32")), - ): ("accumulate", ()), - }, - is_pure=False, - _semantic=_semantic, - ) - -@core.extern -def correct_accumulate(input, partial_sum, full_sum, _semantic=None): - return core.extern_call( - "", # libname - "", # libpath - [input, partial_sum, full_sum], - { - ( - core.pointer_type(core.dtype("int32")), - core.pointer_type(core.dtype("int32")), - core.pointer_type(core.dtype("int32")), - ): ("correct_accumulate", ()), - }, - is_pure=False, - _semantic=_semantic, - ) - -# 05-put-block -@core.extern -def set_and_shift(send_data, recv_data, num_elems, mype, npes, _semantic=None): - return core.extern_call( - "", # libname - "", # libpath - [send_data, recv_data, num_elems, mype, npes], - { - ( - core.pointer_type(core.dtype("fp32")), - core.pointer_type(core.dtype("fp32")), - core.dtype("int32"), - core.dtype("int32"), - core.dtype("int32"), - ): ("set_and_shift", ()), - }, - is_pure=False, - _semantic=_semantic, - ) - -# 00-gemm-allreduce -@core.extern -def tiled_gemm(C, A, B, m, n, k, _semantic=None): - return core.extern_call( - "", - "", - [ - C, A, B, - m, n, k - ], - { - ( - core.pointer_type(core.dtype("fp32")), - core.pointer_type(core.dtype("fp32")), - core.pointer_type(core.dtype("fp32")), - core.dtype("int32"), - core.dtype("int32"), - core.dtype("int32"), - ) : ("tiled_gemm", ()), - }, - is_pure=False, - _semantic=_semantic, - ) - -@core.extern -def per_token_group_quant_8bit(x_q_ptr, x_s_ptr, x_ptr, group_size, num_groups, groups_per_block, eps, fp8_min, fp8_max, - _semantic=None): - return core.extern_call( - "", - "", - [ - tl.cast(x_q_ptr, tl.pointer_type(core.dtype("void")), _semantic=_semantic), - x_s_ptr, - x_ptr, - group_size, - num_groups, - groups_per_block, - eps, - fp8_min, - fp8_max - ], - { - ( - core.pointer_type(core.dtype("void")), - core.pointer_type(core.dtype("fp32")), - core.pointer_type(core.dtype("fp32")), - core.dtype("int32"), - core.dtype("int32"), - core.dtype("int32"), - core.dtype("fp32"), - core.dtype("fp32"), - core.dtype("fp32"), - ): ("per_token_group_quant_8bit", ()), - }, - is_pure=False, - _semantic=_semantic, - ) - -@core.extern -def vector_add(C, A, B, N, _semantic=None): - return core.extern_call( - "", - "", - [C, A, B, N], - { - ( - core.pointer_type(core.dtype("int32")), - core.pointer_type(core.dtype("int32")), - core.pointer_type(core.dtype("int32")), - core.dtype("int32"), - ) : ("vector_add_int", ()), - ( - core.pointer_type(core.dtype("fp32")), - core.pointer_type(core.dtype("fp32")), - core.pointer_type(core.dtype("fp32")), - core.dtype("int32"), - ) : ("vector_add_float", ()), - }, - is_pure=False, - _semantic=_semantic, - ) - -@core.extern -def vector_add_return(C, A, B, N, _semantic=None): - return core.extern_call( - "", - "", - [C, A, B, N], - { - ( - core.pointer_type(core.dtype("fp32")), - core.pointer_type(core.dtype("fp32")), - core.pointer_type(core.dtype("fp32")), - core.dtype("int32"), - ) : ("vector_add_float_return", (core.dtype("uint64"))), - }, - is_pure=False, - _semantic=_semantic, - ) \ No newline at end of file From 3e95f073089fad3af255da05596f0a79b136a01b Mon Sep 17 00:00:00 2001 From: Zhang_kg <1793030808@qq.com> Date: Tue, 23 Jun 2026 06:29:31 +0000 Subject: [PATCH 8/8] Add Triton TLE NVSHMEM MegaMoE intrgration case --- .../test/tle/integration/megamoe/.gitignore | 6 + .../tle/integration/megamoe/RUNBOOK_CN.md | 215 + .../megamoe/megamoe_operator/__init__.py | 13 + .../triton_tle_megamoe_operator.py | 2707 +++++++++++++ .../triton_tle_megamoe_runtime.py | 1547 ++++++++ .../ws_userhopper_dispatch_receiver_device.cu | 1485 +++++++ ...serhopper_dispatch_receiver_extern_call.py | 293 ++ .../ws_userhopper_dispatch_receiver_host.cu | 26 + .../perf/tle_2rank_h128_tokens2_repeats3.log | 2 + ..._masked_tokens1_tile_split_repeats2_w4.log | 8 + .../perf/userhopper_phase5_perf_summary.csv | 29 + .../local_2rank_h128_tokens2_repeats3.log | 2 + ...pk8_masked_tokens1_tile_split_repeats2.log | 8 + .../megamoe/tests/megamoe_local_harness.py | 3518 +++++++++++++++++ .../megamoe/tests/run_isolated_operator.py | 68 + .../megamoe/tools/nvcc_flock_wrapper.sh | 9 + 16 files changed, 9936 insertions(+) create mode 100644 python/test/tle/integration/megamoe/.gitignore create mode 100644 python/test/tle/integration/megamoe/RUNBOOK_CN.md create mode 100644 python/test/tle/integration/megamoe/megamoe_operator/__init__.py create mode 100644 python/test/tle/integration/megamoe/megamoe_operator/triton_tle_megamoe_operator.py create mode 100644 python/test/tle/integration/megamoe/megamoe_operator/triton_tle_megamoe_runtime.py create mode 100644 python/test/tle/integration/megamoe/megamoe_operator/ws_userhopper_dispatch_receiver_device.cu create mode 100644 python/test/tle/integration/megamoe/megamoe_operator/ws_userhopper_dispatch_receiver_extern_call.py create mode 100644 python/test/tle/integration/megamoe/megamoe_operator/ws_userhopper_dispatch_receiver_host.cu create mode 100644 python/test/tle/integration/megamoe/perf/tle_2rank_h128_tokens2_repeats3.log create mode 100644 python/test/tle/integration/megamoe/perf/tle_8rank_h256_topk8_masked_tokens1_tile_split_repeats2_w4.log create mode 100644 python/test/tle/integration/megamoe/perf/userhopper_phase5_perf_summary.csv create mode 100644 python/test/tle/integration/megamoe/test_logs/local_2rank_h128_tokens2_repeats3.log create mode 100644 python/test/tle/integration/megamoe/test_logs/local_8rank_h256_topk8_masked_tokens1_tile_split_repeats2.log create mode 100644 python/test/tle/integration/megamoe/tests/megamoe_local_harness.py create mode 100644 python/test/tle/integration/megamoe/tests/run_isolated_operator.py create mode 100755 python/test/tle/integration/megamoe/tools/nvcc_flock_wrapper.sh diff --git a/python/test/tle/integration/megamoe/.gitignore b/python/test/tle/integration/megamoe/.gitignore new file mode 100644 index 000000000..9c90e7366 --- /dev/null +++ b/python/test/tle/integration/megamoe/.gitignore @@ -0,0 +1,6 @@ +__pycache__/ +*.py[cod] +*.o +*.so +*.lock +*.tmp diff --git a/python/test/tle/integration/megamoe/RUNBOOK_CN.md b/python/test/tle/integration/megamoe/RUNBOOK_CN.md new file mode 100644 index 000000000..46d480c77 --- /dev/null +++ b/python/test/tle/integration/megamoe/RUNBOOK_CN.md @@ -0,0 +1,215 @@ +# Triton TLE MegaMoE 已跑通 Case + +本文档是当前目录唯一说明文档，只记录这个隔离 MegaMoE operator 已跑通的 case 和手动复现命令。 + +当前目录： + +```text +/workspace/megakernel/triton_flagos_support_nvshmem/python/test/tle/integration/megamoe +``` + +核心算子代码： + +```text +megamoe_operator/ +``` + +复测 runner： + +```text +tests/run_isolated_operator.py +``` + +本地 harness： + +```text +tests/megamoe_local_harness.py +``` + +这个 harness 是已复制到当前目录内的 host allocation/reference-check 逻辑。runner 会把 harness 中的 runtime support 和 Triton kernel entry 替换为当前目录里的 `megamoe_operator`。 + +## 基础 Import 自检 + +从 `/workspace/megakernel` 执行： + +```bash +PYTHONDONTWRITEBYTECODE=1 PYTHONNOUSERSITE=1 \ +PYTHONPATH=/workspace/megakernel/triton_flagos_support_nvshmem/python:/workspace/megakernel/triton_flagos_support_nvshmem/python/test/tle/integration/megamoe \ +MEGAMOE_TORCH_SITE_PACKAGES=/workspace/megakernel/.mega-venv/lib/python3.10/site-packages \ +/workspace/megakernel/.triton-flagos-nvshmem-venv/bin/python -c 'import megamoe_operator as op; print(op.__all__)' +``` + +预期输出包含： + +```text +_single_kernel_dispatch_receiver_l1_l2_tile_split_multi_cta_tldot_kernel +``` + +## 公共环境 + +下面两个 case 都需要先设置公共环境： + +```bash +cd /workspace/megakernel + +export CUDA_HOME=/usr/local/cuda-12.8 +export NVSHMEM_HOME=/workspace/megakernel/.mega-venv/lib/python3.10/site-packages/nvidia/nvshmem +export LD_LIBRARY_PATH="$NVSHMEM_HOME/lib:${CUDA_HOME}/lib64:${LD_LIBRARY_PATH:-}" +export CPATH="${CUDA_HOME}/targets/x86_64-linux/include:$NVSHMEM_HOME/include:${CPATH:-}" +export NVSHMEM_BOOTSTRAP=MPI +export PYTHONNOUSERSITE=1 +export PYTHONDONTWRITEBYTECODE=1 +export MEGAMOE_TORCH_SITE_PACKAGES=/workspace/megakernel/.mega-venv/lib/python3.10/site-packages +export MEGAMOE_DIR=/workspace/megakernel/triton_flagos_support_nvshmem/python/test/tle/integration/megamoe +export PYTHONPATH=/workspace/megakernel/triton_flagos_support_nvshmem/python:${MEGAMOE_DIR}:${MEGAMOE_DIR}/tests +export MEGAMOE_RUNNER=${MEGAMOE_DIR}/tests/run_isolated_operator.py +export MEGAMOE_PYTHON=/workspace/megakernel/.triton-flagos-nvshmem-venv/bin/python +``` + +## Case 1：2-rank H128 + +配置： + +```text +world size = 2 +H = 128 +I = 128 +experts = 2 +topk = 1 +tokens/rank = 2 +repeats = 3 +cleanup = 1 +``` + +运行： + +```bash +export USERHOPPER_WS_NUM_RANKS=2 +export USERHOPPER_WS_NUM_EXPERTS=2 +export USERHOPPER_WS_NUM_TOPK=1 +export USERHOPPER_WS_NUM_TOKENS=2 +export USERHOPPER_WS_HIDDEN=128 +export USERHOPPER_WS_INTERMEDIATE_HIDDEN=128 +export USERHOPPER_WS_NUM_WARPS=4 +export USERHOPPER_WS_NUM_DISPATCH_WARPS=1 +export USERHOPPER_WS_CLEANUP=1 +export USERHOPPER_WS_REPEAT_LAUNCHES=3 +export USERHOPPER_WS_L1_I_TILES=2 +export USERHOPPER_WS_L2_H_TILES=2 +export USERHOPPER_WS_L2_BLOCK_N=64 +export USERHOPPER_WS_EXPERT_WAVE_SINGLE_KERNEL=0 +export USERHOPPER_WS_MULTI_CTA_EXPERT_WAVE=0 +export USERHOPPER_WS_MULTI_CTA_L1_ONLY=0 +export USERHOPPER_WS_MULTI_CTA_L1_TILE_SPLIT=0 +export USERHOPPER_WS_MULTI_CTA_L2_TILE_SPLIT=0 + +mkdir -p "${MEGAMOE_DIR}/test_logs" + +timeout 240s mpirun --allow-run-as-root -np 2 \ + -x CUDA_HOME -x NVSHMEM_HOME -x LD_LIBRARY_PATH -x CPATH -x NVSHMEM_BOOTSTRAP \ + -x PYTHONNOUSERSITE -x PYTHONDONTWRITEBYTECODE -x PYTHONPATH \ + -x MEGAMOE_TORCH_SITE_PACKAGES \ + -x USERHOPPER_WS_NUM_RANKS -x USERHOPPER_WS_NUM_EXPERTS -x USERHOPPER_WS_NUM_TOPK -x USERHOPPER_WS_NUM_TOKENS \ + -x USERHOPPER_WS_HIDDEN -x USERHOPPER_WS_INTERMEDIATE_HIDDEN -x USERHOPPER_WS_NUM_WARPS -x USERHOPPER_WS_NUM_DISPATCH_WARPS \ + -x USERHOPPER_WS_CLEANUP -x USERHOPPER_WS_REPEAT_LAUNCHES -x USERHOPPER_WS_L1_I_TILES -x USERHOPPER_WS_L2_H_TILES \ + -x USERHOPPER_WS_L2_BLOCK_N -x USERHOPPER_WS_EXPERT_WAVE_SINGLE_KERNEL -x USERHOPPER_WS_MULTI_CTA_EXPERT_WAVE \ + -x USERHOPPER_WS_MULTI_CTA_L1_ONLY -x USERHOPPER_WS_MULTI_CTA_L1_TILE_SPLIT -x USERHOPPER_WS_MULTI_CTA_L2_TILE_SPLIT \ + "${MEGAMOE_PYTHON}" "${MEGAMOE_RUNNER}" \ + | tee "${MEGAMOE_DIR}/test_logs/local_2rank_h128_tokens2_repeats3.log" +``` + +预期结果： + +```text +rank=0 ... PASS checked=2 counts=[2] ... +rank=1 ... PASS checked=2 counts=[2] ... +``` + +说明：这个 case 会打印 PASS，但 finalize 阶段仍有已知 hang；因此命令最终可能由 `timeout` 结束，退出码可能是 `124`。判断是否跑通以两个 rank 的 PASS 行为准。 + +已有日志： + +```text +test_logs/local_2rank_h128_tokens2_repeats3.log +``` + +## Case 2：8-rank H256 Masked Tile-Split + +配置： + +```text +world size = 8 +H = 256 +I = 128 +experts = 16 +topk = 8 +route mode = masked +tokens/rank = 1 +repeats = 2 +cleanup = 1 +compute_order = expert_wave_multi_cta_l1_l2_tile_split +``` + +运行： + +```bash +export USERHOPPER_WS_NUM_RANKS=8 +export USERHOPPER_WS_NUM_EXPERTS=16 +export USERHOPPER_WS_NUM_TOPK=8 +export USERHOPPER_WS_NUM_TOKENS=1 +export USERHOPPER_WS_ROUTE_MODE=masked +export USERHOPPER_WS_HIDDEN=256 +export USERHOPPER_WS_INTERMEDIATE_HIDDEN=128 +export USERHOPPER_WS_ALLOW_EXPERIMENTAL_H256_TLDOT=1 +export USERHOPPER_WS_NUM_WARPS=4 +export USERHOPPER_WS_NUM_DISPATCH_WARPS=1 +export USERHOPPER_WS_CLEANUP=1 +export USERHOPPER_WS_REPEAT_LAUNCHES=2 +export USERHOPPER_WS_L1_I_TILES=2 +export USERHOPPER_WS_L2_H_TILES=4 +export USERHOPPER_WS_L2_BLOCK_N=64 +export USERHOPPER_WS_EXPERT_WAVE_SINGLE_KERNEL=0 +export USERHOPPER_WS_MULTI_CTA_EXPERT_WAVE=1 +export USERHOPPER_WS_MULTI_CTA_L1_ONLY=1 +export USERHOPPER_WS_MULTI_CTA_L1_TILE_SPLIT=1 +export USERHOPPER_WS_MULTI_CTA_L2_TILE_SPLIT=1 + +mkdir -p "${MEGAMOE_DIR}/test_logs" + +timeout 360s mpirun --allow-run-as-root -np 8 \ + -x CUDA_HOME -x NVSHMEM_HOME -x LD_LIBRARY_PATH -x CPATH -x NVSHMEM_BOOTSTRAP \ + -x PYTHONNOUSERSITE -x PYTHONDONTWRITEBYTECODE -x PYTHONPATH \ + -x MEGAMOE_TORCH_SITE_PACKAGES \ + -x USERHOPPER_WS_NUM_RANKS -x USERHOPPER_WS_NUM_EXPERTS -x USERHOPPER_WS_NUM_TOPK -x USERHOPPER_WS_NUM_TOKENS -x USERHOPPER_WS_ROUTE_MODE \ + -x USERHOPPER_WS_HIDDEN -x USERHOPPER_WS_INTERMEDIATE_HIDDEN -x USERHOPPER_WS_ALLOW_EXPERIMENTAL_H256_TLDOT \ + -x USERHOPPER_WS_NUM_WARPS -x USERHOPPER_WS_NUM_DISPATCH_WARPS \ + -x USERHOPPER_WS_CLEANUP -x USERHOPPER_WS_REPEAT_LAUNCHES -x USERHOPPER_WS_L1_I_TILES -x USERHOPPER_WS_L2_H_TILES \ + -x USERHOPPER_WS_L2_BLOCK_N -x USERHOPPER_WS_EXPERT_WAVE_SINGLE_KERNEL -x USERHOPPER_WS_MULTI_CTA_EXPERT_WAVE \ + -x USERHOPPER_WS_MULTI_CTA_L1_ONLY -x USERHOPPER_WS_MULTI_CTA_L1_TILE_SPLIT -x USERHOPPER_WS_MULTI_CTA_L2_TILE_SPLIT \ + "${MEGAMOE_PYTHON}" "${MEGAMOE_RUNNER}" \ + | tee "${MEGAMOE_DIR}/test_logs/local_8rank_h256_topk8_masked_tokens1_tile_split_repeats2.log" +``` + +预期结果： + +```text +rank=0 ... PASS checked=6 counts=[3, 3] ... +rank=1 ... PASS checked=6 counts=[3, 3] ... +... +rank=7 ... PASS checked=6 counts=[3, 3] ... +``` + +这个 case 已复测为 PASS，并且进程正常退出。 + +已有日志： + +```text +test_logs/local_8rank_h256_topk8_masked_tokens1_tile_split_repeats2.log +``` + +## 当前已跑通结论 + +| case | 状态 | 日志 | +|---|---|---| +| 2-rank H128/I128 experts=2 topk=1 tokens=2 repeats=3 | 两个 rank 打印 PASS；finalize hang，由 timeout 收尾 | `test_logs/local_2rank_h128_tokens2_repeats3.log` | +| 8-rank H256/I128 experts=16 topk=8 masked tokens=1 tile-split repeats=2 | 8 个 rank 全部 PASS，进程正常退出 | `test_logs/local_8rank_h256_topk8_masked_tokens1_tile_split_repeats2.log` | diff --git a/python/test/tle/integration/megamoe/megamoe_operator/__init__.py b/python/test/tle/integration/megamoe/megamoe_operator/__init__.py new file mode 100644 index 000000000..f084cf7d5 --- /dev/null +++ b/python/test/tle/integration/megamoe/megamoe_operator/__init__.py @@ -0,0 +1,13 @@ +"""Triton TLE MegaMoE operator PR payload.""" + +from .triton_tle_megamoe_operator import ( + _single_kernel_dispatch_receiver_l1_l2_expert_wave_tldot_kernel, + _single_kernel_dispatch_receiver_l1_l2_tile_split_multi_cta_tldot_kernel, + _single_kernel_dispatch_receiver_l1_tldot_kernel, +) + +__all__ = [ + "_single_kernel_dispatch_receiver_l1_l2_expert_wave_tldot_kernel", + "_single_kernel_dispatch_receiver_l1_l2_tile_split_multi_cta_tldot_kernel", + "_single_kernel_dispatch_receiver_l1_tldot_kernel", +] diff --git a/python/test/tle/integration/megamoe/megamoe_operator/triton_tle_megamoe_operator.py b/python/test/tle/integration/megamoe/megamoe_operator/triton_tle_megamoe_operator.py new file mode 100644 index 000000000..bbbf06da9 --- /dev/null +++ b/python/test/tle/integration/megamoe/megamoe_operator/triton_tle_megamoe_operator.py @@ -0,0 +1,2707 @@ +"""Triton TLE MegaMoE operator candidate. + +This file is mechanically extracted from the validated smoke implementation, +but excludes the host correctness harness. The current PR payload should treat +``_single_kernel_dispatch_receiver_l1_l2_tile_split_multi_cta_tldot_kernel`` as +the main operator path for the H=256 tile-split merged proof. +""" + +from __future__ import annotations + +import ctypes +import os +import site +from pathlib import Path + +cuda_home = os.environ.get("CUDA_HOME", "/usr/local/cuda-12.8") +os.environ.setdefault("CUDA_HOME", cuda_home) +os.environ.setdefault("USERHOPPER_WS_NUM_RANKS", "2") +os.environ.setdefault("USERHOPPER_WS_NUM_EXPERTS", "2") +os.environ.setdefault("USERHOPPER_WS_NUM_TOPK", "1") +os.environ.setdefault("USERHOPPER_WS_NUM_TOKENS", "4") +os.environ.setdefault("USERHOPPER_WS_HIDDEN", "128") +os.environ.setdefault("USERHOPPER_WS_INTERMEDIATE_HIDDEN", "128") +os.environ.setdefault("USERHOPPER_WS_NUM_WARPS", "4") +os.environ.setdefault("USERHOPPER_WS_NUM_DISPATCH_WARPS", "1") +os.environ.setdefault("USERHOPPER_WS_COMPUTE_FULL_HIDDEN", "1") +os.environ.setdefault("USERHOPPER_WS_ALLOW_EXPERIMENTAL_H256_TLDOT", "0") +os.environ.setdefault("USERHOPPER_WS_CLEANUP", "0") +os.environ.setdefault("USERHOPPER_WS_REPEAT_LAUNCHES", "1") +os.environ.setdefault("USERHOPPER_WS_EXPERT_WAVE_SINGLE_KERNEL", "0") +os.environ.setdefault("USERHOPPER_WS_EXPERT_WAVE_COMPUTE_WARPS", "4") +os.environ.setdefault("USERHOPPER_WS_L2_H_TILES", "0") +os.environ.setdefault("USERHOPPER_WS_L1_I_TILES", "0") +os.environ.setdefault("USERHOPPER_WS_SPLIT_L1_L2_WORKERS", "0") +os.environ.setdefault("USERHOPPER_WS_L1_TILE_LOOP", "0") +os.environ.setdefault("USERHOPPER_WS_L2_BLOCK_N", "64") +os.environ.setdefault("USERHOPPER_WS_L2_SCALAR", "0") +os.environ.setdefault("USERHOPPER_WS_MULTI_CTA_EXPERT_WAVE", "0") +os.environ.setdefault("USERHOPPER_WS_MULTI_CTA_L1_ONLY", "0") +os.environ.setdefault("USERHOPPER_WS_MULTI_CTA_L1_TILE_SPLIT", "0") +os.environ.setdefault("USERHOPPER_WS_MULTI_CTA_L2_TILE_SPLIT", "0") +os.environ.setdefault("USERHOPPER_WS_SKIP_COMBINE", "0") +os.environ.setdefault("USERHOPPER_WS_SKIP_REDUCE", "0") +os.environ["CPATH"] = ( + f"{cuda_home}/targets/x86_64-linux/include:" + os.environ.get("CPATH", "") +) +os.environ["LD_LIBRARY_PATH"] = ( + f"{cuda_home}/lib64:" + os.environ.get("LD_LIBRARY_PATH", "") +) + +import triton +import triton.language as tl +import triton.experimental.tle.language as tle + +torch_site_packages = os.environ.get( + "MEGAMOE_TORCH_SITE_PACKAGES", + "/workspace/megakernel/.mega-venv/lib/python3.10/site-packages", +) +if Path(torch_site_packages).exists(): + site.addsitedir(torch_site_packages) + +import torch +import triton.experimental.tle.language.raw as tle_raw +from triton.experimental.tle.raw import dialect + +try: + from . import triton_tle_megamoe_runtime as uh +except ImportError: + import triton_tle_megamoe_runtime as uh + +__all__ = [ + "_single_kernel_dispatch_receiver_l1_l2_tile_split_multi_cta_tldot_kernel", + "_single_kernel_dispatch_receiver_l1_l2_expert_wave_tldot_kernel", + "_single_kernel_dispatch_receiver_l1_tldot_kernel", +] + + +@dialect( + name="cuda", + compiler="nvcc", + file=uh.HERE / "ws_userhopper_dispatch_receiver_device.cu", + extern=uh.HERE / "ws_userhopper_dispatch_receiver_extern_call.py", + extern_func_name="userhopper_ws_tldot_combine_write_partition", + libs={"nvshmem": uh.NVSHMEM_HOME}, + links=["nvshmem_device"], +) +def edsl_userhopper_ws_tldot_combine_write(*args, **kwargs): + ... + + +DEFAULT_PY_BLOCK_M = 8 if uh.HIDDEN == 256 else 16 +PY_BLOCK_M = int(os.environ.get("USERHOPPER_WS_BLOCK_M", str(DEFAULT_PY_BLOCK_M))) +JIT_BLOCK_M = tl.constexpr(PY_BLOCK_M) +JIT_BLOCK_I = tl.constexpr(64) +JIT_BLOCK_K = tl.constexpr(64) +PY_BLOCK_N = int(os.environ["USERHOPPER_WS_L2_BLOCK_N"]) +JIT_BLOCK_N = tl.constexpr(PY_BLOCK_N) +DEFAULT_PY_L1_I_TILES = uh.INTERMEDIATE_HIDDEN // 64 +PY_L1_I_TILES = int(os.environ["USERHOPPER_WS_L1_I_TILES"]) +if PY_L1_I_TILES == 0: + PY_L1_I_TILES = DEFAULT_PY_L1_I_TILES +JIT_L1_I_TILES = tl.constexpr(PY_L1_I_TILES) +DEFAULT_PY_L2_H_TILES = uh.HIDDEN // PY_BLOCK_N +PY_L2_H_TILES = int(os.environ["USERHOPPER_WS_L2_H_TILES"]) +if PY_L2_H_TILES == 0: + PY_L2_H_TILES = DEFAULT_PY_L2_H_TILES +JIT_L2_H_TILES = tl.constexpr(PY_L2_H_TILES) +CLEANUP_WORKSPACE = int(os.environ["USERHOPPER_WS_CLEANUP"]) +REPEAT_LAUNCHES = int(os.environ["USERHOPPER_WS_REPEAT_LAUNCHES"]) +EXPERT_WAVE_SINGLE_KERNEL = int(os.environ["USERHOPPER_WS_EXPERT_WAVE_SINGLE_KERNEL"]) +EXPERT_WAVE_COMPUTE_WARPS = int(os.environ["USERHOPPER_WS_EXPERT_WAVE_COMPUTE_WARPS"]) +SPLIT_L1_L2_WORKERS = int(os.environ["USERHOPPER_WS_SPLIT_L1_L2_WORKERS"]) +L1_TILE_LOOP = int(os.environ["USERHOPPER_WS_L1_TILE_LOOP"]) +JIT_L1_TILE_LOOP = tl.constexpr(L1_TILE_LOOP) +L2_SCALAR = int(os.environ["USERHOPPER_WS_L2_SCALAR"]) +JIT_L2_SCALAR = tl.constexpr(L2_SCALAR) +MULTI_CTA_EXPERT_WAVE = int(os.environ["USERHOPPER_WS_MULTI_CTA_EXPERT_WAVE"]) +MULTI_CTA_L1_ONLY = int(os.environ["USERHOPPER_WS_MULTI_CTA_L1_ONLY"]) +JIT_MULTI_CTA_L1_ONLY = tl.constexpr(MULTI_CTA_L1_ONLY) +MULTI_CTA_L1_TILE_SPLIT = int(os.environ["USERHOPPER_WS_MULTI_CTA_L1_TILE_SPLIT"]) +JIT_MULTI_CTA_L1_TILE_SPLIT = tl.constexpr(MULTI_CTA_L1_TILE_SPLIT) +MULTI_CTA_L2_TILE_SPLIT = int(os.environ["USERHOPPER_WS_MULTI_CTA_L2_TILE_SPLIT"]) +JIT_MULTI_CTA_L2_TILE_SPLIT = tl.constexpr(MULTI_CTA_L2_TILE_SPLIT) +SKIP_COMBINE = int(os.environ["USERHOPPER_WS_SKIP_COMBINE"]) +JIT_SKIP_COMBINE = tl.constexpr(SKIP_COMBINE) +SKIP_REDUCE = int(os.environ["USERHOPPER_WS_SKIP_REDUCE"]) +JIT_SKIP_REDUCE = tl.constexpr(SKIP_REDUCE) + +if CLEANUP_WORKSPACE not in (0, 1): + raise ValueError(f"USERHOPPER_WS_CLEANUP must be 0 or 1, got {CLEANUP_WORKSPACE}") +if REPEAT_LAUNCHES <= 0: + raise ValueError(f"USERHOPPER_WS_REPEAT_LAUNCHES must be positive, got {REPEAT_LAUNCHES}") +if REPEAT_LAUNCHES > 1 and CLEANUP_WORKSPACE == 0: + raise ValueError("USERHOPPER_WS_REPEAT_LAUNCHES>1 requires USERHOPPER_WS_CLEANUP=1") +if EXPERT_WAVE_SINGLE_KERNEL not in (0, 1): + raise ValueError( + "USERHOPPER_WS_EXPERT_WAVE_SINGLE_KERNEL must be 0 or 1, " + f"got {EXPERT_WAVE_SINGLE_KERNEL}" + ) +if EXPERT_WAVE_COMPUTE_WARPS not in (2, 4): + raise ValueError( + "USERHOPPER_WS_EXPERT_WAVE_COMPUTE_WARPS must be 2 or 4, " + f"got {EXPERT_WAVE_COMPUTE_WARPS}" + ) +if SPLIT_L1_L2_WORKERS not in (0, 1): + raise ValueError( + "USERHOPPER_WS_SPLIT_L1_L2_WORKERS must be 0 or 1, " + f"got {SPLIT_L1_L2_WORKERS}" + ) +if L1_TILE_LOOP not in (0, 1): + raise ValueError(f"USERHOPPER_WS_L1_TILE_LOOP must be 0 or 1, got {L1_TILE_LOOP}") +if L2_SCALAR not in (0, 1): + raise ValueError(f"USERHOPPER_WS_L2_SCALAR must be 0 or 1, got {L2_SCALAR}") +if MULTI_CTA_EXPERT_WAVE not in (0, 1): + raise ValueError( + "USERHOPPER_WS_MULTI_CTA_EXPERT_WAVE must be 0 or 1, " + f"got {MULTI_CTA_EXPERT_WAVE}" + ) +if MULTI_CTA_L1_ONLY not in (0, 1): + raise ValueError( + "USERHOPPER_WS_MULTI_CTA_L1_ONLY must be 0 or 1, " + f"got {MULTI_CTA_L1_ONLY}" + ) +if MULTI_CTA_L1_ONLY != 0 and MULTI_CTA_EXPERT_WAVE == 0: + raise ValueError("USERHOPPER_WS_MULTI_CTA_L1_ONLY requires USERHOPPER_WS_MULTI_CTA_EXPERT_WAVE=1") +if MULTI_CTA_L1_TILE_SPLIT not in (0, 1): + raise ValueError( + "USERHOPPER_WS_MULTI_CTA_L1_TILE_SPLIT must be 0 or 1, " + f"got {MULTI_CTA_L1_TILE_SPLIT}" + ) +if MULTI_CTA_L1_TILE_SPLIT != 0 and (MULTI_CTA_EXPERT_WAVE == 0 or MULTI_CTA_L1_ONLY == 0): + raise ValueError( + "USERHOPPER_WS_MULTI_CTA_L1_TILE_SPLIT requires " + "USERHOPPER_WS_MULTI_CTA_EXPERT_WAVE=1 and USERHOPPER_WS_MULTI_CTA_L1_ONLY=1" + ) +if MULTI_CTA_L2_TILE_SPLIT not in (0, 1): + raise ValueError( + "USERHOPPER_WS_MULTI_CTA_L2_TILE_SPLIT must be 0 or 1, " + f"got {MULTI_CTA_L2_TILE_SPLIT}" + ) +if MULTI_CTA_L2_TILE_SPLIT != 0 and (MULTI_CTA_EXPERT_WAVE == 0 or MULTI_CTA_L1_TILE_SPLIT == 0): + raise ValueError( + "USERHOPPER_WS_MULTI_CTA_L2_TILE_SPLIT requires " + "USERHOPPER_WS_MULTI_CTA_EXPERT_WAVE=1 and USERHOPPER_WS_MULTI_CTA_L1_TILE_SPLIT=1" + ) +if SKIP_COMBINE not in (0, 1): + raise ValueError(f"USERHOPPER_WS_SKIP_COMBINE must be 0 or 1, got {SKIP_COMBINE}") +if SKIP_COMBINE != 0 and MULTI_CTA_L2_TILE_SPLIT == 0: + raise ValueError("USERHOPPER_WS_SKIP_COMBINE requires USERHOPPER_WS_MULTI_CTA_L2_TILE_SPLIT=1") +if SKIP_COMBINE != 0 and CLEANUP_WORKSPACE != 0: + raise ValueError("USERHOPPER_WS_SKIP_COMBINE=1 requires USERHOPPER_WS_CLEANUP=0") +if SKIP_REDUCE not in (0, 1): + raise ValueError(f"USERHOPPER_WS_SKIP_REDUCE must be 0 or 1, got {SKIP_REDUCE}") +if SKIP_REDUCE != 0 and MULTI_CTA_L2_TILE_SPLIT == 0: + raise ValueError("USERHOPPER_WS_SKIP_REDUCE requires USERHOPPER_WS_MULTI_CTA_L2_TILE_SPLIT=1") +if SKIP_REDUCE != 0 and CLEANUP_WORKSPACE != 0: + raise ValueError("USERHOPPER_WS_SKIP_REDUCE=1 requires USERHOPPER_WS_CLEANUP=0") + + +@triton.jit +def _l1_single_cta_two_tile_body( + acts, + acts_sf, + topk_weights, + weights, + weights_sf, + l2_acts, + l2_acts_sf, + marker, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + I_OFFSET_C: tl.constexpr, +): + offs_m = tl.arange(0, JIT_BLOCK_M) + offs_i = I_OFFSET_C + tl.arange(0, JIT_BLOCK_I) + group = offs_i // 8 + lane = offs_i - group * 8 + gate_rows = group * 16 + lane + up_rows = gate_rows + 8 + + gate_acc = tl.zeros((JIT_BLOCK_M, JIT_BLOCK_I), dtype=tl.float32) + up_acc = tl.zeros((JIT_BLOCK_M, JIT_BLOCK_I), dtype=tl.float32) + for k0 in range(0, H_C, JIT_BLOCK_K): + offs_k = k0 + tl.arange(0, JIT_BLOCK_K) + a = tl.load( + acts + offs_m[:, None] * H_C + offs_k[None, :], + mask=offs_m[:, None] < M_C, + other=0.0, + ) + gate_w = tl.load( + weights + gate_rows[None, :] * H_C + offs_k[:, None], + mask=offs_i[None, :] < I_C, + other=0.0, + ) + up_w = tl.load( + weights + up_rows[None, :] * H_C + offs_k[:, None], + mask=offs_i[None, :] < I_C, + other=0.0, + ) + act_scale = tl.load( + acts_sf + (k0 // 128) * NUM_PADDED_M_C + offs_m, + mask=offs_m < M_C, + other=0.0, + ) + gate_scale = tl.load( + weights_sf + (offs_i // 128) * (H_C // 128) + (k0 // 128), + mask=offs_i < I_C, + other=0.0, + ) + up_scale = tl.load( + weights_sf + ((I_C + offs_i) // 128) * (H_C // 128) + (k0 // 128), + mask=offs_i < I_C, + other=0.0, + ) + gate_acc += tl.dot(a, gate_w, out_dtype=tl.float32) * act_scale[:, None] * gate_scale[None, :] + up_acc += tl.dot(a, up_w, out_dtype=tl.float32) * act_scale[:, None] * up_scale[None, :] + + topk = tl.load(topk_weights + offs_m, mask=offs_m < M_C, other=0.0) + swiglu = gate_acc * tl.sigmoid(gate_acc) * up_acc * topk[:, None] + max_abs = tl.max(tl.abs(swiglu), axis=1) + scale = tl.where(max_abs > 0.0, max_abs / 448.0, 1.0) + scaled = swiglu / scale[:, None] + pid_i: tl.constexpr = I_OFFSET_C // JIT_BLOCK_I + + tl.store( + l2_acts_sf + pid_i * NUM_PADDED_M_C + offs_m, + scale, + mask=offs_m < M_C, + ) + tl.store( + l2_acts + offs_m[:, None] * I_C + offs_i[None, :], + scaled, + mask=(offs_m[:, None] < M_C) & (offs_i[None, :] < I_C), + ) + tl.store(marker, 0x4C1107) + + +@triton.jit +def _l1_single_cta_runtime_i_body( + acts, + acts_sf, + topk_weights, + weights, + weights_sf, + l2_acts, + l2_acts_sf, + marker, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + i_offset, +): + offs_m = tl.arange(0, JIT_BLOCK_M) + offs_i = i_offset + tl.arange(0, JIT_BLOCK_I) + group = offs_i // 8 + lane = offs_i - group * 8 + gate_rows = group * 16 + lane + up_rows = gate_rows + 8 + + gate_acc = tl.zeros((JIT_BLOCK_M, JIT_BLOCK_I), dtype=tl.float32) + up_acc = tl.zeros((JIT_BLOCK_M, JIT_BLOCK_I), dtype=tl.float32) + for k0 in range(0, H_C, JIT_BLOCK_K): + offs_k = k0 + tl.arange(0, JIT_BLOCK_K) + a = tl.load( + acts + offs_m[:, None] * H_C + offs_k[None, :], + mask=offs_m[:, None] < M_C, + other=0.0, + ) + gate_w = tl.load( + weights + gate_rows[None, :] * H_C + offs_k[:, None], + mask=offs_i[None, :] < I_C, + other=0.0, + ) + up_w = tl.load( + weights + up_rows[None, :] * H_C + offs_k[:, None], + mask=offs_i[None, :] < I_C, + other=0.0, + ) + act_scale = tl.load( + acts_sf + (k0 // 128) * NUM_PADDED_M_C + offs_m, + mask=offs_m < M_C, + other=0.0, + ) + gate_scale = tl.load( + weights_sf + (offs_i // 128) * (H_C // 128) + (k0 // 128), + mask=offs_i < I_C, + other=0.0, + ) + up_scale = tl.load( + weights_sf + ((I_C + offs_i) // 128) * (H_C // 128) + (k0 // 128), + mask=offs_i < I_C, + other=0.0, + ) + gate_acc += tl.dot(a, gate_w, out_dtype=tl.float32) * act_scale[:, None] * gate_scale[None, :] + up_acc += tl.dot(a, up_w, out_dtype=tl.float32) * act_scale[:, None] * up_scale[None, :] + + topk = tl.load(topk_weights + offs_m, mask=offs_m < M_C, other=0.0) + swiglu = gate_acc * tl.sigmoid(gate_acc) * up_acc * topk[:, None] + max_abs = tl.max(tl.abs(swiglu), axis=1) + scale = tl.where(max_abs > 0.0, max_abs / 448.0, 1.0) + scaled = swiglu / scale[:, None] + pid_i = i_offset // JIT_BLOCK_I + + tl.store( + l2_acts_sf + pid_i * NUM_PADDED_M_C + offs_m, + scale, + mask=offs_m < M_C, + ) + tl.store( + l2_acts + offs_m[:, None] * I_C + offs_i[None, :], + scaled, + mask=(offs_m[:, None] < M_C) & (offs_i[None, :] < I_C), + ) + tl.store(marker, 0x4C1108) + + +@triton.jit +def _l1_single_cta_runtime_m_i_body( + acts, + acts_sf, + topk_weights, + weights, + weights_sf, + l2_acts, + l2_acts_sf, + marker, + M, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + i_offset, +): + offs_m = tl.arange(0, JIT_BLOCK_M) + offs_i = i_offset + tl.arange(0, JIT_BLOCK_I) + group = offs_i // 8 + lane = offs_i - group * 8 + gate_rows = group * 16 + lane + up_rows = gate_rows + 8 + + gate_acc = tl.zeros((JIT_BLOCK_M, JIT_BLOCK_I), dtype=tl.float32) + up_acc = tl.zeros((JIT_BLOCK_M, JIT_BLOCK_I), dtype=tl.float32) + for k0 in range(0, H_C, JIT_BLOCK_K): + offs_k = k0 + tl.arange(0, JIT_BLOCK_K) + a = tl.load( + acts + offs_m[:, None] * H_C + offs_k[None, :], + mask=offs_m[:, None] < M, + other=0.0, + ) + gate_w = tl.load( + weights + gate_rows[None, :] * H_C + offs_k[:, None], + mask=offs_i[None, :] < I_C, + other=0.0, + ) + up_w = tl.load( + weights + up_rows[None, :] * H_C + offs_k[:, None], + mask=offs_i[None, :] < I_C, + other=0.0, + ) + act_scale = tl.load( + acts_sf + (k0 // 128) * NUM_PADDED_M_C + offs_m, + mask=offs_m < M, + other=0.0, + ) + gate_scale = tl.load( + weights_sf + (offs_i // 128) * (H_C // 128) + (k0 // 128), + mask=offs_i < I_C, + other=0.0, + ) + up_scale = tl.load( + weights_sf + ((I_C + offs_i) // 128) * (H_C // 128) + (k0 // 128), + mask=offs_i < I_C, + other=0.0, + ) + gate_acc += tl.dot(a, gate_w, out_dtype=tl.float32) * act_scale[:, None] * gate_scale[None, :] + up_acc += tl.dot(a, up_w, out_dtype=tl.float32) * act_scale[:, None] * up_scale[None, :] + + topk = tl.load(topk_weights + offs_m, mask=offs_m < M, other=0.0) + swiglu = gate_acc * tl.sigmoid(gate_acc) * up_acc * topk[:, None] + max_abs = tl.max(tl.abs(swiglu), axis=1) + scale = tl.where(max_abs > 0.0, max_abs / 448.0, 1.0) + scaled = swiglu / scale[:, None] + pid_i = i_offset // JIT_BLOCK_I + + tl.store( + l2_acts_sf + pid_i * NUM_PADDED_M_C + offs_m, + scale, + mask=offs_m < M, + ) + tl.store( + l2_acts + offs_m[:, None] * I_C + offs_i[None, :], + scaled, + mask=(offs_m[:, None] < M) & (offs_i[None, :] < I_C), + ) + tl.store(marker, 0x4C1109) + + +@triton.jit +def _l2_single_cta_tile_body( + acts, + acts_sf, + weights, + weights_sf, + out, + marker, + M_C: tl.constexpr, + N_C: tl.constexpr, + K_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + N_OFFSET_C: tl.constexpr, +): + offs_m = tl.arange(0, JIT_BLOCK_M) + offs_n = N_OFFSET_C + tl.arange(0, JIT_BLOCK_N) + + acc = tl.zeros((JIT_BLOCK_M, JIT_BLOCK_N), dtype=tl.float32) + for k0 in range(0, K_C, JIT_BLOCK_K): + offs_k = k0 + tl.arange(0, JIT_BLOCK_K) + a = tl.load( + acts + offs_m[:, None] * K_C + offs_k[None, :], + mask=offs_m[:, None] < M_C, + other=0.0, + ) + b = tl.load( + weights + offs_n[None, :] * K_C + offs_k[:, None], + mask=offs_n[None, :] < N_C, + other=0.0, + ) + partial = tl.dot(a, b, out_dtype=tl.float32) + act_scale = tl.load( + acts_sf + (k0 // 64) * NUM_PADDED_M_C + offs_m, + mask=offs_m < M_C, + other=0.0, + ) + weight_scale = tl.load( + weights_sf + (offs_n // 128) * (K_C // 128) + (k0 // 128), + mask=offs_n < N_C, + other=0.0, + ) + acc += partial * act_scale[:, None] * weight_scale[None, :] + + tl.store( + out + offs_m[:, None] * N_C + offs_n[None, :], + acc, + mask=(offs_m[:, None] < M_C) & (offs_n[None, :] < N_C), + ) + tl.store(marker, 0x4C2207) + + +@triton.jit +def _l2_single_cta_scalar_tile_body( + acts, + acts_sf, + weights, + weights_sf, + out, + marker, + M_C: tl.constexpr, + N_C: tl.constexpr, + K_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + N_OFFSET_C: tl.constexpr, +): + offs_m = tl.arange(0, JIT_BLOCK_M) + offs_n = N_OFFSET_C + tl.arange(0, JIT_BLOCK_N) + acc = tl.zeros((JIT_BLOCK_M, JIT_BLOCK_N), dtype=tl.float32) + + for kk in range(0, K_C): + a = tl.load( + acts + offs_m * K_C + kk, + mask=offs_m < M_C, + other=0.0, + ).to(tl.float32) + b = tl.load( + weights + offs_n * K_C + kk, + mask=offs_n < N_C, + other=0.0, + ).to(tl.float32) + act_scale = tl.load( + acts_sf + (kk // 64) * NUM_PADDED_M_C + offs_m, + mask=offs_m < M_C, + other=0.0, + ) + weight_scale = tl.load( + weights_sf + (offs_n // 128) * (K_C // 128) + (kk // 128), + mask=offs_n < N_C, + other=0.0, + ) + acc += (a * act_scale)[:, None] * (b * weight_scale)[None, :] + + tl.store( + out + offs_m[:, None] * N_C + offs_n[None, :], + acc, + mask=(offs_m[:, None] < M_C) & (offs_n[None, :] < N_C), + ) + tl.store(marker, 0x4C2307) + + +@triton.jit +def _l2_single_cta_runtime_m_n_body( + acts, + acts_sf, + weights, + weights_sf, + out, + marker, + M, + N_C: tl.constexpr, + K_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + n_offset, +): + offs_m = tl.arange(0, JIT_BLOCK_M) + offs_n = n_offset + tl.arange(0, JIT_BLOCK_N) + + acc = tl.zeros((JIT_BLOCK_M, JIT_BLOCK_N), dtype=tl.float32) + for k0 in range(0, K_C, JIT_BLOCK_K): + offs_k = k0 + tl.arange(0, JIT_BLOCK_K) + a = tl.load( + acts + offs_m[:, None] * K_C + offs_k[None, :], + mask=offs_m[:, None] < M, + other=0.0, + ) + b = tl.load( + weights + offs_n[None, :] * K_C + offs_k[:, None], + mask=offs_n[None, :] < N_C, + other=0.0, + ) + partial = tl.dot(a, b, out_dtype=tl.float32) + act_scale = tl.load( + acts_sf + (k0 // 64) * NUM_PADDED_M_C + offs_m, + mask=offs_m < M, + other=0.0, + ) + weight_scale = tl.load( + weights_sf + (offs_n // 128) * (K_C // 128) + (k0 // 128), + mask=offs_n < N_C, + other=0.0, + ) + acc += partial * act_scale[:, None] * weight_scale[None, :] + + tl.store( + out + offs_m[:, None] * N_C + offs_n[None, :], + acc, + mask=(offs_m[:, None] < M) & (offs_n[None, :] < N_C), + ) + tl.store(marker, 0x4C2208) + + +@triton.jit +def _single_expert_l1_body( + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + POOL_BASE_C: tl.constexpr, + LOCAL_EXPERT_C: tl.constexpr, +): + l1_acts_e = l1_acts + POOL_BASE_C * H_C + l1_acts_sf_e = l1_acts_sf + POOL_BASE_C + l1_topk_weights_e = l1_topk_weights + POOL_BASE_C + l2_acts_e = l2_acts + POOL_BASE_C * I_C + l2_acts_sf_e = l2_acts_sf + POOL_BASE_C + l1_weights_e = l1_weights + LOCAL_EXPERT_C * (2 * I_C * H_C) + l1_weights_sf_e = l1_weights_sf + LOCAL_EXPERT_C * ((2 * I_C // 128) * (H_C // 128)) + + if JIT_L1_TILE_LOOP != 0: + for i_offset in tl.range(0, JIT_L1_I_TILES * JIT_BLOCK_I, JIT_BLOCK_I): + _l1_single_cta_runtime_i_body( + l1_acts_e, + l1_acts_sf_e, + l1_topk_weights_e, + l1_weights_e, + l1_weights_sf_e, + l2_acts_e, + l2_acts_sf_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + i_offset, + ) + else: + if JIT_L1_I_TILES >= 1: + _l1_single_cta_two_tile_body( + l1_acts_e, + l1_acts_sf_e, + l1_topk_weights_e, + l1_weights_e, + l1_weights_sf_e, + l2_acts_e, + l2_acts_sf_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + 0, + ) + if JIT_L1_I_TILES >= 2: + _l1_single_cta_two_tile_body( + l1_acts_e, + l1_acts_sf_e, + l1_topk_weights_e, + l1_weights_e, + l1_weights_sf_e, + l2_acts_e, + l2_acts_sf_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + 64, + ) + + +@triton.jit +def _single_expert_l2_body( + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + marker, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + POOL_BASE_C: tl.constexpr, + LOCAL_EXPERT_C: tl.constexpr, +): + l2_acts_e = l2_acts + POOL_BASE_C * I_C + l2_acts_sf_e = l2_acts_sf + POOL_BASE_C + l2_out_e = l2_out + POOL_BASE_C * H_C + l2_weights_e = l2_weights + LOCAL_EXPERT_C * (H_C * I_C) + l2_weights_sf_e = l2_weights_sf + LOCAL_EXPERT_C * ((H_C // 128) * (I_C // 128)) + + if JIT_L2_H_TILES >= 1: + if JIT_L2_SCALAR != 0: + _l2_single_cta_scalar_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + 0, + ) + else: + _l2_single_cta_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + 0, + ) + if JIT_L2_H_TILES >= 2: + if JIT_L2_SCALAR != 0: + _l2_single_cta_scalar_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N, + ) + else: + _l2_single_cta_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N, + ) + if JIT_L2_H_TILES >= 3: + if JIT_L2_SCALAR != 0: + _l2_single_cta_scalar_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 2, + ) + else: + _l2_single_cta_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 2, + ) + if JIT_L2_H_TILES >= 4: + if JIT_L2_SCALAR != 0: + _l2_single_cta_scalar_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 3, + ) + else: + _l2_single_cta_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 3, + ) + if JIT_L2_H_TILES >= 5: + if JIT_L2_SCALAR != 0: + _l2_single_cta_scalar_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 4, + ) + else: + _l2_single_cta_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 4, + ) + if JIT_L2_H_TILES >= 6: + if JIT_L2_SCALAR != 0: + _l2_single_cta_scalar_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 5, + ) + else: + _l2_single_cta_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 5, + ) + if JIT_L2_H_TILES >= 7: + if JIT_L2_SCALAR != 0: + _l2_single_cta_scalar_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 6, + ) + else: + _l2_single_cta_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 6, + ) + if JIT_L2_H_TILES >= 8: + if JIT_L2_SCALAR != 0: + _l2_single_cta_scalar_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 7, + ) + else: + _l2_single_cta_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 7, + ) + + +@triton.jit +def _single_expert_l1_l2_body( + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + marker, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + POOL_BASE_C: tl.constexpr, + LOCAL_EXPERT_C: tl.constexpr, +): + _single_expert_l1_body( + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + POOL_BASE_C, + LOCAL_EXPERT_C, + ) + tl.debug_barrier() + _single_expert_l2_body( + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + POOL_BASE_C, + LOCAL_EXPERT_C, + ) + + +@triton.jit +def _l1_single_kernel_worker( + compute_reader, + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, +): + wait_result = compute_reader.wait(0) + _ = tl.load(tle.gpu.local_ptr(wait_result.slot.done, (0,))) + + _single_expert_l1_l2_body( + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + marker, + EXPERT0_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + 0, + 0, + ) + if NUM_EXPERTS_PER_RANK_C >= 2: + tl.debug_barrier() + _single_expert_l1_l2_body( + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + marker, + EXPERT1_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + EXPERT1_POOL_BASE_C, + 1, + ) + tl.debug_barrier() + tle_raw.call( + edsl_userhopper_ws_tldot_combine_write, + [ + symm_buffer, + l2_out, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + H_C, + I_C, + NUM_PADDED_M_C, + ], + ) + tle_raw.call( + uh.edsl_userhopper_ws_combine_reduce, + [ + symm_buffer, + y, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + H_C, + I_C, + NUM_PADDED_M_C, + CLEANUP_WORKSPACE_C, + ], + ) + tl.store(marker, 0x4C3307) + compute_reader.release(0) + + +@triton.jit +def _l1_l2_expert_wave_single_kernel_worker( + compute_reader, + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, +): + wait_result = compute_reader.wait(0) + _ = tl.load(tle.gpu.local_ptr(wait_result.slot.done, (0,))) + + _single_expert_l1_body( + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + EXPERT0_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + 0, + 0, + ) + if NUM_EXPERTS_PER_RANK_C >= 2: + tl.debug_barrier() + _single_expert_l1_body( + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + EXPERT1_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + EXPERT1_POOL_BASE_C, + 1, + ) + + tl.debug_barrier() + _single_expert_l2_body( + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + marker, + EXPERT0_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + 0, + 0, + ) + if NUM_EXPERTS_PER_RANK_C >= 2: + tl.debug_barrier() + _single_expert_l2_body( + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + marker, + EXPERT1_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + EXPERT1_POOL_BASE_C, + 1, + ) + + tl.debug_barrier() + tle_raw.call( + edsl_userhopper_ws_tldot_combine_write, + [ + symm_buffer, + l2_out, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + H_C, + I_C, + NUM_PADDED_M_C, + ], + ) + tle_raw.call( + uh.edsl_userhopper_ws_combine_reduce, + [ + symm_buffer, + y, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + H_C, + I_C, + NUM_PADDED_M_C, + CLEANUP_WORKSPACE_C, + ], + ) + tl.store(marker, 0x4C3507) + compute_reader.release(0) + + +@triton.jit +def _wait_l1_arrival_one_block( + symm_buffer, + M_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + POOL_BASE_C: tl.constexpr, +): + if M_C > 0: + arrival_offset: tl.constexpr = 32 + NUM_EXPERTS_C * 8 * 2 + NUM_EXPERTS_PER_RANK_C * 8 + arrival_u32 = (symm_buffer + arrival_offset).to(tl.pointer_type(tl.uint32)) + first_block: tl.constexpr = POOL_BASE_C // 64 + while tl.load(arrival_u32 + first_block, volatile=True) < M_C: + pass + + +@triton.jit +def _wait_l1_arrival_one_block_dynamic( + symm_buffer, + M, + NUM_EXPERTS_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + pool_base, +): + if M > 0: + arrival_offset: tl.constexpr = 32 + NUM_EXPERTS_C * 8 * 2 + NUM_EXPERTS_PER_RANK_C * 8 + arrival_u32 = (symm_buffer + arrival_offset).to(tl.pointer_type(tl.uint32)) + first_block = pool_base // 64 + while tl.load(arrival_u32 + first_block, volatile=True) < M: + pass + + +@triton.jit +def _l1_l2_multi_cta_expert_wave_worker( + compute_reader, + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + cta_done, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, +): + wait_result = compute_reader.wait(0) + _ = tl.load(tle.gpu.local_ptr(wait_result.slot.done, (0,))) + pid = tl.program_id(0) + + if pid == 1: + _wait_l1_arrival_one_block( + symm_buffer, + EXPERT0_COUNT_C, + NUM_EXPERTS_C, + NUM_EXPERTS_PER_RANK_C, + 0, + ) + _single_expert_l1_body( + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + EXPERT0_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + 0, + 0, + ) + if JIT_MULTI_CTA_L1_ONLY == 0: + tl.debug_barrier() + _single_expert_l2_body( + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + marker, + EXPERT0_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + 0, + 0, + ) + tl.atomic_add(cta_done, 1, sem="release") + if pid == 2 and NUM_EXPERTS_PER_RANK_C >= 2: + _wait_l1_arrival_one_block( + symm_buffer, + EXPERT1_COUNT_C, + NUM_EXPERTS_C, + NUM_EXPERTS_PER_RANK_C, + EXPERT1_POOL_BASE_C, + ) + _single_expert_l1_body( + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + EXPERT1_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + EXPERT1_POOL_BASE_C, + 1, + ) + if JIT_MULTI_CTA_L1_ONLY == 0: + tl.debug_barrier() + _single_expert_l2_body( + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + marker, + EXPERT1_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + EXPERT1_POOL_BASE_C, + 1, + ) + tl.atomic_add(cta_done, 1, sem="release") + + if pid == 0: + while tl.load(cta_done, volatile=True) < NUM_EXPERTS_PER_RANK_C: + pass + if JIT_MULTI_CTA_L1_ONLY == 0: + tle_raw.call( + edsl_userhopper_ws_tldot_combine_write, + [ + symm_buffer, + l2_out, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + H_C, + I_C, + NUM_PADDED_M_C, + ], + ) + tle_raw.call( + uh.edsl_userhopper_ws_combine_reduce, + [ + symm_buffer, + y, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + H_C, + I_C, + NUM_PADDED_M_C, + CLEANUP_WORKSPACE_C, + ], + ) + tl.store(marker, 0x4C3707) + else: + tl.store(marker, 0x4C3711) + compute_reader.release(0) + + +@triton.jit +def _l1_tile_split_multi_cta_worker( + compute_reader, + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + cta_done, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, +): + wait_result = compute_reader.wait(0) + _ = tl.load(tle.gpu.local_ptr(wait_result.slot.done, (0,))) + pid = tl.program_id(0) + + tile_cta = pid - 1 + active_tile_ctas: tl.constexpr = NUM_EXPERTS_PER_RANK_C * JIT_L1_I_TILES + if tile_cta >= 0 and tile_cta < active_tile_ctas: + local_expert = tile_cta // JIT_L1_I_TILES + tile_idx = tile_cta - local_expert * JIT_L1_I_TILES + count = tl.where(local_expert == 0, EXPERT0_COUNT_C, EXPERT1_COUNT_C) + pool_base = tl.where(local_expert == 0, 0, EXPERT1_POOL_BASE_C) + + _wait_l1_arrival_one_block_dynamic( + symm_buffer, + count, + NUM_EXPERTS_C, + NUM_EXPERTS_PER_RANK_C, + pool_base, + ) + + l1_acts_e = l1_acts + pool_base * H_C + l1_acts_sf_e = l1_acts_sf + pool_base + l1_topk_weights_e = l1_topk_weights + pool_base + l2_acts_e = l2_acts + pool_base * I_C + l2_acts_sf_e = l2_acts_sf + pool_base + l1_weights_e = l1_weights + local_expert * (2 * I_C * H_C) + l1_weights_sf_e = l1_weights_sf + local_expert * ((2 * I_C // 128) * (H_C // 128)) + i_offset = tile_idx * JIT_BLOCK_I + + _l1_single_cta_runtime_m_i_body( + l1_acts_e, + l1_acts_sf_e, + l1_topk_weights_e, + l1_weights_e, + l1_weights_sf_e, + l2_acts_e, + l2_acts_sf_e, + marker, + count, + H_C, + I_C, + NUM_PADDED_M_C, + i_offset, + ) + tl.atomic_add(cta_done, 1, sem="release") + + if pid == 0: + while tl.load(cta_done, volatile=True) < NUM_EXPERTS_PER_RANK_C * JIT_L1_I_TILES: + pass + tl.store(marker, 0x4C3712) + compute_reader.release(0) + + +@triton.jit +def _l1_l2_tile_split_multi_cta_worker( + compute_reader, + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + cta_done, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, + SKIP_COMBINE_C: tl.constexpr, + SKIP_REDUCE_C: tl.constexpr, +): + wait_result = compute_reader.wait(0) + _ = tl.load(tle.gpu.local_ptr(wait_result.slot.done, (0,))) + pid = tl.program_id(0) + + l1_tile_ctas: tl.constexpr = NUM_EXPERTS_PER_RANK_C * JIT_L1_I_TILES + l2_tile_ctas: tl.constexpr = NUM_EXPERTS_PER_RANK_C * JIT_L2_H_TILES + tile_cta = pid - 1 + + if tile_cta >= 0 and tile_cta < l1_tile_ctas: + local_expert = tile_cta // JIT_L1_I_TILES + tile_idx = tile_cta - local_expert * JIT_L1_I_TILES + count = tl.where(local_expert == 0, EXPERT0_COUNT_C, EXPERT1_COUNT_C) + pool_base = tl.where(local_expert == 0, 0, EXPERT1_POOL_BASE_C) + + _wait_l1_arrival_one_block_dynamic( + symm_buffer, + count, + NUM_EXPERTS_C, + NUM_EXPERTS_PER_RANK_C, + pool_base, + ) + + l1_acts_e = l1_acts + pool_base * H_C + l1_acts_sf_e = l1_acts_sf + pool_base + l1_topk_weights_e = l1_topk_weights + pool_base + l2_acts_e = l2_acts + pool_base * I_C + l2_acts_sf_e = l2_acts_sf + pool_base + l1_weights_e = l1_weights + local_expert * (2 * I_C * H_C) + l1_weights_sf_e = l1_weights_sf + local_expert * ((2 * I_C // 128) * (H_C // 128)) + i_offset = tile_idx * JIT_BLOCK_I + + _l1_single_cta_runtime_m_i_body( + l1_acts_e, + l1_acts_sf_e, + l1_topk_weights_e, + l1_weights_e, + l1_weights_sf_e, + l2_acts_e, + l2_acts_sf_e, + marker, + count, + H_C, + I_C, + NUM_PADDED_M_C, + i_offset, + ) + tl.atomic_add(cta_done, 1, sem="release") + + l2_tile_cta = tile_cta - l1_tile_ctas + if l2_tile_cta >= 0 and l2_tile_cta < l2_tile_ctas: + while tl.load(cta_done, volatile=True) < NUM_EXPERTS_PER_RANK_C * JIT_L1_I_TILES: + pass + + local_expert = l2_tile_cta // JIT_L2_H_TILES + hidden_tile = l2_tile_cta - local_expert * JIT_L2_H_TILES + count = tl.where(local_expert == 0, EXPERT0_COUNT_C, EXPERT1_COUNT_C) + pool_base = tl.where(local_expert == 0, 0, EXPERT1_POOL_BASE_C) + + l2_acts_e = l2_acts + pool_base * I_C + l2_acts_sf_e = l2_acts_sf + pool_base + l2_out_e = l2_out + pool_base * H_C + l2_weights_e = l2_weights + local_expert * (H_C * I_C) + l2_weights_sf_e = l2_weights_sf + local_expert * ((H_C // 128) * (I_C // 128)) + n_offset = hidden_tile * JIT_BLOCK_N + + _l2_single_cta_runtime_m_n_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + count, + H_C, + I_C, + NUM_PADDED_M_C, + n_offset, + ) + tl.atomic_add(cta_done + 1, 1, sem="release") + + if pid == 0: + while tl.load(cta_done + 1, volatile=True) < NUM_EXPERTS_PER_RANK_C * JIT_L2_H_TILES: + pass + if SKIP_COMBINE_C == 0: + tle_raw.call( + edsl_userhopper_ws_tldot_combine_write, + [ + symm_buffer, + l2_out, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + H_C, + I_C, + NUM_PADDED_M_C, + ], + ) + if SKIP_REDUCE_C == 0: + tle_raw.call( + uh.edsl_userhopper_ws_combine_reduce, + [ + symm_buffer, + y, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + H_C, + I_C, + NUM_PADDED_M_C, + CLEANUP_WORKSPACE_C, + ], + ) + tl.store(marker, 0x4C3713) + compute_reader.release(0) + + +@triton.jit +def _l1_expert_wave_split_worker( + compute_reader, + l2_writer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, +): + wait_result = compute_reader.wait(0) + _ = tl.load(tle.gpu.local_ptr(wait_result.slot.done, (0,))) + + _single_expert_l1_body( + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + EXPERT0_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + 0, + 0, + ) + if NUM_EXPERTS_PER_RANK_C >= 2: + tl.debug_barrier() + _single_expert_l1_body( + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + EXPERT1_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + EXPERT1_POOL_BASE_C, + 1, + ) + + tl.debug_barrier() + l2_slot = l2_writer.acquire(0) + tl.store(tle.gpu.local_ptr(l2_slot.done, (0,)), 1) + l2_writer.commit(0) + tl.store(marker, 0x4C3601) + compute_reader.release(0) + + +@triton.jit +def _l2_combine_expert_wave_split_worker( + l2_reader, + symm_buffer, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, +): + wait_result = l2_reader.wait(0) + _ = tl.load(tle.gpu.local_ptr(wait_result.slot.done, (0,))) + + _single_expert_l2_body( + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + marker, + EXPERT0_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + 0, + 0, + ) + if NUM_EXPERTS_PER_RANK_C >= 2: + tl.debug_barrier() + _single_expert_l2_body( + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + marker, + EXPERT1_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + EXPERT1_POOL_BASE_C, + 1, + ) + + tl.debug_barrier() + tle_raw.call( + edsl_userhopper_ws_tldot_combine_write, + [ + symm_buffer, + l2_out, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + H_C, + I_C, + NUM_PADDED_M_C, + ], + ) + tle_raw.call( + uh.edsl_userhopper_ws_combine_reduce, + [ + symm_buffer, + y, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + H_C, + I_C, + NUM_PADDED_M_C, + CLEANUP_WORKSPACE_C, + ], + ) + tl.store(marker, 0x4C3607) + l2_reader.release(0) + + +@triton.jit +def _single_kernel_dispatch_receiver_l1_tldot_kernel( + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + NUM_TOKENS_C: tl.constexpr, + EXPECTED_LOCAL_RECV_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + INTERMEDIATE_HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, + NUM_DISPATCH_WARPS_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, +): + dispatch_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + compute_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + dispatch_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_dispatch_receiver_sync", + done=dispatch_done, + ) + compute_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_receiver_compute_sync", + done=compute_done, + ) + tle.gpu.warp_specialize( + [ + ( + uh._dispatch_pipe_partition, + ( + dispatch_pipe.writer(), + symm_buffer, + marker, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_DISPATCH_WARPS_C, + ), + ), + ( + uh._receiver_pipe_to_compute_partition, + ( + dispatch_pipe.reader(), + compute_pipe.writer(), + symm_buffer, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + ), + ), + ( + _l1_single_kernel_worker, + ( + compute_pipe.reader(), + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + EXPECTED_LOCAL_RECV_TOKENS_C, + HIDDEN_C, + INTERMEDIATE_HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + NUM_EXPERTS_PER_RANK_C, + EXPERT0_COUNT_C, + EXPERT1_COUNT_C, + EXPERT1_POOL_BASE_C, + CLEANUP_WORKSPACE_C, + ), + ), + ], + [1, 4], + [80, 180], + ) + + +@triton.jit +def _single_kernel_dispatch_receiver_l1_l2_expert_wave_tldot_kernel( + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + NUM_TOKENS_C: tl.constexpr, + EXPECTED_LOCAL_RECV_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + INTERMEDIATE_HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, + NUM_DISPATCH_WARPS_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, +): + dispatch_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + compute_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + dispatch_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_expert_wave_dispatch_receiver_sync", + done=dispatch_done, + ) + compute_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_expert_wave_receiver_compute_sync", + done=compute_done, + ) + tle.gpu.warp_specialize( + [ + ( + uh._dispatch_pipe_partition, + ( + dispatch_pipe.writer(), + symm_buffer, + marker, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_DISPATCH_WARPS_C, + ), + ), + ( + uh._receiver_pipe_to_compute_partition, + ( + dispatch_pipe.reader(), + compute_pipe.writer(), + symm_buffer, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + ), + ), + ( + _l1_l2_expert_wave_single_kernel_worker, + ( + compute_pipe.reader(), + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + EXPECTED_LOCAL_RECV_TOKENS_C, + HIDDEN_C, + INTERMEDIATE_HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + NUM_EXPERTS_PER_RANK_C, + EXPERT0_COUNT_C, + EXPERT1_COUNT_C, + EXPERT1_POOL_BASE_C, + CLEANUP_WORKSPACE_C, + ), + ), + ], + [1, 4], + [80, 180], + ) + + +@triton.jit +def _single_kernel_dispatch_receiver_l1_l2_split_workers_tldot_kernel( + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + NUM_TOKENS_C: tl.constexpr, + EXPECTED_LOCAL_RECV_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + INTERMEDIATE_HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, + NUM_DISPATCH_WARPS_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, +): + dispatch_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + compute_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + l2_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + dispatch_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_split_dispatch_receiver_sync", + done=dispatch_done, + ) + compute_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_split_receiver_l1_sync", + done=compute_done, + ) + l2_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_split_l1_l2_sync", + done=l2_done, + ) + tle.gpu.warp_specialize( + [ + ( + uh._dispatch_pipe_partition, + ( + dispatch_pipe.writer(), + symm_buffer, + marker, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_DISPATCH_WARPS_C, + ), + ), + ( + uh._receiver_pipe_to_compute_partition, + ( + dispatch_pipe.reader(), + compute_pipe.writer(), + symm_buffer, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + ), + ), + ( + _l1_expert_wave_split_worker, + ( + compute_pipe.reader(), + l2_pipe.writer(), + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + EXPECTED_LOCAL_RECV_TOKENS_C, + HIDDEN_C, + INTERMEDIATE_HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + NUM_EXPERTS_PER_RANK_C, + EXPERT0_COUNT_C, + EXPERT1_COUNT_C, + EXPERT1_POOL_BASE_C, + ), + ), + ( + _l2_combine_expert_wave_split_worker, + ( + l2_pipe.reader(), + symm_buffer, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + EXPECTED_LOCAL_RECV_TOKENS_C, + HIDDEN_C, + INTERMEDIATE_HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + NUM_EXPERTS_PER_RANK_C, + EXPERT0_COUNT_C, + EXPERT1_COUNT_C, + EXPERT1_POOL_BASE_C, + CLEANUP_WORKSPACE_C, + ), + ), + ], + [1, 4, 4], + [80, 180, 180], + ) + + +@triton.jit +def _single_kernel_dispatch_receiver_l1_l2_multi_cta_expert_wave_tldot_kernel( + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + cta_done, + NUM_TOKENS_C: tl.constexpr, + EXPECTED_LOCAL_RECV_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + INTERMEDIATE_HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, + NUM_DISPATCH_WARPS_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, +): + dispatch_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + compute_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + dispatch_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_multi_cta_dispatch_receiver_sync", + done=dispatch_done, + ) + compute_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_multi_cta_receiver_compute_sync", + done=compute_done, + ) + tle.gpu.warp_specialize( + [ + ( + uh._dispatch_pipe_partition, + ( + dispatch_pipe.writer(), + symm_buffer, + marker, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_DISPATCH_WARPS_C, + ), + ), + ( + uh._receiver_pipe_to_compute_partition, + ( + dispatch_pipe.reader(), + compute_pipe.writer(), + symm_buffer, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + ), + ), + ( + _l1_l2_multi_cta_expert_wave_worker, + ( + compute_pipe.reader(), + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + cta_done, + EXPECTED_LOCAL_RECV_TOKENS_C, + HIDDEN_C, + INTERMEDIATE_HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + NUM_EXPERTS_PER_RANK_C, + EXPERT0_COUNT_C, + EXPERT1_COUNT_C, + EXPERT1_POOL_BASE_C, + CLEANUP_WORKSPACE_C, + ), + ), + ], + [1, 4], + [80, 180], + ) + + +@triton.jit +def _single_kernel_dispatch_receiver_l1_tile_split_multi_cta_tldot_kernel( + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + cta_done, + NUM_TOKENS_C: tl.constexpr, + EXPECTED_LOCAL_RECV_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + INTERMEDIATE_HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, + NUM_DISPATCH_WARPS_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, +): + dispatch_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + compute_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + dispatch_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_l1_tile_split_dispatch_receiver_sync", + done=dispatch_done, + ) + compute_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_l1_tile_split_receiver_compute_sync", + done=compute_done, + ) + tle.gpu.warp_specialize( + [ + ( + uh._dispatch_pipe_partition, + ( + dispatch_pipe.writer(), + symm_buffer, + marker, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_DISPATCH_WARPS_C, + ), + ), + ( + uh._receiver_pipe_to_compute_partition, + ( + dispatch_pipe.reader(), + compute_pipe.writer(), + symm_buffer, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + ), + ), + ( + _l1_tile_split_multi_cta_worker, + ( + compute_pipe.reader(), + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + cta_done, + EXPECTED_LOCAL_RECV_TOKENS_C, + HIDDEN_C, + INTERMEDIATE_HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + NUM_EXPERTS_C, + NUM_EXPERTS_PER_RANK_C, + EXPERT0_COUNT_C, + EXPERT1_COUNT_C, + EXPERT1_POOL_BASE_C, + ), + ), + ], + [1, 4], + [80, 180], + ) + + +@triton.jit +def _single_kernel_dispatch_receiver_l1_l2_tile_split_multi_cta_tldot_kernel( + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + cta_done, + NUM_TOKENS_C: tl.constexpr, + EXPECTED_LOCAL_RECV_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + INTERMEDIATE_HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, + NUM_DISPATCH_WARPS_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, + SKIP_COMBINE_C: tl.constexpr, + SKIP_REDUCE_C: tl.constexpr, +): + dispatch_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + compute_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + dispatch_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_l1_l2_tile_split_dispatch_receiver_sync", + done=dispatch_done, + ) + compute_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_l1_l2_tile_split_receiver_compute_sync", + done=compute_done, + ) + tle.gpu.warp_specialize( + [ + ( + uh._dispatch_pipe_partition, + ( + dispatch_pipe.writer(), + symm_buffer, + marker, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_DISPATCH_WARPS_C, + ), + ), + ( + uh._receiver_pipe_to_compute_partition, + ( + dispatch_pipe.reader(), + compute_pipe.writer(), + symm_buffer, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + ), + ), + ( + _l1_l2_tile_split_multi_cta_worker, + ( + compute_pipe.reader(), + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + cta_done, + EXPECTED_LOCAL_RECV_TOKENS_C, + HIDDEN_C, + INTERMEDIATE_HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + NUM_EXPERTS_PER_RANK_C, + EXPERT0_COUNT_C, + EXPERT1_COUNT_C, + EXPERT1_POOL_BASE_C, + CLEANUP_WORKSPACE_C, + SKIP_COMBINE_C, + SKIP_REDUCE_C, + ), + ), + ], + [1, 4], + [80, 180], + ) + + +@triton.jit +def _single_kernel_dispatch_receiver_l1_l2_expert_wave_tldot_w2_kernel( + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + NUM_TOKENS_C: tl.constexpr, + EXPECTED_LOCAL_RECV_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + INTERMEDIATE_HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, + NUM_DISPATCH_WARPS_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, +): + dispatch_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + compute_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + dispatch_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_expert_wave_w2_dispatch_receiver_sync", + done=dispatch_done, + ) + compute_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_expert_wave_w2_receiver_compute_sync", + done=compute_done, + ) + tle.gpu.warp_specialize( + [ + ( + uh._dispatch_pipe_partition, + ( + dispatch_pipe.writer(), + symm_buffer, + marker, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_DISPATCH_WARPS_C, + ), + ), + ( + uh._receiver_pipe_to_compute_partition, + ( + dispatch_pipe.reader(), + compute_pipe.writer(), + symm_buffer, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + ), + ), + ( + _l1_l2_expert_wave_single_kernel_worker, + ( + compute_pipe.reader(), + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + EXPECTED_LOCAL_RECV_TOKENS_C, + HIDDEN_C, + INTERMEDIATE_HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + NUM_EXPERTS_PER_RANK_C, + EXPERT0_COUNT_C, + EXPERT1_COUNT_C, + EXPERT1_POOL_BASE_C, + CLEANUP_WORKSPACE_C, + ), + ), + ], + [1, 2], + [80, 180], + ) diff --git a/python/test/tle/integration/megamoe/megamoe_operator/triton_tle_megamoe_runtime.py b/python/test/tle/integration/megamoe/megamoe_operator/triton_tle_megamoe_runtime.py new file mode 100644 index 000000000..f93456ae3 --- /dev/null +++ b/python/test/tle/integration/megamoe/megamoe_operator/triton_tle_megamoe_runtime.py @@ -0,0 +1,1547 @@ +"""Triton TLE MegaMoE runtime support for the operator candidate. + +This module is mechanically extracted from the validated dispatch/receiver +smoke and keeps the pieces required by ``triton_tle_megamoe_operator.py``. +It keeps the public symmetric-buffer byte layout aligned with UserHopper +MegaMoE and provides explicit TLE WS worker partitions: + +* dispatch: build UserHopper-style per-expert send counts, write remote + src_token_topk queues, and publish remote recv counts through raw NVSHMEM. +* receiver: wait on local recv counts, select source ranks with the same + round-robin policy, and pull token bytes/scale/weight into the local L1 pool. + +The PR payload should treat this file as operator support code, not as a +standalone benchmark or production runtime. +""" + +from __future__ import annotations + +import ctypes +import fcntl +import math +import os +import subprocess +import time +from functools import lru_cache +from pathlib import Path + +import torch +import triton +import triton.knobs as knobs +import triton.language as tl +import triton.experimental.tle.language as tle +import triton.experimental.tle.language.raw as tle_raw +from triton.experimental.tle.raw import dialect + + +HERE = Path(__file__).resolve().parent +NVSHMEM_HOME = os.environ.get( + "NVSHMEM_HOME", + "/workspace/megakernel/.mega-venv/lib/python3.10/site-packages/nvidia/nvshmem", +) + +def _env_int(name: str, default: int) -> int: + return int(os.environ.get(name, str(default))) + + +NUM_RANKS = _env_int("USERHOPPER_WS_NUM_RANKS", 2) +NUM_EXPERTS = _env_int("USERHOPPER_WS_NUM_EXPERTS", NUM_RANKS) +NUM_TOPK = _env_int("USERHOPPER_WS_NUM_TOPK", 1) +NUM_TOKENS = _env_int("USERHOPPER_WS_NUM_TOKENS", 4) +NUM_MAX_TOKENS_PER_RANK = _env_int("USERHOPPER_WS_NUM_MAX_TOKENS_PER_RANK", 384) +HIDDEN = _env_int("USERHOPPER_WS_HIDDEN", 128) +INTERMEDIATE_HIDDEN = _env_int("USERHOPPER_WS_INTERMEDIATE_HIDDEN", 128) +ROUTE_MODE = os.environ.get("USERHOPPER_WS_ROUTE_MODE", "uniform") +PRINT_LIMIT = _env_int("USERHOPPER_WS_PRINT_LIMIT", 32) +MAXNREG = _env_int("USERHOPPER_WS_MAXNREG", 240) +CLEANUP_WORKSPACE = _env_int("USERHOPPER_WS_CLEANUP", 0) +REPEAT_LAUNCHES = _env_int("USERHOPPER_WS_REPEAT_LAUNCHES", 1) +NUM_WARPS = _env_int("USERHOPPER_WS_NUM_WARPS", 4) +NUM_DISPATCH_WARPS = _env_int("USERHOPPER_WS_NUM_DISPATCH_WARPS", 1) +COMPUTE_FULL_HIDDEN = _env_int("USERHOPPER_WS_COMPUTE_FULL_HIDDEN", 0) +COMPUTE_PARALLEL = _env_int("USERHOPPER_WS_COMPUTE_PARALLEL", 0) +COMPUTE_WORKER_WARPS = _env_int("USERHOPPER_WS_COMPUTE_WORKER_WARPS", 1) +MAX_DISPATCH_WARPS = 8 + +if NUM_EXPERTS % NUM_RANKS != 0: + raise ValueError(f"NUM_EXPERTS must be divisible by NUM_RANKS, got {NUM_EXPERTS}/{NUM_RANKS}") +if NUM_TOPK <= 0 or NUM_TOPK > 32: + raise ValueError(f"NUM_TOPK must be in [1, 32], got {NUM_TOPK}") +if NUM_TOPK > NUM_EXPERTS: + raise ValueError(f"NUM_TOPK must be <= NUM_EXPERTS for unique topk experts, got {NUM_TOPK}>{NUM_EXPERTS}") +if ROUTE_MODE not in {"uniform", "skew", "masked"}: + raise ValueError(f"unknown USERHOPPER_WS_ROUTE_MODE={ROUTE_MODE!r}; expected uniform, skew, or masked") +if HIDDEN % 128 != 0: + raise ValueError(f"HIDDEN must be divisible by 128 for SM90 SF layout, got {HIDDEN}") +if INTERMEDIATE_HIDDEN % 128 != 0: + raise ValueError( + f"INTERMEDIATE_HIDDEN must be divisible by 128 for UserHopper layout, got {INTERMEDIATE_HIDDEN}" + ) +if NUM_TOKENS > NUM_MAX_TOKENS_PER_RANK: + raise ValueError( + f"NUM_TOKENS must be <= NUM_MAX_TOKENS_PER_RANK, got {NUM_TOKENS}>{NUM_MAX_TOKENS_PER_RANK}" + ) +if PRINT_LIMIT <= 0: + raise ValueError(f"USERHOPPER_WS_PRINT_LIMIT must be positive, got {PRINT_LIMIT}") +if MAXNREG <= 0: + raise ValueError(f"USERHOPPER_WS_MAXNREG must be positive, got {MAXNREG}") +if CLEANUP_WORKSPACE not in (0, 1): + raise ValueError(f"USERHOPPER_WS_CLEANUP must be 0 or 1, got {CLEANUP_WORKSPACE}") +if REPEAT_LAUNCHES <= 0: + raise ValueError(f"USERHOPPER_WS_REPEAT_LAUNCHES must be positive, got {REPEAT_LAUNCHES}") +if NUM_WARPS <= 0 or NUM_WARPS % 4 != 0: + raise ValueError(f"USERHOPPER_WS_NUM_WARPS must be a positive multiple of 4, got {NUM_WARPS}") +if NUM_DISPATCH_WARPS <= 0: + raise ValueError(f"USERHOPPER_WS_NUM_DISPATCH_WARPS must be positive, got {NUM_DISPATCH_WARPS}") +if NUM_DISPATCH_WARPS > MAX_DISPATCH_WARPS: + raise ValueError( + f"USERHOPPER_WS_NUM_DISPATCH_WARPS must be <= {MAX_DISPATCH_WARPS}, got {NUM_DISPATCH_WARPS}" + ) +if NUM_DISPATCH_WARPS > NUM_WARPS: + raise ValueError( + f"USERHOPPER_WS_NUM_DISPATCH_WARPS must be <= USERHOPPER_WS_NUM_WARPS, " + f"got {NUM_DISPATCH_WARPS}>{NUM_WARPS}" + ) +if COMPUTE_FULL_HIDDEN not in (0, 1): + raise ValueError(f"USERHOPPER_WS_COMPUTE_FULL_HIDDEN must be 0 or 1, got {COMPUTE_FULL_HIDDEN}") +if COMPUTE_PARALLEL not in (0, 1): + raise ValueError(f"USERHOPPER_WS_COMPUTE_PARALLEL must be 0 or 1, got {COMPUTE_PARALLEL}") +if COMPUTE_WORKER_WARPS <= 0: + raise ValueError(f"USERHOPPER_WS_COMPUTE_WORKER_WARPS must be positive, got {COMPUTE_WORKER_WARPS}") +EXPECTED_LOCAL_RECV_TOKENS = NUM_TOKENS * NUM_TOPK +NUM_EXPERTS_PER_RANK = NUM_EXPERTS // NUM_RANKS + +K_LCM_CANDIDATE_BLOCK_M = 384 +K_MAX_CANDIDATE_BLOCK_M = 192 +K_MIN_CANDIDATE_BLOCK_M = 8 +K_BLOCK_M = 64 +K_CANDIDATE_BLOCK_M = (8, 16, 32, 64, 96, 128, 192) + + +def _align(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +def _num_max_pool_tokens(num_ranks: int, num_max_tokens_per_rank: int, num_topk: int, + num_experts_per_rank: int) -> int: + num_max_recv_tokens = num_ranks * num_max_tokens_per_rank + num_max_experts_per_token = min(num_topk, num_experts_per_rank) + return _align( + num_max_recv_tokens * num_max_experts_per_token + + num_experts_per_rank * (K_MAX_CANDIDATE_BLOCK_M - 1), + K_LCM_CANDIDATE_BLOCK_M, + ) + + +def _layout() -> dict[str, int]: + num_experts_per_rank = NUM_EXPERTS // NUM_RANKS + num_max_recv_tokens_per_expert = NUM_RANKS * NUM_MAX_TOKENS_PER_RANK + num_max_pool_tokens = _num_max_pool_tokens( + NUM_RANKS, NUM_MAX_TOKENS_PER_RANK, NUM_TOPK, num_experts_per_rank + ) + num_max_pool_blocks = num_max_pool_tokens // K_MIN_CANDIDATE_BLOCK_M + workspace = 0 + workspace += 32 + workspace += NUM_EXPERTS * 8 * 2 + workspace += num_experts_per_rank * 8 + workspace += _align(num_max_pool_blocks, 2) * 4 + workspace += num_max_pool_blocks * 8 + workspace += num_experts_per_rank * NUM_RANKS * num_max_recv_tokens_per_expert * 4 + workspace += num_max_pool_tokens * 12 + workspace = _align(workspace, 16) + + input_token = workspace + input_sf = input_token + NUM_MAX_TOKENS_PER_RANK * HIDDEN + input_topk_idx = input_sf + NUM_MAX_TOKENS_PER_RANK * (HIDDEN // 32) + input_topk_weight = input_topk_idx + NUM_MAX_TOKENS_PER_RANK * NUM_TOPK * 8 + l1_token = input_topk_weight + NUM_MAX_TOKENS_PER_RANK * NUM_TOPK * 4 + + num_max_padded_sf_pool_tokens = max( + (num_max_pool_tokens // block_m) * _align(block_m, 128) + for block_m in K_CANDIDATE_BLOCK_M + ) + l1_sf = l1_token + num_max_pool_tokens * HIDDEN + l1_topk_weight = l1_sf + num_max_padded_sf_pool_tokens * (HIDDEN // 32) + l2_token = l1_topk_weight + num_max_pool_tokens * 4 + l2_sf = l2_token + num_max_pool_tokens * INTERMEDIATE_HIDDEN + l2_sf_bytes_per_token = INTERMEDIATE_HIDDEN // 16 + combine_token = l2_sf + num_max_padded_sf_pool_tokens * l2_sf_bytes_per_token + total_bytes = combine_token + NUM_TOPK * NUM_MAX_TOKENS_PER_RANK * HIDDEN * 2 + + return { + "workspace": workspace, + "input_token": input_token, + "input_sf": input_sf, + "input_topk_idx": input_topk_idx, + "input_topk_weight": input_topk_weight, + "l1_token": l1_token, + "l1_sf": l1_sf, + "l1_topk_weight": l1_topk_weight, + "l2_token": l2_token, + "l2_sf": l2_sf, + "combine_token": combine_token, + "num_max_pool_tokens": num_max_pool_tokens, + "num_max_pool_blocks": num_max_pool_blocks, + "num_max_padded_sf_pool_tokens": num_max_padded_sf_pool_tokens, + "total_bytes": total_bytes, + } + + +LAYOUT = _layout() + + +def _uniform_route_expert(rank: int, topk: int) -> int: + dst_rank = (rank + 1 + topk) % NUM_RANKS + dst_local_expert = topk % NUM_EXPERTS_PER_RANK + return dst_rank * NUM_EXPERTS_PER_RANK + dst_local_expert + + +def _route_experts_for_rank(rank: int) -> list[int]: + if ROUTE_MODE == "uniform": + candidates = [_uniform_route_expert(rank, topk) for topk in range(NUM_TOPK)] + elif ROUTE_MODE == "skew": + rank_offsets = (1, 1, 2, 4, 4, 7, 3, 6, 0, 5, 2, 7, 1, 4, 6, 3) + local_offsets = (0, 1, 0, 1, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 0, 1) + candidates = [] + for topk in range(NUM_TOPK): + dst_rank = (rank + rank_offsets[topk % len(rank_offsets)]) % NUM_RANKS + dst_local = (local_offsets[topk % len(local_offsets)] + topk // len(local_offsets)) % NUM_EXPERTS_PER_RANK + candidates.append(dst_rank * NUM_EXPERTS_PER_RANK + dst_local) + else: + candidates = [] + for topk in range(NUM_TOPK): + candidates.append(-1 if topk % 3 == 2 else _uniform_route_expert(rank, topk)) + + seen = set() + unique = [] + for expert in candidates: + if expert < 0: + unique.append(expert) + continue + while expert in seen: + expert = (expert + 1) % NUM_EXPERTS + seen.add(expert) + unique.append(expert) + return unique + + +def _route_expert(rank: int, topk: int) -> int: + return _route_experts_for_rank(rank)[topk] + + +def _topks_for_destination(src_rank: int, dst_rank: int, dst_local_expert: int) -> list[int]: + topks = [] + for topk in range(NUM_TOPK): + expert = _route_expert(src_rank, topk) + if expert >= 0 and expert // NUM_EXPERTS_PER_RANK == dst_rank and expert % NUM_EXPERTS_PER_RANK == dst_local_expert: + topks.append(topk) + return topks + + +def _rank_counts_for_destination(dst_rank: int, dst_local_expert: int) -> list[int]: + return [ + NUM_TOKENS * len(_topks_for_destination(src_rank, dst_rank, dst_local_expert)) + for src_rank in range(NUM_RANKS) + ] + + +def _expected_counts_for_rank(rank: int) -> list[int]: + return [ + sum(_rank_counts_for_destination(rank, local_expert)) + for local_expert in range(NUM_EXPERTS_PER_RANK) + ] + + +def _choose_rank_round_robin(token_idx_in_expert: int, rank_counts: list[int]) -> tuple[int, int]: + remaining = rank_counts.copy() + offset = 0 + slot_idx = token_idx_in_expert + while True: + active_ranks = [rank for rank, count in enumerate(remaining) if count > 0] + if not active_ranks: + break + length = min(remaining[rank] for rank in active_ranks) + num_round_tokens = length * len(active_ranks) + if slot_idx < num_round_tokens: + return active_ranks[slot_idx % len(active_ranks)], offset + slot_idx // len(active_ranks) + slot_idx -= num_round_tokens + offset += length + for rank in active_ranks: + remaining[rank] -= length + return 0, 0 + + +def _input_fp8_value(rank: int, token: int) -> float: + return 0.25 + 0.25 * rank + 0.015625 * token + + +@lru_cache(maxsize=None) +def _input_fp8_byte(rank: int, token: int) -> int: + value = torch.tensor([_input_fp8_value(rank, token)], dtype=torch.float32) + return int(value.to(torch.float8_e4m3fn).view(torch.uint8)[0].item()) + + +def _expected_send_counts_for_rank(rank: int) -> list[int]: + counts = [0 for _ in range(NUM_EXPERTS)] + for _token in range(NUM_TOKENS): + for topk in range(NUM_TOPK): + expert = _route_expert(rank, topk) + if expert >= 0: + counts[expert] += 1 + return counts + + +def _expected_queue_for_destination(dst_rank: int, dst_local_expert: int, src_rank: int) -> list[int]: + entries: list[int] = [] + for token in range(NUM_TOKENS): + for topk in range(NUM_TOPK): + expert = _route_expert(src_rank, topk) + if expert >= 0 and expert // NUM_EXPERTS_PER_RANK == dst_rank and expert % NUM_EXPERTS_PER_RANK == dst_local_expert: + entries.append(token * NUM_TOPK + topk) + return entries + + +def _validate_dispatch_queue(rank: int, queue: torch.Tensor) -> None: + for local_expert in range(NUM_EXPERTS_PER_RANK): + for src_rank in range(NUM_RANKS): + expected_entries = _expected_queue_for_destination(rank, local_expert, src_rank) + if not expected_entries: + continue + got_entries = queue[local_expert, src_rank, :len(expected_entries)] + expected_tensor = torch.tensor(expected_entries, dtype=torch.uint32) + if not torch.equal(got_entries, expected_tensor): + raise SystemExit( + "dispatch queue mismatch: local_expert={} src_rank={} got={} expected={}".format( + local_expert, + src_rank, + got_entries.tolist(), + expected_tensor.tolist(), + ) + ) + + +def _validate_dispatch_workspace( + rank: int, + send_and_recv_count: torch.Tensor, + recv_sum: torch.Tensor, + queue: torch.Tensor, +) -> None: + send_counts = _expected_send_counts_for_rank(rank) + expected_send = torch.tensor( + [(1 << 32) | count for count in send_counts], + dtype=torch.uint64, + ) + got_send = send_and_recv_count[:NUM_EXPERTS] + if not torch.equal(got_send, expected_send): + raise SystemExit( + "dispatch send_count mismatch: got={} expected={}".format( + got_send.tolist(), expected_send.tolist() + ) + ) + + expected_recv_rows = torch.tensor( + [ + _rank_counts_for_destination(rank, local_expert) + for local_expert in range(NUM_EXPERTS_PER_RANK) + ], + dtype=torch.uint64, + ).T.contiguous() + got_recv_rows = send_and_recv_count[NUM_EXPERTS:NUM_EXPERTS * 2].reshape(NUM_RANKS, NUM_EXPERTS_PER_RANK) + if not torch.equal(got_recv_rows, expected_recv_rows): + raise SystemExit( + "dispatch recv_count mismatch: got={} expected={}".format( + got_recv_rows.tolist(), expected_recv_rows.tolist() + ) + ) + + expected_recv_sum = torch.tensor( + [ + (NUM_RANKS << 32) | sum(_rank_counts_for_destination(rank, local_expert)) + for local_expert in range(NUM_EXPERTS_PER_RANK) + ], + dtype=torch.uint64, + ) + if not torch.equal(recv_sum, expected_recv_sum): + raise SystemExit( + "dispatch recv_sum mismatch: got={} expected={}".format( + recv_sum.tolist(), expected_recv_sum.tolist() + ) + ) + + _validate_dispatch_queue(rank, queue) + + +def _validate_workspace_cleanup( + send_and_recv_count: torch.Tensor, + recv_sum: torch.Tensor, + arrival: torch.Tensor, + l2_arrival_mask: torch.Tensor, +) -> None: + if bool(torch.any(send_and_recv_count != 0).item()): + raise SystemExit(f"cleanup send/recv count not zero: {send_and_recv_count.tolist()}") + if bool(torch.any(recv_sum != 0).item()): + raise SystemExit(f"cleanup recv_sum not zero: {recv_sum.tolist()}") + if bool(torch.any(arrival != 0).item()): + raise SystemExit(f"cleanup arrival not zero: nonzero={torch.nonzero(arrival).flatten().tolist()}") + if bool(torch.any(l2_arrival_mask != 0).item()): + raise SystemExit(f"cleanup l2_arrival_mask not zero: nonzero={torch.nonzero(l2_arrival_mask).flatten().tolist()}") + + +def _expected_receive(rank: int) -> tuple[list[int], torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict[int, int]]: + indices: list[int] = [] + rows: list[torch.Tensor] = [] + sf_rows: list[torch.Tensor] = [] + weight_values: list[float] = [] + meta_values: list[list[int]] = [] + arrival_counts: dict[int, int] = {} + pool_token_base = 0 + + for local_expert in range(NUM_EXPERTS_PER_RANK): + rank_counts = _rank_counts_for_destination(rank, local_expert) + total = sum(rank_counts) + block_offset = pool_token_base // K_BLOCK_M + for token_idx_in_expert in range(total): + src_rank, token_idx_in_rank = _choose_rank_round_robin(token_idx_in_expert, rank_counts) + topks = _topks_for_destination(src_rank, rank, local_expert) + topk = topks[token_idx_in_rank % len(topks)] + src_token = token_idx_in_rank // len(topks) + pool_idx = pool_token_base + token_idx_in_expert + row = torch.empty((HIDDEN,), dtype=torch.uint8) + row.fill_(_input_fp8_byte(src_rank, src_token)) + indices.append(pool_idx) + rows.append(row) + sf_base = src_rank + 0.125 * (src_token + 1) + sf_row = torch.tensor( + [sf_base + 0.01 * sf_idx for sf_idx in range(HIDDEN // 128)], + dtype=torch.float32, + ) + sf_rows.append(sf_row) + weight_values.append(src_rank + 0.25 * (src_token + 1) + 0.03125 * topk) + meta_values.append([src_rank, src_token, topk]) + arrival_idx = block_offset + token_idx_in_expert // K_BLOCK_M + arrival_counts[arrival_idx] = arrival_counts.get(arrival_idx, 0) + 1 + pool_token_base += _align(total, K_BLOCK_M) + + expected_rows = torch.stack(rows) if rows else torch.empty((0, HIDDEN), dtype=torch.uint8) + expected_sf = torch.stack(sf_rows) if sf_rows else torch.empty((0, HIDDEN // 128), dtype=torch.float32) + expected_weight = torch.tensor(weight_values, dtype=torch.float32) + expected_meta = torch.tensor(meta_values, dtype=torch.uint32) + return indices, expected_rows, expected_sf, expected_weight, expected_meta, arrival_counts + + +def _expected_l1_checksum(expected_rows: torch.Tensor, expected_weight: torch.Tensor) -> int: + checksum = 2166136261 + first_bytes = expected_rows[:, 0].to(torch.int64).tolist() + weight_bits = expected_weight.view(torch.int32).to(torch.int64).tolist() + for byte, bits in zip(first_bytes, weight_bits): + checksum ^= int(byte) & 0xffffffff + checksum = (checksum * 16777619) & 0xffffffff + checksum ^= int(bits) & 0xffffffff + checksum = (checksum * 16777619) & 0xffffffff + return checksum + + +def _expected_l1_weight_checksum( + rank: int, + l1_compute_w: torch.Tensor, + l1_compute_sf: torch.Tensor, +) -> int: + checksum = 2166136261 + compute_h = _compute_h() + for local_expert in range(NUM_EXPERTS_PER_RANK): + rank_counts = _rank_counts_for_destination(rank, local_expert) + total = sum(rank_counts) + scale_bits = int(l1_compute_sf[local_expert, 0, 0].view(torch.int32).item()) & 0xffffffff + for token_idx_in_expert in range(total): + src_rank, token_idx_in_rank = _choose_rank_round_robin(token_idx_in_expert, rank_counts) + topks = _topks_for_destination(src_rank, rank, local_expert) + src_token = token_idx_in_rank // len(topks) + token_byte = _input_fp8_byte(src_rank, src_token) + weighted = 0 + for h in range(compute_h): + weighted = (weighted + token_byte * int(l1_compute_w[local_expert, 0, h].item())) & 0xffffffff + checksum ^= weighted + checksum = (checksum * 16777619) & 0xffffffff + checksum ^= scale_bits + checksum = (checksum * 16777619) & 0xffffffff + return checksum + + +def _fp8_e4m3fn_byte_to_float(value: int) -> float: + value &= 0xff + if value == 0: + return 0.0 + sign = -1.0 if value & 0x80 else 1.0 + exp = (value >> 3) & 0x0f + mant = value & 0x07 + if exp == 0: + mag = math.ldexp(mant / 8.0, -6) + else: + mag = math.ldexp(1.0 + mant / 8.0, exp - 7) + return sign * mag + + +def _topk_weight_value(src_rank: int, src_token: int, src_topk: int) -> float: + return src_rank + 0.25 * (src_token + 1) + 0.03125 * src_topk + + +def _compute_h() -> int: + return HIDDEN if COMPUTE_FULL_HIDDEN else min(HIDDEN, 32) + + +def _input_sf_value(src_rank: int, src_token: int, h: int) -> float: + return src_rank + 0.125 * (src_token + 1) + 0.01 * (h // 128) + + +def _l1_weight_sf_value(l1_compute_sf: torch.Tensor, local_expert: int, row: int, h: int) -> float: + return float(l1_compute_sf[local_expert, row // 128, h // 128].item()) + + +def _expected_l1_scalar_sum( + rank: int, + l1_compute_w: torch.Tensor, + l1_compute_sf: torch.Tensor, +) -> float: + scalar_sum = 0.0 + compute_h = _compute_h() + compute_i = INTERMEDIATE_HIDDEN + for local_expert in range(NUM_EXPERTS_PER_RANK): + rank_counts = _rank_counts_for_destination(rank, local_expert) + total = sum(rank_counts) + for token_idx_in_expert in range(total): + src_rank, token_idx_in_rank = _choose_rank_round_robin(token_idx_in_expert, rank_counts) + topks = _topks_for_destination(src_rank, rank, local_expert) + src_token = token_idx_in_rank // len(topks) + src_topk = topks[token_idx_in_rank % len(topks)] + topk_weight = _topk_weight_value(src_rank, src_token, src_topk) + token_raw = _input_fp8_byte(src_rank, src_token) + for ii in range(compute_i): + group = ii // 8 + lane = ii - group * 8 + gate_row = group * 16 + lane + up_row = gate_row + 8 + gate_sf = _l1_weight_sf_value(l1_compute_sf, local_expert, ii, 0) + up_sf = _l1_weight_sf_value(l1_compute_sf, local_expert, INTERMEDIATE_HIDDEN + ii, 0) + gate_acc = 0.0 + up_acc = 0.0 + for h in range(compute_h): + token_value = _fp8_e4m3fn_byte_to_float(token_raw) * _input_sf_value(src_rank, src_token, h) + gate_sf_h = ( + _l1_weight_sf_value(l1_compute_sf, local_expert, ii, h) + if COMPUTE_FULL_HIDDEN + else gate_sf + ) + up_sf_h = ( + _l1_weight_sf_value(l1_compute_sf, local_expert, INTERMEDIATE_HIDDEN + ii, h) + if COMPUTE_FULL_HIDDEN + else up_sf + ) + gate_acc += ( + token_value + * _fp8_e4m3fn_byte_to_float(int(l1_compute_w[local_expert, gate_row, h].item())) + * gate_sf_h + ) + up_acc += ( + token_value + * _fp8_e4m3fn_byte_to_float(int(l1_compute_w[local_expert, up_row, h].item())) + * up_sf_h + ) + scalar_sum += gate_acc * (1.0 / (1.0 + math.exp(-gate_acc))) * up_acc * topk_weight + return scalar_sum + + +def _expected_l2_token_floats( + rank: int, + l1_compute_w: torch.Tensor, + l1_compute_sf: torch.Tensor, +) -> torch.Tensor: + rows: list[list[float]] = [] + compute_h = _compute_h() + compute_i = INTERMEDIATE_HIDDEN + for local_expert in range(NUM_EXPERTS_PER_RANK): + rank_counts = _rank_counts_for_destination(rank, local_expert) + total = sum(rank_counts) + for token_idx_in_expert in range(total): + src_rank, token_idx_in_rank = _choose_rank_round_robin(token_idx_in_expert, rank_counts) + topks = _topks_for_destination(src_rank, rank, local_expert) + src_token = token_idx_in_rank // len(topks) + src_topk = topks[token_idx_in_rank % len(topks)] + topk_weight = _topk_weight_value(src_rank, src_token, src_topk) + token_raw = _input_fp8_byte(src_rank, src_token) + row: list[float] = [] + for ii in range(compute_i): + group = ii // 8 + lane = ii - group * 8 + gate_row = group * 16 + lane + up_row = gate_row + 8 + gate_sf = _l1_weight_sf_value(l1_compute_sf, local_expert, ii, 0) + up_sf = _l1_weight_sf_value(l1_compute_sf, local_expert, INTERMEDIATE_HIDDEN + ii, 0) + gate_acc = 0.0 + up_acc = 0.0 + for h in range(compute_h): + token_value = _fp8_e4m3fn_byte_to_float(token_raw) * _input_sf_value(src_rank, src_token, h) + gate_sf_h = ( + _l1_weight_sf_value(l1_compute_sf, local_expert, ii, h) + if COMPUTE_FULL_HIDDEN + else gate_sf + ) + up_sf_h = ( + _l1_weight_sf_value(l1_compute_sf, local_expert, INTERMEDIATE_HIDDEN + ii, h) + if COMPUTE_FULL_HIDDEN + else up_sf + ) + gate_acc += ( + token_value + * _fp8_e4m3fn_byte_to_float(int(l1_compute_w[local_expert, gate_row, h].item())) + * gate_sf_h + ) + up_acc += ( + token_value + * _fp8_e4m3fn_byte_to_float(int(l1_compute_w[local_expert, up_row, h].item())) + * up_sf_h + ) + row.append(gate_acc * (1.0 / (1.0 + math.exp(-gate_acc))) * up_acc * topk_weight) + rows.append(row) + if not rows: + return torch.empty((0, INTERMEDIATE_HIDDEN), dtype=torch.float32) + return torch.tensor(rows, dtype=torch.float32) + + +def _expected_l2_sf_and_scaled_floats(l2_float: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + if l2_float.numel() == 0: + return ( + torch.empty((0, INTERMEDIATE_HIDDEN // 64), dtype=torch.float32), + torch.empty((0, INTERMEDIATE_HIDDEN), dtype=torch.float32), + ) + grouped = l2_float.reshape(l2_float.shape[0], INTERMEDIATE_HIDDEN // 64, 64) + max_abs = grouped.abs().amax(dim=2) + scale = torch.where(max_abs > 0.0, max_abs / 448.0, torch.ones_like(max_abs)) + scaled = (grouped / scale[:, :, None]).reshape_as(l2_float) + return scale.contiguous(), scaled.contiguous() + + +def _float_to_cuda_satfinite_e4m3_bytes(values: torch.Tensor, device: torch.device) -> torch.Tensor: + bytes_ = values.to(device=device).to(torch.float8_e4m3fn).view(torch.uint8).detach().cpu() + bytes_[bytes_ == 0x7F] = 0x7E + bytes_[bytes_ == 0xFF] = 0xFE + return bytes_ + + +def _expected_l2_float_for_route( + src_rank: int, + src_token: int, + src_topk: int, + local_expert: int, + l1_compute_w: torch.Tensor, + l1_compute_sf: torch.Tensor, +) -> torch.Tensor: + compute_h = _compute_h() + topk_weight = _topk_weight_value(src_rank, src_token, src_topk) + token_raw = _input_fp8_byte(src_rank, src_token) + row: list[float] = [] + for ii in range(INTERMEDIATE_HIDDEN): + group = ii // 8 + lane = ii - group * 8 + gate_row = group * 16 + lane + up_row = gate_row + 8 + gate_sf = _l1_weight_sf_value(l1_compute_sf, local_expert, ii, 0) + up_sf = _l1_weight_sf_value(l1_compute_sf, local_expert, INTERMEDIATE_HIDDEN + ii, 0) + gate_acc = 0.0 + up_acc = 0.0 + for h in range(compute_h): + token_value = _fp8_e4m3fn_byte_to_float(token_raw) * _input_sf_value(src_rank, src_token, h) + gate_sf_h = ( + _l1_weight_sf_value(l1_compute_sf, local_expert, ii, h) + if COMPUTE_FULL_HIDDEN + else gate_sf + ) + up_sf_h = ( + _l1_weight_sf_value(l1_compute_sf, local_expert, INTERMEDIATE_HIDDEN + ii, h) + if COMPUTE_FULL_HIDDEN + else up_sf + ) + gate_acc += ( + token_value + * _fp8_e4m3fn_byte_to_float(int(l1_compute_w[local_expert, gate_row, h].item())) + * gate_sf_h + ) + up_acc += ( + token_value + * _fp8_e4m3fn_byte_to_float(int(l1_compute_w[local_expert, up_row, h].item())) + * up_sf_h + ) + row.append(gate_acc * (1.0 / (1.0 + math.exp(-gate_acc))) * up_acc * topk_weight) + return torch.tensor(row, dtype=torch.float32) + + +def _expected_combine_float( + rank: int, + l1_compute_w: torch.Tensor, + l1_compute_sf: torch.Tensor, + l2_compute_w: torch.Tensor, + l2_compute_sf: torch.Tensor, + device: torch.device, +) -> torch.Tensor: + expected = torch.zeros((NUM_TOPK, NUM_MAX_TOKENS_PER_RANK, HIDDEN), dtype=torch.float32) + for token in range(NUM_TOKENS): + for topk in range(NUM_TOPK): + expert = _route_expert(rank, topk) + if expert < 0: + continue + local_expert = expert % NUM_EXPERTS_PER_RANK + l2_float = _expected_l2_float_for_route( + rank, + token, + topk, + local_expert, + l1_compute_w, + l1_compute_sf, + ).unsqueeze(0) + l2_sf_row, l2_scaled = _expected_l2_sf_and_scaled_floats(l2_float) + l2_bytes = _float_to_cuda_satfinite_e4m3_bytes(l2_scaled, device)[0] + l2_sf_values = l2_sf_row[0] + for h in range(HIDDEN): + acc = 0.0 + for ii in range(INTERMEDIATE_HIDDEN): + act = _fp8_e4m3fn_byte_to_float(int(l2_bytes[ii].item())) * float( + l2_sf_values[ii // 64].item() + ) + w = _fp8_e4m3fn_byte_to_float(int(l2_compute_w[local_expert, h, ii].item())) + sf = float(l2_compute_sf[local_expert, h // 128, ii // 128].item()) + acc += act * w * sf + expected[topk, token, h] = acc + return expected + + +@dialect( + name="cuda", + compiler="nvcc", + file=HERE / "ws_userhopper_dispatch_receiver_device.cu", + extern=HERE / "ws_userhopper_dispatch_receiver_extern_call.py", + extern_func_name="userhopper_ws_dispatch_partition", + libs={"nvshmem": NVSHMEM_HOME}, + links=["nvshmem_device"], +) +def edsl_userhopper_ws_dispatch(*args, **kwargs): + ... + + +@dialect( + name="cuda", + compiler="nvcc", + file=HERE / "ws_userhopper_dispatch_receiver_device.cu", + extern=HERE / "ws_userhopper_dispatch_receiver_extern_call.py", + extern_func_name="userhopper_ws_dispatch_partition_cta_warp0", + libs={"nvshmem": NVSHMEM_HOME}, + links=["nvshmem_device"], +) +def edsl_userhopper_ws_dispatch_cta_warp0(*args, **kwargs): + ... + + +@dialect( + name="cuda", + compiler="nvcc", + file=HERE / "ws_userhopper_dispatch_receiver_device.cu", + extern=HERE / "ws_userhopper_dispatch_receiver_extern_call.py", + extern_func_name="userhopper_ws_dispatch_partition_cta_multiwarp", + libs={"nvshmem": NVSHMEM_HOME}, + links=["nvshmem_device"], +) +def edsl_userhopper_ws_dispatch_cta_multiwarp(*args, **kwargs): + ... + + +@dialect( + name="cuda", + compiler="nvcc", + file=HERE / "ws_userhopper_dispatch_receiver_device.cu", + extern=HERE / "ws_userhopper_dispatch_receiver_extern_call.py", + extern_func_name="userhopper_ws_receiver_partition", + libs={"nvshmem": NVSHMEM_HOME}, + links=["nvshmem_device"], +) +def edsl_userhopper_ws_receiver(*args, **kwargs): + ... + + +@dialect( + name="cuda", + compiler="nvcc", + file=HERE / "ws_userhopper_dispatch_receiver_device.cu", + extern=HERE / "ws_userhopper_dispatch_receiver_extern_call.py", + extern_func_name="userhopper_ws_receiver_partition_bounded", + libs={"nvshmem": NVSHMEM_HOME}, + links=["nvshmem_device"], +) +def edsl_userhopper_ws_receiver_bounded(*args, **kwargs): + ... + + +@dialect( + name="cuda", + compiler="nvcc", + file=HERE / "ws_userhopper_dispatch_receiver_device.cu", + extern=HERE / "ws_userhopper_dispatch_receiver_extern_call.py", + extern_func_name="userhopper_ws_compute_stub_partition", + libs={"nvshmem": NVSHMEM_HOME}, + links=["nvshmem_device"], +) +def edsl_userhopper_ws_compute_stub(*args, **kwargs): + ... + + +@dialect( + name="cuda", + compiler="nvcc", + file=HERE / "ws_userhopper_dispatch_receiver_device.cu", + extern=HERE / "ws_userhopper_dispatch_receiver_extern_call.py", + extern_func_name="userhopper_ws_combine_reduce_partition", + libs={"nvshmem": NVSHMEM_HOME}, + links=["nvshmem_device"], +) +def edsl_userhopper_ws_combine_reduce(*args, **kwargs): + ... + + +@triton.jit +def _stage_default_partition(marker, VALUE: tl.constexpr): + tl.store(marker, VALUE) + + +@triton.jit +def _dispatch_default_partition( + symm_buffer, + marker, + NUM_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, +): + tle_raw.call( + edsl_userhopper_ws_dispatch, + [ + symm_buffer, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + ], + ) + tl.store(marker, 0x4452) + + +@triton.jit +def _dispatch_pipe_partition( + sync_writer, + symm_buffer, + marker, + NUM_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + NUM_DISPATCH_WARPS_C: tl.constexpr, +): + sync_slot = sync_writer.acquire(0) + if NUM_DISPATCH_WARPS_C == 1: + tle_raw.call( + edsl_userhopper_ws_dispatch_cta_warp0, + [ + symm_buffer, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + ], + ) + else: + tle_raw.call( + edsl_userhopper_ws_dispatch_cta_multiwarp, + [ + symm_buffer, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_DISPATCH_WARPS_C, + ], + ) + tl.store(tle.gpu.local_ptr(sync_slot.done, (0,)), 1) + tl.store(marker, 0x5100) + sync_writer.commit(0) + + +@triton.jit +def _receiver_partition( + symm_buffer, + NUM_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, +): + tle_raw.call( + edsl_userhopper_ws_receiver, + [ + symm_buffer, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + ], + ) + + +@triton.jit +def _receiver_pipe_partition( + sync_reader, + symm_buffer, + NUM_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, +): + wait_result = sync_reader.wait(0) + _ = tl.load(tle.gpu.local_ptr(wait_result.slot.done, (0,))) + tle_raw.call( + edsl_userhopper_ws_receiver, + [ + symm_buffer, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + ], + ) + sync_reader.release(0) + + +@triton.jit +def _receiver_pipe_to_compute_partition( + dispatch_reader, + compute_writer, + symm_buffer, + NUM_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, +): + wait_result = dispatch_reader.wait(0) + _ = tl.load(tle.gpu.local_ptr(wait_result.slot.done, (0,))) + tle_raw.call( + edsl_userhopper_ws_receiver, + [ + symm_buffer, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + ], + ) + compute_slot = compute_writer.acquire(0) + tl.store(tle.gpu.local_ptr(compute_slot.done, (0,)), 1) + compute_writer.commit(0) + dispatch_reader.release(0) + + +@triton.jit +def _compute_stub_pipe_partition( + compute_reader, + combine_writer, + symm_buffer, + l1_weights, + l1_weights_sf, + l2_weights, + l2_weights_sf, + marker, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + INTERMEDIATE_HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, + COMPUTE_FULL_HIDDEN_C: tl.constexpr, + COMPUTE_PARALLEL_C: tl.constexpr, + COMPUTE_WORKER_WARPS_C: tl.constexpr, +): + wait_result = compute_reader.wait(0) + _ = tl.load(tle.gpu.local_ptr(wait_result.slot.done, (0,))) + tle_raw.call( + edsl_userhopper_ws_compute_stub, + [ + symm_buffer, + l1_weights, + l1_weights_sf, + l2_weights, + l2_weights_sf, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + INTERMEDIATE_HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + COMPUTE_FULL_HIDDEN_C, + COMPUTE_PARALLEL_C, + COMPUTE_WORKER_WARPS_C, + ], + ) + combine_slot = combine_writer.acquire(0) + tl.store(tle.gpu.local_ptr(combine_slot.done, (0,)), 1) + combine_writer.commit(0) + tl.store(marker, 0x5300) + compute_reader.release(0) + + +@triton.jit +def _combine_reduce_pipe_partition( + combine_reader, + symm_buffer, + y, + marker, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + INTERMEDIATE_HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, +): + wait_result = combine_reader.wait(0) + _ = tl.load(tle.gpu.local_ptr(wait_result.slot.done, (0,))) + tle_raw.call( + edsl_userhopper_ws_combine_reduce, + [ + symm_buffer, + y, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + INTERMEDIATE_HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + CLEANUP_WORKSPACE_C, + ], + ) + tl.store(marker, 0x5400) + combine_reader.release(0) + + +@triton.jit +def _receiver_pipe_debug_partition( + sync_reader, + symm_buffer, + marker, + NUM_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, +): + wait_result = sync_reader.wait(0) + _ = tl.load(tle.gpu.local_ptr(wait_result.slot.done, (0,))) + tle_raw.call( + edsl_userhopper_ws_receiver_bounded, + [ + symm_buffer, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + ], + ) + tl.store(marker, 0x510D) + sync_reader.release(0) + + +@triton.jit +def _ws_dispatch_stage_kernel( + symm_buffer, + marker, + NUM_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, +): + tle.gpu.warp_specialize( + [ + (_stage_default_partition, (marker, 0xD100)), + ( + _dispatch_default_partition, + ( + symm_buffer, + marker, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_DISPATCH_WARPS_C, + ), + ), + ], + [1], + [40], + ) + + +@triton.jit +def _ws_receiver_stage_kernel( + symm_buffer, + marker, + NUM_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, +): + tle.gpu.warp_specialize( + [ + (_stage_default_partition, (marker, 0xD200)), + ( + _receiver_partition, + ( + symm_buffer, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + ), + ), + ], + [1], + [40], + ) + + +@triton.jit +def _ws_single_pipe_dispatch_receiver_kernel( + symm_buffer, + marker, + NUM_TOKENS_C: tl.constexpr, + EXPECTED_LOCAL_RECV_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, + NUM_DISPATCH_WARPS_C: tl.constexpr, +): + sync_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + sync_pipe = tle.pipe( + capacity=1, + scope="cta", + name="userhopper_dispatch_receiver_sync", + done=sync_done, + ) + sync_writer = sync_pipe.writer() + sync_reader = sync_pipe.reader() + tle.gpu.warp_specialize( + [ + ( + _dispatch_pipe_partition, + ( + sync_writer, + symm_buffer, + marker, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_DISPATCH_WARPS_C, + ), + ), + ( + _receiver_pipe_partition, + ( + sync_reader, + symm_buffer, + EXPECTED_LOCAL_RECV_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + ), + ), + ], + [1], + [40], + ) + + +@triton.jit +def _ws_single_pipe_compute_stub_dispatch_receiver_kernel( + symm_buffer, + y, + l1_weights, + l1_weights_sf, + l2_weights, + l2_weights_sf, + marker, + NUM_TOKENS_C: tl.constexpr, + EXPECTED_LOCAL_RECV_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + INTERMEDIATE_HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, + NUM_DISPATCH_WARPS_C: tl.constexpr, + COMPUTE_FULL_HIDDEN_C: tl.constexpr, + COMPUTE_PARALLEL_C: tl.constexpr, + COMPUTE_WORKER_WARPS_C: tl.constexpr, +): + dispatch_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + compute_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + combine_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + dispatch_pipe = tle.pipe( + capacity=1, + scope="cta", + name="userhopper_dispatch_receiver_compute_dispatch_sync", + done=dispatch_done, + ) + compute_pipe = tle.pipe( + capacity=1, + scope="cta", + name="userhopper_dispatch_receiver_compute_stub_sync", + done=compute_done, + ) + combine_pipe = tle.pipe( + capacity=1, + scope="cta", + name="userhopper_dispatch_receiver_combine_sync", + done=combine_done, + ) + dispatch_writer = dispatch_pipe.writer() + dispatch_reader = dispatch_pipe.reader() + compute_writer = compute_pipe.writer() + compute_reader = compute_pipe.reader() + combine_writer = combine_pipe.writer() + combine_reader = combine_pipe.reader() + tle.gpu.warp_specialize( + [ + ( + _dispatch_pipe_partition, + ( + dispatch_writer, + symm_buffer, + marker, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_DISPATCH_WARPS_C, + ), + ), + ( + _receiver_pipe_to_compute_partition, + ( + dispatch_reader, + compute_writer, + symm_buffer, + EXPECTED_LOCAL_RECV_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + ), + ), + ( + _compute_stub_pipe_partition, + ( + compute_reader, + combine_writer, + symm_buffer, + l1_weights, + l1_weights_sf, + l2_weights, + l2_weights_sf, + marker, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + INTERMEDIATE_HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + COMPUTE_FULL_HIDDEN_C, + COMPUTE_PARALLEL_C, + COMPUTE_WORKER_WARPS_C, + ), + ), + ( + _combine_reduce_pipe_partition, + ( + combine_reader, + symm_buffer, + y, + marker, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + INTERMEDIATE_HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + CLEANUP_WORKSPACE_C, + ), + ), + ], + [1, COMPUTE_WORKER_WARPS_C, 1], + [40, 40, 40], + ) + + +@triton.jit +def _ws_single_pipe_debug_dispatch_receiver_kernel( + symm_buffer, + marker, + NUM_TOKENS_C: tl.constexpr, + EXPECTED_LOCAL_RECV_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, + NUM_DISPATCH_WARPS_C: tl.constexpr, +): + sync_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + sync_pipe = tle.pipe( + capacity=1, + scope="cta", + name="userhopper_dispatch_receiver_debug_sync", + done=sync_done, + ) + sync_writer = sync_pipe.writer() + sync_reader = sync_pipe.reader() + tle.gpu.warp_specialize( + [ + ( + _dispatch_pipe_partition, + ( + sync_writer, + symm_buffer, + marker, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_DISPATCH_WARPS_C, + ), + ), + ( + _receiver_pipe_debug_partition, + ( + sync_reader, + symm_buffer, + marker, + EXPECTED_LOCAL_RECV_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + ), + ), + ], + [1], + [40], + ) + + +def _nvcc_path() -> str: + explicit = os.environ.get("NVCC") + if explicit: + return explicit + cuda_home = os.environ.get("CUDA_HOME", "/usr/local/cuda-12.8") + candidate = Path(cuda_home) / "bin" / "nvcc" + return str(candidate if candidate.exists() else "nvcc") + + +def _compile_host_library(arch: str) -> Path: + src = HERE / "ws_userhopper_dispatch_receiver_host.cu" + out = src.with_suffix(".so") + lock = out.with_suffix(".so.lock") + with open(lock, "w") as lock_file: + while True: + try: + fcntl.flock(lock_file, fcntl.LOCK_EX) + break + except BlockingIOError: + time.sleep(0.1) + if out.exists() and out.stat().st_mtime_ns >= src.stat().st_mtime_ns: + return out + tmp_out = out.with_suffix(f".so.{os.getpid()}.tmp") + cmd = [ + _nvcc_path(), + "-shared", + "-Xcompiler", + "-fPIC", + "-rdc=true", + f"-arch={arch}", + f"-I{NVSHMEM_HOME}/include", + f"-L{NVSHMEM_HOME}/lib", + "-lnvshmem_host", + "-lnvshmem_device", + "-o", + str(tmp_out), + str(src), + ] + try: + subprocess.run(cmd, check=True) + os.replace(tmp_out, out) + finally: + if tmp_out.exists(): + tmp_out.unlink() + return out + + +def _tensor_from_pointer(ptr: int, shape: tuple[int, ...], dtype: torch.dtype, device: torch.device): + elem_size = torch.empty((), dtype=dtype).element_size() + numel = 1 + for dim in shape: + numel *= dim + storage = torch._C._construct_storage_from_data_pointer(ptr, device, elem_size * numel) + return torch.empty(0, dtype=dtype, device=device).set_(storage).view(*shape) + + +def _view(ptr: ctypes.c_void_p, offset: int, shape: tuple[int, ...], dtype: torch.dtype, device: torch.device): + return _tensor_from_pointer(ptr.value + offset, shape, dtype, device) + + +def _view_strided( + ptr: ctypes.c_void_p, + offset: int, + shape: tuple[int, ...], + stride: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, +): + if len(shape) != len(stride): + raise ValueError(f"shape/stride rank mismatch: shape={shape}, stride={stride}") + storage_numel = 1 + if shape: + storage_numel = 1 + sum((dim - 1) * step for dim, step in zip(shape, stride)) + base = _tensor_from_pointer(ptr.value + offset, (storage_numel,), dtype, device) + return torch.as_strided(base, shape, stride) + + +def _install_cumodule_hook(lib) -> None: + def cumodule_init_hook(*args, **kwargs): + key = kwargs["key"] + jit_function = kwargs["fn"].jit_function + dev = kwargs["compile"]["device"] + kernel_cache = jit_function.device_caches[dev][0] + kernel = kernel_cache.get(key, None) + assert kernel is not None + kernel._init_handles() + ret = lib.userhopper_ws_nvshmemx_cumodule_init_wrapper(ctypes.c_void_p(kernel.module)) + assert ret == 0, f"nvshmemx_cumodule_init_wrapper failed: {ret}" + + knobs.runtime.jit_post_compile_hook = cumodule_init_hook + + +def _setup_lib(): + lib_path = _compile_host_library(os.environ.get("TRITON_WS_NVSHMEM_ARCH", "sm_90")) + lib = ctypes.CDLL(str(lib_path)) + lib.userhopper_ws_nvshmem_init_wrapper.argtypes = [] + lib.userhopper_ws_nvshmem_init_wrapper.restype = None + lib.userhopper_ws_nvshmem_team_mype_wrapper.argtypes = [] + lib.userhopper_ws_nvshmem_team_mype_wrapper.restype = ctypes.c_int + lib.userhopper_ws_nvshmem_n_pes_wrapper.argtypes = [] + lib.userhopper_ws_nvshmem_n_pes_wrapper.restype = ctypes.c_int + lib.userhopper_ws_nvshmem_alloc_bytes_wrapper.argtypes = [ctypes.c_longlong] + lib.userhopper_ws_nvshmem_alloc_bytes_wrapper.restype = ctypes.c_void_p + lib.userhopper_ws_nvshmem_barrier_all_wrapper.argtypes = [] + lib.userhopper_ws_nvshmem_barrier_all_wrapper.restype = None + lib.userhopper_ws_nvshmem_finalize_wrapper.argtypes = [ctypes.c_void_p] + lib.userhopper_ws_nvshmem_finalize_wrapper.restype = None + return lib diff --git a/python/test/tle/integration/megamoe/megamoe_operator/ws_userhopper_dispatch_receiver_device.cu b/python/test/tle/integration/megamoe/megamoe_operator/ws_userhopper_dispatch_receiver_device.cu new file mode 100644 index 000000000..7e86ce8a4 --- /dev/null +++ b/python/test/tle/integration/megamoe/megamoe_operator/ws_userhopper_dispatch_receiver_device.cu @@ -0,0 +1,1485 @@ +#include +#include +#include + +#include +#include +#include + +namespace { + +constexpr int kNumExpertsLimit = 16; +constexpr int kMaxDispatchWarps = 8; +constexpr uint64_t kNumBarrierSignalBytes = 32; +constexpr int kLCMCandidateBlockM = 384; +constexpr int kMaxCandidateBlockM = 192; +constexpr int kMinCandidateBlockM = 8; +constexpr int kBlockM = 64; + +__device__ __forceinline__ uint64_t align_u64(uint64_t value, uint64_t alignment) { + return (value + alignment - 1) / alignment * alignment; +} + +__device__ __forceinline__ uint32_t align_u32(uint32_t value, uint32_t alignment) { + return (value + alignment - 1) / alignment * alignment; +} + +__device__ __forceinline__ uint32_t num_max_pool_tokens( + uint32_t num_ranks, + uint32_t num_max_tokens_per_rank, + uint32_t num_topk, + uint32_t num_experts_per_rank) { + const uint32_t num_max_recv_tokens = num_ranks * num_max_tokens_per_rank; + const uint32_t num_max_experts_per_token = + num_topk < num_experts_per_rank ? num_topk : num_experts_per_rank; + return align_u32( + num_max_recv_tokens * num_max_experts_per_token + + num_experts_per_rank * (kMaxCandidateBlockM - 1), + kLCMCandidateBlockM); +} + +__device__ __forceinline__ uint32_t num_max_pool_blocks( + uint32_t num_ranks, + uint32_t num_max_tokens_per_rank, + uint32_t num_topk, + uint32_t num_experts_per_rank) { + return num_max_pool_tokens(num_ranks, num_max_tokens_per_rank, num_topk, + num_experts_per_rank) / + kMinCandidateBlockM; +} + +__device__ __forceinline__ uint64_t workspace_bytes( + uint32_t num_ranks, + uint32_t num_experts, + uint32_t num_max_tokens_per_rank, + uint32_t num_topk) { + const uint32_t num_experts_per_rank = num_experts / num_ranks; + const uint32_t max_recv = num_ranks * num_max_tokens_per_rank; + const uint32_t pool_tokens = + num_max_pool_tokens(num_ranks, num_max_tokens_per_rank, num_topk, + num_experts_per_rank); + const uint32_t pool_blocks = + num_max_pool_blocks(num_ranks, num_max_tokens_per_rank, num_topk, + num_experts_per_rank); + + uint64_t bytes = 0; + bytes += kNumBarrierSignalBytes; + bytes += static_cast(num_experts) * sizeof(uint64_t) * 2; + bytes += static_cast(num_experts_per_rank) * sizeof(uint64_t); + bytes += static_cast(align_u32(pool_blocks, 2)) * sizeof(uint32_t); + bytes += static_cast(pool_blocks) * sizeof(uint64_t); + bytes += static_cast(num_experts_per_rank) * num_ranks * max_recv * + sizeof(uint32_t); + bytes += static_cast(pool_tokens) * 3 * sizeof(uint32_t); + return align_u64(bytes, 16); +} + +__device__ __forceinline__ uint64_t *expert_send_count_ptr( + uint8_t *base, + uint32_t expert_idx) { + return reinterpret_cast(base + kNumBarrierSignalBytes) + expert_idx; +} + +__device__ __forceinline__ uint64_t *expert_recv_count_ptr( + uint8_t *base, + uint32_t num_experts, + uint32_t num_experts_per_rank, + uint32_t rank_idx, + uint32_t expert_idx) { + return expert_send_count_ptr(base, num_experts) + + rank_idx * num_experts_per_rank + expert_idx; +} + +__device__ __forceinline__ uint64_t *expert_recv_count_sum_ptr( + uint8_t *base, + uint32_t num_experts, + uint32_t num_experts_per_rank, + uint32_t expert_idx) { + return expert_send_count_ptr(base, num_experts * 2) + expert_idx; +} + +__device__ __forceinline__ uint32_t *l1_arrival_count_ptr( + uint8_t *base, + uint32_t num_experts, + uint32_t num_experts_per_rank, + uint32_t pool_block_idx) { + return reinterpret_cast( + expert_recv_count_sum_ptr(base, num_experts, num_experts_per_rank, + num_experts_per_rank)) + + pool_block_idx; +} + +__device__ __forceinline__ uint64_t *l2_arrival_mask_ptr( + uint8_t *base, + uint32_t num_ranks, + uint32_t num_experts, + uint32_t num_max_tokens_per_rank, + uint32_t num_topk, + uint32_t pool_block_idx) { + const uint32_t num_experts_per_rank = num_experts / num_ranks; + const uint32_t pool_blocks = + num_max_pool_blocks(num_ranks, num_max_tokens_per_rank, num_topk, + num_experts_per_rank); + return reinterpret_cast( + l1_arrival_count_ptr(base, num_experts, num_experts_per_rank, + align_u32(pool_blocks, 2))) + + pool_block_idx; +} + +__device__ __forceinline__ uint32_t *src_token_topk_idx_ptr( + uint8_t *base, + uint32_t num_ranks, + uint32_t num_experts, + uint32_t num_max_tokens_per_rank, + uint32_t num_topk, + uint32_t expert_idx, + uint32_t rank_idx, + uint32_t token_idx) { + const uint32_t num_experts_per_rank = num_experts / num_ranks; + const uint32_t max_recv = num_ranks * num_max_tokens_per_rank; + const uint32_t pool_blocks = + num_max_pool_blocks(num_ranks, num_max_tokens_per_rank, num_topk, + num_experts_per_rank); + return reinterpret_cast( + l2_arrival_mask_ptr(base, num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, + pool_blocks)) + + expert_idx * (num_ranks * max_recv) + rank_idx * max_recv + token_idx; +} + +__device__ __forceinline__ uint32_t *token_src_metadata_ptr( + uint8_t *base, + uint32_t num_ranks, + uint32_t num_experts, + uint32_t num_max_tokens_per_rank, + uint32_t num_topk, + uint32_t pool_token_idx) { + const uint32_t num_experts_per_rank = num_experts / num_ranks; + return reinterpret_cast(src_token_topk_idx_ptr( + base, num_ranks, num_experts, num_max_tokens_per_rank, num_topk, + num_experts_per_rank, 0, 0)) + + pool_token_idx * 3; +} + +__device__ __forceinline__ uint64_t input_token_offset( + uint32_t num_ranks, + uint32_t num_experts, + uint32_t num_max_tokens_per_rank, + uint32_t num_topk) { + return workspace_bytes(num_ranks, num_experts, num_max_tokens_per_rank, + num_topk); +} + +__device__ __forceinline__ uint64_t input_sf_offset( + uint32_t num_ranks, + uint32_t num_experts, + uint32_t num_max_tokens_per_rank, + uint32_t num_topk, + uint32_t hidden) { + return input_token_offset(num_ranks, num_experts, num_max_tokens_per_rank, + num_topk) + + static_cast(num_max_tokens_per_rank) * hidden; +} + +__device__ __forceinline__ uint64_t input_topk_idx_offset( + uint32_t num_ranks, + uint32_t num_experts, + uint32_t num_max_tokens_per_rank, + uint32_t num_topk, + uint32_t hidden) { + return input_sf_offset(num_ranks, num_experts, num_max_tokens_per_rank, + num_topk, hidden) + + static_cast(num_max_tokens_per_rank) * (hidden / 32); +} + +__device__ __forceinline__ uint64_t input_topk_weight_offset( + uint32_t num_ranks, + uint32_t num_experts, + uint32_t num_max_tokens_per_rank, + uint32_t num_topk, + uint32_t hidden) { + return input_topk_idx_offset(num_ranks, num_experts, num_max_tokens_per_rank, + num_topk, hidden) + + static_cast(num_max_tokens_per_rank) * num_topk * + sizeof(int64_t); +} + +__device__ __forceinline__ uint64_t l1_token_offset( + uint32_t num_ranks, + uint32_t num_experts, + uint32_t num_max_tokens_per_rank, + uint32_t num_topk, + uint32_t hidden) { + return input_topk_weight_offset(num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, hidden) + + static_cast(num_max_tokens_per_rank) * num_topk * + sizeof(float); +} + +__device__ __forceinline__ uint64_t l1_sf_offset( + uint32_t num_ranks, + uint32_t num_experts, + uint32_t num_max_tokens_per_rank, + uint32_t num_topk, + uint32_t hidden) { + const uint32_t epr = num_experts / num_ranks; + const uint32_t pool_tokens = + num_max_pool_tokens(num_ranks, num_max_tokens_per_rank, num_topk, epr); + return l1_token_offset(num_ranks, num_experts, num_max_tokens_per_rank, + num_topk, hidden) + + static_cast(pool_tokens) * hidden; +} + +__device__ __forceinline__ uint64_t l1_topk_weight_offset( + uint32_t num_ranks, + uint32_t num_experts, + uint32_t num_max_tokens_per_rank, + uint32_t num_topk, + uint32_t hidden, + uint32_t num_padded_sf_pool_tokens) { + return l1_sf_offset(num_ranks, num_experts, num_max_tokens_per_rank, num_topk, + hidden) + + static_cast(num_padded_sf_pool_tokens) * (hidden / 32); +} + +__device__ __forceinline__ uint64_t l2_token_offset( + uint32_t num_ranks, + uint32_t num_experts, + uint32_t num_max_tokens_per_rank, + uint32_t num_topk, + uint32_t hidden, + uint32_t num_padded_sf_pool_tokens) { + const uint32_t epr = num_experts / num_ranks; + const uint32_t pool_tokens = + num_max_pool_tokens(num_ranks, num_max_tokens_per_rank, num_topk, epr); + return l1_topk_weight_offset(num_ranks, num_experts, num_max_tokens_per_rank, + num_topk, hidden, + num_padded_sf_pool_tokens) + + static_cast(pool_tokens) * sizeof(float); +} + +__device__ __forceinline__ uint64_t l2_sf_offset( + uint32_t num_ranks, + uint32_t num_experts, + uint32_t num_max_tokens_per_rank, + uint32_t num_topk, + uint32_t hidden, + uint32_t intermediate_hidden, + uint32_t num_padded_sf_pool_tokens) { + const uint32_t epr = num_experts / num_ranks; + const uint32_t pool_tokens = + num_max_pool_tokens(num_ranks, num_max_tokens_per_rank, num_topk, epr); + return l2_token_offset(num_ranks, num_experts, num_max_tokens_per_rank, + num_topk, hidden, num_padded_sf_pool_tokens) + + static_cast(pool_tokens) * intermediate_hidden; +} + +__device__ __forceinline__ uint64_t combine_token_offset( + uint32_t num_ranks, + uint32_t num_experts, + uint32_t num_max_tokens_per_rank, + uint32_t num_topk, + uint32_t hidden, + uint32_t intermediate_hidden, + uint32_t num_padded_sf_pool_tokens) { + return l2_sf_offset(num_ranks, num_experts, num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, + num_padded_sf_pool_tokens) + + static_cast(num_padded_sf_pool_tokens) * + (intermediate_hidden / 16u); +} + +__device__ __forceinline__ uint16_t float_to_bf16_bits(float value) { + const uint32_t bits = __float_as_uint(value); + const uint32_t lsb = (bits >> 16) & 1u; + return static_cast((bits + 0x7fffu + lsb) >> 16); +} + +__device__ __forceinline__ float fp8_e4m3fn_to_float(uint8_t value) { + if (value == 0) { + return 0.0f; + } + const int sign = (value & 0x80u) ? -1 : 1; + const int exp = (value >> 3) & 0x0f; + const int mant = value & 0x07; + float mag; + if (exp == 0) { + mag = ldexpf(static_cast(mant) / 8.0f, -6); + } else { + mag = ldexpf(1.0f + static_cast(mant) / 8.0f, exp - 7); + } + return sign < 0 ? -mag : mag; +} + +__device__ __forceinline__ void choose_rank_round_robin( + uint32_t token_idx_in_expert, + const uint32_t *rank_counts, + uint32_t num_ranks, + uint32_t *rank_out, + uint32_t *token_idx_in_rank_out) { + uint32_t remaining[8]; + for (uint32_t r = 0; r < num_ranks; ++r) { + remaining[r] = rank_counts[r]; + } + + uint32_t offset = 0; + uint32_t slot_idx = token_idx_in_expert; + while (true) { + uint32_t active = 0; + uint32_t length = 0xffffffffu; + for (uint32_t r = 0; r < num_ranks; ++r) { + if (remaining[r] > 0) { + ++active; + length = remaining[r] < length ? remaining[r] : length; + } + } + if (active == 0) { + break; + } + const uint32_t num_round_tokens = length * active; + if (slot_idx < num_round_tokens) { + const uint32_t slot_idx_in_round = slot_idx % active; + uint32_t seen = 0; + for (uint32_t r = 0; r < num_ranks; ++r) { + if (remaining[r] == 0) { + continue; + } + if (slot_idx_in_round == seen) { + *rank_out = r; + *token_idx_in_rank_out = offset + slot_idx / active; + return; + } + ++seen; + } + } + slot_idx -= num_round_tokens; + offset += length; + for (uint32_t r = 0; r < num_ranks; ++r) { + remaining[r] -= remaining[r] < length ? remaining[r] : length; + } + } + *rank_out = 0; + *token_idx_in_rank_out = 0; +} + +} // namespace + +extern "C" __device__ void userhopper_ws_dispatch_partition( + uint8_t *symm_buffer, + int num_tokens, + int num_ranks, + int num_experts, + int num_max_tokens_per_rank, + int num_topk, + int hidden) { + if (blockIdx.x != 0) { + return; + } + + const uint32_t rank = static_cast(nvshmem_my_pe()); + const uint32_t epr = static_cast(num_experts / num_ranks); + const uint32_t lane_idx = static_cast(threadIdx.x & 31); + const uint32_t tokens_per_warp = 32u / static_cast(num_topk); + const uint32_t activate_lanes = tokens_per_warp * static_cast(num_topk); + __shared__ uint32_t local_counts[kNumExpertsLimit]; + __shared__ uint32_t local_slots[kNumExpertsLimit]; + for (uint32_t i = lane_idx; i < kNumExpertsLimit; i += 32) { + local_counts[i] = 0; + local_slots[i] = 0; + } + __syncwarp(); + + int64_t *topk_idx = reinterpret_cast( + symm_buffer + input_topk_idx_offset(num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, + hidden)); + + for (uint32_t i = 0; i < static_cast(num_tokens); + i += tokens_per_warp) { + int expert = -1; + const uint32_t token_idx = i + lane_idx / static_cast(num_topk); + const bool active = + lane_idx < activate_lanes && token_idx < static_cast(num_tokens); + const uint32_t token_topk_idx = i * static_cast(num_topk) + lane_idx; + if (active) { + expert = static_cast(topk_idx[token_topk_idx]); + if (expert >= 0 && expert < num_experts && expert < kNumExpertsLimit) { + atomicAdd(local_counts + expert, 1u); + } + } + __syncwarp(); + } + + __syncwarp(); + for (uint32_t expert = lane_idx; + expert < static_cast(num_experts) && expert < kNumExpertsLimit; + expert += 32) { + const uint64_t send_value = + (static_cast(1) << 32) | local_counts[expert]; + local_slots[expert] = static_cast( + nvshmem_uint64_atomic_fetch_add(expert_send_count_ptr(symm_buffer, expert), + send_value, rank)); + } + __syncwarp(); + + for (uint32_t i = 0; i < static_cast(num_tokens); + i += tokens_per_warp) { + int expert_idx = -1; + const uint32_t token_idx = i + lane_idx / static_cast(num_topk); + const bool active = + lane_idx < activate_lanes && token_idx < static_cast(num_tokens); + const uint32_t token_topk_idx = i * static_cast(num_topk) + lane_idx; + if (active) { + expert_idx = static_cast(topk_idx[token_topk_idx]); + } + if (expert_idx >= 0 && expert_idx < num_experts && + expert_idx < kNumExpertsLimit) { + const uint32_t expert = static_cast(expert_idx); + const uint32_t dst_rank = expert / epr; + const uint32_t dst_local_expert = expert - dst_rank * epr; + const uint32_t dst_slot = atomicAdd(local_slots + expert, 1u); + uint32_t *remote_queue = src_token_topk_idx_ptr( + symm_buffer, num_ranks, num_experts, num_max_tokens_per_rank, num_topk, + dst_local_expert, rank, dst_slot); + nvshmem_uint_p(remote_queue, token_topk_idx, dst_rank); + } + __syncwarp(); + } + + __syncwarp(); + nvshmem_quiet(); + + for (uint32_t expert = lane_idx; + expert < static_cast(num_experts) && expert < kNumExpertsLimit; + expert += 32) { + const uint64_t expert_status = *expert_send_count_ptr(symm_buffer, expert); + const uint32_t dst_rank = expert / epr; + const uint32_t dst_local_expert = expert - dst_rank * epr; + nvshmem_uint64_p( + expert_recv_count_ptr(symm_buffer, num_experts, epr, rank, + dst_local_expert), + expert_status & 0xffffffffu, dst_rank); + nvshmem_quiet(); + nvshmem_uint64_atomic_add( + expert_recv_count_sum_ptr(symm_buffer, num_experts, epr, + dst_local_expert), + expert_status, dst_rank); + } + __syncwarp(); + nvshmem_quiet(); +} + +extern "C" __device__ void userhopper_ws_dispatch_partition_cta_warp0( + uint8_t *symm_buffer, + int num_tokens, + int num_ranks, + int num_experts, + int num_max_tokens_per_rank, + int num_topk, + int hidden) { + if (blockIdx.x != 0 || (threadIdx.x >> 5) != 0) { + return; + } + userhopper_ws_dispatch_partition(symm_buffer, num_tokens, num_ranks, + num_experts, num_max_tokens_per_rank, + num_topk, hidden); +} + +extern "C" __device__ void userhopper_ws_dispatch_partition_cta_multiwarp( + uint8_t *symm_buffer, + int num_tokens, + int num_ranks, + int num_experts, + int num_max_tokens_per_rank, + int num_topk, + int hidden, + int num_dispatch_warps) { + if (blockIdx.x != 0) { + return; + } + const uint32_t warp_idx = static_cast(threadIdx.x >> 5); + const uint32_t lane_idx = static_cast(threadIdx.x & 31); + uint32_t active_dispatch_warps = static_cast(num_dispatch_warps); + active_dispatch_warps = + active_dispatch_warps < static_cast(kMaxDispatchWarps) + ? active_dispatch_warps + : static_cast(kMaxDispatchWarps); + if (warp_idx >= active_dispatch_warps) { + return; + } + + const uint32_t rank = static_cast(nvshmem_my_pe()); + const uint32_t epr = static_cast(num_experts / num_ranks); + const uint32_t tokens_per_warp = 32u / static_cast(num_topk); + const uint32_t activate_lanes = tokens_per_warp * static_cast(num_topk); + __shared__ uint32_t local_counts[kNumExpertsLimit]; + __shared__ uint32_t local_slots[kNumExpertsLimit]; + __shared__ uint32_t round_counts[kMaxDispatchWarps * kNumExpertsLimit]; + __shared__ uint32_t round_slots[kMaxDispatchWarps * kNumExpertsLimit]; + __shared__ volatile uint32_t sync_state[6]; + + if (warp_idx == 0) { + for (uint32_t i = lane_idx; i < kNumExpertsLimit; i += 32) { + local_counts[i] = 0; + local_slots[i] = 0; + } + for (uint32_t i = lane_idx; + i < static_cast(kMaxDispatchWarps * kNumExpertsLimit); + i += 32) { + round_counts[i] = 0; + round_slots[i] = 0; + } + if (lane_idx < 6u) { + sync_state[lane_idx] = 0; + } + __syncwarp(); + if (lane_idx == 0) { + sync_state[0] = 1; + } + } else { + while (sync_state[0] == 0) { + } + } + __syncwarp(); + + int64_t *topk_idx = reinterpret_cast( + symm_buffer + input_topk_idx_offset(num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, + hidden)); + + const uint32_t stride = + active_dispatch_warps * tokens_per_warp; + for (uint32_t i = warp_idx * tokens_per_warp; + i < static_cast(num_tokens); i += stride) { + int expert = -1; + const uint32_t token_idx = i + lane_idx / static_cast(num_topk); + const bool active = + lane_idx < activate_lanes && token_idx < static_cast(num_tokens); + const uint32_t token_topk_idx = i * static_cast(num_topk) + lane_idx; + if (active) { + expert = static_cast(topk_idx[token_topk_idx]); + if (expert >= 0 && expert < num_experts && expert < kNumExpertsLimit) { + atomicAdd(local_counts + expert, 1u); + } + } + __syncwarp(); + } + if (lane_idx == 0) { + atomicAdd(const_cast(&sync_state[1]), 1u); + } + while (sync_state[1] < active_dispatch_warps) { + } + __syncwarp(); + + if (warp_idx == 0) { + for (uint32_t expert = lane_idx; + expert < static_cast(num_experts) && expert < kNumExpertsLimit; + expert += 32) { + const uint64_t send_value = + (static_cast(1) << 32) | local_counts[expert]; + local_slots[expert] = static_cast( + nvshmem_uint64_atomic_fetch_add(expert_send_count_ptr(symm_buffer, expert), + send_value, rank)); + } + __syncwarp(); + if (lane_idx == 0) { + sync_state[2] = 1; + } + } else { + while (sync_state[2] == 0) { + } + } + __syncwarp(); + + uint32_t round_idx = 0; + for (uint32_t round_base = 0; round_base < static_cast(num_tokens); + round_base += stride, ++round_idx) { + uint32_t *my_round_counts = + round_counts + warp_idx * static_cast(kNumExpertsLimit); + uint32_t *my_round_slots = + round_slots + warp_idx * static_cast(kNumExpertsLimit); + for (uint32_t expert = lane_idx; expert < kNumExpertsLimit; expert += 32) { + my_round_counts[expert] = 0; + my_round_slots[expert] = 0; + } + __syncwarp(); + + const uint32_t i = round_base + warp_idx * tokens_per_warp; + int expert_idx = -1; + const uint32_t token_idx = i + lane_idx / static_cast(num_topk); + const bool active = + lane_idx < activate_lanes && token_idx < static_cast(num_tokens); + const uint32_t token_topk_idx = i * static_cast(num_topk) + lane_idx; + if (active) { + expert_idx = static_cast(topk_idx[token_topk_idx]); + if (expert_idx >= 0 && expert_idx < num_experts && + expert_idx < kNumExpertsLimit) { + atomicAdd(my_round_counts + static_cast(expert_idx), 1u); + } + } + __syncwarp(); + + const uint32_t count_ready_target = (round_idx + 1u) * active_dispatch_warps; + if (lane_idx == 0) { + atomicAdd(const_cast(&sync_state[3]), 1u); + } + while (sync_state[3] < count_ready_target) { + } + __syncwarp(); + + if (warp_idx == 0) { + for (uint32_t expert = lane_idx; + expert < static_cast(num_experts) && expert < kNumExpertsLimit; + expert += 32) { + uint32_t slot = local_slots[expert]; + for (uint32_t w = 0; w < active_dispatch_warps; ++w) { + const uint32_t idx = w * static_cast(kNumExpertsLimit) + expert; + round_slots[idx] = slot; + slot += round_counts[idx]; + } + local_slots[expert] = slot; + } + __syncwarp(); + if (lane_idx == 0) { + sync_state[4] = round_idx + 1u; + } + } else { + while (sync_state[4] < round_idx + 1u) { + } + } + __syncwarp(); + + if (expert_idx >= 0 && expert_idx < num_experts && + expert_idx < kNumExpertsLimit) { + const uint32_t expert = static_cast(expert_idx); + const uint32_t dst_rank = expert / epr; + const uint32_t dst_local_expert = expert - dst_rank * epr; + const uint32_t dst_slot = atomicAdd(my_round_slots + expert, 1u); + uint32_t *remote_queue = src_token_topk_idx_ptr( + symm_buffer, num_ranks, num_experts, num_max_tokens_per_rank, num_topk, + dst_local_expert, rank, dst_slot); + nvshmem_uint_p(remote_queue, token_topk_idx, dst_rank); + } + __syncwarp(); + + const uint32_t write_done_target = (round_idx + 1u) * active_dispatch_warps; + if (lane_idx == 0) { + atomicAdd(const_cast(&sync_state[5]), 1u); + } + while (sync_state[5] < write_done_target) { + } + __syncwarp(); + } + nvshmem_quiet(); + + if (warp_idx == 0) { + for (uint32_t expert = lane_idx; + expert < static_cast(num_experts) && expert < kNumExpertsLimit; + expert += 32) { + const uint64_t expert_status = *expert_send_count_ptr(symm_buffer, expert); + const uint32_t dst_rank = expert / epr; + const uint32_t dst_local_expert = expert - dst_rank * epr; + nvshmem_uint64_p( + expert_recv_count_ptr(symm_buffer, num_experts, epr, rank, + dst_local_expert), + expert_status & 0xffffffffu, dst_rank); + nvshmem_quiet(); + nvshmem_uint64_atomic_add( + expert_recv_count_sum_ptr(symm_buffer, num_experts, epr, + dst_local_expert), + expert_status, dst_rank); + } + __syncwarp(); + nvshmem_quiet(); + } +} + +extern "C" __device__ void userhopper_ws_receiver_partition( + uint8_t *symm_buffer, + int expected_local_recv_tokens, + int num_ranks, + int num_experts, + int num_max_tokens_per_rank, + int num_topk, + int hidden, + int num_padded_sf_pool_tokens) { + if ((threadIdx.x & 31) != 0 || blockIdx.x != 0) { + return; + } + + const uint32_t rank = static_cast(nvshmem_my_pe()); + const uint32_t epr = static_cast(num_experts / num_ranks); + uint8_t *l1_tokens = + symm_buffer + l1_token_offset(num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, hidden); + uint8_t *l1_sf = + symm_buffer + l1_sf_offset(num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, hidden); + float *l1_topk_weights = reinterpret_cast( + symm_buffer + l1_topk_weight_offset(num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, + hidden, num_padded_sf_pool_tokens)); + const uint64_t remote_x_off = + input_token_offset(num_ranks, num_experts, num_max_tokens_per_rank, + num_topk); + const uint64_t remote_sf_off = + input_sf_offset(num_ranks, num_experts, num_max_tokens_per_rank, num_topk, + hidden); + const uint64_t remote_weight_off = input_topk_weight_offset( + num_ranks, num_experts, num_max_tokens_per_rank, num_topk, hidden); + const uint32_t sf_count = hidden / 128; + const uint32_t input_sf_bytes = hidden / 32; + uint32_t pool_token_base = 0; + volatile uint32_t *configured_expected = + reinterpret_cast(symm_buffer); + const bool use_configured_expected = configured_expected[7] == 0x45585043u; + + for (uint32_t local_expert = 0; local_expert < epr; ++local_expert) { + uint32_t expected_for_expert = + use_configured_expected ? configured_expected[local_expert] : 0; + if (!use_configured_expected) { + for (uint32_t topk = 0; topk < static_cast(num_topk); ++topk) { + if ((topk % epr) == local_expert) { + expected_for_expert += + static_cast(expected_local_recv_tokens) / + static_cast(num_topk); + } + } + } + if (expected_for_expert == 0) { + continue; + } + + volatile uint64_t *sum_ptr = reinterpret_cast( + expert_recv_count_sum_ptr(symm_buffer, num_experts, epr, + local_expert)); + uint64_t observed = 0; + while (true) { + observed = *sum_ptr; + if ((observed & 0xffffffffu) >= expected_for_expert) { + break; + } + } + + uint32_t rank_counts[8]; + uint32_t total = 0; + while (true) { + total = 0; + for (uint32_t r = 0; r < static_cast(num_ranks); ++r) { + const uint64_t count = *reinterpret_cast( + expert_recv_count_ptr(symm_buffer, num_experts, epr, r, + local_expert)); + rank_counts[r] = static_cast(count & 0xffffffffu); + total += rank_counts[r]; + } + if (total >= expected_for_expert) { + break; + } + } + + const uint32_t expert_block_offset = pool_token_base / kBlockM; + for (uint32_t token_idx_in_expert = 0; token_idx_in_expert < total; + ++token_idx_in_expert) { + uint32_t src_rank = 0; + uint32_t token_idx_in_rank = 0; + choose_rank_round_robin(token_idx_in_expert, rank_counts, num_ranks, + &src_rank, &token_idx_in_rank); + + volatile uint32_t *queue_ptr = reinterpret_cast( + src_token_topk_idx_ptr( + symm_buffer, num_ranks, num_experts, num_max_tokens_per_rank, + num_topk, local_expert, src_rank, token_idx_in_rank)); + const uint32_t src_token_topk_idx = *queue_ptr; + const uint32_t src_token = src_token_topk_idx / num_topk; + const uint32_t src_topk = src_token_topk_idx - src_token * num_topk; + const uint32_t pool_idx = pool_token_base + token_idx_in_expert; + + nvshmem_getmem(l1_tokens + static_cast(pool_idx) * hidden, + symm_buffer + remote_x_off + + static_cast(src_token) * hidden, + static_cast(hidden), src_rank); + for (uint32_t sf_idx = 0; sf_idx < sf_count; ++sf_idx) { + nvshmem_getmem( + l1_sf + (static_cast(sf_idx) * num_padded_sf_pool_tokens + + pool_idx) * + sizeof(float), + symm_buffer + remote_sf_off + + static_cast(src_token) * input_sf_bytes + + static_cast(sf_idx) * sizeof(float), + sizeof(float), src_rank); + } + + float *remote_weight = + reinterpret_cast(symm_buffer + remote_weight_off) + + src_token_topk_idx; + l1_topk_weights[pool_idx] = nvshmem_float_g(remote_weight, src_rank); + + uint32_t *metadata = token_src_metadata_ptr( + symm_buffer, num_ranks, num_experts, num_max_tokens_per_rank, + num_topk, pool_idx); + metadata[0] = src_rank; + metadata[1] = src_token; + metadata[2] = src_topk; + atomicAdd(l1_arrival_count_ptr( + symm_buffer, num_experts, epr, + expert_block_offset + token_idx_in_expert / kBlockM), + 1u); + } + pool_token_base += align_u32(total, kBlockM); + } +} + +extern "C" __device__ void userhopper_ws_compute_stub_partition( + uint8_t *symm_buffer, + uint8_t *l1_weights, + float *l1_weights_sf, + uint8_t *l2_weights, + float *l2_weights_sf, + int num_ranks, + int num_experts, + int num_max_tokens_per_rank, + int num_topk, + int hidden, + int intermediate_hidden, + int num_padded_sf_pool_tokens, + int compute_full_hidden, + int compute_parallel, + int compute_worker_warps) { + const uint32_t lane_idx = static_cast(threadIdx.x & 31); + const uint32_t active_compute_warps = + compute_worker_warps > 0 ? static_cast(compute_worker_warps) : 1u; + const uint32_t worker_warp_idx = + static_cast(threadIdx.x >> 5) % active_compute_warps; + if (blockIdx.x != 0 || + (compute_parallel == 0 && (worker_warp_idx != 0 || lane_idx != 0))) { + return; + } + + const uint32_t epr = static_cast(num_experts / num_ranks); + volatile uint32_t *debug = reinterpret_cast(symm_buffer); + const bool use_configured_expected = debug[7] == 0x45585043u; + if (!use_configured_expected) { + if (worker_warp_idx == 0 && lane_idx == 0) { + debug[4] = 0xffffffffu; + debug[5] = 0xffffffffu; + debug[6] = 0xBAD00001u; + } + return; + } + + uint8_t *l1_tokens = + symm_buffer + l1_token_offset(num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, hidden); + float *l1_sf = reinterpret_cast( + symm_buffer + l1_sf_offset(num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, hidden)); + float *l1_topk_weights = reinterpret_cast( + symm_buffer + l1_topk_weight_offset(num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, + hidden, num_padded_sf_pool_tokens)); + uint8_t *l2_tokens = + symm_buffer + l2_token_offset(num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, hidden, + num_padded_sf_pool_tokens); + float *l2_sf = reinterpret_cast( + symm_buffer + + l2_sf_offset(num_ranks, num_experts, num_max_tokens_per_rank, num_topk, + hidden, static_cast(intermediate_hidden), + num_padded_sf_pool_tokens)); + const uint64_t l2_token_bytes = static_cast( + num_max_pool_tokens(static_cast(num_ranks), + static_cast(num_max_tokens_per_rank), + static_cast(num_topk), epr)) * + static_cast(intermediate_hidden); + uint32_t *l2_debug = reinterpret_cast( + l2_tokens + l2_token_bytes - 8u * sizeof(uint32_t)); + const uint64_t combine_base_off = combine_token_offset( + static_cast(num_ranks), static_cast(num_experts), + static_cast(num_max_tokens_per_rank), + static_cast(num_topk), static_cast(hidden), + static_cast(intermediate_hidden), + static_cast(num_padded_sf_pool_tokens)); + __shared__ float parallel_l1_values[64]; + __shared__ float parallel_l2_scale; + __shared__ volatile uint32_t compute_sync[5]; + if (compute_parallel != 0) { + if (worker_warp_idx == 0 && lane_idx < 5u) { + compute_sync[lane_idx] = 0; + } + __syncwarp(); + if (worker_warp_idx == 0 && lane_idx == 0) { + compute_sync[0] = 1; + } else { + while (compute_sync[0] == 0) { + } + } + __syncwarp(); + } + + uint32_t expected_total = 0; + uint32_t observed_total = 0; + uint32_t checksum = 2166136261u; + uint32_t l1_weight_checksum = 2166136261u; + float l1_scalar_sum = 0.0f; + uint32_t pool_token_base = 0; + uint32_t compute_barrier_epoch = 0; + uint32_t compute_token_epoch = 0; + for (uint32_t local_expert = 0; local_expert < epr; ++local_expert) { + const uint32_t expected_for_expert = debug[local_expert]; + if (worker_warp_idx == 0 && lane_idx == 0) { + expected_total += expected_for_expert; + } + const uint32_t first_block = pool_token_base / kBlockM; + const uint32_t num_blocks = + expected_for_expert == 0 ? 0 : ((expected_for_expert + kBlockM - 1) / kBlockM); + if (worker_warp_idx == 0 && lane_idx == 0) { + for (uint32_t block = 0; block < num_blocks; ++block) { + observed_total += *reinterpret_cast( + l1_arrival_count_ptr(symm_buffer, num_experts, epr, first_block + block)); + } + } + for (uint32_t token_idx_in_expert = 0; token_idx_in_expert < expected_for_expert; + ++token_idx_in_expert) { + const uint32_t pool_idx = pool_token_base + token_idx_in_expert; + const uint32_t *metadata = token_src_metadata_ptr( + symm_buffer, num_ranks, num_experts, num_max_tokens_per_rank, num_topk, + pool_idx); + const uint32_t src_rank_for_combine = metadata[0]; + const uint32_t src_token_for_combine = metadata[1]; + const uint32_t src_topk_for_combine = metadata[2]; + const float topk_weight_for_l2 = l1_topk_weights[pool_idx]; + const uint32_t compute_h = + compute_full_hidden != 0 ? static_cast(hidden) + : (hidden < 32 ? hidden : 32); + const uint32_t compute_i = static_cast(intermediate_hidden); + const uint64_t expert_weight_base = + static_cast(local_expert) * + (2u * static_cast(intermediate_hidden)) * + static_cast(hidden); + const uint64_t gate_weight_base = expert_weight_base; + const uint32_t sf_stride = + (2u * static_cast(intermediate_hidden) / 128u) * + (static_cast(hidden) / 128u); + const uint64_t sf_base = static_cast(local_expert) * sf_stride; + if (worker_warp_idx == 0 && lane_idx == 0) { + checksum ^= static_cast( + l1_tokens[static_cast(pool_idx) * hidden]); + checksum *= 16777619u; + checksum ^= __float_as_uint(l1_topk_weights[pool_idx]); + checksum *= 16777619u; + uint32_t weighted = 0; + for (uint32_t h = 0; h < compute_h; ++h) { + const uint8_t token_raw = + l1_tokens[static_cast(pool_idx) * hidden + h]; + const uint32_t token_byte = static_cast(token_raw); + const uint32_t weight_byte = static_cast( + l1_weights[gate_weight_base + h]); + weighted += token_byte * weight_byte; + } + const float sf = l1_weights_sf + [static_cast(local_expert) * sf_stride]; + l1_weight_checksum ^= weighted; + l1_weight_checksum *= 16777619u; + l1_weight_checksum ^= __float_as_uint(sf); + l1_weight_checksum *= 16777619u; + } + const uint32_t num_l2_sf_groups = compute_i / 64u; + for (uint32_t sf_group = 0; sf_group < num_l2_sf_groups; ++sf_group) { + const uint32_t first_in_group = + compute_parallel != 0 ? worker_warp_idx * 32u + lane_idx : 0u; + const uint32_t in_group_stride = + compute_parallel != 0 ? active_compute_warps * 32u : 1u; + for (uint32_t in_group = first_in_group; in_group < 64u; + in_group += in_group_stride) { + const uint32_t ii = sf_group * 64u + in_group; + const uint32_t group = ii / 8u; + const uint32_t lane = ii - group * 8u; + const uint32_t gate_row = group * 16u + lane; + const uint32_t up_row = gate_row + 8u; + const uint64_t gate_row_base = + expert_weight_base + static_cast(gate_row) * hidden; + const uint64_t up_row_base = + expert_weight_base + static_cast(up_row) * hidden; + const float gate_sf_for_i = l1_weights_sf + [sf_base + (ii / 128u) * (static_cast(hidden) / 128u)]; + const float up_sf_for_i = l1_weights_sf + [sf_base + + ((static_cast(intermediate_hidden) + ii) / 128u) * + (static_cast(hidden) / 128u)]; + float gate_acc = 0.0f; + float up_acc = 0.0f; + for (uint32_t h = 0; h < compute_h; ++h) { + const uint8_t token_raw = + l1_tokens[static_cast(pool_idx) * hidden + h]; + const float token_sf_for_h = + l1_sf[static_cast(h / 128u) * + static_cast(num_padded_sf_pool_tokens) + + pool_idx]; + const float token_value = + fp8_e4m3fn_to_float(token_raw) * token_sf_for_h; + const uint32_t h_sf_block = h / 128u; + const float gate_sf_for_h = l1_weights_sf + [sf_base + + (ii / 128u) * (static_cast(hidden) / 128u) + + h_sf_block]; + const float up_sf_for_h = l1_weights_sf + [sf_base + + ((static_cast(intermediate_hidden) + ii) / 128u) * + (static_cast(hidden) / 128u) + + h_sf_block]; + gate_acc += + token_value * fp8_e4m3fn_to_float(l1_weights[gate_row_base + h]) * + (compute_full_hidden != 0 ? gate_sf_for_h : gate_sf_for_i); + up_acc += + token_value * fp8_e4m3fn_to_float(l1_weights[up_row_base + h]) * + (compute_full_hidden != 0 ? up_sf_for_h : up_sf_for_i); + } + const float sigmoid_gate = 1.0f / (1.0f + expf(-gate_acc)); + const float swiglu = gate_acc * sigmoid_gate * up_acc; + const float weighted_swiglu = swiglu * topk_weight_for_l2; + if (compute_parallel != 0) { + parallel_l1_values[in_group] = weighted_swiglu; + } else { + const float abs_value = fabsf(weighted_swiglu); + parallel_l1_values[in_group] = weighted_swiglu; + parallel_l2_scale = + in_group == 0 || abs_value > parallel_l2_scale ? abs_value : parallel_l2_scale; + } + } + if (compute_parallel != 0) { + __syncwarp(); + if (lane_idx == 0) { + atomicAdd(const_cast(&compute_sync[1]), 1u); + } + const uint32_t target = (compute_barrier_epoch + 1u) * active_compute_warps; + while (compute_sync[1] < target) { + } + } + if (worker_warp_idx == 0 && lane_idx == 0) { + float max_abs = 0.0f; + for (uint32_t in_group = 0; in_group < 64u; ++in_group) { + l1_scalar_sum += parallel_l1_values[in_group]; + const float abs_value = fabsf(parallel_l1_values[in_group]); + max_abs = abs_value > max_abs ? abs_value : max_abs; + } + parallel_l2_scale = max_abs > 0.0f ? max_abs / 448.0f : 1.0f; + l2_sf[static_cast(sf_group) * + static_cast(num_padded_sf_pool_tokens) + + pool_idx] = parallel_l2_scale; + } + if (compute_parallel != 0) { + __syncwarp(); + if (lane_idx == 0) { + atomicAdd(const_cast(&compute_sync[2]), 1u); + } + const uint32_t target = (compute_barrier_epoch + 1u) * active_compute_warps; + while (compute_sync[2] < target) { + } + } + for (uint32_t in_group = first_in_group; in_group < 64u; + in_group += in_group_stride) { + const uint32_t ii = sf_group * 64u + in_group; + l2_tokens[static_cast(pool_idx) * + static_cast(intermediate_hidden) + + ii] = __nv_cvt_float_to_fp8( + parallel_l1_values[in_group] / parallel_l2_scale, + __NV_SATFINITE, __NV_E4M3); + } + if (compute_parallel != 0) { + __syncwarp(); + if (lane_idx == 0) { + atomicAdd(const_cast(&compute_sync[3]), 1u); + } + const uint32_t target = (compute_barrier_epoch + 1u) * active_compute_warps; + while (compute_sync[3] < target) { + } + } + ++compute_barrier_epoch; + } + const uint64_t l2_weight_expert_base = + static_cast(local_expert) * static_cast(hidden) * + static_cast(intermediate_hidden); + const uint32_t l2_sf_stride = + (static_cast(hidden) / 128u) * + (static_cast(intermediate_hidden) / 128u); + const uint64_t l2_sf_base = static_cast(local_expert) * l2_sf_stride; + if (compute_parallel != 0) { + __syncwarp(); + if (lane_idx == 0) { + atomicAdd(const_cast(&compute_sync[4]), 1u); + } + const uint32_t target = (compute_token_epoch + 1u) * active_compute_warps; + while (compute_sync[4] < target) { + } + ++compute_token_epoch; + } + const uint32_t first_h_pair = + compute_parallel != 0 ? (worker_warp_idx * 32u + lane_idx) * 2u : 0u; + const uint32_t h_pair_stride = + compute_parallel != 0 ? active_compute_warps * 64u : 2u; + for (uint32_t h_pair = first_h_pair; + h_pair < static_cast(hidden); h_pair += h_pair_stride) { + float accum0 = 0.0f; + float accum1 = 0.0f; + for (uint32_t ii = 0; ii < compute_i; ++ii) { + const uint32_t sf_group = ii / 64u; + const float act_scale = + l2_sf[static_cast(sf_group) * + static_cast(num_padded_sf_pool_tokens) + + pool_idx]; + const float act_value = + fp8_e4m3fn_to_float(l2_tokens[static_cast(pool_idx) * + static_cast(intermediate_hidden) + + ii]) * + act_scale; + const uint32_t sf_i_block = ii / 128u; + const uint32_t sf_h_block0 = h_pair / 128u; + const uint32_t sf_h_block1 = (h_pair + 1u) / 128u; + const float weight_sf0 = + l2_weights_sf[l2_sf_base + + sf_h_block0 * + (static_cast(intermediate_hidden) / 128u) + + sf_i_block]; + const float weight_sf1 = + l2_weights_sf[l2_sf_base + + sf_h_block1 * + (static_cast(intermediate_hidden) / 128u) + + sf_i_block]; + accum0 += act_value * + fp8_e4m3fn_to_float( + l2_weights[l2_weight_expert_base + + static_cast(h_pair) * + static_cast(intermediate_hidden) + + ii]) * + weight_sf0; + accum1 += act_value * + fp8_e4m3fn_to_float( + l2_weights[l2_weight_expert_base + + static_cast(h_pair + 1u) * + static_cast(intermediate_hidden) + + ii]) * + weight_sf1; + } + const uint32_t packed = + static_cast(float_to_bf16_bits(accum0)) | + (static_cast(float_to_bf16_bits(accum1)) << 16); + uint32_t *remote_combine = reinterpret_cast( + symm_buffer + combine_base_off + + (static_cast(src_topk_for_combine) * + static_cast(num_max_tokens_per_rank) + + src_token_for_combine) * + static_cast(hidden) * sizeof(uint16_t) + + static_cast(h_pair) * sizeof(uint16_t)); + nvshmem_uint_p(remote_combine, packed, static_cast(src_rank_for_combine)); + } + if (compute_parallel != 0) { + __syncwarp(); + } + } + pool_token_base += align_u32(expected_for_expert, kBlockM); + } + + nvshmem_quiet(); + if (worker_warp_idx == 0 && lane_idx == 0) { + l2_debug[0] = checksum; + l2_debug[1] = observed_total; + l2_debug[2] = expected_total; + l2_debug[3] = observed_total == expected_total ? 0xC0DEC0DEu : 0xBAD00002u; + l2_debug[4] = l1_weight_checksum; + l2_debug[5] = 0x1A10C001u; + l2_debug[6] = __float_as_uint(l1_scalar_sum); + l2_debug[7] = 0x51A10F32u; + debug[4] = observed_total; + debug[5] = expected_total; + debug[6] = observed_total == expected_total ? 0xC0DEC0DEu : 0xBAD00002u; + } +} + +extern "C" __device__ void userhopper_ws_combine_reduce_partition( + uint8_t *symm_buffer, + uint8_t *y, + int num_ranks, + int num_experts, + int num_max_tokens_per_rank, + int num_topk, + int hidden, + int intermediate_hidden, + int num_padded_sf_pool_tokens, + int cleanup_workspace) { + if ((threadIdx.x & 31) != 0 || blockIdx.x != 0) { + return; + } + + nvshmem_barrier_all(); + + const uint64_t combine_base_off = combine_token_offset( + static_cast(num_ranks), static_cast(num_experts), + static_cast(num_max_tokens_per_rank), + static_cast(num_topk), static_cast(hidden), + static_cast(intermediate_hidden), + static_cast(num_padded_sf_pool_tokens)); + uint16_t *combine_tokens = reinterpret_cast( + symm_buffer + combine_base_off); + uint16_t *y_bf16 = reinterpret_cast(y); + for (uint32_t token = 0; token < static_cast(num_max_tokens_per_rank); + ++token) { + for (uint32_t h_pair = 0; h_pair < static_cast(hidden); h_pair += 2u) { + float accum0 = 0.0f; + float accum1 = 0.0f; + for (uint32_t topk = 0; topk < static_cast(num_topk); ++topk) { + const uint64_t combine_idx = + (static_cast(topk) * + static_cast(num_max_tokens_per_rank) + + token) * + static_cast(hidden) + + h_pair; + accum0 += __bfloat162float(*reinterpret_cast<__nv_bfloat16 *>( + combine_tokens + combine_idx)); + accum1 += __bfloat162float(*reinterpret_cast<__nv_bfloat16 *>( + combine_tokens + combine_idx + 1u)); + } + y_bf16[static_cast(token) * static_cast(hidden) + + h_pair] = float_to_bf16_bits(accum0); + y_bf16[static_cast(token) * static_cast(hidden) + + h_pair + 1u] = float_to_bf16_bits(accum1); + } + } + + if (cleanup_workspace == 0) { + return; + } + + const uint32_t epr = static_cast(num_experts / num_ranks); + const uint32_t pool_blocks = num_max_pool_blocks( + static_cast(num_ranks), + static_cast(num_max_tokens_per_rank), + static_cast(num_topk), + epr); + + for (uint32_t expert = 0; expert < static_cast(num_experts); ++expert) { + *expert_send_count_ptr(symm_buffer, expert) = 0; + } + for (uint32_t rank_idx = 0; rank_idx < static_cast(num_ranks); ++rank_idx) { + for (uint32_t local_expert = 0; local_expert < epr; ++local_expert) { + *expert_recv_count_ptr(symm_buffer, static_cast(num_experts), epr, + rank_idx, local_expert) = 0; + } + } + for (uint32_t local_expert = 0; local_expert < epr; ++local_expert) { + *expert_recv_count_sum_ptr(symm_buffer, static_cast(num_experts), epr, + local_expert) = 0; + } + for (uint32_t block = 0; block < pool_blocks; ++block) { + *l1_arrival_count_ptr(symm_buffer, static_cast(num_experts), epr, + block) = 0; + *l2_arrival_mask_ptr(symm_buffer, static_cast(num_ranks), + static_cast(num_experts), + static_cast(num_max_tokens_per_rank), + static_cast(num_topk), block) = 0; + } +} + +extern "C" __device__ void userhopper_ws_tldot_combine_write_partition( + uint8_t *symm_buffer, + float *l2_out, + int num_ranks, + int num_experts, + int num_max_tokens_per_rank, + int num_topk, + int hidden, + int intermediate_hidden, + int num_padded_sf_pool_tokens) { + if (blockIdx.x != 0) { + return; + } + + const uint32_t lane_idx = static_cast(threadIdx.x & 31); + const uint32_t epr = static_cast(num_experts / num_ranks); + volatile uint32_t *debug = reinterpret_cast(symm_buffer); + if (debug[7] != 0x45585043u) { + return; + } + + const uint64_t combine_base_off = combine_token_offset( + static_cast(num_ranks), static_cast(num_experts), + static_cast(num_max_tokens_per_rank), + static_cast(num_topk), static_cast(hidden), + static_cast(intermediate_hidden), + static_cast(num_padded_sf_pool_tokens)); + + uint32_t pool_token_base = 0; + for (uint32_t local_expert = 0; local_expert < epr; ++local_expert) { + const uint32_t expected_for_expert = debug[local_expert]; + for (uint32_t token_idx_in_expert = 0; token_idx_in_expert < expected_for_expert; + ++token_idx_in_expert) { + const uint32_t pool_idx = pool_token_base + token_idx_in_expert; + const uint32_t *metadata = token_src_metadata_ptr( + symm_buffer, static_cast(num_ranks), + static_cast(num_experts), + static_cast(num_max_tokens_per_rank), + static_cast(num_topk), pool_idx); + const uint32_t src_rank = metadata[0]; + const uint32_t src_token = metadata[1]; + const uint32_t src_topk = metadata[2]; + for (uint32_t h_pair = lane_idx * 2u; h_pair < static_cast(hidden); + h_pair += 64u) { + const float value0 = + l2_out[static_cast(pool_idx) * static_cast(hidden) + + h_pair]; + const float value1 = + l2_out[static_cast(pool_idx) * static_cast(hidden) + + h_pair + 1u]; + const uint32_t packed = + static_cast(float_to_bf16_bits(value0)) | + (static_cast(float_to_bf16_bits(value1)) << 16); + uint32_t *remote_combine = reinterpret_cast( + symm_buffer + combine_base_off + + (static_cast(src_topk) * + static_cast(num_max_tokens_per_rank) + + src_token) * + static_cast(hidden) * sizeof(uint16_t) + + static_cast(h_pair) * sizeof(uint16_t)); + nvshmem_uint_p(remote_combine, packed, static_cast(src_rank)); + } + } + pool_token_base += align_u32(expected_for_expert, kBlockM); + } + + nvshmem_quiet(); +} + +extern "C" __device__ void userhopper_ws_receiver_partition_bounded( + uint8_t *symm_buffer, + int expected_local_recv_tokens, + int num_ranks, + int num_experts, + int num_max_tokens_per_rank, + int num_topk, + int hidden, + int num_padded_sf_pool_tokens) { + if ((threadIdx.x & 31) != 0 || blockIdx.x != 0) { + return; + } + + volatile uint32_t *debug = reinterpret_cast(symm_buffer); + debug[0] = 0xB001u; + debug[6] = static_cast(threadIdx.x); + debug[7] = static_cast(nvshmem_my_pe()); + + const uint32_t rank = static_cast(nvshmem_my_pe()); + const uint32_t epr = static_cast(num_experts / num_ranks); + const uint32_t local_expert = 0; + volatile uint64_t *sum_ptr = reinterpret_cast( + expert_recv_count_sum_ptr(symm_buffer, num_experts, epr, local_expert)); + + uint64_t observed = 0; + uint32_t spins = 0; + constexpr uint32_t kMaxSpins = 1u << 24; + while (spins < kMaxSpins) { + observed = *sum_ptr; + debug[1] = static_cast(observed & 0xffffffffu); + debug[2] = static_cast(observed >> 32); + if ((observed & 0xffffffffu) >= + static_cast(expected_local_recv_tokens)) { + break; + } + ++spins; + } + if (spins >= kMaxSpins) { + debug[0] = 0xB0E1u; + return; + } + + uint32_t rank_counts[8]; + uint32_t total = 0; + spins = 0; + while (spins < kMaxSpins) { + total = 0; + for (uint32_t r = 0; r < static_cast(num_ranks); ++r) { + const uint64_t count = *reinterpret_cast( + expert_recv_count_ptr(symm_buffer, num_experts, epr, r, + local_expert)); + rank_counts[r] = static_cast(count & 0xffffffffu); + total += rank_counts[r]; + if (r < 2) { + debug[3 + r] = rank_counts[r]; + } + } + debug[5] = total; + if (total >= static_cast(expected_local_recv_tokens)) { + break; + } + ++spins; + } + if (spins >= kMaxSpins) { + debug[0] = 0xB0E2u; + return; + } + + uint8_t *l1_tokens = + symm_buffer + l1_token_offset(num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, hidden); + uint8_t *l1_sf = + symm_buffer + l1_sf_offset(num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, hidden); + float *l1_topk_weights = reinterpret_cast( + symm_buffer + l1_topk_weight_offset(num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, + hidden, num_padded_sf_pool_tokens)); + const uint64_t remote_x_off = + input_token_offset(num_ranks, num_experts, num_max_tokens_per_rank, + num_topk); + const uint64_t remote_sf_off = + input_sf_offset(num_ranks, num_experts, num_max_tokens_per_rank, num_topk, + hidden); + const uint64_t remote_weight_off = input_topk_weight_offset( + num_ranks, num_experts, num_max_tokens_per_rank, num_topk, hidden); + const uint32_t sf_count = hidden / 128; + const uint32_t input_sf_bytes = hidden / 32; + + for (uint32_t pool_idx = 0; pool_idx < total; ++pool_idx) { + uint32_t src_rank = 0; + uint32_t token_idx_in_rank = 0; + choose_rank_round_robin(pool_idx, rank_counts, num_ranks, &src_rank, + &token_idx_in_rank); + + volatile uint32_t *queue_ptr = reinterpret_cast( + src_token_topk_idx_ptr( + symm_buffer, num_ranks, num_experts, num_max_tokens_per_rank, num_topk, + local_expert, src_rank, token_idx_in_rank)); + const uint32_t src_token_topk_idx = *queue_ptr; + const uint32_t src_token = src_token_topk_idx / num_topk; + const uint32_t src_topk = src_token_topk_idx - src_token * num_topk; + + nvshmem_getmem(l1_tokens + static_cast(pool_idx) * hidden, + symm_buffer + remote_x_off + + static_cast(src_token) * hidden, + static_cast(hidden), src_rank); + for (uint32_t sf_idx = 0; sf_idx < sf_count; ++sf_idx) { + nvshmem_getmem( + l1_sf + (static_cast(sf_idx) * num_padded_sf_pool_tokens + + pool_idx) * + sizeof(float), + symm_buffer + remote_sf_off + + static_cast(src_token) * input_sf_bytes + + static_cast(sf_idx) * sizeof(float), + sizeof(float), src_rank); + } + + float *remote_weight = + reinterpret_cast(symm_buffer + remote_weight_off) + + src_token_topk_idx; + l1_topk_weights[pool_idx] = nvshmem_float_g(remote_weight, src_rank); + + uint32_t *metadata = token_src_metadata_ptr( + symm_buffer, num_ranks, num_experts, num_max_tokens_per_rank, num_topk, + pool_idx); + metadata[0] = src_rank; + metadata[1] = src_token; + metadata[2] = src_topk; + atomicAdd(l1_arrival_count_ptr(symm_buffer, num_experts, epr, + pool_idx / kBlockM), + 1u); + } + debug[0] = 0xB0FFu; +} diff --git a/python/test/tle/integration/megamoe/megamoe_operator/ws_userhopper_dispatch_receiver_extern_call.py b/python/test/tle/integration/megamoe/megamoe_operator/ws_userhopper_dispatch_receiver_extern_call.py new file mode 100644 index 000000000..4dbb04278 --- /dev/null +++ b/python/test/tle/integration/megamoe/megamoe_operator/ws_userhopper_dispatch_receiver_extern_call.py @@ -0,0 +1,293 @@ +import triton.language as tl +import triton.language.core as core + + +@core.extern +def userhopper_ws_dispatch_partition(symm_buffer, num_tokens, num_ranks, num_experts, num_max_tokens_per_rank, + num_topk, hidden, _semantic=None): + return core.extern_call( + "", + "", + [ + symm_buffer, + tl.cast(num_tokens, tl.int32, _semantic=_semantic), + tl.cast(num_ranks, tl.int32, _semantic=_semantic), + tl.cast(num_experts, tl.int32, _semantic=_semantic), + tl.cast(num_max_tokens_per_rank, tl.int32, _semantic=_semantic), + tl.cast(num_topk, tl.int32, _semantic=_semantic), + tl.cast(hidden, tl.int32, _semantic=_semantic), + ], + { + ( + core.pointer_type(core.dtype("uint8")), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + ): ("userhopper_ws_dispatch_partition", ()), + }, + is_pure=False, + _semantic=_semantic, + ) + + +@core.extern +def userhopper_ws_dispatch_partition_cta_warp0(symm_buffer, num_tokens, num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, hidden, _semantic=None): + return core.extern_call( + "", + "", + [ + symm_buffer, + tl.cast(num_tokens, tl.int32, _semantic=_semantic), + tl.cast(num_ranks, tl.int32, _semantic=_semantic), + tl.cast(num_experts, tl.int32, _semantic=_semantic), + tl.cast(num_max_tokens_per_rank, tl.int32, _semantic=_semantic), + tl.cast(num_topk, tl.int32, _semantic=_semantic), + tl.cast(hidden, tl.int32, _semantic=_semantic), + ], + { + ( + core.pointer_type(core.dtype("uint8")), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + ): ("userhopper_ws_dispatch_partition_cta_warp0", ()), + }, + is_pure=False, + _semantic=_semantic, + ) + + +@core.extern +def userhopper_ws_dispatch_partition_cta_multiwarp(symm_buffer, num_tokens, num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, hidden, + num_dispatch_warps, _semantic=None): + return core.extern_call( + "", + "", + [ + symm_buffer, + tl.cast(num_tokens, tl.int32, _semantic=_semantic), + tl.cast(num_ranks, tl.int32, _semantic=_semantic), + tl.cast(num_experts, tl.int32, _semantic=_semantic), + tl.cast(num_max_tokens_per_rank, tl.int32, _semantic=_semantic), + tl.cast(num_topk, tl.int32, _semantic=_semantic), + tl.cast(hidden, tl.int32, _semantic=_semantic), + tl.cast(num_dispatch_warps, tl.int32, _semantic=_semantic), + ], + { + ( + core.pointer_type(core.dtype("uint8")), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + ): ("userhopper_ws_dispatch_partition_cta_multiwarp", ()), + }, + is_pure=False, + _semantic=_semantic, + ) + + +@core.extern +def userhopper_ws_receiver_partition(symm_buffer, expected_local_recv_tokens, num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, hidden, num_padded_sf_pool_tokens, + _semantic=None): + return core.extern_call( + "", + "", + [ + symm_buffer, + tl.cast(expected_local_recv_tokens, tl.int32, _semantic=_semantic), + tl.cast(num_ranks, tl.int32, _semantic=_semantic), + tl.cast(num_experts, tl.int32, _semantic=_semantic), + tl.cast(num_max_tokens_per_rank, tl.int32, _semantic=_semantic), + tl.cast(num_topk, tl.int32, _semantic=_semantic), + tl.cast(hidden, tl.int32, _semantic=_semantic), + tl.cast(num_padded_sf_pool_tokens, tl.int32, _semantic=_semantic), + ], + { + ( + core.pointer_type(core.dtype("uint8")), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + ): ("userhopper_ws_receiver_partition", ()), + }, + is_pure=False, + _semantic=_semantic, + ) + + +@core.extern +def userhopper_ws_receiver_partition_bounded(symm_buffer, expected_local_recv_tokens, num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, hidden, + num_padded_sf_pool_tokens, _semantic=None): + return core.extern_call( + "", + "", + [ + symm_buffer, + tl.cast(expected_local_recv_tokens, tl.int32, _semantic=_semantic), + tl.cast(num_ranks, tl.int32, _semantic=_semantic), + tl.cast(num_experts, tl.int32, _semantic=_semantic), + tl.cast(num_max_tokens_per_rank, tl.int32, _semantic=_semantic), + tl.cast(num_topk, tl.int32, _semantic=_semantic), + tl.cast(hidden, tl.int32, _semantic=_semantic), + tl.cast(num_padded_sf_pool_tokens, tl.int32, _semantic=_semantic), + ], + { + ( + core.pointer_type(core.dtype("uint8")), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + ): ("userhopper_ws_receiver_partition_bounded", ()), + }, + is_pure=False, + _semantic=_semantic, + ) + + +@core.extern +def userhopper_ws_compute_stub_partition(symm_buffer, l1_weights, l1_weights_sf, l2_weights, l2_weights_sf, + num_ranks, num_experts, num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, num_padded_sf_pool_tokens, + compute_full_hidden, + compute_parallel, + compute_worker_warps, + _semantic=None): + return core.extern_call( + "", + "", + [ + symm_buffer, + l1_weights, + l1_weights_sf, + l2_weights, + l2_weights_sf, + tl.cast(num_ranks, tl.int32, _semantic=_semantic), + tl.cast(num_experts, tl.int32, _semantic=_semantic), + tl.cast(num_max_tokens_per_rank, tl.int32, _semantic=_semantic), + tl.cast(num_topk, tl.int32, _semantic=_semantic), + tl.cast(hidden, tl.int32, _semantic=_semantic), + tl.cast(intermediate_hidden, tl.int32, _semantic=_semantic), + tl.cast(num_padded_sf_pool_tokens, tl.int32, _semantic=_semantic), + tl.cast(compute_full_hidden, tl.int32, _semantic=_semantic), + tl.cast(compute_parallel, tl.int32, _semantic=_semantic), + tl.cast(compute_worker_warps, tl.int32, _semantic=_semantic), + ], + { + ( + core.pointer_type(core.dtype("uint8")), + core.pointer_type(core.dtype("uint8")), + core.pointer_type(core.dtype("fp32")), + core.pointer_type(core.dtype("uint8")), + core.pointer_type(core.dtype("fp32")), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + ): ("userhopper_ws_compute_stub_partition", ()), + }, + is_pure=False, + _semantic=_semantic, + ) + + +@core.extern +def userhopper_ws_combine_reduce_partition(symm_buffer, y, num_ranks, num_experts, num_max_tokens_per_rank, + num_topk, hidden, intermediate_hidden, num_padded_sf_pool_tokens, + cleanup_workspace, _semantic=None): + return core.extern_call( + "", + "", + [ + symm_buffer, + y, + tl.cast(num_ranks, tl.int32, _semantic=_semantic), + tl.cast(num_experts, tl.int32, _semantic=_semantic), + tl.cast(num_max_tokens_per_rank, tl.int32, _semantic=_semantic), + tl.cast(num_topk, tl.int32, _semantic=_semantic), + tl.cast(hidden, tl.int32, _semantic=_semantic), + tl.cast(intermediate_hidden, tl.int32, _semantic=_semantic), + tl.cast(num_padded_sf_pool_tokens, tl.int32, _semantic=_semantic), + tl.cast(cleanup_workspace, tl.int32, _semantic=_semantic), + ], + { + ( + core.pointer_type(core.dtype("uint8")), + core.pointer_type(core.dtype("uint8")), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + ): ("userhopper_ws_combine_reduce_partition", ()), + }, + is_pure=False, + _semantic=_semantic, + ) + + +@core.extern +def userhopper_ws_tldot_combine_write_partition(symm_buffer, l2_out, num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, hidden, + intermediate_hidden, num_padded_sf_pool_tokens, + _semantic=None): + return core.extern_call( + "", + "", + [ + symm_buffer, + l2_out, + tl.cast(num_ranks, tl.int32, _semantic=_semantic), + tl.cast(num_experts, tl.int32, _semantic=_semantic), + tl.cast(num_max_tokens_per_rank, tl.int32, _semantic=_semantic), + tl.cast(num_topk, tl.int32, _semantic=_semantic), + tl.cast(hidden, tl.int32, _semantic=_semantic), + tl.cast(intermediate_hidden, tl.int32, _semantic=_semantic), + tl.cast(num_padded_sf_pool_tokens, tl.int32, _semantic=_semantic), + ], + { + ( + core.pointer_type(core.dtype("uint8")), + core.pointer_type(core.dtype("fp32")), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + core.dtype("int32"), + ): ("userhopper_ws_tldot_combine_write_partition", ()), + }, + is_pure=False, + _semantic=_semantic, + ) diff --git a/python/test/tle/integration/megamoe/megamoe_operator/ws_userhopper_dispatch_receiver_host.cu b/python/test/tle/integration/megamoe/megamoe_operator/ws_userhopper_dispatch_receiver_host.cu new file mode 100644 index 000000000..472e7cdc1 --- /dev/null +++ b/python/test/tle/integration/megamoe/megamoe_operator/ws_userhopper_dispatch_receiver_host.cu @@ -0,0 +1,26 @@ +#include +#include +#include + +extern "C" void userhopper_ws_nvshmem_init_wrapper() { nvshmem_init(); } + +extern "C" int userhopper_ws_nvshmemx_cumodule_init_wrapper(CUmodule module) { + return nvshmemx_cumodule_init(module); +} + +extern "C" int userhopper_ws_nvshmem_team_mype_wrapper() { + return nvshmem_team_my_pe(NVSHMEMX_TEAM_NODE); +} + +extern "C" int userhopper_ws_nvshmem_n_pes_wrapper() { return nvshmem_n_pes(); } + +extern "C" unsigned char *userhopper_ws_nvshmem_alloc_bytes_wrapper(long long bytes) { + return reinterpret_cast(nvshmem_malloc(static_cast(bytes))); +} + +extern "C" void userhopper_ws_nvshmem_barrier_all_wrapper() { nvshmem_barrier_all(); } + +extern "C" void userhopper_ws_nvshmem_finalize_wrapper(unsigned char *ptr) { + nvshmem_free(ptr); + nvshmem_finalize(); +} diff --git a/python/test/tle/integration/megamoe/perf/tle_2rank_h128_tokens2_repeats3.log b/python/test/tle/integration/megamoe/perf/tle_2rank_h128_tokens2_repeats3.log new file mode 100644 index 000000000..15495cd93 --- /dev/null +++ b/python/test/tle/integration/megamoe/perf/tle_2rank_h128_tokens2_repeats3.log @@ -0,0 +1,2 @@ +rank=0 userhopper_single_kernel_l1_l2_combine_tldot_smoke=PASS checked=2 counts=[2] ws=checked raw_nvshmem=checked tensor_core=checked combine=checked y=checked compute_order=sequential compute_warps=4 l1_i_tiles=2/2 l2_h_tiles=2/2 l2_block_n=64 l2_mode=tensorcore cleanup=1 repeats=3 launch_avg_us=16311756.652 launch_min_us=3235.424 launch_max_us=48928007.812 launch_steady_avg_us=3631.072 launch_steady_min_us=3235.424 launch_steady_max_us=4026.720 l2_max_abs=0.00012207 l2_tol=5.21506 +rank=1 userhopper_single_kernel_l1_l2_combine_tldot_smoke=PASS checked=2 counts=[2] ws=checked raw_nvshmem=checked tensor_core=checked combine=checked y=checked compute_order=sequential compute_warps=4 l1_i_tiles=2/2 l2_h_tiles=2/2 l2_block_n=64 l2_mode=tensorcore cleanup=1 repeats=3 launch_avg_us=16312680.060 launch_min_us=4005.088 launch_max_us=48928035.156 launch_steady_avg_us=5002.512 launch_steady_min_us=4005.088 launch_steady_max_us=5999.936 l2_max_abs=2.38419e-07 l2_tol=0.0127156 diff --git a/python/test/tle/integration/megamoe/perf/tle_8rank_h256_topk8_masked_tokens1_tile_split_repeats2_w4.log b/python/test/tle/integration/megamoe/perf/tle_8rank_h256_topk8_masked_tokens1_tile_split_repeats2_w4.log new file mode 100644 index 000000000..6c66952d7 --- /dev/null +++ b/python/test/tle/integration/megamoe/perf/tle_8rank_h256_topk8_masked_tokens1_tile_split_repeats2_w4.log @@ -0,0 +1,8 @@ +rank=3 userhopper_single_kernel_l1_l2_combine_tldot_smoke=PASS checked=6 counts=[3, 3] ws=checked raw_nvshmem=checked tensor_core=checked combine=checked y=checked compute_order=expert_wave_multi_cta_l1_l2_tile_split compute_warps=4 l1_i_tiles=2/2 l2_h_tiles=4/4 l2_block_n=64 l2_mode=tensorcore cleanup=1 repeats=2 launch_avg_us=58912527.696 launch_min_us=16555.391 launch_max_us=117808500.000 launch_steady_avg_us=16555.391 launch_steady_min_us=16555.391 launch_steady_max_us=16555.391 l2_max_abs=8 l2_tol=182439 +rank=5 userhopper_single_kernel_l1_l2_combine_tldot_smoke=PASS checked=6 counts=[3, 3] ws=checked raw_nvshmem=checked tensor_core=checked combine=checked y=checked compute_order=expert_wave_multi_cta_l1_l2_tile_split compute_warps=4 l1_i_tiles=2/2 l2_h_tiles=4/4 l2_block_n=64 l2_mode=tensorcore cleanup=1 repeats=2 launch_avg_us=58914472.631 launch_min_us=18531.200 launch_max_us=117810414.062 launch_steady_avg_us=18531.200 launch_steady_min_us=18531.200 launch_steady_max_us=18531.200 l2_max_abs=1 l2_tol=39561.7 +rank=2 userhopper_single_kernel_l1_l2_combine_tldot_smoke=PASS checked=6 counts=[3, 3] ws=checked raw_nvshmem=checked tensor_core=checked combine=checked y=checked compute_order=expert_wave_multi_cta_l1_l2_tile_split compute_warps=4 l1_i_tiles=2/2 l2_h_tiles=4/4 l2_block_n=64 l2_mode=tensorcore cleanup=1 repeats=2 launch_avg_us=58913101.111 launch_min_us=18397.535 launch_max_us=117807804.688 launch_steady_avg_us=18397.535 launch_steady_min_us=18397.535 launch_steady_max_us=18397.535 l2_max_abs=4 l2_tol=89188.2 +rank=1 userhopper_single_kernel_l1_l2_combine_tldot_smoke=PASS checked=6 counts=[3, 3] ws=checked raw_nvshmem=checked tensor_core=checked combine=checked y=checked compute_order=expert_wave_multi_cta_l1_l2_tile_split compute_warps=4 l1_i_tiles=2/2 l2_h_tiles=4/4 l2_block_n=64 l2_mode=tensorcore cleanup=1 repeats=2 launch_avg_us=58923113.928 launch_min_us=18423.168 launch_max_us=117827804.688 launch_steady_avg_us=18423.168 launch_steady_min_us=18423.168 launch_steady_max_us=18423.168 l2_max_abs=8 l2_tol=180886 +rank=7 userhopper_single_kernel_l1_l2_combine_tldot_smoke=PASS checked=6 counts=[3, 3] ws=checked raw_nvshmem=checked tensor_core=checked combine=checked y=checked compute_order=expert_wave_multi_cta_l1_l2_tile_split compute_warps=4 l1_i_tiles=2/2 l2_h_tiles=4/4 l2_block_n=64 l2_mode=tensorcore cleanup=1 repeats=2 launch_avg_us=58912976.476 launch_min_us=17968.576 launch_max_us=117807984.375 launch_steady_avg_us=17968.576 launch_steady_min_us=17968.576 launch_steady_max_us=17968.576 l2_max_abs=4 l2_tol=185544 +rank=0 userhopper_single_kernel_l1_l2_combine_tldot_smoke=PASS checked=6 counts=[3, 3] ws=checked raw_nvshmem=checked tensor_core=checked combine=checked y=checked compute_order=expert_wave_multi_cta_l1_l2_tile_split compute_warps=4 l1_i_tiles=2/2 l2_h_tiles=4/4 l2_block_n=64 l2_mode=tensorcore cleanup=1 repeats=2 launch_avg_us=58913241.586 launch_min_us=18311.296 launch_max_us=117808171.875 launch_steady_avg_us=18311.296 launch_steady_min_us=18311.296 launch_steady_max_us=18311.296 l2_max_abs=2 l2_tol=88309.5 +rank=6 userhopper_single_kernel_l1_l2_combine_tldot_smoke=PASS checked=6 counts=[3, 3] ws=checked raw_nvshmem=checked tensor_core=checked combine=checked y=checked compute_order=expert_wave_multi_cta_l1_l2_tile_split compute_warps=4 l1_i_tiles=2/2 l2_h_tiles=4/4 l2_block_n=64 l2_mode=tensorcore cleanup=1 repeats=2 launch_avg_us=58922711.125 launch_min_us=18570.688 launch_max_us=117826851.562 launch_steady_avg_us=18570.688 launch_steady_min_us=18570.688 launch_steady_max_us=18570.688 l2_max_abs=2 l2_tol=90945.6 +rank=4 userhopper_single_kernel_l1_l2_combine_tldot_smoke=PASS checked=6 counts=[3, 3] ws=checked raw_nvshmem=checked tensor_core=checked combine=checked y=checked compute_order=expert_wave_multi_cta_l1_l2_tile_split compute_warps=4 l1_i_tiles=2/2 l2_h_tiles=4/4 l2_block_n=64 l2_mode=tensorcore cleanup=1 repeats=2 launch_avg_us=58913215.096 launch_min_us=18180.193 launch_max_us=117808250.000 launch_steady_avg_us=18180.193 launch_steady_min_us=18180.193 launch_steady_max_us=18180.193 l2_max_abs=2 l2_tol=73815.9 diff --git a/python/test/tle/integration/megamoe/perf/userhopper_phase5_perf_summary.csv b/python/test/tle/integration/megamoe/perf/userhopper_phase5_perf_summary.csv new file mode 100644 index 000000000..42e92605b --- /dev/null +++ b/python/test/tle/integration/megamoe/perf/userhopper_phase5_perf_summary.csv @@ -0,0 +1,29 @@ +scope,implementation,kernel_or_path,shape,tokens_per_rank,world_size,status,avg_us,max_us,min_us,speedup_vs_baseline,source,notes +tle_prototype,Triton TLE merged,sequential_single_cta,h128_i128_e2_k1,2,2,pass_output_finalize_hang,4316.792,5002.512,3235.424,,tle_2rank_h128_tokens2_repeats3.log,"steady CUDA-event timing excludes first JIT launch; PASS lines printed, then run was manually killed after finalize hang" +tle_prototype,Triton TLE merged,multi_cta_tile_split,h256_i128_e16_k8_masked,1,8,pass,18117.256,18570.688,16555.391,,tle_8rank_h256_topk8_masked_tokens1_tile_split_repeats2_w4.log,"steady CUDA-event timing excludes first JIT launch; checked=6 counts=[3,3]" +direct_shape_probe,PR360,auto,h256_i128_e16_k8,1,8,failed_layout_assert,,,,,pr360_auto_h256_i128_e16_k8_b1.log,"DeepGEMM MegaMoE layout assertion: num_bytes % 16 == 0 or not require_tma_alignment" +production_existing,SplitHopper,split,h7168_i2048_e256_k8,512,8,pass,976.113,997.832,948.974,1.000,pr360_vs_split_rank_summary_official_fp8_shapes_20260617.csv,"DeepSeek-V3/R1 official-FP8 shape" +production_existing,PR360,auto,h7168_i2048_e256_k8,512,8,pass,858.010,858.058,857.958,1.163,pr360_vs_split_rank_summary_official_fp8_shapes_20260617.csv,"best PR360 path for this shape" +production_existing,SplitHopper,split,h7168_i2048_e256_k8,1024,8,pass,1619.321,1648.287,1520.706,1.000,pr360_vs_split_rank_summary_official_fp8_shapes_20260617.csv,"DeepSeek-V3/R1 official-FP8 shape" +production_existing,PR360,auto,h7168_i2048_e256_k8,1024,8,pass,1426.368,1426.387,1426.334,1.156,pr360_vs_split_rank_summary_official_fp8_shapes_20260617.csv,"best PR360 path for this shape" +production_existing,SplitHopper,split,h6144_i2048_e256_k8,512,8,pass,933.292,972.271,857.497,1.000,pr360_vs_split_rank_summary_official_fp8_shapes_20260617.csv,"GLM-5-FP8 official-FP8 shape" +production_existing,PR360,auto,h6144_i2048_e256_k8,512,8,pass,750.884,750.906,750.864,1.295,pr360_vs_split_rank_summary_official_fp8_shapes_20260617.csv,"best PR360 path for this shape" +production_existing,SplitHopper,split,h6144_i2048_e256_k8,1024,8,pass,1359.335,1380.092,1293.277,1.000,pr360_vs_split_rank_summary_official_fp8_shapes_20260617.csv,"GLM-5-FP8 official-FP8 shape" +production_existing,PR360,auto,h6144_i2048_e256_k8,1024,8,pass,1223.466,1223.606,1223.363,1.128,pr360_vs_split_rank_summary_official_fp8_shapes_20260617.csv,"best PR360 path for this shape" +production_existing,SplitHopper,split,h7168_i2048_e384_k8,512,8,pass,1367.106,1398.987,1288.131,1.000,pr360_vs_split_rank_summary_official_fp8_shapes_20260617.csv,"Kimi-K2 official-FP8 shape" +production_existing,PR360,cooperative,h7168_i2048_e384_k8,512,8,pass,927.747,927.763,927.710,1.508,pr360_vs_split_rank_summary_official_fp8_shapes_20260617.csv,"best PR360 path for this shape" +production_existing,SplitHopper,split,h7168_i2048_e384_k8,1024,8,pass,1742.677,1752.606,1726.092,1.000,pr360_vs_split_rank_summary_official_fp8_shapes_20260617.csv,"Kimi-K2 official-FP8 shape" +production_existing,PR360,auto,h7168_i2048_e384_k8,1024,8,pass,1592.736,1592.795,1592.694,1.100,pr360_vs_split_rank_summary_official_fp8_shapes_20260617.csv,"best PR360 path for this shape" +production_existing,SplitHopper,split,h6144_i2560_e160_k8,512,8,pass,1016.077,1042.815,938.165,1.000,pr360_vs_split_rank_summary_official_fp8_shapes_20260617.csv,"Qwen3-Coder official-FP8 shape" +production_existing,PR360,cooperative,h6144_i2560_e160_k8,512,8,pass,730.488,730.512,730.437,1.428,pr360_vs_split_rank_summary_official_fp8_shapes_20260617.csv,"best PR360 path for this shape" +production_existing,SplitHopper,split,h6144_i2560_e160_k8,1024,8,pass,1573.009,1585.388,1539.723,1.000,pr360_vs_split_rank_summary_official_fp8_shapes_20260617.csv,"Qwen3-Coder official-FP8 shape" +production_existing,PR360,cooperative,h6144_i2560_e160_k8,1024,8,pass,1357.076,1357.123,1357.014,1.168,pr360_vs_split_rank_summary_official_fp8_shapes_20260617.csv,"best PR360 path for this shape" +production_existing,SplitHopper,split,h2048_i512_e512_k10,512,8,pass,315.767,340.658,272.628,1.000,pr360_vs_split_rank_summary_official_fp8_shapes_20260617.csv,"Qwen3-Next official-FP8 shape" +production_existing,PR360,pingpong,h2048_i512_e512_k10,512,8,pass,199.870,199.883,199.848,1.704,pr360_vs_split_rank_summary_official_fp8_shapes_20260617.csv,"best PR360 path for this shape" +production_existing,SplitHopper,split,h2048_i512_e512_k10,1024,8,pass,356.479,370.444,344.103,1.000,pr360_vs_split_rank_summary_official_fp8_shapes_20260617.csv,"Qwen3-Next official-FP8 shape" +production_existing,PR360,pingpong,h2048_i512_e512_k10,1024,8,pass,316.471,316.502,316.458,1.170,pr360_vs_split_rank_summary_official_fp8_shapes_20260617.csv,"best PR360 path for this shape" +production_existing,UserHopperMegaMoE,add_hopper_mega,h5120_i1536_e1280_k6,4096,8,pass,,3005.000,,,user_hopper_moe_bench_results.csv,"MoE-4 rank0 timing, not same matrix as PR360-vs-Split rank-summary table" +production_existing,UserHopperMegaMoE,add_hopper_mega,h2048_i768_e1024_k8,4096,8,pass,,985.400,,,user_hopper_moe_bench_results.csv,"MoE-6 rank0 timing, not same matrix as PR360-vs-Split rank-summary table" +production_existing,UserHopperMegaMoE,add_hopper_mega,h4096_i1536_e1024_k8,4096,8,pass,,2481.000,,,user_hopper_moe_bench_results.csv,"MoE-7 rank0 timing, not same matrix as PR360-vs-Split rank-summary table" +production_existing,UserHopperMegaMoE,add_hopper_mega,h6144_i2560_e1280_k8,4096,8,pass,,5980.000,,,user_hopper_moe_bench_results.csv,"MoE-8 rank0 timing, not same matrix as PR360-vs-Split rank-summary table" +production_existing,UserHopperMegaMoE,add_hopper_mega,h2048_i512_e4096_k10,4096,8,pass,,1454.000,,,user_hopper_moe_bench_results.csv,"MoE-9 rank0 timing, not same matrix as PR360-vs-Split rank-summary table" diff --git a/python/test/tle/integration/megamoe/test_logs/local_2rank_h128_tokens2_repeats3.log b/python/test/tle/integration/megamoe/test_logs/local_2rank_h128_tokens2_repeats3.log new file mode 100644 index 000000000..41fa5c2c4 --- /dev/null +++ b/python/test/tle/integration/megamoe/test_logs/local_2rank_h128_tokens2_repeats3.log @@ -0,0 +1,2 @@ +rank=1 userhopper_single_kernel_l1_l2_combine_tldot_smoke=PASS checked=2 counts=[2] ws=checked raw_nvshmem=checked tensor_core=checked combine=checked y=checked compute_order=sequential compute_warps=4 l1_i_tiles=2/2 l2_h_tiles=2/2 l2_block_n=64 l2_mode=tensorcore cleanup=1 repeats=3 launch_avg_us=13731456.018 launch_min_us=3234.944 launch_max_us=41187207.031 launch_steady_avg_us=3580.512 launch_steady_min_us=3234.944 launch_steady_max_us=3926.080 l2_max_abs=2.38419e-07 l2_tol=0.0127156 +rank=0 userhopper_single_kernel_l1_l2_combine_tldot_smoke=PASS checked=2 counts=[2] ws=checked raw_nvshmem=checked tensor_core=checked combine=checked y=checked compute_order=sequential compute_warps=4 l1_i_tiles=2/2 l2_h_tiles=2/2 l2_block_n=64 l2_mode=tensorcore cleanup=1 repeats=3 launch_avg_us=13732348.124 launch_min_us=3809.312 launch_max_us=41187203.125 launch_steady_avg_us=4920.624 launch_steady_min_us=3809.312 launch_steady_max_us=6031.936 l2_max_abs=0.00012207 l2_tol=5.21506 diff --git a/python/test/tle/integration/megamoe/test_logs/local_8rank_h256_topk8_masked_tokens1_tile_split_repeats2.log b/python/test/tle/integration/megamoe/test_logs/local_8rank_h256_topk8_masked_tokens1_tile_split_repeats2.log new file mode 100644 index 000000000..80c8ce603 --- /dev/null +++ b/python/test/tle/integration/megamoe/test_logs/local_8rank_h256_topk8_masked_tokens1_tile_split_repeats2.log @@ -0,0 +1,8 @@ +rank=7 userhopper_single_kernel_l1_l2_combine_tldot_smoke=PASS checked=6 counts=[3, 3] ws=checked raw_nvshmem=checked tensor_core=checked combine=checked y=checked compute_order=expert_wave_multi_cta_l1_l2_tile_split compute_warps=4 l1_i_tiles=2/2 l2_h_tiles=4/4 l2_block_n=64 l2_mode=tensorcore cleanup=1 repeats=2 launch_avg_us=15411598.269 launch_min_us=16427.008 launch_max_us=30806769.531 launch_steady_avg_us=16427.008 launch_steady_min_us=16427.008 launch_steady_max_us=16427.008 l2_max_abs=4 l2_tol=185544 +rank=1 userhopper_single_kernel_l1_l2_combine_tldot_smoke=PASS checked=6 counts=[3, 3] ws=checked raw_nvshmem=checked tensor_core=checked combine=checked y=checked compute_order=expert_wave_multi_cta_l1_l2_tile_split compute_warps=4 l1_i_tiles=2/2 l2_h_tiles=4/4 l2_block_n=64 l2_mode=tensorcore cleanup=1 repeats=2 launch_avg_us=15411565.987 launch_min_us=16399.551 launch_max_us=30806732.422 launch_steady_avg_us=16399.551 launch_steady_min_us=16399.551 launch_steady_max_us=16399.551 l2_max_abs=8 l2_tol=180886 +rank=5 userhopper_single_kernel_l1_l2_combine_tldot_smoke=PASS checked=6 counts=[3, 3] ws=checked raw_nvshmem=checked tensor_core=checked combine=checked y=checked compute_order=expert_wave_multi_cta_l1_l2_tile_split compute_warps=4 l1_i_tiles=2/2 l2_h_tiles=4/4 l2_block_n=64 l2_mode=tensorcore cleanup=1 repeats=2 launch_avg_us=15411638.280 launch_min_us=16471.872 launch_max_us=30806804.688 launch_steady_avg_us=16471.872 launch_steady_min_us=16471.872 launch_steady_max_us=16471.872 l2_max_abs=1 l2_tol=39561.7 +rank=4 userhopper_single_kernel_l1_l2_combine_tldot_smoke=PASS checked=6 counts=[3, 3] ws=checked raw_nvshmem=checked tensor_core=checked combine=checked y=checked compute_order=expert_wave_multi_cta_l1_l2_tile_split compute_warps=4 l1_i_tiles=2/2 l2_h_tiles=4/4 l2_block_n=64 l2_mode=tensorcore cleanup=1 repeats=2 launch_avg_us=15411462.614 launch_min_us=16294.369 launch_max_us=30806630.859 launch_steady_avg_us=16294.369 launch_steady_min_us=16294.369 launch_steady_max_us=16294.369 l2_max_abs=2 l2_tol=73815.9 +rank=6 userhopper_single_kernel_l1_l2_combine_tldot_smoke=PASS checked=6 counts=[3, 3] ws=checked raw_nvshmem=checked tensor_core=checked combine=checked y=checked compute_order=expert_wave_multi_cta_l1_l2_tile_split compute_warps=4 l1_i_tiles=2/2 l2_h_tiles=4/4 l2_block_n=64 l2_mode=tensorcore cleanup=1 repeats=2 launch_avg_us=15411536.393 launch_min_us=16373.568 launch_max_us=30806699.219 launch_steady_avg_us=16373.568 launch_steady_min_us=16373.568 launch_steady_max_us=16373.568 l2_max_abs=2 l2_tol=90945.6 +rank=3 userhopper_single_kernel_l1_l2_combine_tldot_smoke=PASS checked=6 counts=[3, 3] ws=checked raw_nvshmem=checked tensor_core=checked combine=checked y=checked compute_order=expert_wave_multi_cta_l1_l2_tile_split compute_warps=4 l1_i_tiles=2/2 l2_h_tiles=4/4 l2_block_n=64 l2_mode=tensorcore cleanup=1 repeats=2 launch_avg_us=15411574.425 launch_min_us=16396.896 launch_max_us=30806751.953 launch_steady_avg_us=16396.896 launch_steady_min_us=16396.896 launch_steady_max_us=16396.896 l2_max_abs=8 l2_tol=182439 +rank=2 userhopper_single_kernel_l1_l2_combine_tldot_smoke=PASS checked=6 counts=[3, 3] ws=checked raw_nvshmem=checked tensor_core=checked combine=checked y=checked compute_order=expert_wave_multi_cta_l1_l2_tile_split compute_warps=4 l1_i_tiles=2/2 l2_h_tiles=4/4 l2_block_n=64 l2_mode=tensorcore cleanup=1 repeats=2 launch_avg_us=15411391.862 launch_min_us=16207.552 launch_max_us=30806576.172 launch_steady_avg_us=16207.552 launch_steady_min_us=16207.552 launch_steady_max_us=16207.552 l2_max_abs=4 l2_tol=89188.2 +rank=0 userhopper_single_kernel_l1_l2_combine_tldot_smoke=PASS checked=6 counts=[3, 3] ws=checked raw_nvshmem=checked tensor_core=checked combine=checked y=checked compute_order=expert_wave_multi_cta_l1_l2_tile_split compute_warps=4 l1_i_tiles=2/2 l2_h_tiles=4/4 l2_block_n=64 l2_mode=tensorcore cleanup=1 repeats=2 launch_avg_us=15411480.759 launch_min_us=16328.705 launch_max_us=30806632.812 launch_steady_avg_us=16328.705 launch_steady_min_us=16328.705 launch_steady_max_us=16328.705 l2_max_abs=2 l2_tol=88309.5 diff --git a/python/test/tle/integration/megamoe/tests/megamoe_local_harness.py b/python/test/tle/integration/megamoe/tests/megamoe_local_harness.py new file mode 100644 index 000000000..2152a7cdf --- /dev/null +++ b/python/test/tle/integration/megamoe/tests/megamoe_local_harness.py @@ -0,0 +1,3518 @@ +"""Single-kernel raw-NVSHMEM + Tensor Core L1 smoke. + +This is a narrow stage-4 capability test: one TLE warp-specialized kernel runs +UserHopper-like raw-NVSHMEM dispatch, raw-NVSHMEM receiver, and an L1/L2 FP8 +``tl.dot`` worker in the same ``warp_specialize`` region. + +The test intentionally fixes the local shape to the smallest useful cases: +up to 2 experts per rank, I=128, H=128 with <=16 received token-topk entries +per local expert, or H=256 with <=8 entries per local expert. That lets one +CTA prove the communication-to-TensorCore handoff without introducing a +multi-CTA scheduler. +""" + +from __future__ import annotations + +import ctypes +import os +import site +from pathlib import Path + +cuda_home = os.environ.get("CUDA_HOME", "/usr/local/cuda-12.8") +os.environ.setdefault("CUDA_HOME", cuda_home) +os.environ.setdefault("USERHOPPER_WS_NUM_RANKS", "2") +os.environ.setdefault("USERHOPPER_WS_NUM_EXPERTS", "2") +os.environ.setdefault("USERHOPPER_WS_NUM_TOPK", "1") +os.environ.setdefault("USERHOPPER_WS_NUM_TOKENS", "4") +os.environ.setdefault("USERHOPPER_WS_HIDDEN", "128") +os.environ.setdefault("USERHOPPER_WS_INTERMEDIATE_HIDDEN", "128") +os.environ.setdefault("USERHOPPER_WS_NUM_WARPS", "4") +os.environ.setdefault("USERHOPPER_WS_NUM_DISPATCH_WARPS", "1") +os.environ.setdefault("USERHOPPER_WS_COMPUTE_FULL_HIDDEN", "1") +os.environ.setdefault("USERHOPPER_WS_ALLOW_EXPERIMENTAL_H256_TLDOT", "0") +os.environ.setdefault("USERHOPPER_WS_CLEANUP", "0") +os.environ.setdefault("USERHOPPER_WS_REPEAT_LAUNCHES", "1") +os.environ.setdefault("USERHOPPER_WS_EXPERT_WAVE_SINGLE_KERNEL", "0") +os.environ.setdefault("USERHOPPER_WS_EXPERT_WAVE_COMPUTE_WARPS", "4") +os.environ.setdefault("USERHOPPER_WS_L2_H_TILES", "0") +os.environ.setdefault("USERHOPPER_WS_L1_I_TILES", "0") +os.environ.setdefault("USERHOPPER_WS_SPLIT_L1_L2_WORKERS", "0") +os.environ.setdefault("USERHOPPER_WS_L1_TILE_LOOP", "0") +os.environ.setdefault("USERHOPPER_WS_L2_BLOCK_N", "64") +os.environ.setdefault("USERHOPPER_WS_L2_SCALAR", "0") +os.environ.setdefault("USERHOPPER_WS_MULTI_CTA_EXPERT_WAVE", "0") +os.environ.setdefault("USERHOPPER_WS_MULTI_CTA_L1_ONLY", "0") +os.environ.setdefault("USERHOPPER_WS_MULTI_CTA_L1_TILE_SPLIT", "0") +os.environ.setdefault("USERHOPPER_WS_MULTI_CTA_L2_TILE_SPLIT", "0") +os.environ.setdefault("USERHOPPER_WS_SKIP_COMBINE", "0") +os.environ.setdefault("USERHOPPER_WS_SKIP_REDUCE", "0") +os.environ["CPATH"] = ( + f"{cuda_home}/targets/x86_64-linux/include:" + os.environ.get("CPATH", "") +) +os.environ["LD_LIBRARY_PATH"] = ( + f"{cuda_home}/lib64:" + os.environ.get("LD_LIBRARY_PATH", "") +) + +import triton +import triton.language as tl +import triton.experimental.tle.language as tle + +torch_site_packages = os.environ.get( + "MEGAMOE_TORCH_SITE_PACKAGES", + "/workspace/megakernel/.mega-venv/lib/python3.10/site-packages", +) +if Path(torch_site_packages).exists(): + site.addsitedir(torch_site_packages) + +import torch +import triton.experimental.tle.language.raw as tle_raw +from triton.experimental.tle.raw import dialect +import megamoe_operator.triton_tle_megamoe_runtime as uh + + +@dialect( + name="cuda", + compiler="nvcc", + file=uh.HERE / "ws_userhopper_dispatch_receiver_device.cu", + extern=uh.HERE / "ws_userhopper_dispatch_receiver_extern_call.py", + extern_func_name="userhopper_ws_tldot_combine_write_partition", + libs={"nvshmem": uh.NVSHMEM_HOME}, + links=["nvshmem_device"], +) +def edsl_userhopper_ws_tldot_combine_write(*args, **kwargs): + ... + + +DEFAULT_PY_BLOCK_M = 8 if uh.HIDDEN == 256 else 16 +PY_BLOCK_M = int(os.environ.get("USERHOPPER_WS_BLOCK_M", str(DEFAULT_PY_BLOCK_M))) +JIT_BLOCK_M = tl.constexpr(PY_BLOCK_M) +JIT_BLOCK_I = tl.constexpr(64) +JIT_BLOCK_K = tl.constexpr(64) +PY_BLOCK_N = int(os.environ["USERHOPPER_WS_L2_BLOCK_N"]) +JIT_BLOCK_N = tl.constexpr(PY_BLOCK_N) +DEFAULT_PY_L1_I_TILES = uh.INTERMEDIATE_HIDDEN // 64 +PY_L1_I_TILES = int(os.environ["USERHOPPER_WS_L1_I_TILES"]) +if PY_L1_I_TILES == 0: + PY_L1_I_TILES = DEFAULT_PY_L1_I_TILES +JIT_L1_I_TILES = tl.constexpr(PY_L1_I_TILES) +DEFAULT_PY_L2_H_TILES = uh.HIDDEN // PY_BLOCK_N +PY_L2_H_TILES = int(os.environ["USERHOPPER_WS_L2_H_TILES"]) +if PY_L2_H_TILES == 0: + PY_L2_H_TILES = DEFAULT_PY_L2_H_TILES +JIT_L2_H_TILES = tl.constexpr(PY_L2_H_TILES) +CLEANUP_WORKSPACE = int(os.environ["USERHOPPER_WS_CLEANUP"]) +REPEAT_LAUNCHES = int(os.environ["USERHOPPER_WS_REPEAT_LAUNCHES"]) +EXPERT_WAVE_SINGLE_KERNEL = int(os.environ["USERHOPPER_WS_EXPERT_WAVE_SINGLE_KERNEL"]) +EXPERT_WAVE_COMPUTE_WARPS = int(os.environ["USERHOPPER_WS_EXPERT_WAVE_COMPUTE_WARPS"]) +SPLIT_L1_L2_WORKERS = int(os.environ["USERHOPPER_WS_SPLIT_L1_L2_WORKERS"]) +L1_TILE_LOOP = int(os.environ["USERHOPPER_WS_L1_TILE_LOOP"]) +JIT_L1_TILE_LOOP = tl.constexpr(L1_TILE_LOOP) +L2_SCALAR = int(os.environ["USERHOPPER_WS_L2_SCALAR"]) +JIT_L2_SCALAR = tl.constexpr(L2_SCALAR) +MULTI_CTA_EXPERT_WAVE = int(os.environ["USERHOPPER_WS_MULTI_CTA_EXPERT_WAVE"]) +MULTI_CTA_L1_ONLY = int(os.environ["USERHOPPER_WS_MULTI_CTA_L1_ONLY"]) +JIT_MULTI_CTA_L1_ONLY = tl.constexpr(MULTI_CTA_L1_ONLY) +MULTI_CTA_L1_TILE_SPLIT = int(os.environ["USERHOPPER_WS_MULTI_CTA_L1_TILE_SPLIT"]) +JIT_MULTI_CTA_L1_TILE_SPLIT = tl.constexpr(MULTI_CTA_L1_TILE_SPLIT) +MULTI_CTA_L2_TILE_SPLIT = int(os.environ["USERHOPPER_WS_MULTI_CTA_L2_TILE_SPLIT"]) +JIT_MULTI_CTA_L2_TILE_SPLIT = tl.constexpr(MULTI_CTA_L2_TILE_SPLIT) +SKIP_COMBINE = int(os.environ["USERHOPPER_WS_SKIP_COMBINE"]) +JIT_SKIP_COMBINE = tl.constexpr(SKIP_COMBINE) +SKIP_REDUCE = int(os.environ["USERHOPPER_WS_SKIP_REDUCE"]) +JIT_SKIP_REDUCE = tl.constexpr(SKIP_REDUCE) + +if CLEANUP_WORKSPACE not in (0, 1): + raise ValueError(f"USERHOPPER_WS_CLEANUP must be 0 or 1, got {CLEANUP_WORKSPACE}") +if REPEAT_LAUNCHES <= 0: + raise ValueError(f"USERHOPPER_WS_REPEAT_LAUNCHES must be positive, got {REPEAT_LAUNCHES}") +if REPEAT_LAUNCHES > 1 and CLEANUP_WORKSPACE == 0: + raise ValueError("USERHOPPER_WS_REPEAT_LAUNCHES>1 requires USERHOPPER_WS_CLEANUP=1") +if EXPERT_WAVE_SINGLE_KERNEL not in (0, 1): + raise ValueError( + "USERHOPPER_WS_EXPERT_WAVE_SINGLE_KERNEL must be 0 or 1, " + f"got {EXPERT_WAVE_SINGLE_KERNEL}" + ) +if EXPERT_WAVE_COMPUTE_WARPS not in (2, 4): + raise ValueError( + "USERHOPPER_WS_EXPERT_WAVE_COMPUTE_WARPS must be 2 or 4, " + f"got {EXPERT_WAVE_COMPUTE_WARPS}" + ) +if SPLIT_L1_L2_WORKERS not in (0, 1): + raise ValueError( + "USERHOPPER_WS_SPLIT_L1_L2_WORKERS must be 0 or 1, " + f"got {SPLIT_L1_L2_WORKERS}" + ) +if L1_TILE_LOOP not in (0, 1): + raise ValueError(f"USERHOPPER_WS_L1_TILE_LOOP must be 0 or 1, got {L1_TILE_LOOP}") +if L2_SCALAR not in (0, 1): + raise ValueError(f"USERHOPPER_WS_L2_SCALAR must be 0 or 1, got {L2_SCALAR}") +if MULTI_CTA_EXPERT_WAVE not in (0, 1): + raise ValueError( + "USERHOPPER_WS_MULTI_CTA_EXPERT_WAVE must be 0 or 1, " + f"got {MULTI_CTA_EXPERT_WAVE}" + ) +if MULTI_CTA_L1_ONLY not in (0, 1): + raise ValueError( + "USERHOPPER_WS_MULTI_CTA_L1_ONLY must be 0 or 1, " + f"got {MULTI_CTA_L1_ONLY}" + ) +if MULTI_CTA_L1_ONLY != 0 and MULTI_CTA_EXPERT_WAVE == 0: + raise ValueError("USERHOPPER_WS_MULTI_CTA_L1_ONLY requires USERHOPPER_WS_MULTI_CTA_EXPERT_WAVE=1") +if MULTI_CTA_L1_TILE_SPLIT not in (0, 1): + raise ValueError( + "USERHOPPER_WS_MULTI_CTA_L1_TILE_SPLIT must be 0 or 1, " + f"got {MULTI_CTA_L1_TILE_SPLIT}" + ) +if MULTI_CTA_L1_TILE_SPLIT != 0 and (MULTI_CTA_EXPERT_WAVE == 0 or MULTI_CTA_L1_ONLY == 0): + raise ValueError( + "USERHOPPER_WS_MULTI_CTA_L1_TILE_SPLIT requires " + "USERHOPPER_WS_MULTI_CTA_EXPERT_WAVE=1 and USERHOPPER_WS_MULTI_CTA_L1_ONLY=1" + ) +if MULTI_CTA_L2_TILE_SPLIT not in (0, 1): + raise ValueError( + "USERHOPPER_WS_MULTI_CTA_L2_TILE_SPLIT must be 0 or 1, " + f"got {MULTI_CTA_L2_TILE_SPLIT}" + ) +if MULTI_CTA_L2_TILE_SPLIT != 0 and (MULTI_CTA_EXPERT_WAVE == 0 or MULTI_CTA_L1_TILE_SPLIT == 0): + raise ValueError( + "USERHOPPER_WS_MULTI_CTA_L2_TILE_SPLIT requires " + "USERHOPPER_WS_MULTI_CTA_EXPERT_WAVE=1 and USERHOPPER_WS_MULTI_CTA_L1_TILE_SPLIT=1" + ) +if SKIP_COMBINE not in (0, 1): + raise ValueError(f"USERHOPPER_WS_SKIP_COMBINE must be 0 or 1, got {SKIP_COMBINE}") +if SKIP_COMBINE != 0 and MULTI_CTA_L2_TILE_SPLIT == 0: + raise ValueError("USERHOPPER_WS_SKIP_COMBINE requires USERHOPPER_WS_MULTI_CTA_L2_TILE_SPLIT=1") +if SKIP_COMBINE != 0 and CLEANUP_WORKSPACE != 0: + raise ValueError("USERHOPPER_WS_SKIP_COMBINE=1 requires USERHOPPER_WS_CLEANUP=0") +if SKIP_REDUCE not in (0, 1): + raise ValueError(f"USERHOPPER_WS_SKIP_REDUCE must be 0 or 1, got {SKIP_REDUCE}") +if SKIP_REDUCE != 0 and MULTI_CTA_L2_TILE_SPLIT == 0: + raise ValueError("USERHOPPER_WS_SKIP_REDUCE requires USERHOPPER_WS_MULTI_CTA_L2_TILE_SPLIT=1") +if SKIP_REDUCE != 0 and CLEANUP_WORKSPACE != 0: + raise ValueError("USERHOPPER_WS_SKIP_REDUCE=1 requires USERHOPPER_WS_CLEANUP=0") + + +@triton.jit +def _l1_single_cta_two_tile_body( + acts, + acts_sf, + topk_weights, + weights, + weights_sf, + l2_acts, + l2_acts_sf, + marker, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + I_OFFSET_C: tl.constexpr, +): + offs_m = tl.arange(0, JIT_BLOCK_M) + offs_i = I_OFFSET_C + tl.arange(0, JIT_BLOCK_I) + group = offs_i // 8 + lane = offs_i - group * 8 + gate_rows = group * 16 + lane + up_rows = gate_rows + 8 + + gate_acc = tl.zeros((JIT_BLOCK_M, JIT_BLOCK_I), dtype=tl.float32) + up_acc = tl.zeros((JIT_BLOCK_M, JIT_BLOCK_I), dtype=tl.float32) + for k0 in range(0, H_C, JIT_BLOCK_K): + offs_k = k0 + tl.arange(0, JIT_BLOCK_K) + a = tl.load( + acts + offs_m[:, None] * H_C + offs_k[None, :], + mask=offs_m[:, None] < M_C, + other=0.0, + ) + gate_w = tl.load( + weights + gate_rows[None, :] * H_C + offs_k[:, None], + mask=offs_i[None, :] < I_C, + other=0.0, + ) + up_w = tl.load( + weights + up_rows[None, :] * H_C + offs_k[:, None], + mask=offs_i[None, :] < I_C, + other=0.0, + ) + act_scale = tl.load( + acts_sf + (k0 // 128) * NUM_PADDED_M_C + offs_m, + mask=offs_m < M_C, + other=0.0, + ) + gate_scale = tl.load( + weights_sf + (offs_i // 128) * (H_C // 128) + (k0 // 128), + mask=offs_i < I_C, + other=0.0, + ) + up_scale = tl.load( + weights_sf + ((I_C + offs_i) // 128) * (H_C // 128) + (k0 // 128), + mask=offs_i < I_C, + other=0.0, + ) + gate_acc += tl.dot(a, gate_w, out_dtype=tl.float32) * act_scale[:, None] * gate_scale[None, :] + up_acc += tl.dot(a, up_w, out_dtype=tl.float32) * act_scale[:, None] * up_scale[None, :] + + topk = tl.load(topk_weights + offs_m, mask=offs_m < M_C, other=0.0) + swiglu = gate_acc * tl.sigmoid(gate_acc) * up_acc * topk[:, None] + max_abs = tl.max(tl.abs(swiglu), axis=1) + scale = tl.where(max_abs > 0.0, max_abs / 448.0, 1.0) + scaled = swiglu / scale[:, None] + pid_i: tl.constexpr = I_OFFSET_C // JIT_BLOCK_I + + tl.store( + l2_acts_sf + pid_i * NUM_PADDED_M_C + offs_m, + scale, + mask=offs_m < M_C, + ) + tl.store( + l2_acts + offs_m[:, None] * I_C + offs_i[None, :], + scaled, + mask=(offs_m[:, None] < M_C) & (offs_i[None, :] < I_C), + ) + tl.store(marker, 0x4C1107) + + +@triton.jit +def _l1_single_cta_runtime_i_body( + acts, + acts_sf, + topk_weights, + weights, + weights_sf, + l2_acts, + l2_acts_sf, + marker, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + i_offset, +): + offs_m = tl.arange(0, JIT_BLOCK_M) + offs_i = i_offset + tl.arange(0, JIT_BLOCK_I) + group = offs_i // 8 + lane = offs_i - group * 8 + gate_rows = group * 16 + lane + up_rows = gate_rows + 8 + + gate_acc = tl.zeros((JIT_BLOCK_M, JIT_BLOCK_I), dtype=tl.float32) + up_acc = tl.zeros((JIT_BLOCK_M, JIT_BLOCK_I), dtype=tl.float32) + for k0 in range(0, H_C, JIT_BLOCK_K): + offs_k = k0 + tl.arange(0, JIT_BLOCK_K) + a = tl.load( + acts + offs_m[:, None] * H_C + offs_k[None, :], + mask=offs_m[:, None] < M_C, + other=0.0, + ) + gate_w = tl.load( + weights + gate_rows[None, :] * H_C + offs_k[:, None], + mask=offs_i[None, :] < I_C, + other=0.0, + ) + up_w = tl.load( + weights + up_rows[None, :] * H_C + offs_k[:, None], + mask=offs_i[None, :] < I_C, + other=0.0, + ) + act_scale = tl.load( + acts_sf + (k0 // 128) * NUM_PADDED_M_C + offs_m, + mask=offs_m < M_C, + other=0.0, + ) + gate_scale = tl.load( + weights_sf + (offs_i // 128) * (H_C // 128) + (k0 // 128), + mask=offs_i < I_C, + other=0.0, + ) + up_scale = tl.load( + weights_sf + ((I_C + offs_i) // 128) * (H_C // 128) + (k0 // 128), + mask=offs_i < I_C, + other=0.0, + ) + gate_acc += tl.dot(a, gate_w, out_dtype=tl.float32) * act_scale[:, None] * gate_scale[None, :] + up_acc += tl.dot(a, up_w, out_dtype=tl.float32) * act_scale[:, None] * up_scale[None, :] + + topk = tl.load(topk_weights + offs_m, mask=offs_m < M_C, other=0.0) + swiglu = gate_acc * tl.sigmoid(gate_acc) * up_acc * topk[:, None] + max_abs = tl.max(tl.abs(swiglu), axis=1) + scale = tl.where(max_abs > 0.0, max_abs / 448.0, 1.0) + scaled = swiglu / scale[:, None] + pid_i = i_offset // JIT_BLOCK_I + + tl.store( + l2_acts_sf + pid_i * NUM_PADDED_M_C + offs_m, + scale, + mask=offs_m < M_C, + ) + tl.store( + l2_acts + offs_m[:, None] * I_C + offs_i[None, :], + scaled, + mask=(offs_m[:, None] < M_C) & (offs_i[None, :] < I_C), + ) + tl.store(marker, 0x4C1108) + + +@triton.jit +def _l1_single_cta_runtime_m_i_body( + acts, + acts_sf, + topk_weights, + weights, + weights_sf, + l2_acts, + l2_acts_sf, + marker, + M, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + i_offset, +): + offs_m = tl.arange(0, JIT_BLOCK_M) + offs_i = i_offset + tl.arange(0, JIT_BLOCK_I) + group = offs_i // 8 + lane = offs_i - group * 8 + gate_rows = group * 16 + lane + up_rows = gate_rows + 8 + + gate_acc = tl.zeros((JIT_BLOCK_M, JIT_BLOCK_I), dtype=tl.float32) + up_acc = tl.zeros((JIT_BLOCK_M, JIT_BLOCK_I), dtype=tl.float32) + for k0 in range(0, H_C, JIT_BLOCK_K): + offs_k = k0 + tl.arange(0, JIT_BLOCK_K) + a = tl.load( + acts + offs_m[:, None] * H_C + offs_k[None, :], + mask=offs_m[:, None] < M, + other=0.0, + ) + gate_w = tl.load( + weights + gate_rows[None, :] * H_C + offs_k[:, None], + mask=offs_i[None, :] < I_C, + other=0.0, + ) + up_w = tl.load( + weights + up_rows[None, :] * H_C + offs_k[:, None], + mask=offs_i[None, :] < I_C, + other=0.0, + ) + act_scale = tl.load( + acts_sf + (k0 // 128) * NUM_PADDED_M_C + offs_m, + mask=offs_m < M, + other=0.0, + ) + gate_scale = tl.load( + weights_sf + (offs_i // 128) * (H_C // 128) + (k0 // 128), + mask=offs_i < I_C, + other=0.0, + ) + up_scale = tl.load( + weights_sf + ((I_C + offs_i) // 128) * (H_C // 128) + (k0 // 128), + mask=offs_i < I_C, + other=0.0, + ) + gate_acc += tl.dot(a, gate_w, out_dtype=tl.float32) * act_scale[:, None] * gate_scale[None, :] + up_acc += tl.dot(a, up_w, out_dtype=tl.float32) * act_scale[:, None] * up_scale[None, :] + + topk = tl.load(topk_weights + offs_m, mask=offs_m < M, other=0.0) + swiglu = gate_acc * tl.sigmoid(gate_acc) * up_acc * topk[:, None] + max_abs = tl.max(tl.abs(swiglu), axis=1) + scale = tl.where(max_abs > 0.0, max_abs / 448.0, 1.0) + scaled = swiglu / scale[:, None] + pid_i = i_offset // JIT_BLOCK_I + + tl.store( + l2_acts_sf + pid_i * NUM_PADDED_M_C + offs_m, + scale, + mask=offs_m < M, + ) + tl.store( + l2_acts + offs_m[:, None] * I_C + offs_i[None, :], + scaled, + mask=(offs_m[:, None] < M) & (offs_i[None, :] < I_C), + ) + tl.store(marker, 0x4C1109) + + +@triton.jit +def _l2_single_cta_tile_body( + acts, + acts_sf, + weights, + weights_sf, + out, + marker, + M_C: tl.constexpr, + N_C: tl.constexpr, + K_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + N_OFFSET_C: tl.constexpr, +): + offs_m = tl.arange(0, JIT_BLOCK_M) + offs_n = N_OFFSET_C + tl.arange(0, JIT_BLOCK_N) + + acc = tl.zeros((JIT_BLOCK_M, JIT_BLOCK_N), dtype=tl.float32) + for k0 in range(0, K_C, JIT_BLOCK_K): + offs_k = k0 + tl.arange(0, JIT_BLOCK_K) + a = tl.load( + acts + offs_m[:, None] * K_C + offs_k[None, :], + mask=offs_m[:, None] < M_C, + other=0.0, + ) + b = tl.load( + weights + offs_n[None, :] * K_C + offs_k[:, None], + mask=offs_n[None, :] < N_C, + other=0.0, + ) + partial = tl.dot(a, b, out_dtype=tl.float32) + act_scale = tl.load( + acts_sf + (k0 // 64) * NUM_PADDED_M_C + offs_m, + mask=offs_m < M_C, + other=0.0, + ) + weight_scale = tl.load( + weights_sf + (offs_n // 128) * (K_C // 128) + (k0 // 128), + mask=offs_n < N_C, + other=0.0, + ) + acc += partial * act_scale[:, None] * weight_scale[None, :] + + tl.store( + out + offs_m[:, None] * N_C + offs_n[None, :], + acc, + mask=(offs_m[:, None] < M_C) & (offs_n[None, :] < N_C), + ) + tl.store(marker, 0x4C2207) + + +@triton.jit +def _l2_single_cta_scalar_tile_body( + acts, + acts_sf, + weights, + weights_sf, + out, + marker, + M_C: tl.constexpr, + N_C: tl.constexpr, + K_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + N_OFFSET_C: tl.constexpr, +): + offs_m = tl.arange(0, JIT_BLOCK_M) + offs_n = N_OFFSET_C + tl.arange(0, JIT_BLOCK_N) + acc = tl.zeros((JIT_BLOCK_M, JIT_BLOCK_N), dtype=tl.float32) + + for kk in range(0, K_C): + a = tl.load( + acts + offs_m * K_C + kk, + mask=offs_m < M_C, + other=0.0, + ).to(tl.float32) + b = tl.load( + weights + offs_n * K_C + kk, + mask=offs_n < N_C, + other=0.0, + ).to(tl.float32) + act_scale = tl.load( + acts_sf + (kk // 64) * NUM_PADDED_M_C + offs_m, + mask=offs_m < M_C, + other=0.0, + ) + weight_scale = tl.load( + weights_sf + (offs_n // 128) * (K_C // 128) + (kk // 128), + mask=offs_n < N_C, + other=0.0, + ) + acc += (a * act_scale)[:, None] * (b * weight_scale)[None, :] + + tl.store( + out + offs_m[:, None] * N_C + offs_n[None, :], + acc, + mask=(offs_m[:, None] < M_C) & (offs_n[None, :] < N_C), + ) + tl.store(marker, 0x4C2307) + + +@triton.jit +def _l2_single_cta_runtime_m_n_body( + acts, + acts_sf, + weights, + weights_sf, + out, + marker, + M, + N_C: tl.constexpr, + K_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + n_offset, +): + offs_m = tl.arange(0, JIT_BLOCK_M) + offs_n = n_offset + tl.arange(0, JIT_BLOCK_N) + + acc = tl.zeros((JIT_BLOCK_M, JIT_BLOCK_N), dtype=tl.float32) + for k0 in range(0, K_C, JIT_BLOCK_K): + offs_k = k0 + tl.arange(0, JIT_BLOCK_K) + a = tl.load( + acts + offs_m[:, None] * K_C + offs_k[None, :], + mask=offs_m[:, None] < M, + other=0.0, + ) + b = tl.load( + weights + offs_n[None, :] * K_C + offs_k[:, None], + mask=offs_n[None, :] < N_C, + other=0.0, + ) + partial = tl.dot(a, b, out_dtype=tl.float32) + act_scale = tl.load( + acts_sf + (k0 // 64) * NUM_PADDED_M_C + offs_m, + mask=offs_m < M, + other=0.0, + ) + weight_scale = tl.load( + weights_sf + (offs_n // 128) * (K_C // 128) + (k0 // 128), + mask=offs_n < N_C, + other=0.0, + ) + acc += partial * act_scale[:, None] * weight_scale[None, :] + + tl.store( + out + offs_m[:, None] * N_C + offs_n[None, :], + acc, + mask=(offs_m[:, None] < M) & (offs_n[None, :] < N_C), + ) + tl.store(marker, 0x4C2208) + + +@triton.jit +def _single_expert_l1_body( + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + POOL_BASE_C: tl.constexpr, + LOCAL_EXPERT_C: tl.constexpr, +): + l1_acts_e = l1_acts + POOL_BASE_C * H_C + l1_acts_sf_e = l1_acts_sf + POOL_BASE_C + l1_topk_weights_e = l1_topk_weights + POOL_BASE_C + l2_acts_e = l2_acts + POOL_BASE_C * I_C + l2_acts_sf_e = l2_acts_sf + POOL_BASE_C + l1_weights_e = l1_weights + LOCAL_EXPERT_C * (2 * I_C * H_C) + l1_weights_sf_e = l1_weights_sf + LOCAL_EXPERT_C * ((2 * I_C // 128) * (H_C // 128)) + + if JIT_L1_TILE_LOOP != 0: + for i_offset in tl.range(0, JIT_L1_I_TILES * JIT_BLOCK_I, JIT_BLOCK_I): + _l1_single_cta_runtime_i_body( + l1_acts_e, + l1_acts_sf_e, + l1_topk_weights_e, + l1_weights_e, + l1_weights_sf_e, + l2_acts_e, + l2_acts_sf_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + i_offset, + ) + else: + if JIT_L1_I_TILES >= 1: + _l1_single_cta_two_tile_body( + l1_acts_e, + l1_acts_sf_e, + l1_topk_weights_e, + l1_weights_e, + l1_weights_sf_e, + l2_acts_e, + l2_acts_sf_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + 0, + ) + if JIT_L1_I_TILES >= 2: + _l1_single_cta_two_tile_body( + l1_acts_e, + l1_acts_sf_e, + l1_topk_weights_e, + l1_weights_e, + l1_weights_sf_e, + l2_acts_e, + l2_acts_sf_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + 64, + ) + + +@triton.jit +def _single_expert_l2_body( + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + marker, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + POOL_BASE_C: tl.constexpr, + LOCAL_EXPERT_C: tl.constexpr, +): + l2_acts_e = l2_acts + POOL_BASE_C * I_C + l2_acts_sf_e = l2_acts_sf + POOL_BASE_C + l2_out_e = l2_out + POOL_BASE_C * H_C + l2_weights_e = l2_weights + LOCAL_EXPERT_C * (H_C * I_C) + l2_weights_sf_e = l2_weights_sf + LOCAL_EXPERT_C * ((H_C // 128) * (I_C // 128)) + + if JIT_L2_H_TILES >= 1: + if JIT_L2_SCALAR != 0: + _l2_single_cta_scalar_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + 0, + ) + else: + _l2_single_cta_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + 0, + ) + if JIT_L2_H_TILES >= 2: + if JIT_L2_SCALAR != 0: + _l2_single_cta_scalar_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N, + ) + else: + _l2_single_cta_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N, + ) + if JIT_L2_H_TILES >= 3: + if JIT_L2_SCALAR != 0: + _l2_single_cta_scalar_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 2, + ) + else: + _l2_single_cta_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 2, + ) + if JIT_L2_H_TILES >= 4: + if JIT_L2_SCALAR != 0: + _l2_single_cta_scalar_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 3, + ) + else: + _l2_single_cta_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 3, + ) + if JIT_L2_H_TILES >= 5: + if JIT_L2_SCALAR != 0: + _l2_single_cta_scalar_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 4, + ) + else: + _l2_single_cta_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 4, + ) + if JIT_L2_H_TILES >= 6: + if JIT_L2_SCALAR != 0: + _l2_single_cta_scalar_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 5, + ) + else: + _l2_single_cta_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 5, + ) + if JIT_L2_H_TILES >= 7: + if JIT_L2_SCALAR != 0: + _l2_single_cta_scalar_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 6, + ) + else: + _l2_single_cta_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 6, + ) + if JIT_L2_H_TILES >= 8: + if JIT_L2_SCALAR != 0: + _l2_single_cta_scalar_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 7, + ) + else: + _l2_single_cta_tile_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + JIT_BLOCK_N * 7, + ) + + +@triton.jit +def _single_expert_l1_l2_body( + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + marker, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + POOL_BASE_C: tl.constexpr, + LOCAL_EXPERT_C: tl.constexpr, +): + _single_expert_l1_body( + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + POOL_BASE_C, + LOCAL_EXPERT_C, + ) + tl.debug_barrier() + _single_expert_l2_body( + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + marker, + M_C, + H_C, + I_C, + NUM_PADDED_M_C, + POOL_BASE_C, + LOCAL_EXPERT_C, + ) + + +@triton.jit +def _l1_single_kernel_worker( + compute_reader, + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, +): + wait_result = compute_reader.wait(0) + _ = tl.load(tle.gpu.local_ptr(wait_result.slot.done, (0,))) + + _single_expert_l1_l2_body( + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + marker, + EXPERT0_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + 0, + 0, + ) + if NUM_EXPERTS_PER_RANK_C >= 2: + tl.debug_barrier() + _single_expert_l1_l2_body( + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + marker, + EXPERT1_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + EXPERT1_POOL_BASE_C, + 1, + ) + tl.debug_barrier() + tle_raw.call( + edsl_userhopper_ws_tldot_combine_write, + [ + symm_buffer, + l2_out, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + H_C, + I_C, + NUM_PADDED_M_C, + ], + ) + tle_raw.call( + uh.edsl_userhopper_ws_combine_reduce, + [ + symm_buffer, + y, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + H_C, + I_C, + NUM_PADDED_M_C, + CLEANUP_WORKSPACE_C, + ], + ) + tl.store(marker, 0x4C3307) + compute_reader.release(0) + + +@triton.jit +def _l1_l2_expert_wave_single_kernel_worker( + compute_reader, + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, +): + wait_result = compute_reader.wait(0) + _ = tl.load(tle.gpu.local_ptr(wait_result.slot.done, (0,))) + + _single_expert_l1_body( + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + EXPERT0_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + 0, + 0, + ) + if NUM_EXPERTS_PER_RANK_C >= 2: + tl.debug_barrier() + _single_expert_l1_body( + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + EXPERT1_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + EXPERT1_POOL_BASE_C, + 1, + ) + + tl.debug_barrier() + _single_expert_l2_body( + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + marker, + EXPERT0_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + 0, + 0, + ) + if NUM_EXPERTS_PER_RANK_C >= 2: + tl.debug_barrier() + _single_expert_l2_body( + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + marker, + EXPERT1_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + EXPERT1_POOL_BASE_C, + 1, + ) + + tl.debug_barrier() + tle_raw.call( + edsl_userhopper_ws_tldot_combine_write, + [ + symm_buffer, + l2_out, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + H_C, + I_C, + NUM_PADDED_M_C, + ], + ) + tle_raw.call( + uh.edsl_userhopper_ws_combine_reduce, + [ + symm_buffer, + y, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + H_C, + I_C, + NUM_PADDED_M_C, + CLEANUP_WORKSPACE_C, + ], + ) + tl.store(marker, 0x4C3507) + compute_reader.release(0) + + +@triton.jit +def _wait_l1_arrival_one_block( + symm_buffer, + M_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + POOL_BASE_C: tl.constexpr, +): + if M_C > 0: + arrival_offset: tl.constexpr = 32 + NUM_EXPERTS_C * 8 * 2 + NUM_EXPERTS_PER_RANK_C * 8 + arrival_u32 = (symm_buffer + arrival_offset).to(tl.pointer_type(tl.uint32)) + first_block: tl.constexpr = POOL_BASE_C // 64 + while tl.load(arrival_u32 + first_block, volatile=True) < M_C: + pass + + +@triton.jit +def _wait_l1_arrival_one_block_dynamic( + symm_buffer, + M, + NUM_EXPERTS_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + pool_base, +): + if M > 0: + arrival_offset: tl.constexpr = 32 + NUM_EXPERTS_C * 8 * 2 + NUM_EXPERTS_PER_RANK_C * 8 + arrival_u32 = (symm_buffer + arrival_offset).to(tl.pointer_type(tl.uint32)) + first_block = pool_base // 64 + while tl.load(arrival_u32 + first_block, volatile=True) < M: + pass + + +@triton.jit +def _l1_l2_multi_cta_expert_wave_worker( + compute_reader, + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + cta_done, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, +): + wait_result = compute_reader.wait(0) + _ = tl.load(tle.gpu.local_ptr(wait_result.slot.done, (0,))) + pid = tl.program_id(0) + + if pid == 1: + _wait_l1_arrival_one_block( + symm_buffer, + EXPERT0_COUNT_C, + NUM_EXPERTS_C, + NUM_EXPERTS_PER_RANK_C, + 0, + ) + _single_expert_l1_body( + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + EXPERT0_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + 0, + 0, + ) + if JIT_MULTI_CTA_L1_ONLY == 0: + tl.debug_barrier() + _single_expert_l2_body( + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + marker, + EXPERT0_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + 0, + 0, + ) + tl.atomic_add(cta_done, 1, sem="release") + if pid == 2 and NUM_EXPERTS_PER_RANK_C >= 2: + _wait_l1_arrival_one_block( + symm_buffer, + EXPERT1_COUNT_C, + NUM_EXPERTS_C, + NUM_EXPERTS_PER_RANK_C, + EXPERT1_POOL_BASE_C, + ) + _single_expert_l1_body( + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + EXPERT1_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + EXPERT1_POOL_BASE_C, + 1, + ) + if JIT_MULTI_CTA_L1_ONLY == 0: + tl.debug_barrier() + _single_expert_l2_body( + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + marker, + EXPERT1_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + EXPERT1_POOL_BASE_C, + 1, + ) + tl.atomic_add(cta_done, 1, sem="release") + + if pid == 0: + while tl.load(cta_done, volatile=True) < NUM_EXPERTS_PER_RANK_C: + pass + if JIT_MULTI_CTA_L1_ONLY == 0: + tle_raw.call( + edsl_userhopper_ws_tldot_combine_write, + [ + symm_buffer, + l2_out, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + H_C, + I_C, + NUM_PADDED_M_C, + ], + ) + tle_raw.call( + uh.edsl_userhopper_ws_combine_reduce, + [ + symm_buffer, + y, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + H_C, + I_C, + NUM_PADDED_M_C, + CLEANUP_WORKSPACE_C, + ], + ) + tl.store(marker, 0x4C3707) + else: + tl.store(marker, 0x4C3711) + compute_reader.release(0) + + +@triton.jit +def _l1_tile_split_multi_cta_worker( + compute_reader, + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + cta_done, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, +): + wait_result = compute_reader.wait(0) + _ = tl.load(tle.gpu.local_ptr(wait_result.slot.done, (0,))) + pid = tl.program_id(0) + + tile_cta = pid - 1 + active_tile_ctas: tl.constexpr = NUM_EXPERTS_PER_RANK_C * JIT_L1_I_TILES + if tile_cta >= 0 and tile_cta < active_tile_ctas: + local_expert = tile_cta // JIT_L1_I_TILES + tile_idx = tile_cta - local_expert * JIT_L1_I_TILES + count = tl.where(local_expert == 0, EXPERT0_COUNT_C, EXPERT1_COUNT_C) + pool_base = tl.where(local_expert == 0, 0, EXPERT1_POOL_BASE_C) + + _wait_l1_arrival_one_block_dynamic( + symm_buffer, + count, + NUM_EXPERTS_C, + NUM_EXPERTS_PER_RANK_C, + pool_base, + ) + + l1_acts_e = l1_acts + pool_base * H_C + l1_acts_sf_e = l1_acts_sf + pool_base + l1_topk_weights_e = l1_topk_weights + pool_base + l2_acts_e = l2_acts + pool_base * I_C + l2_acts_sf_e = l2_acts_sf + pool_base + l1_weights_e = l1_weights + local_expert * (2 * I_C * H_C) + l1_weights_sf_e = l1_weights_sf + local_expert * ((2 * I_C // 128) * (H_C // 128)) + i_offset = tile_idx * JIT_BLOCK_I + + _l1_single_cta_runtime_m_i_body( + l1_acts_e, + l1_acts_sf_e, + l1_topk_weights_e, + l1_weights_e, + l1_weights_sf_e, + l2_acts_e, + l2_acts_sf_e, + marker, + count, + H_C, + I_C, + NUM_PADDED_M_C, + i_offset, + ) + tl.atomic_add(cta_done, 1, sem="release") + + if pid == 0: + while tl.load(cta_done, volatile=True) < NUM_EXPERTS_PER_RANK_C * JIT_L1_I_TILES: + pass + tl.store(marker, 0x4C3712) + compute_reader.release(0) + + +@triton.jit +def _l1_l2_tile_split_multi_cta_worker( + compute_reader, + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + cta_done, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, + SKIP_COMBINE_C: tl.constexpr, + SKIP_REDUCE_C: tl.constexpr, +): + wait_result = compute_reader.wait(0) + _ = tl.load(tle.gpu.local_ptr(wait_result.slot.done, (0,))) + pid = tl.program_id(0) + + l1_tile_ctas: tl.constexpr = NUM_EXPERTS_PER_RANK_C * JIT_L1_I_TILES + l2_tile_ctas: tl.constexpr = NUM_EXPERTS_PER_RANK_C * JIT_L2_H_TILES + tile_cta = pid - 1 + + if tile_cta >= 0 and tile_cta < l1_tile_ctas: + local_expert = tile_cta // JIT_L1_I_TILES + tile_idx = tile_cta - local_expert * JIT_L1_I_TILES + count = tl.where(local_expert == 0, EXPERT0_COUNT_C, EXPERT1_COUNT_C) + pool_base = tl.where(local_expert == 0, 0, EXPERT1_POOL_BASE_C) + + _wait_l1_arrival_one_block_dynamic( + symm_buffer, + count, + NUM_EXPERTS_C, + NUM_EXPERTS_PER_RANK_C, + pool_base, + ) + + l1_acts_e = l1_acts + pool_base * H_C + l1_acts_sf_e = l1_acts_sf + pool_base + l1_topk_weights_e = l1_topk_weights + pool_base + l2_acts_e = l2_acts + pool_base * I_C + l2_acts_sf_e = l2_acts_sf + pool_base + l1_weights_e = l1_weights + local_expert * (2 * I_C * H_C) + l1_weights_sf_e = l1_weights_sf + local_expert * ((2 * I_C // 128) * (H_C // 128)) + i_offset = tile_idx * JIT_BLOCK_I + + _l1_single_cta_runtime_m_i_body( + l1_acts_e, + l1_acts_sf_e, + l1_topk_weights_e, + l1_weights_e, + l1_weights_sf_e, + l2_acts_e, + l2_acts_sf_e, + marker, + count, + H_C, + I_C, + NUM_PADDED_M_C, + i_offset, + ) + tl.atomic_add(cta_done, 1, sem="release") + + l2_tile_cta = tile_cta - l1_tile_ctas + if l2_tile_cta >= 0 and l2_tile_cta < l2_tile_ctas: + while tl.load(cta_done, volatile=True) < NUM_EXPERTS_PER_RANK_C * JIT_L1_I_TILES: + pass + + local_expert = l2_tile_cta // JIT_L2_H_TILES + hidden_tile = l2_tile_cta - local_expert * JIT_L2_H_TILES + count = tl.where(local_expert == 0, EXPERT0_COUNT_C, EXPERT1_COUNT_C) + pool_base = tl.where(local_expert == 0, 0, EXPERT1_POOL_BASE_C) + + l2_acts_e = l2_acts + pool_base * I_C + l2_acts_sf_e = l2_acts_sf + pool_base + l2_out_e = l2_out + pool_base * H_C + l2_weights_e = l2_weights + local_expert * (H_C * I_C) + l2_weights_sf_e = l2_weights_sf + local_expert * ((H_C // 128) * (I_C // 128)) + n_offset = hidden_tile * JIT_BLOCK_N + + _l2_single_cta_runtime_m_n_body( + l2_acts_e, + l2_acts_sf_e, + l2_weights_e, + l2_weights_sf_e, + l2_out_e, + marker, + count, + H_C, + I_C, + NUM_PADDED_M_C, + n_offset, + ) + tl.atomic_add(cta_done + 1, 1, sem="release") + + if pid == 0: + while tl.load(cta_done + 1, volatile=True) < NUM_EXPERTS_PER_RANK_C * JIT_L2_H_TILES: + pass + if SKIP_COMBINE_C == 0: + tle_raw.call( + edsl_userhopper_ws_tldot_combine_write, + [ + symm_buffer, + l2_out, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + H_C, + I_C, + NUM_PADDED_M_C, + ], + ) + if SKIP_REDUCE_C == 0: + tle_raw.call( + uh.edsl_userhopper_ws_combine_reduce, + [ + symm_buffer, + y, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + H_C, + I_C, + NUM_PADDED_M_C, + CLEANUP_WORKSPACE_C, + ], + ) + tl.store(marker, 0x4C3713) + compute_reader.release(0) + + +@triton.jit +def _l1_expert_wave_split_worker( + compute_reader, + l2_writer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, +): + wait_result = compute_reader.wait(0) + _ = tl.load(tle.gpu.local_ptr(wait_result.slot.done, (0,))) + + _single_expert_l1_body( + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + EXPERT0_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + 0, + 0, + ) + if NUM_EXPERTS_PER_RANK_C >= 2: + tl.debug_barrier() + _single_expert_l1_body( + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + EXPERT1_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + EXPERT1_POOL_BASE_C, + 1, + ) + + tl.debug_barrier() + l2_slot = l2_writer.acquire(0) + tl.store(tle.gpu.local_ptr(l2_slot.done, (0,)), 1) + l2_writer.commit(0) + tl.store(marker, 0x4C3601) + compute_reader.release(0) + + +@triton.jit +def _l2_combine_expert_wave_split_worker( + l2_reader, + symm_buffer, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + M_C: tl.constexpr, + H_C: tl.constexpr, + I_C: tl.constexpr, + NUM_PADDED_M_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, +): + wait_result = l2_reader.wait(0) + _ = tl.load(tle.gpu.local_ptr(wait_result.slot.done, (0,))) + + _single_expert_l2_body( + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + marker, + EXPERT0_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + 0, + 0, + ) + if NUM_EXPERTS_PER_RANK_C >= 2: + tl.debug_barrier() + _single_expert_l2_body( + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + marker, + EXPERT1_COUNT_C, + H_C, + I_C, + NUM_PADDED_M_C, + EXPERT1_POOL_BASE_C, + 1, + ) + + tl.debug_barrier() + tle_raw.call( + edsl_userhopper_ws_tldot_combine_write, + [ + symm_buffer, + l2_out, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + H_C, + I_C, + NUM_PADDED_M_C, + ], + ) + tle_raw.call( + uh.edsl_userhopper_ws_combine_reduce, + [ + symm_buffer, + y, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + H_C, + I_C, + NUM_PADDED_M_C, + CLEANUP_WORKSPACE_C, + ], + ) + tl.store(marker, 0x4C3607) + l2_reader.release(0) + + +@triton.jit +def _single_kernel_dispatch_receiver_l1_tldot_kernel( + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + NUM_TOKENS_C: tl.constexpr, + EXPECTED_LOCAL_RECV_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + INTERMEDIATE_HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, + NUM_DISPATCH_WARPS_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, +): + dispatch_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + compute_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + dispatch_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_dispatch_receiver_sync", + done=dispatch_done, + ) + compute_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_receiver_compute_sync", + done=compute_done, + ) + tle.gpu.warp_specialize( + [ + ( + uh._dispatch_pipe_partition, + ( + dispatch_pipe.writer(), + symm_buffer, + marker, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_DISPATCH_WARPS_C, + ), + ), + ( + uh._receiver_pipe_to_compute_partition, + ( + dispatch_pipe.reader(), + compute_pipe.writer(), + symm_buffer, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + ), + ), + ( + _l1_single_kernel_worker, + ( + compute_pipe.reader(), + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + EXPECTED_LOCAL_RECV_TOKENS_C, + HIDDEN_C, + INTERMEDIATE_HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + NUM_EXPERTS_PER_RANK_C, + EXPERT0_COUNT_C, + EXPERT1_COUNT_C, + EXPERT1_POOL_BASE_C, + CLEANUP_WORKSPACE_C, + ), + ), + ], + [1, 4], + [80, 180], + ) + + +@triton.jit +def _single_kernel_dispatch_receiver_l1_l2_expert_wave_tldot_kernel( + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + NUM_TOKENS_C: tl.constexpr, + EXPECTED_LOCAL_RECV_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + INTERMEDIATE_HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, + NUM_DISPATCH_WARPS_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, +): + dispatch_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + compute_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + dispatch_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_expert_wave_dispatch_receiver_sync", + done=dispatch_done, + ) + compute_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_expert_wave_receiver_compute_sync", + done=compute_done, + ) + tle.gpu.warp_specialize( + [ + ( + uh._dispatch_pipe_partition, + ( + dispatch_pipe.writer(), + symm_buffer, + marker, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_DISPATCH_WARPS_C, + ), + ), + ( + uh._receiver_pipe_to_compute_partition, + ( + dispatch_pipe.reader(), + compute_pipe.writer(), + symm_buffer, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + ), + ), + ( + _l1_l2_expert_wave_single_kernel_worker, + ( + compute_pipe.reader(), + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + EXPECTED_LOCAL_RECV_TOKENS_C, + HIDDEN_C, + INTERMEDIATE_HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + NUM_EXPERTS_PER_RANK_C, + EXPERT0_COUNT_C, + EXPERT1_COUNT_C, + EXPERT1_POOL_BASE_C, + CLEANUP_WORKSPACE_C, + ), + ), + ], + [1, 4], + [80, 180], + ) + + +@triton.jit +def _single_kernel_dispatch_receiver_l1_l2_split_workers_tldot_kernel( + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + NUM_TOKENS_C: tl.constexpr, + EXPECTED_LOCAL_RECV_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + INTERMEDIATE_HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, + NUM_DISPATCH_WARPS_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, +): + dispatch_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + compute_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + l2_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + dispatch_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_split_dispatch_receiver_sync", + done=dispatch_done, + ) + compute_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_split_receiver_l1_sync", + done=compute_done, + ) + l2_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_split_l1_l2_sync", + done=l2_done, + ) + tle.gpu.warp_specialize( + [ + ( + uh._dispatch_pipe_partition, + ( + dispatch_pipe.writer(), + symm_buffer, + marker, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_DISPATCH_WARPS_C, + ), + ), + ( + uh._receiver_pipe_to_compute_partition, + ( + dispatch_pipe.reader(), + compute_pipe.writer(), + symm_buffer, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + ), + ), + ( + _l1_expert_wave_split_worker, + ( + compute_pipe.reader(), + l2_pipe.writer(), + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + EXPECTED_LOCAL_RECV_TOKENS_C, + HIDDEN_C, + INTERMEDIATE_HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + NUM_EXPERTS_PER_RANK_C, + EXPERT0_COUNT_C, + EXPERT1_COUNT_C, + EXPERT1_POOL_BASE_C, + ), + ), + ( + _l2_combine_expert_wave_split_worker, + ( + l2_pipe.reader(), + symm_buffer, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + EXPECTED_LOCAL_RECV_TOKENS_C, + HIDDEN_C, + INTERMEDIATE_HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + NUM_EXPERTS_PER_RANK_C, + EXPERT0_COUNT_C, + EXPERT1_COUNT_C, + EXPERT1_POOL_BASE_C, + CLEANUP_WORKSPACE_C, + ), + ), + ], + [1, 4, 4], + [80, 180, 180], + ) + + +@triton.jit +def _single_kernel_dispatch_receiver_l1_l2_multi_cta_expert_wave_tldot_kernel( + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + cta_done, + NUM_TOKENS_C: tl.constexpr, + EXPECTED_LOCAL_RECV_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + INTERMEDIATE_HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, + NUM_DISPATCH_WARPS_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, +): + dispatch_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + compute_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + dispatch_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_multi_cta_dispatch_receiver_sync", + done=dispatch_done, + ) + compute_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_multi_cta_receiver_compute_sync", + done=compute_done, + ) + tle.gpu.warp_specialize( + [ + ( + uh._dispatch_pipe_partition, + ( + dispatch_pipe.writer(), + symm_buffer, + marker, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_DISPATCH_WARPS_C, + ), + ), + ( + uh._receiver_pipe_to_compute_partition, + ( + dispatch_pipe.reader(), + compute_pipe.writer(), + symm_buffer, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + ), + ), + ( + _l1_l2_multi_cta_expert_wave_worker, + ( + compute_pipe.reader(), + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + cta_done, + EXPECTED_LOCAL_RECV_TOKENS_C, + HIDDEN_C, + INTERMEDIATE_HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + NUM_EXPERTS_PER_RANK_C, + EXPERT0_COUNT_C, + EXPERT1_COUNT_C, + EXPERT1_POOL_BASE_C, + CLEANUP_WORKSPACE_C, + ), + ), + ], + [1, 4], + [80, 180], + ) + + +@triton.jit +def _single_kernel_dispatch_receiver_l1_tile_split_multi_cta_tldot_kernel( + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + cta_done, + NUM_TOKENS_C: tl.constexpr, + EXPECTED_LOCAL_RECV_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + INTERMEDIATE_HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, + NUM_DISPATCH_WARPS_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, +): + dispatch_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + compute_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + dispatch_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_l1_tile_split_dispatch_receiver_sync", + done=dispatch_done, + ) + compute_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_l1_tile_split_receiver_compute_sync", + done=compute_done, + ) + tle.gpu.warp_specialize( + [ + ( + uh._dispatch_pipe_partition, + ( + dispatch_pipe.writer(), + symm_buffer, + marker, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_DISPATCH_WARPS_C, + ), + ), + ( + uh._receiver_pipe_to_compute_partition, + ( + dispatch_pipe.reader(), + compute_pipe.writer(), + symm_buffer, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + ), + ), + ( + _l1_tile_split_multi_cta_worker, + ( + compute_pipe.reader(), + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + cta_done, + EXPECTED_LOCAL_RECV_TOKENS_C, + HIDDEN_C, + INTERMEDIATE_HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + NUM_EXPERTS_C, + NUM_EXPERTS_PER_RANK_C, + EXPERT0_COUNT_C, + EXPERT1_COUNT_C, + EXPERT1_POOL_BASE_C, + ), + ), + ], + [1, 4], + [80, 180], + ) + + +@triton.jit +def _single_kernel_dispatch_receiver_l1_l2_tile_split_multi_cta_tldot_kernel( + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + cta_done, + NUM_TOKENS_C: tl.constexpr, + EXPECTED_LOCAL_RECV_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + INTERMEDIATE_HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, + NUM_DISPATCH_WARPS_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, + SKIP_COMBINE_C: tl.constexpr, + SKIP_REDUCE_C: tl.constexpr, +): + dispatch_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + compute_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + dispatch_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_l1_l2_tile_split_dispatch_receiver_sync", + done=dispatch_done, + ) + compute_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_l1_l2_tile_split_receiver_compute_sync", + done=compute_done, + ) + tle.gpu.warp_specialize( + [ + ( + uh._dispatch_pipe_partition, + ( + dispatch_pipe.writer(), + symm_buffer, + marker, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_DISPATCH_WARPS_C, + ), + ), + ( + uh._receiver_pipe_to_compute_partition, + ( + dispatch_pipe.reader(), + compute_pipe.writer(), + symm_buffer, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + ), + ), + ( + _l1_l2_tile_split_multi_cta_worker, + ( + compute_pipe.reader(), + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + cta_done, + EXPECTED_LOCAL_RECV_TOKENS_C, + HIDDEN_C, + INTERMEDIATE_HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + NUM_EXPERTS_PER_RANK_C, + EXPERT0_COUNT_C, + EXPERT1_COUNT_C, + EXPERT1_POOL_BASE_C, + CLEANUP_WORKSPACE_C, + SKIP_COMBINE_C, + SKIP_REDUCE_C, + ), + ), + ], + [1, 4], + [80, 180], + ) + + +@triton.jit +def _single_kernel_dispatch_receiver_l1_l2_expert_wave_tldot_w2_kernel( + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + NUM_TOKENS_C: tl.constexpr, + EXPECTED_LOCAL_RECV_TOKENS_C: tl.constexpr, + NUM_RANKS_C: tl.constexpr, + NUM_EXPERTS_C: tl.constexpr, + NUM_MAX_TOKENS_PER_RANK_C: tl.constexpr, + NUM_TOPK_C: tl.constexpr, + HIDDEN_C: tl.constexpr, + INTERMEDIATE_HIDDEN_C: tl.constexpr, + NUM_PADDED_SF_POOL_TOKENS_C: tl.constexpr, + NUM_DISPATCH_WARPS_C: tl.constexpr, + NUM_EXPERTS_PER_RANK_C: tl.constexpr, + EXPERT0_COUNT_C: tl.constexpr, + EXPERT1_COUNT_C: tl.constexpr, + EXPERT1_POOL_BASE_C: tl.constexpr, + CLEANUP_WORKSPACE_C: tl.constexpr, +): + dispatch_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + compute_done = tle.gpu.alloc( + [1, 1], + dtype=tl.int32, + layout=None, + scope=tle.gpu.smem, + nv_mma_shared_layout=False, + ) + dispatch_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_expert_wave_w2_dispatch_receiver_sync", + done=dispatch_done, + ) + compute_pipe = tle.pipe( + capacity=1, + scope="cta", + name="single_kernel_expert_wave_w2_receiver_compute_sync", + done=compute_done, + ) + tle.gpu.warp_specialize( + [ + ( + uh._dispatch_pipe_partition, + ( + dispatch_pipe.writer(), + symm_buffer, + marker, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_DISPATCH_WARPS_C, + ), + ), + ( + uh._receiver_pipe_to_compute_partition, + ( + dispatch_pipe.reader(), + compute_pipe.writer(), + symm_buffer, + NUM_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + ), + ), + ( + _l1_l2_expert_wave_single_kernel_worker, + ( + compute_pipe.reader(), + symm_buffer, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y, + marker, + EXPECTED_LOCAL_RECV_TOKENS_C, + HIDDEN_C, + INTERMEDIATE_HIDDEN_C, + NUM_PADDED_SF_POOL_TOKENS_C, + NUM_RANKS_C, + NUM_EXPERTS_C, + NUM_MAX_TOKENS_PER_RANK_C, + NUM_TOPK_C, + NUM_EXPERTS_PER_RANK_C, + EXPERT0_COUNT_C, + EXPERT1_COUNT_C, + EXPERT1_POOL_BASE_C, + CLEANUP_WORKSPACE_C, + ), + ), + ], + [1, 2], + [80, 180], + ) + + +def _check_supported_shape() -> None: + if uh.NUM_EXPERTS_PER_RANK <= 0 or uh.NUM_EXPERTS_PER_RANK > 2: + raise SystemExit("this smoke currently supports 1 or 2 local experts per rank") + if uh.NUM_TOPK <= 0 or uh.NUM_TOPK > uh.NUM_RANKS: + raise SystemExit("this smoke requires 0 < NUM_TOPK <= NUM_RANKS") + if uh.HIDDEN not in (128, 256) or uh.INTERMEDIATE_HIDDEN != 128: + raise SystemExit("this smoke currently requires H in {128, 256} and I=128") + if uh.HIDDEN == 256 and os.environ.get("USERHOPPER_WS_ALLOW_EXPERIMENTAL_H256_TLDOT") != "1": + raise SystemExit( + "H=256 tl.dot path is experimental and currently expected to hit " + "OutOfResources in the full single-kernel L1+L2 worker; set " + "USERHOPPER_WS_ALLOW_EXPERIMENTAL_H256_TLDOT=1 to reproduce" + ) + if PY_BLOCK_M <= 0 or PY_BLOCK_M > DEFAULT_PY_BLOCK_M: + raise SystemExit( + f"USERHOPPER_WS_BLOCK_M must be in [1, {DEFAULT_PY_BLOCK_M}], got {PY_BLOCK_M}" + ) + if PY_BLOCK_N not in (32, 64): + raise SystemExit(f"USERHOPPER_WS_L2_BLOCK_N must be 32 or 64, got {PY_BLOCK_N}") + if uh.HIDDEN % PY_BLOCK_N != 0: + raise SystemExit(f"USERHOPPER_WS_L2_BLOCK_N={PY_BLOCK_N} must divide HIDDEN={uh.HIDDEN}") + if PY_L1_I_TILES <= 0 or PY_L1_I_TILES > DEFAULT_PY_L1_I_TILES: + raise SystemExit( + f"USERHOPPER_WS_L1_I_TILES must be 0 or in [1, {DEFAULT_PY_L1_I_TILES}], " + f"got {os.environ['USERHOPPER_WS_L1_I_TILES']}" + ) + if PY_L2_H_TILES <= 0 or PY_L2_H_TILES > DEFAULT_PY_L2_H_TILES: + raise SystemExit( + f"USERHOPPER_WS_L2_H_TILES must be 0 or in [1, {DEFAULT_PY_L2_H_TILES}], " + f"got {os.environ['USERHOPPER_WS_L2_H_TILES']}" + ) + counts = uh._expected_counts_for_rank(0) + if any(count > PY_BLOCK_M for count in counts): + raise SystemExit(f"this smoke requires received token-topk entries per expert <= {PY_BLOCK_M}") + if any(uh._expected_counts_for_rank(rank) != counts for rank in range(uh.NUM_RANKS)): + raise SystemExit("this smoke currently requires rank-invariant per-local-expert receive counts") + + +def _expected_l2_out(acts, acts_sf, weights, weights_sf, count: int): + acts_f32 = acts[:count].to(torch.float32) + weights_f32 = weights.to(torch.float32) + out = torch.zeros((count, uh.HIDDEN), device=acts.device, dtype=torch.float32) + for k0 in range(0, uh.INTERMEDIATE_HIDDEN, 64): + a = acts_f32[:, k0:k0 + 64] * acts_sf[:count, k0 // 64][:, None] + b = weights_f32[:, k0:k0 + 64] + b = b * weights_sf[:, k0 // 128].repeat_interleave(128)[: uh.HIDDEN][:, None] + out += a @ b.T + return out + + +def main() -> None: + _check_supported_shape() + lib = uh._setup_lib() + lib.userhopper_ws_nvshmem_init_wrapper() + rank = lib.userhopper_ws_nvshmem_team_mype_wrapper() + npes = lib.userhopper_ws_nvshmem_n_pes_wrapper() + if npes != uh.NUM_RANKS: + raise RuntimeError(f"this smoke expects exactly {uh.NUM_RANKS} PEs, got {npes}") + per_expert_counts = uh._expected_counts_for_rank(rank) + expert1_count = per_expert_counts[1] if uh.NUM_EXPERTS_PER_RANK > 1 else 0 + expert1_pool_base = uh._align(per_expert_counts[0], 64) + + torch.cuda.set_device(rank) + device = triton.runtime.driver.active.get_active_torch_device() + stream = torch.cuda.Stream(device=device) + ptr = ctypes.c_void_p(lib.userhopper_ws_nvshmem_alloc_bytes_wrapper(uh.LAYOUT["total_bytes"])) + + whole = uh._view(ptr, 0, (uh.LAYOUT["total_bytes"],), torch.uint8, device) + x_fp8 = uh._view( + ptr, + uh.LAYOUT["input_token"], + (uh.NUM_MAX_TOKENS_PER_RANK, uh.HIDDEN), + torch.float8_e4m3fn, + device, + ) + x_sf = uh._view( + ptr, + uh.LAYOUT["input_sf"], + (uh.NUM_MAX_TOKENS_PER_RANK, uh.HIDDEN // 128), + torch.float32, + device, + ) + topk_idx = uh._view( + ptr, + uh.LAYOUT["input_topk_idx"], + (uh.NUM_MAX_TOKENS_PER_RANK, uh.NUM_TOPK), + torch.int64, + device, + ) + topk_weights = uh._view( + ptr, + uh.LAYOUT["input_topk_weight"], + (uh.NUM_MAX_TOKENS_PER_RANK, uh.NUM_TOPK), + torch.float32, + device, + ) + l1_acts_u8 = uh._view( + ptr, + uh.LAYOUT["l1_token"], + (uh.LAYOUT["num_max_pool_tokens"], uh.HIDDEN), + torch.uint8, + device, + ) + l1_acts = uh._view( + ptr, + uh.LAYOUT["l1_token"], + (uh.LAYOUT["num_max_pool_tokens"], uh.HIDDEN), + torch.float8_e4m3fn, + device, + ) + l1_acts_sf = uh._view_strided( + ptr, + uh.LAYOUT["l1_sf"], + (uh.LAYOUT["num_max_padded_sf_pool_tokens"], uh.HIDDEN // 128), + (1, uh.LAYOUT["num_max_padded_sf_pool_tokens"]), + torch.float32, + device, + ) + l1_topk_weights = uh._view( + ptr, + uh.LAYOUT["l1_topk_weight"], + (uh.LAYOUT["num_max_pool_tokens"],), + torch.float32, + device, + ) + l2_acts_u8 = uh._view( + ptr, + uh.LAYOUT["l2_token"], + (uh.LAYOUT["num_max_pool_tokens"], uh.INTERMEDIATE_HIDDEN), + torch.uint8, + device, + ) + l2_acts = uh._view( + ptr, + uh.LAYOUT["l2_token"], + (uh.LAYOUT["num_max_pool_tokens"], uh.INTERMEDIATE_HIDDEN), + torch.float8_e4m3fn, + device, + ) + l2_acts_sf = uh._view_strided( + ptr, + uh.LAYOUT["l2_sf"], + (uh.LAYOUT["num_max_padded_sf_pool_tokens"], uh.INTERMEDIATE_HIDDEN // 64), + (1, uh.LAYOUT["num_max_padded_sf_pool_tokens"]), + torch.float32, + device, + ) + debug_words = uh._view(ptr, 0, (8,), torch.uint32, device) + expected_count_words = uh._view(ptr, 0, (uh.NUM_EXPERTS_PER_RANK,), torch.uint32, device) + send_count = uh._view(ptr, 32, (uh.NUM_EXPERTS * 2,), torch.uint64, device) + recv_sum = uh._view( + ptr, + 32 + uh.NUM_EXPERTS * 8 * 2, + (uh.NUM_EXPERTS_PER_RANK,), + torch.uint64, + device, + ) + arrival_offset = 32 + uh.NUM_EXPERTS * 8 * 2 + uh.NUM_EXPERTS_PER_RANK * 8 + arrival = uh._view( + ptr, + arrival_offset, + (uh._align(uh.LAYOUT["num_max_pool_blocks"], 2),), + torch.uint32, + device, + ) + l2_arrival_mask = uh._view( + ptr, + arrival_offset + uh._align(uh.LAYOUT["num_max_pool_blocks"], 2) * 4, + (uh.LAYOUT["num_max_pool_blocks"],), + torch.uint64, + device, + ) + marker = torch.empty((1,), dtype=torch.int32, device=device) + cta_done = torch.empty((2,), dtype=torch.int32, device=device) + l1_weights_u8 = torch.empty( + (uh.NUM_EXPERTS_PER_RANK, 2 * uh.INTERMEDIATE_HIDDEN, uh.HIDDEN), + dtype=torch.uint8, + device=device, + ) + l1_weights = l1_weights_u8.view(torch.float8_e4m3fn) + l1_weights_sf = torch.empty( + (uh.NUM_EXPERTS_PER_RANK, 2 * uh.INTERMEDIATE_HIDDEN // 128, uh.HIDDEN // 128), + dtype=torch.float32, + device=device, + ) + l2_weights_u8 = torch.empty( + (uh.NUM_EXPERTS_PER_RANK, uh.HIDDEN, uh.INTERMEDIATE_HIDDEN), + dtype=torch.uint8, + device=device, + ) + l2_weights = l2_weights_u8.view(torch.float8_e4m3fn) + l2_weights_sf = torch.empty( + (uh.NUM_EXPERTS_PER_RANK, uh.HIDDEN // 128, uh.INTERMEDIATE_HIDDEN // 128), + dtype=torch.float32, + device=device, + ) + l2_out = torch.empty( + (uh.LAYOUT["num_max_pool_tokens"], uh.HIDDEN), + dtype=torch.float32, + device=device, + ) + combine = uh._view( + ptr, + uh.LAYOUT["combine_token"], + (uh.NUM_TOPK, uh.NUM_MAX_TOKENS_PER_RANK, uh.HIDDEN), + torch.uint16, + device, + ) + y = torch.empty((uh.NUM_MAX_TOKENS_PER_RANK, uh.HIDDEN), dtype=torch.bfloat16, device=device) + y_u8 = y.view(torch.uint8) + + try: + with torch.cuda.stream(stream): + whole.zero_() + marker.zero_() + l1_weight_values = ( + 0.03125 + * ( + 1.0 + + ( + torch.arange(l1_weights_u8.numel(), dtype=torch.float32, device=device) + % 7.0 + ) + ) + ).reshape_as(l1_weights) + l1_weights.copy_(l1_weight_values) + l1_weights_sf.copy_( + ( + 1.0 + + 0.125 + * torch.arange(l1_weights_sf.numel(), dtype=torch.float32, device=device) + ).reshape_as(l1_weights_sf) + ) + l2_weight_values = ( + 0.015625 + * ( + 1.0 + + ( + torch.arange(l2_weights_u8.numel(), dtype=torch.float32, device=device) + % 11.0 + ) + ) + ).reshape_as(l2_weights) + l2_weights.copy_(l2_weight_values) + l2_weights_sf.copy_( + ( + 0.5 + + 0.0625 + * torch.arange(l2_weights_sf.numel(), dtype=torch.float32, device=device) + ).reshape_as(l2_weights_sf) + ) + l2_out.zero_() + y.zero_() + expected_count_words.copy_( + torch.tensor(per_expert_counts, dtype=torch.uint32, device=device) + ) + debug_words[7] = 0x45585043 + for token in range(uh.NUM_TOKENS): + x_fp8[token].fill_(uh._input_fp8_value(rank, token)) + for sf_idx in range(uh.HIDDEN // 128): + x_sf[token, sf_idx] = rank + 0.125 * (token + 1) + 0.01 * sf_idx + for topk in range(uh.NUM_TOPK): + topk_idx[token, topk] = uh._route_expert(rank, topk) + topk_weights[token, topk] = rank + 0.25 * (token + 1) + 0.03125 * topk + stream.synchronize() + lib.userhopper_ws_nvshmem_barrier_all_wrapper() + + uh._install_cumodule_hook(lib) + if MULTI_CTA_L2_TILE_SPLIT != 0: + kernel = _single_kernel_dispatch_receiver_l1_l2_tile_split_multi_cta_tldot_kernel + elif MULTI_CTA_L1_TILE_SPLIT != 0: + kernel = _single_kernel_dispatch_receiver_l1_tile_split_multi_cta_tldot_kernel + elif MULTI_CTA_EXPERT_WAVE != 0: + kernel = _single_kernel_dispatch_receiver_l1_l2_multi_cta_expert_wave_tldot_kernel + elif SPLIT_L1_L2_WORKERS != 0: + kernel = _single_kernel_dispatch_receiver_l1_l2_split_workers_tldot_kernel + elif EXPERT_WAVE_SINGLE_KERNEL != 0 and EXPERT_WAVE_COMPUTE_WARPS == 2: + kernel = _single_kernel_dispatch_receiver_l1_l2_expert_wave_tldot_w2_kernel + elif EXPERT_WAVE_SINGLE_KERNEL != 0: + kernel = _single_kernel_dispatch_receiver_l1_l2_expert_wave_tldot_kernel + else: + kernel = _single_kernel_dispatch_receiver_l1_tldot_kernel + expected_marker = ( + ( + 0x4C3713 + if MULTI_CTA_L2_TILE_SPLIT != 0 + else ( + 0x4C3712 + if MULTI_CTA_L1_TILE_SPLIT != 0 + else (0x4C3711 if MULTI_CTA_L1_ONLY != 0 else 0x4C3707) + ) + ) + if MULTI_CTA_EXPERT_WAVE != 0 + else ( + 0x4C3607 + if SPLIT_L1_L2_WORKERS != 0 + else (0x4C3507 if EXPERT_WAVE_SINGLE_KERNEL != 0 else 0x4C3307) + ) + ) + compute_order = ( + ( + "expert_wave_multi_cta_l1_l2_tile_split" + if MULTI_CTA_L2_TILE_SPLIT != 0 + else ( + "expert_wave_multi_cta_l1_tile_split" + if MULTI_CTA_L1_TILE_SPLIT != 0 + else ("expert_wave_multi_cta_l1_only" if MULTI_CTA_L1_ONLY != 0 else "expert_wave_multi_cta") + ) + ) + if MULTI_CTA_EXPERT_WAVE != 0 + else ( + "expert_wave_split_workers" + if SPLIT_L1_L2_WORKERS != 0 + else ("expert_wave" if EXPERT_WAVE_SINGLE_KERNEL != 0 else "sequential") + ) + ) + grid = ( + (1 + uh.NUM_EXPERTS_PER_RANK * PY_L1_I_TILES + uh.NUM_EXPERTS_PER_RANK * PY_L2_H_TILES,) + if MULTI_CTA_L2_TILE_SPLIT != 0 + else ( + (1 + uh.NUM_EXPERTS_PER_RANK * PY_L1_I_TILES,) + if MULTI_CTA_L1_TILE_SPLIT != 0 + else ((1 + uh.NUM_EXPERTS_PER_RANK,) if MULTI_CTA_EXPERT_WAVE != 0 else (1,)) + ) + ) + launch_times_us = [] + launch_start = torch.cuda.Event(enable_timing=True) + launch_end = torch.cuda.Event(enable_timing=True) + with torch.cuda.stream(stream): + cta_done.zero_() + marker.zero_() + launch_start.record(stream) + if MULTI_CTA_L2_TILE_SPLIT != 0: + compiled = kernel[grid]( + whole, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y_u8, + marker, + cta_done, + uh.NUM_TOKENS, + uh.EXPECTED_LOCAL_RECV_TOKENS, + uh.NUM_RANKS, + uh.NUM_EXPERTS, + uh.NUM_MAX_TOKENS_PER_RANK, + uh.NUM_TOPK, + uh.HIDDEN, + uh.INTERMEDIATE_HIDDEN, + uh.LAYOUT["num_max_padded_sf_pool_tokens"], + uh.NUM_DISPATCH_WARPS, + uh.NUM_EXPERTS_PER_RANK, + per_expert_counts[0], + expert1_count, + expert1_pool_base, + CLEANUP_WORKSPACE, + SKIP_COMBINE, + SKIP_REDUCE, + num_warps=uh.NUM_WARPS, + maxnreg=uh.MAXNREG, + ) + elif MULTI_CTA_L1_TILE_SPLIT != 0: + compiled = kernel[grid]( + whole, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + cta_done, + uh.NUM_TOKENS, + uh.EXPECTED_LOCAL_RECV_TOKENS, + uh.NUM_RANKS, + uh.NUM_EXPERTS, + uh.NUM_MAX_TOKENS_PER_RANK, + uh.NUM_TOPK, + uh.HIDDEN, + uh.INTERMEDIATE_HIDDEN, + uh.LAYOUT["num_max_padded_sf_pool_tokens"], + uh.NUM_DISPATCH_WARPS, + uh.NUM_EXPERTS_PER_RANK, + per_expert_counts[0], + expert1_count, + expert1_pool_base, + num_warps=uh.NUM_WARPS, + maxnreg=uh.MAXNREG, + ) + elif MULTI_CTA_EXPERT_WAVE != 0: + compiled = kernel[grid]( + whole, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y_u8, + marker, + cta_done, + uh.NUM_TOKENS, + uh.EXPECTED_LOCAL_RECV_TOKENS, + uh.NUM_RANKS, + uh.NUM_EXPERTS, + uh.NUM_MAX_TOKENS_PER_RANK, + uh.NUM_TOPK, + uh.HIDDEN, + uh.INTERMEDIATE_HIDDEN, + uh.LAYOUT["num_max_padded_sf_pool_tokens"], + uh.NUM_DISPATCH_WARPS, + uh.NUM_EXPERTS_PER_RANK, + per_expert_counts[0], + expert1_count, + expert1_pool_base, + CLEANUP_WORKSPACE, + num_warps=uh.NUM_WARPS, + maxnreg=uh.MAXNREG, + ) + else: + compiled = kernel[grid]( + whole, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y_u8, + marker, + uh.NUM_TOKENS, + uh.EXPECTED_LOCAL_RECV_TOKENS, + uh.NUM_RANKS, + uh.NUM_EXPERTS, + uh.NUM_MAX_TOKENS_PER_RANK, + uh.NUM_TOPK, + uh.HIDDEN, + uh.INTERMEDIATE_HIDDEN, + uh.LAYOUT["num_max_padded_sf_pool_tokens"], + uh.NUM_DISPATCH_WARPS, + uh.NUM_EXPERTS_PER_RANK, + per_expert_counts[0], + expert1_count, + expert1_pool_base, + CLEANUP_WORKSPACE, + num_warps=uh.NUM_WARPS, + maxnreg=uh.MAXNREG, + ) + launch_end.record(stream) + stream.synchronize() + launch_times_us.append(launch_start.elapsed_time(launch_end) * 1000.0) + lib.userhopper_ws_nvshmem_barrier_all_wrapper() + + if int(marker.cpu()[0]) != expected_marker: + raise SystemExit(f"single-kernel L1 marker mismatch: {int(marker.cpu()[0])}") + ttgir = compiled.asm.get("ttgir", "") + ptx = compiled.asm.get("ptx", "") + if "ttg.warp_specialize" not in ttgir: + raise SystemExit("single-kernel TTGIR missing ttg.warp_specialize") + if "wgmma.mma_async" not in ptx and "mma.sync" not in ptx: + raise SystemExit("single-kernel PTX missing Tensor Core mma instruction") + + expected_indices, expected_rows, expected_sf, expected_weight, _, _ = uh._expected_receive(rank) + got_l1 = l1_acts_u8.detach().cpu()[expected_indices] + got_l1_sf = l1_acts_sf.detach().cpu()[expected_indices, :] + got_l1_weight = l1_topk_weights.detach().cpu()[expected_indices] + if not torch.equal(got_l1, expected_rows): + raise SystemExit("single-kernel receiver L1 token mismatch") + if not torch.allclose(got_l1_sf, expected_sf): + raise SystemExit("single-kernel receiver L1 scale mismatch") + if not torch.allclose(got_l1_weight, expected_weight): + raise SystemExit("single-kernel receiver L1 topk weight mismatch") + + expected_l2_float = uh._expected_l2_token_floats( + rank, + l1_weights_u8.detach().cpu(), + l1_weights_sf.detach().cpu(), + ) + expected_l2_sf, expected_l2_scaled = uh._expected_l2_sf_and_scaled_floats(expected_l2_float) + expected_l2_tokens = uh._float_to_cuda_satfinite_e4m3_bytes(expected_l2_scaled, device) + got_l2_tokens = l2_acts_u8.detach().cpu()[expected_indices] + got_l2_sf = l2_acts_sf.detach().cpu()[expected_indices, :] + computed_intermediate = PY_L1_I_TILES * 64 + full_l1_validation = computed_intermediate >= uh.INTERMEDIATE_HIDDEN + if not torch.equal( + got_l2_tokens[:, :computed_intermediate], + expected_l2_tokens[:, :computed_intermediate], + ): + raise SystemExit( + "single-kernel L1 Tensor Core l2_acts mismatch: " + f"checked_intermediate={computed_intermediate} " + f"got={got_l2_tokens[:, :16].tolist()} expected={expected_l2_tokens[:, :16].tolist()}" + ) + expected_l2_sf_checked = expected_l2_sf[:, :PY_L1_I_TILES] + got_l2_sf_checked = got_l2_sf[:, :PY_L1_I_TILES] + l2_sf_tol = max(1e-5, 1e-4 * float(expected_l2_sf_checked.abs().max().item())) + if not torch.allclose(got_l2_sf_checked, expected_l2_sf_checked, atol=l2_sf_tol, rtol=1e-4): + max_abs = float((got_l2_sf_checked - expected_l2_sf_checked).abs().max().item()) + raise SystemExit( + "single-kernel L1 Tensor Core l2_acts_sf mismatch: " + f"checked_i_tiles={PY_L1_I_TILES} max_abs={max_abs} tol={l2_sf_tol}" + ) + + pool_base = 0 + max_l2_abs = 0.0 + l2_out_tol = 0.0 + computed_hidden = PY_L2_H_TILES * PY_BLOCK_N + l2_is_computed = MULTI_CTA_L2_TILE_SPLIT != 0 or MULTI_CTA_L1_ONLY == 0 + full_l2_validation = l2_is_computed and computed_hidden >= uh.HIDDEN + full_output_validation = full_l1_validation and full_l2_validation + combine_validation = full_output_validation and SKIP_COMBINE == 0 + y_validation = combine_validation and SKIP_REDUCE == 0 + if l2_is_computed: + for local_expert, count in enumerate(per_expert_counts): + if count > 0: + expected_out = _expected_l2_out( + l2_acts[pool_base:pool_base + count], + l2_acts_sf[pool_base:], + l2_weights[local_expert], + l2_weights_sf[local_expert], + count, + ) + got_out = l2_out[pool_base:pool_base + count] + expected_checked = expected_out[:, :computed_hidden] + got_checked = got_out[:, :computed_hidden] + expert_max_abs = float((got_checked - expected_checked).abs().max().item()) + expert_max_ref = float(expected_checked.abs().max().item()) + expert_tol = max(1e-2, 5e-3 * expert_max_ref) + max_l2_abs = max(max_l2_abs, expert_max_abs) + l2_out_tol = max(l2_out_tol, expert_tol) + if not torch.allclose(got_checked, expected_checked, atol=expert_tol, rtol=5e-3): + raise SystemExit( + "single-kernel L2 Tensor Core output mismatch: " + f"local_expert={local_expert} checked_hidden={computed_hidden} " + f"max_abs={expert_max_abs} tol={expert_tol}" + ) + pool_base += uh._align(count, 64) + + combine_expected_mask = None + expected_combine_valid = None + expected_y = None + combine_tol = 0.0 + y_tol = 0.0 + if combine_validation: + expected_combine = uh._expected_combine_float( + rank, + l1_weights_u8.detach().cpu(), + l1_weights_sf.detach().cpu(), + l2_weights_u8.detach().cpu(), + l2_weights_sf.detach().cpu(), + device, + ).to(torch.bfloat16).float() + got_combine = combine.detach().cpu().contiguous().view(torch.bfloat16).float() + combine_expected_mask = torch.zeros( + (uh.NUM_TOPK, uh.NUM_MAX_TOKENS_PER_RANK), dtype=torch.bool + ) + for token in range(uh.NUM_TOKENS): + for topk in range(uh.NUM_TOPK): + if uh._route_expert(rank, topk) >= 0: + combine_expected_mask[topk, token] = True + got_combine_valid = got_combine[combine_expected_mask] + expected_combine_valid = expected_combine[combine_expected_mask] + combine_tol = max( + 1e-1, + 5e-3 * float(expected_combine_valid.abs().max().item() if expected_combine_valid.numel() else 1.0), + ) + if not torch.allclose(got_combine_valid, expected_combine_valid, atol=combine_tol, rtol=5e-3): + raise SystemExit(f"single-kernel combine mismatch: tol={combine_tol}") + expected_y = got_combine.sum(dim=0).to(torch.bfloat16) + if y_validation: + got_y = y.detach().cpu() + y_tol = max(1e-1, 5e-3 * float(expected_y[: uh.NUM_TOKENS].float().abs().max().item())) + if not torch.allclose(got_y[: uh.NUM_TOKENS].float(), expected_y[: uh.NUM_TOKENS].float(), atol=y_tol, rtol=5e-3): + raise SystemExit(f"single-kernel y mismatch: tol={y_tol}") + if CLEANUP_WORKSPACE != 0: + uh._validate_workspace_cleanup( + send_count.detach().cpu(), + recv_sum.detach().cpu(), + arrival.detach().cpu(), + l2_arrival_mask.detach().cpu(), + ) + + for repeat_idx in range(1, REPEAT_LAUNCHES): + lib.userhopper_ws_nvshmem_barrier_all_wrapper() + launch_start = torch.cuda.Event(enable_timing=True) + launch_end = torch.cuda.Event(enable_timing=True) + with torch.cuda.stream(stream): + marker.zero_() + cta_done.zero_() + launch_start.record(stream) + if MULTI_CTA_L2_TILE_SPLIT != 0: + compiled = kernel[grid]( + whole, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y_u8, + marker, + cta_done, + uh.NUM_TOKENS, + uh.EXPECTED_LOCAL_RECV_TOKENS, + uh.NUM_RANKS, + uh.NUM_EXPERTS, + uh.NUM_MAX_TOKENS_PER_RANK, + uh.NUM_TOPK, + uh.HIDDEN, + uh.INTERMEDIATE_HIDDEN, + uh.LAYOUT["num_max_padded_sf_pool_tokens"], + uh.NUM_DISPATCH_WARPS, + uh.NUM_EXPERTS_PER_RANK, + per_expert_counts[0], + expert1_count, + expert1_pool_base, + CLEANUP_WORKSPACE, + SKIP_COMBINE, + SKIP_REDUCE, + num_warps=uh.NUM_WARPS, + maxnreg=uh.MAXNREG, + ) + elif MULTI_CTA_L1_TILE_SPLIT != 0: + compiled = kernel[grid]( + whole, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + marker, + cta_done, + uh.NUM_TOKENS, + uh.EXPECTED_LOCAL_RECV_TOKENS, + uh.NUM_RANKS, + uh.NUM_EXPERTS, + uh.NUM_MAX_TOKENS_PER_RANK, + uh.NUM_TOPK, + uh.HIDDEN, + uh.INTERMEDIATE_HIDDEN, + uh.LAYOUT["num_max_padded_sf_pool_tokens"], + uh.NUM_DISPATCH_WARPS, + uh.NUM_EXPERTS_PER_RANK, + per_expert_counts[0], + expert1_count, + expert1_pool_base, + num_warps=uh.NUM_WARPS, + maxnreg=uh.MAXNREG, + ) + elif MULTI_CTA_EXPERT_WAVE != 0: + compiled = kernel[grid]( + whole, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y_u8, + marker, + cta_done, + uh.NUM_TOKENS, + uh.EXPECTED_LOCAL_RECV_TOKENS, + uh.NUM_RANKS, + uh.NUM_EXPERTS, + uh.NUM_MAX_TOKENS_PER_RANK, + uh.NUM_TOPK, + uh.HIDDEN, + uh.INTERMEDIATE_HIDDEN, + uh.LAYOUT["num_max_padded_sf_pool_tokens"], + uh.NUM_DISPATCH_WARPS, + uh.NUM_EXPERTS_PER_RANK, + per_expert_counts[0], + expert1_count, + expert1_pool_base, + CLEANUP_WORKSPACE, + num_warps=uh.NUM_WARPS, + maxnreg=uh.MAXNREG, + ) + else: + compiled = kernel[grid]( + whole, + l1_acts, + l1_acts_sf, + l1_topk_weights, + l1_weights, + l1_weights_sf, + l2_acts, + l2_acts_sf, + l2_weights, + l2_weights_sf, + l2_out, + y_u8, + marker, + uh.NUM_TOKENS, + uh.EXPECTED_LOCAL_RECV_TOKENS, + uh.NUM_RANKS, + uh.NUM_EXPERTS, + uh.NUM_MAX_TOKENS_PER_RANK, + uh.NUM_TOPK, + uh.HIDDEN, + uh.INTERMEDIATE_HIDDEN, + uh.LAYOUT["num_max_padded_sf_pool_tokens"], + uh.NUM_DISPATCH_WARPS, + uh.NUM_EXPERTS_PER_RANK, + per_expert_counts[0], + expert1_count, + expert1_pool_base, + CLEANUP_WORKSPACE, + num_warps=uh.NUM_WARPS, + maxnreg=uh.MAXNREG, + ) + launch_end.record(stream) + stream.synchronize() + launch_times_us.append(launch_start.elapsed_time(launch_end) * 1000.0) + lib.userhopper_ws_nvshmem_barrier_all_wrapper() + if int(marker.cpu()[0]) != expected_marker: + raise SystemExit( + f"repeat launch {repeat_idx} marker mismatch: {int(marker.cpu()[0])}" + ) + uh._validate_workspace_cleanup( + send_count.detach().cpu(), + recv_sum.detach().cpu(), + arrival.detach().cpu(), + l2_arrival_mask.detach().cpu(), + ) + if combine_validation: + repeat_combine = combine.detach().cpu().contiguous().view(torch.bfloat16).float() + repeat_combine_valid = repeat_combine[combine_expected_mask] + if not torch.allclose( + repeat_combine_valid, + expected_combine_valid, + atol=combine_tol, + rtol=5e-3, + ): + raise SystemExit(f"repeat launch {repeat_idx} combine mismatch: tol={combine_tol}") + if y_validation: + repeat_y = y.detach().cpu() + if not torch.allclose( + repeat_y[: uh.NUM_TOKENS].float(), + expected_y[: uh.NUM_TOKENS].float(), + atol=y_tol, + rtol=5e-3, + ): + raise SystemExit(f"repeat launch {repeat_idx} y mismatch: tol={y_tol}") + + steady_launch_times_us = launch_times_us[1:] if len(launch_times_us) > 1 else launch_times_us + print( + "rank={} userhopper_single_kernel_l1_l2_combine_tldot_smoke=PASS checked={} " + "counts={} ws=checked raw_nvshmem=checked tensor_core=checked " + "combine={} y={} compute_order={} compute_warps={} l1_i_tiles={}/{} " + "l2_h_tiles={}/{} l2_block_n={} l2_mode={} cleanup={} repeats={} " + "launch_avg_us={:.3f} launch_min_us={:.3f} launch_max_us={:.3f} " + "launch_steady_avg_us={:.3f} launch_steady_min_us={:.3f} launch_steady_max_us={:.3f} " + "l2_max_abs={:.6g} l2_tol={:.6g}".format( + rank, + len(expected_indices), + uh._expected_counts_for_rank(rank), + "skipped" if SKIP_COMBINE != 0 else ("checked" if full_output_validation else "partial-skip"), + "skipped" if (SKIP_COMBINE != 0 or SKIP_REDUCE != 0) else ("checked" if full_output_validation else "partial-skip"), + compute_order, + EXPERT_WAVE_COMPUTE_WARPS if EXPERT_WAVE_SINGLE_KERNEL != 0 else 4, + PY_L1_I_TILES, + DEFAULT_PY_L1_I_TILES, + PY_L2_H_TILES, + DEFAULT_PY_L2_H_TILES, + PY_BLOCK_N, + "scalar" if L2_SCALAR != 0 else "tensorcore", + CLEANUP_WORKSPACE, + REPEAT_LAUNCHES, + sum(launch_times_us) / len(launch_times_us), + min(launch_times_us), + max(launch_times_us), + sum(steady_launch_times_us) / len(steady_launch_times_us), + min(steady_launch_times_us), + max(steady_launch_times_us), + max_l2_abs, + l2_out_tol, + ), + flush=True, + ) + finally: + lib.userhopper_ws_nvshmem_barrier_all_wrapper() + lib.userhopper_ws_nvshmem_finalize_wrapper(ptr) + + +if __name__ == "__main__": + main() diff --git a/python/test/tle/integration/megamoe/tests/run_isolated_operator.py b/python/test/tle/integration/megamoe/tests/run_isolated_operator.py new file mode 100644 index 000000000..6ca017c84 --- /dev/null +++ b/python/test/tle/integration/megamoe/tests/run_isolated_operator.py @@ -0,0 +1,68 @@ +"""Run the isolated operator package with the local correctness harness. + +This script reuses the copied host allocation/reference-check harness in this +directory, but redirects the runtime support and Triton kernel entries to this +isolated package. +""" + +from __future__ import annotations + +import importlib +import pathlib +import sys + + +ROOT = pathlib.Path(__file__).resolve().parents[1] +TRITON_PYTHON_DIR = ROOT.parents[3] + +sys.path.insert(0, str(ROOT)) +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) +sys.path.insert(0, str(TRITON_PYTHON_DIR)) + +import megamoe_operator.triton_tle_megamoe_operator as isolated_op +import megamoe_operator.triton_tle_megamoe_runtime as isolated_runtime + +harness = importlib.import_module("megamoe_local_harness") + +harness.uh = isolated_runtime + +for name in ( + "_single_kernel_dispatch_receiver_l1_tldot_kernel", + "_single_kernel_dispatch_receiver_l1_l2_expert_wave_tldot_kernel", + "_single_kernel_dispatch_receiver_l1_l2_expert_wave_tldot_w2_kernel", + "_single_kernel_dispatch_receiver_l1_l2_split_workers_tldot_kernel", + "_single_kernel_dispatch_receiver_l1_l2_multi_cta_expert_wave_tldot_kernel", + "_single_kernel_dispatch_receiver_l1_tile_split_multi_cta_tldot_kernel", + "_single_kernel_dispatch_receiver_l1_l2_tile_split_multi_cta_tldot_kernel", +): + if hasattr(isolated_op, name): + setattr(harness, name, getattr(isolated_op, name)) + +for name in ( + "PY_BLOCK_M", + "PY_BLOCK_N", + "PY_L1_I_TILES", + "PY_L2_H_TILES", + "DEFAULT_PY_BLOCK_M", + "DEFAULT_PY_L1_I_TILES", + "DEFAULT_PY_L2_H_TILES", + "CLEANUP_WORKSPACE", + "REPEAT_LAUNCHES", + "EXPERT_WAVE_SINGLE_KERNEL", + "EXPERT_WAVE_COMPUTE_WARPS", + "SPLIT_L1_L2_WORKERS", + "L1_TILE_LOOP", + "L2_SCALAR", + "MULTI_CTA_EXPERT_WAVE", + "MULTI_CTA_L1_ONLY", + "MULTI_CTA_L1_TILE_SPLIT", + "MULTI_CTA_L2_TILE_SPLIT", + "SKIP_COMBINE", + "SKIP_REDUCE", +): + if hasattr(isolated_op, name): + setattr(harness, name, getattr(isolated_op, name)) + + +if __name__ == "__main__": + harness.main() diff --git a/python/test/tle/integration/megamoe/tools/nvcc_flock_wrapper.sh b/python/test/tle/integration/megamoe/tools/nvcc_flock_wrapper.sh new file mode 100755 index 000000000..93e49bbac --- /dev/null +++ b/python/test/tle/integration/megamoe/tools/nvcc_flock_wrapper.sh @@ -0,0 +1,9 @@ +#!/usr/bin/env bash +set -euo pipefail + +lock_file="${TRITON_TLE_NVCC_LOCK:-/tmp/triton_tle_nvcc.lock}" +real_nvcc="${REAL_NVCC:-/usr/local/cuda-12.8/bin/nvcc}" + +exec 9>"${lock_file}" +flock 9 +exec "${real_nvcc}" "$@"