[TLERaw][Mega] Add Triton TLE NVSHMEM MegaMoE correctness case

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,

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<MemoryEffectsOpInterface>,
+  ConditionallySpeculatable,
+]> {
+
+    let description = [{
+        call an external function $symbol implemented in $libpath/$libname with $args
+        return $libpath/$libname:$symbol($args...)
+    }];
+
+    let arguments = (ins Variadic<TT_Type>:$srcs, StrAttr:$libname, StrAttr:$libpath, StrAttr:$symbol, BoolAttr:$pure);
+
+    let results = (outs Variadic<TT_Type>:$result);
+
+    let assemblyFormat = "operands attr-dict `:` functional-type(operands, $result)";
+
+    let extraClassDeclaration = [{
+      // Interface method for ConditionallySpeculatable.
+      Speculation::Speculatability getSpeculatability();
+    }];
+
+}
+
 //
 // External Elementwise op
 //

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

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<triton::ExternCallOp> {
+public:
+  ExternCallOpConversion(const LLVMTypeConverter &converter,
+                         const PatternBenefit &benefit)
+      : ConvertOpToLLVMPattern<triton::ExternCallOp>(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<ExternCallOpConversion>(typeConverter, benefit);
+}

lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp

@@ -618,6 +618,7 @@ void populateTritonPatterns(TritonGPUTypeConverter &typeConverter,
       GenericOpPattern<triton::HistogramOp>,
       GenericOpPattern<triton::GatherOp>,
       GenericOpPattern<triton::ExternElementwiseOp>,
+      GenericOpPattern<triton::ExternCallOp>,
       GenericOpPattern<triton::PrintOp>,
       GenericOpPattern<triton::AssertOp>,
       GenericOpPattern<triton::AtomicCASOp>,

lib/Dialect/Triton/IR/Ops.cpp

@@ -1313,6 +1313,24 @@ Speculation::Speculatability ExternElementwiseOp::getSpeculatability() {
   return Speculation::NotSpeculatable;
 }
 
+// -- ExternCallOp --
+void ExternCallOp::getEffects(
+    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
+        &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();

python/src/ir.cc

@@ -1673,6 +1673,14 @@ void init_triton_ir(py::module &&m) {
              return self.create<ExternElementwiseOp>(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<Value> &argList, const std::vector<Type> &retTypes,
+              bool isPure) -> OpState {
+             return self.create<ExternCallOp>(retTypes, argList, libName,
+                                              libPath, symbol, isPure);
+           })
       // Built-in instruction
       .def("create_get_program_id",
            [](TritonOpBuilder &self, int axis) -> Value {

python/test/tle/integration/megamoe/.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+*.py[cod]
+*.o
+*.so
+*.lock
+*.tmp

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` |

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",
+]