feat(amd): EP intra-node normal and low-latency kernels with mori shmem (#164)

* feat(amd): EP intra-node normal and low-latency kernels with mori shmem

- Implement EP intra-node dispatch/combine kernels using mori shmem P2P (putmem_signal_warp) on AMD MI325X
- Add Low Latency EP v1 (raw all-to-all) and v2 (online FP8 quant + combine with topk weighted reduce)
- Fix shfl_up/shfl_down_sync implementation and golden reference calculation in test_language_extra.py
- Fix mixed-bitwidth ld/st implementation and add kernel test coverage
- Update mori submodule to main with JIT bitcode compilation, replacing manual hipcc/llvm-link build
- Simplify `build_mori_shmem.sh` to use mori JIT (`mori.ir.bitcode.find_bitcode()`)
- Add AlgoBW and BusBW metrics to EP A2A benchmark output
- Add CI tests for EP A2A (correctness + perf), LL v2 (correctness + perf M=64/128)

---------

Co-authored-by: Wu, Yutong <yutong.wu@amd.com>

* fix(ci): use non-recursive submodule checkout to avoid pulling mori's nested submodules

---------

Co-authored-by: Wu, Yutong <yutong.wu@amd.com>

Author: jhchouuu

.github/workflows/amd-ci.yml

@@ -32,7 +32,7 @@ jobs:
       - name: Checkout
         uses: actions/checkout@v4
         with:
-          submodules: 'recursive'
+          submodules: 'true'
       - name: Build rocmshmem bind
         run: |
           bash ./shmem/rocshmem_bind/build.sh
@@ -78,7 +78,7 @@ jobs:
       - name: Checkout
         uses: actions/checkout@v4
         with:
-          submodules: 'recursive'
+          submodules: 'true'
       - name: Build rocmshmem bind
         run: |
           bash ./shmem/rocshmem_bind/build.sh
@@ -124,14 +124,26 @@ jobs:
       - name: Checkout
         uses: actions/checkout@v4
         with:
-          submodules: 'recursive'
+          submodules: 'true'
       - name: Build triton-distributed
         run: |
           pip3 install -e python --verbose --no-build-isolation --use-pep517
       - name: Build mori shmem
         run: |
-          bash ./scripts/build_mori_shmem.sh 
-      - name: Mori SHMEM API tests
+          bash ./scripts/build_mori_shmem.sh
+      - name: MoRI SHMEM API tests
         run: |
           bash ./scripts/launch_amd.sh ./python/triton_dist/test/amd/test_mori_shmem_api.py
-          bash ./scripts/launch_amd.sh ./python/triton_dist/test/amd/test_mori_shmem_bw.py
+          bash ./scripts/launch_amd.sh ./python/triton_dist/test/amd/test_mori_shmem_bw.py
+      - name: EP A2A intra-node tests
+        run: |
+          bash ./scripts/launch_amd.sh ./python/triton_dist/test/amd/test_ep_a2a.py --check
+          bash ./scripts/launch_amd.sh ./python/triton_dist/test/amd/test_ep_a2a.py --check --with-scatter-indices
+          bash ./scripts/launch_amd.sh ./python/triton_dist/test/amd/test_ep_a2a.py --check --enable-local-combine
+          bash ./scripts/launch_amd.sh ./python/triton_dist/test/amd/test_ep_a2a.py --rounds 3 --bench_iters 10
+      - name: EP Low Latency v2 tests
+        run: |
+          bash ./scripts/launch_amd.sh ./python/triton_dist/test/amd/test_ep_ll_a2a.py --check -M 64 -N 7168 -G 256 --topk 8
+          bash ./scripts/launch_amd.sh ./python/triton_dist/test/amd/test_ep_ll_a2a.py -M 64 -N 7168 -G 256 --topk 8 --rounds 3
+          bash ./scripts/launch_amd.sh ./python/triton_dist/test/amd/test_ep_ll_a2a.py -M 128 -N 7168 -G 256 --topk 8 --rounds 3
+

.gitmodules

@@ -4,7 +4,6 @@
 [submodule "3rdparty/mori"]
 	path = 3rdparty/mori
 	url = https://github.com/ROCm/mori.git
-	branch = jiahzhou/triton_dis_support
 [submodule "3rdparty/cutlass"]
 	path = 3rdparty/cutlass
 	url = https://github.com/NVIDIA/cutlass.git

3rdparty/mori

@@ -118,6 +118,14 @@ class CallOpConversion : public OpRewritePattern<LLVM::CallOp> {
         llvm_unreachable("unknown scope string");
     };
 
+    auto skipBitwidthPrefix =
+        [](const SmallVector<StringRef> &parts) -> size_t {
+      if (!parts.empty() &&
+          llvm::all_of(parts[0], [](char c) { return std::isdigit(c); }))
+        return 1;
+      return 0;
+    };
+
     auto operands = callOp.getOperands();
     auto result = callOp.getResult();
 
@@ -193,10 +201,12 @@ class CallOpConversion : public OpRewritePattern<LLVM::CallOp> {
     if (auto maybeParts =
             matchPrefixAndSplitRemainder(calleeName, "__triton_hip_load_")) {
       auto parts = maybeParts.value();
-      assert(parts.size() == 2 &&
-             "expected load function to have 2 parts after prefix");
-      LLVM::AtomicOrdering memOrder = strToMemoryOrder(parts[0]);
-      auto scopeStr = strToScope(parts[1]);
+      size_t idx = skipBitwidthPrefix(parts);
+      assert(parts.size() - idx == 2 &&
+             "expected load function to have 2 parts (memOrder, scope) after "
+             "optional bitwidth prefix");
+      LLVM::AtomicOrdering memOrder = strToMemoryOrder(parts[idx]);
+      auto scopeStr = strToScope(parts[idx + 1]);
       assert(operands.size() == 1 && "expected load to have 1 operand");
 
       replacementOp = buildAtomicLoad(operands[0], memOrder, scopeStr);
@@ -206,10 +216,12 @@ class CallOpConversion : public OpRewritePattern<LLVM::CallOp> {
     else if (auto maybeParts = matchPrefixAndSplitRemainder(
                  calleeName, "__triton_hip_store_")) {
       auto parts = maybeParts.value();
-      assert(parts.size() == 2 &&
-             "expected store function to have 2 parts after prefix");
-      LLVM::AtomicOrdering memOrder = strToMemoryOrder(parts[0]);
-      auto scopeStr = strToScope(parts[1]);
+      size_t idx = skipBitwidthPrefix(parts);
+      assert(parts.size() - idx == 2 &&
+             "expected store function to have 2 parts (memOrder, scope) after "
+             "optional bitwidth prefix");
+      LLVM::AtomicOrdering memOrder = strToMemoryOrder(parts[idx]);
+      auto scopeStr = strToScope(parts[idx + 1]);
       assert(operands.size() == 2 && "expected store to have 2 operands");
       buildAtomicStore(operands[1], operands[0], memOrder, scopeStr);
       rewriter.eraseOp(callOp);
@@ -220,11 +232,13 @@ class CallOpConversion : public OpRewritePattern<LLVM::CallOp> {
     else if (auto maybeParts = matchPrefixAndSplitRemainder(
                  calleeName, "__triton_hip_atom_add_")) {
       auto parts = maybeParts.value();
-      assert(parts.size() == 2 &&
-             "expected atomic add function to have 2 parts after prefix");
+      size_t idx = skipBitwidthPrefix(parts);
+      assert(parts.size() - idx == 2 &&
+             "expected atomic add function to have 2 parts (memOrder, scope) "
+             "after optional bitwidth prefix");
       assert(operands.size() == 2 && "expected atomic add to have 2 operands");
-      LLVM::AtomicOrdering memOrder = strToMemoryOrder(parts[0]);
-      auto scopeStr = strToScope(parts[1]);
+      LLVM::AtomicOrdering memOrder = strToMemoryOrder(parts[idx]);
+      auto scopeStr = strToScope(parts[idx + 1]);
       replacementOp =
           buildAtomicFetchAdd(operands[0], operands[1], memOrder, scopeStr);
     }
@@ -233,12 +247,15 @@ class CallOpConversion : public OpRewritePattern<LLVM::CallOp> {
     else if (auto maybeParts = matchPrefixAndSplitRemainder(
                  calleeName, "__triton_hip_atom_cas_")) {
       auto parts = maybeParts.value();
-      assert(parts.size() == 3 &&
-             "expected atomic cas function to have 3 parts after prefix");
+      size_t idx = skipBitwidthPrefix(parts);
+      assert(parts.size() - idx == 3 &&
+             "expected atomic cas function to have 3 parts "
+             "(successOrder, failureOrder, scope) after optional bitwidth "
+             "prefix");
       assert(operands.size() == 3 && "expected atomic cas to have 3 operands");
-      LLVM::AtomicOrdering successOrdering = strToMemoryOrder(parts[0]);
-      LLVM::AtomicOrdering failureOrdering = strToMemoryOrder(parts[1]);
-      auto scopeStr = strToScope(parts[2]);
+      LLVM::AtomicOrdering successOrdering = strToMemoryOrder(parts[idx]);
+      LLVM::AtomicOrdering failureOrdering = strToMemoryOrder(parts[idx + 1]);
+      auto scopeStr = strToScope(parts[idx + 2]);
       replacementOp = buildAtomicCompareExchangeStrong(
           operands[0], operands[1], operands[2], successOrdering,
           failureOrdering, scopeStr);

lib/Conversion/TritonDistributedToLLVM/AMD/BuiltinFuncToLLVMExt.cpp

@@ -276,7 +276,14 @@ def _get_amdsmi_device_index(device_id: int | None):
     uuid = _get_gpu_uuid(device_id)
 
     uuid_map = {get_uuid_by_physical_device_id(i)[-12:]: i for i in range(get_physical_device_count())}
-    return uuid_map[uuid[-12:]]
+    # TODO-rocm fix error
+    uuid_tail = uuid[-12:]
+    if uuid_tail not in uuid_map:
+        warnings.warn(f"UUID mapping miss in _get_amdsmi_device_index: device_id={device_id}, "
+                      f"uuid_tail={uuid_tail}, available_tails={sorted(uuid_map.keys())}. "
+                      f"Fallback to logical device_id.")
+        return device_id
+    return uuid_map[uuid_tail]
 
 
 def get_physical_device_count():

python/triton_dist/amd_utils.py

@@ -88,9 +88,9 @@ def shmem_kernel_module_init_hook(*args, **kwargs) -> None:
             else:
                 hip.hipGetLastError()  # Discard the last error
         elif backend == 'mori_shmem':
-            # Initialize mori_shmem device symbols in this kernel module
-            import mori.shmem as mori_shmem
-            mori_shmem.shmem_module_init(kernel_module)
+            if "mori_shmem" in kernel.asm.get('llir', ''):
+                import mori.shmem as mori_shmem
+                mori_shmem.shmem_module_init(kernel_module)
     elif is_maca():
         if "mxshmem" in kernel.asm['ttir']:
             import pymxshmem

python/triton_dist/jit.py

@@ -22,9 +22,33 @@
 # SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 #
 ################################################################################
-from .allgather_gemm import ag_gemm_intra_node, create_ag_gemm_intra_node_context
-from .gemm_reduce_scatter import gemm_rs_intra_node, create_gemm_rs_intra_node_context
+from .ep_a2a_intra_node import (
+    kernel_dispatch_token_intra_node,
+    kernel_skipped_token_local_dispatch_intra_node,
+    kernel_skipped_token_inplace_local_combine_intra_node,
+    kernel_combine_token_intra_node,
+    get_ag_splits_and_recv_offset_for_dispatch_intra_node,
+)
+from .low_latency_all_to_all import create_all_to_all_context, fast_all_to_all, all_to_all_post_process
+
+try:
+    from .allgather_gemm import ag_gemm_intra_node, create_ag_gemm_intra_node_context
+    from .gemm_reduce_scatter import gemm_rs_intra_node, create_gemm_rs_intra_node_context
+except ImportError as e:
+    import warnings
+    warnings.warn(f"allgather_gemm/gemm_reduce_scatter unavailable (pyrocshmem not installed): {e}")
 
 __all__ = [
-    "ag_gemm_intra_node", "create_ag_gemm_intra_node_context", "gemm_rs_intra_node", "create_gemm_rs_intra_node_context"
+    "ag_gemm_intra_node",
+    "create_ag_gemm_intra_node_context",
+    "gemm_rs_intra_node",
+    "create_gemm_rs_intra_node_context",
+    "kernel_dispatch_token_intra_node",
+    "kernel_skipped_token_local_dispatch_intra_node",
+    "kernel_skipped_token_inplace_local_combine_intra_node",
+    "kernel_combine_token_intra_node",
+    "get_ag_splits_and_recv_offset_for_dispatch_intra_node",
+    "create_all_to_all_context",
+    "fast_all_to_all",
+    "all_to_all_post_process",
 ]

python/triton_dist/kernels/amd/__init__.py

@@ -31,7 +31,12 @@
 
 from triton_dist.language.extra.hip.language_extra import load, atomic_add, sync_grid, atomic_cas, tid, __syncthreads
 from hip import hip
-from triton_dist.utils import HIP_CHECK, rocshmem_barrier_all_on_stream
+from triton_dist.utils import (
+    HIP_CHECK,
+    get_shmem_backend,
+    mori_shmem_barrier_all_on_stream,
+    rocshmem_barrier_all_on_stream,
+)
 
 
 @triton.jit
@@ -175,9 +180,9 @@ def barrier_all_kernel(rank, num_ranks, comm_buf_ptr):
 
 
 def barrier_all_on_stream(stream: Optional[torch.cuda.Stream] = None):
-    '''
-    call rocshmem barrier api
-    '''
+    """Call shmem barrier on stream: mori_shmem when backend is mori_shmem, else rocshmem."""
+    if get_shmem_backend() == "mori_shmem":
+        return mori_shmem_barrier_all_on_stream(stream)
     return rocshmem_barrier_all_on_stream(stream)
 
 

python/triton_dist/kernels/amd/common_ops.py

@@ -0,0 +1,82 @@
+################################################################################
+#
+# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates
+#
+# Permission is hereby granted, free of charge, to any person obtaining
+# a copy of this software and associated documentation files
+# (the "Software"), to deal in the Software without restriction,
+# including without limitation the rights to use, copy, modify, merge,
+# publish, distribute, sublicense, and/or sell copies of the Software,
+# and to permit persons to whom the Software is furnished to do so,
+# subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be
+# included in all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
+# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
+# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
+# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+#
+################################################################################
+"""
+AMD EP A2A kernels and helpers.
+Provides bincount and re-exports intra-node kernels from ep_a2a_intra_node.
+"""
+
+import torch
+import triton.language as tl
+import triton_dist
+from triton_dist.language.extra.hip.language_extra import tid, ld, atomic_add
+from triton_dist.language.extra.language_extra import threads_per_warp
+
+
+@triton_dist.jit(do_not_specialize=["n", "length", "num_sm"])
+def kernel_bincount(n, input, output, length, num_sm, num_warps: tl.constexpr):
+    """
+    GPU bincount: count occurrences of each index in [0, length). AMD version using tid(0)
+    and fixed threads_per_block (no simt_exec_region). Same semantics as nvidia/ep_a2a.py.
+    """
+    pid = tl.program_id(0)
+    num_pid = tl.num_programs(0)
+    thread_idx = tid(0)
+    threads_per_block = num_warps * threads_per_warp()
+    for i in range(pid * threads_per_block + thread_idx, n, num_pid * threads_per_block):
+        val = ld(input + i)
+        if val < length:
+            atomic_add(output + val, 1, scope="agent", semantic="relaxed")
+
+
+def bincount(input_tensor, length, output=None, output_dtype=torch.int32, num_sm=16, num_warps=8):
+    """GPU bincount for AMD (no AOT). input_tensor: 1D int32 on device; output: length elements."""
+    if output is None:
+        output = torch.zeros(length, dtype=output_dtype, device=input_tensor.device)
+    assert input_tensor.dim() == 1 and input_tensor.is_contiguous()
+    assert output.size(0) >= length and output.dtype == output_dtype
+    n = input_tensor.size(0)
+    grid = (num_sm, )
+    kernel_bincount[grid](n, input_tensor, output, length, num_sm, num_warps=num_warps)
+    return output
+
+
+# Re-export intra-node kernels and helpers so layer can import from this module only.
+from triton_dist.kernels.amd.ep_a2a_intra_node import (
+    kernel_combine_token_intra_node,
+    kernel_dispatch_token_intra_node,
+    get_ag_splits_and_recv_offset_for_dispatch_intra_node,
+    kernel_skipped_token_local_dispatch_intra_node,
+    kernel_skipped_token_inplace_local_combine_intra_node,
+)
+
+__all__ = [
+    "kernel_bincount",
+    "bincount",
+    "kernel_combine_token_intra_node",
+    "kernel_dispatch_token_intra_node",
+    "get_ag_splits_and_recv_offset_for_dispatch_intra_node",
+    "kernel_skipped_token_local_dispatch_intra_node",
+    "kernel_skipped_token_inplace_local_combine_intra_node",
+]