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/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}" "$@" diff --git a/python/triton/experimental/tle/language/raw/core.py b/python/triton/experimental/tle/language/raw/core.py index d7029380e..51836be76 100644 --- a/python/triton/experimental/tle/language/raw/core.py +++ b/python/triton/experimental/tle/language/raw/core.py @@ -2,9 +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.compiler).lower() == "nvcc": + func.make_cubin() + 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 9015f2bf7..c3f21bc89 100644 --- a/python/triton/experimental/tle/raw/cuda/runtime.py +++ b/python/triton/experimental/tle/raw/cuda/runtime.py @@ -3,29 +3,163 @@ from pathlib import Path import subprocess 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] # 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") +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: - super().__init__(*args, **kwargs) + super().__init__() 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("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" 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 def make_llvm(self, mlir_context) -> str: + 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( [ CLANG, "-x", "cuda", "--cuda-device-only", + f"--cuda-gpu-arch={arch}", "-emit-llvm", "-O2", "-S", @@ -40,3 +174,6 @@ 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): + ... \ No newline at end of file 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..5e4448e65 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py @@ -0,0 +1,89 @@ +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" + +@dialect( + name="cuda", + 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): + ... + + +@triton.jit +def simple_shift_kernel(destination_ptr, ): + tle_raw.call(edsl, [destination_ptr]) + + +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("PMI_RANK", "0")) + + 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/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..041b44ca6 --- /dev/null +++ b/python/tutorials/tle/raw/nvshmem/02-ring-reduce/ring-reduce.py @@ -0,0 +1,110 @@ +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" +@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": NVSHMEM_HOME}, + links=["nvshmem_device"] +) +def edsl(*args, **kwargs): + ... + + +@triton.jit +def ring_reduce_kernel( + dst, + src, + nreduce, + signal, + chunk_size, +): + tle_raw.call(edsl, [dst, src, nreduce, signal, chunk_size]) + +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"]) + + 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 + ) + + 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) + + 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, ) + + 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): + chunk_size = int(M // num_blocks * dtype.itemsize) + ring_reduce_kernel[(num_blocks, )]( + dst_tensor, + src_tensor, + M * N, + signal_tensor, + chunk_size + ) + + print(f"PE {mype.value}: {npes.value}") + 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-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/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/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/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,