Codestin Search App

jhchouuu · 2025-12-19T00:46:26Z

Description

Add mori-shmem backend support for AMD GPUs

Changes

Integrated mori-shmem library as a submodule
Added backend switching via TRITON_DIST_SHMEM_BACKEND environment variable
Added test suite for mori-shmem functionality

- Add libmori_shmem_device Python bindings for Triton kernels - Add comprehensive test suite for mori-shmem APIs (PyTorch/UniqueID init) - Add build script (build_mori_shmem.sh) for mori and device BC

CLAassistant · 2025-12-19T00:46:32Z

All committers have signed the CLA.

Copilot

Pull request overview

This PR adds support for the mori-shmem backend as an alternative to rocshmem for AMD GPU distributed computing. The implementation introduces a new backend selection mechanism via the TRITON_DIST_SHMEM_BACKEND environment variable and integrates the mori library as a git submodule.

Adds mori-shmem as a selectable backend alongside rocshmem for AMD GPUs
Implements dynamic backend switching through environment variable configuration
Provides comprehensive test coverage for mori-shmem APIs

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 19 comments.

Show a summary per file

File	Description
scripts/build_mori_shmem.sh	Build script for compiling mori library and linking bitcode files for device execution
python/triton_dist/utils.py	Core utilities for backend selection, initialization, and version/hash management
python/triton_dist/test/amd/test_mori_shmem_api.py	Test suite validating mori-shmem basic operations and device-level APIs
python/triton_dist/language/extra/libshmem_device.py	Updated module proxy to route to mori-shmem device library based on backend
python/triton_dist/language/extra/hip/libmori_shmem_device.py	Device-level API bindings for mori-shmem operations (my_pe, n_pes, int_p)
python/triton_dist/jit.py	JIT compilation integration with backend-specific module initialization
.gitmodules	Adds mori library as a git submodule dependency

Comments suppressed due to low confidence (3)

python/triton_dist/test/amd/test_mori_shmem_api.py:28

Import of 'time' is not used.

import time

python/triton_dist/test/amd/test_mori_shmem_api.py:29

Import of 'shutil' is not used.

import shutil

python/triton_dist/test/amd/test_mori_shmem_api.py:31

Import of 'dist' is not used.

import torch.distributed as dist

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated 3 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

wenlei-bao · 2026-01-06T19:54:21Z

+
+
+# Simple helper to wrap mori shmem pointer as torch tensor
+class MoriShmemBuffer:


This part of code and below shouldn't be here, can you please make these consistent with other shmem component? @jhchouuu

Understood, it was just a temporary placement before. Do you prefer to maintain it in triton-dist like rocshmem, or directly in mori_shmem like nvshmem4py?

@jhchouuu right now we can follow rocshmem, and can be refactor like nvshmem4py later.

Now we wrap it into mori_shmem similar to nvshmem4py

wenlei-bao · 2026-01-06T20:16:39Z

+
+
+@core.extern
+def my_pe(_semantic=None):


where are other shmem APIs ? Ideally it should be consistent compared to other shmem. @jhchouuu

In another branch that doesn't involve mr, we need to discuss the API encapsulation method, as this will add a parameter qp_id for selecting the QP during RDMA communication. Maybe I should talk with @XG-zheng offline?

And the CI has encountered a build error due to failure in pulling the submodules of mori. I will fix this issue.

* Feature: add more mori_shmem bitcode wrappers && add mori_shmem bandwith test * Feature: mori_shmem supoort dl op call && small refactor * fix ci && move tensor create to mori library * Refine setup.py

Copilot

Pull request overview

Copilot reviewed 16 out of 16 changed files in this pull request and generated 7 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-01-08T11:10:03Z

 def get_rocshmem_home():
    return os.getenv("ROCSHMEM_HOME",
-                     Path(__file__).parent.parent / "shmem" / "rocshmem_bind" / "rocshmem_build" / "install")
+                     Path(__file__).parent.parent.parent / "shmem" / "rocshmem_bind" / "rocshmem_build" / "install")


The path traverses up three parent directories (parent.parent.parent) which appears incorrect. This changes the rocshmem default path from the original two parents to three, which would break existing rocshmem installations. The original path used Path(__file__).parent.parent, and this should remain unchanged.

Path corrected to use three parents - the shmem directory is at project root, not inside python/.

Copilot

Pull request overview

Copilot reviewed 16 out of 16 changed files in this pull request and generated 4 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot

Pull request overview

Copilot reviewed 16 out of 16 changed files in this pull request and generated 6 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-01-27T05:29:25Z

+[submodule "3rdparty/mori"]
+	path = 3rdparty/mori
+	url = https://github.com/ROCm/mori.git
+	branch = jiahzhou/triton_dis_support


Corrected spelling of 'jiahzhou' to 'jiahzhou'. However, note that this appears to be referencing a development branch which may not be appropriate for production. Consider using a stable release branch instead.

Suggested change

branch = jiahzhou/triton_dis_support

Copilot · 2026-01-27T05:29:25Z

@@ -82,8 +88,8 @@ jobs:
      - name: E2E tests
        run: |
          bash ./scripts/build_e2e_env.sh --download_model


The change from Qwen3-32B to Qwen3-0.6B significantly reduces the model size for testing. While this may speed up CI, ensure this smaller model still provides adequate coverage for the test scenarios. Consider documenting why this change was made (e.g., CI resource constraints) in the PR description or commit message.

Suggested change

bash ./scripts/build_e2e_env.sh --download_model

bash ./scripts/build_e2e_env.sh --download_model

# Use the smaller Qwen3-0.6B model in CI to keep AMD E2E tests within runtime and memory limits.

# The larger Qwen3-32B variants below remain commented out until the CI image supports them (e.g., flash-attention).

Copilot · 2026-01-27T05:29:25Z

 def get_rocshmem_home():
    return os.getenv("ROCSHMEM_HOME",
-                     Path(__file__).parent.parent / "shmem" / "rocshmem_bind" / "rocshmem_build" / "install")
+                     Path(__file__).parent.parent.parent / "shmem" / "rocshmem_bind" / "rocshmem_build" / "install")


The path has been changed from parent.parent to parent.parent.parent, adding an extra level of directory traversal. This breaks the relative path resolution for rocshmem. The original path structure should be maintained to ensure rocshmem can be located correctly.

Suggested change

Path(__file__).parent.parent.parent / "shmem" / "rocshmem_bind" / "rocshmem_build" / "install")

Path(__file__).parent.parent / "shmem" / "rocshmem_bind" / "rocshmem_build" / "install")

Copilot · 2026-01-27T05:29:26Z

+            return get_rocshmem_version()
+        elif backend == 'mori_shmem':
+            return get_mori_version()
+    return "unknown"


When neither CUDA nor HIP is available, or when backend detection fails, returning 'unknown' could mask configuration issues. Consider raising an exception or logging a warning to make debugging easier.

Copilot · 2026-01-27T05:29:26Z

+            return get_rocshmem_hash()
+        elif backend == 'mori_shmem':
+            return get_mori_shmem_hash()
+    return "unknown"


Similar to get_shmem_version, returning 'unknown' silently when neither CUDA nor HIP is detected could hide configuration problems. Consider raising an exception or logging a warning.

Copilot · 2026-01-27T05:29:26Z

+        raise RuntimeError(f"Unknown TRITON_DIST_SHMEM_BACKEND: {shmem_backend}. Must be 'mori_shmem' or 'rocshmem'")
+
+    # Also build if explicitly requested via env var (for backward compatibility)
+    if check_env_flag("TRITON_DISTRIBUTED_BUILD_PYROCSHMEM", "0") and shmem_backend != "rocshmem":


The logic for building rocshmem when TRITON_DISTRIBUTED_BUILD_PYROCSHMEM is set seems redundant since rocshmem is already built based on TRITON_DIST_SHMEM_BACKEND. This could lead to building rocshmem twice in some scenarios. Consider simplifying this logic or clarifying the intended behavior in a comment.

jhchouuu added 7 commits December 17, 2025 22:00

Feature: Add mori-shmem host api support

99ead3f

Rebase: Integrate mori-shmem device library into Triton

a385fd9

- Add libmori_shmem_device Python bindings for Triton kernels - Add comprehensive test suite for mori-shmem APIs (PyTorch/UniqueID init) - Add build script (build_mori_shmem.sh) for mori and device BC

Chore: Update triton submodule URL to fork

b4e5187

Milestone: Add mori shmem ring put test

76e9f59

Refactor: Support mori-shmem in the refactored triton-dist

297b5ed

remove unused log

353722c

update gitmodules

40fb4f0

Copilot AI review requested due to automatic review settings December 19, 2025 00:46

Copilot started reviewing on behalf of jhchouuu December 19, 2025 00:46 View session

Copilot AI reviewed Dec 19, 2025

View reviewed changes

XG-zheng reviewed Dec 19, 2025

View reviewed changes

Comment thread python/triton_dist/test/amd/test_mori_shmem_api.py

XG-zheng reviewed Dec 19, 2025

View reviewed changes

Comment thread scripts/build_mori_shmem.sh

preminstrel self-requested a review December 19, 2025 18:35

jhchouuu added 2 commits December 22, 2025 19:58

Add mori submodule

913a906

Refine the script && workaroud for hipIpcOpenMemHandle on rocm7.1.0

eff4cce

Copilot AI review requested due to automatic review settings December 22, 2025 12:32

Copilot AI reviewed Dec 22, 2025

View reviewed changes

Comment thread .gitmodules

jhchouuu added 2 commits December 23, 2025 17:19

add ci amd_mori_test scripts && fix ci

301a11d

apply copilot's suggestions

5c5d54a

Copilot AI review requested due to automatic review settings December 26, 2025 06:08

Copilot AI reviewed Dec 26, 2025

View reviewed changes

Comment thread scripts/build_mori_shmem.sh Outdated

Comment thread python/triton_dist/utils.py

Comment thread python/triton_dist/test/amd/test_mori_shmem_api.py Outdated

wenlei-bao self-requested a review December 29, 2025 19:41

wenlei-bao reviewed Jan 6, 2026

View reviewed changes

jhchouuu added 2 commits January 8, 2026 18:28

supoort dl op call && add more mori_shmem bitcode wrappers && bw test

68ac15a

* Feature: add more mori_shmem bitcode wrappers && add mori_shmem bandwith test * Feature: mori_shmem supoort dl op call && small refactor * fix ci && move tensor create to mori library * Refine setup.py

delete unused code && explicitly specify rocshmem bc code version

ab634c0

Copilot AI review requested due to automatic review settings January 8, 2026 11:06

Copilot AI reviewed Jan 8, 2026

View reviewed changes

fix ci

0200e2d

enhance amd ci

7647bf4

Copilot AI review requested due to automatic review settings January 8, 2026 12:13

Copilot AI reviewed Jan 8, 2026

View reviewed changes

Comment thread python/triton_dist/test/amd/test_mori_shmem_bw.py Outdated

Comment thread python/triton_dist/test/amd/test_mori_shmem_bw.py Outdated

Comment thread lib/Conversion/TritonDistributedToLLVM/AMD/DistributedOpToLLVM.cpp Outdated

Comment thread python/setup.py

jhchouuu added 3 commits January 8, 2026 21:31

fix ci

89a5ef6

fix ci && apply copilot suggestion

bc5fca4

format using code-format.sh

5cb77fa

Copilot AI review requested due to automatic review settings January 27, 2026 05:25

Copilot AI reviewed Jan 27, 2026

View reviewed changes

wenlei-bao approved these changes Jan 27, 2026

View reviewed changes

wenlei-bao merged commit 1242224 into ByteDance-Seed:main Jan 27, 2026
9 of 11 checks passed



		# Simple helper to wrap mori shmem pointer as torch tensor
		class MoriShmemBuffer:

	Path(__file__).parent.parent.parent / "shmem" / "rocshmem_bind" / "rocshmem_build" / "install")
	Path(__file__).parent.parent / "shmem" / "rocshmem_bind" / "rocshmem_build" / "install")

Conversation

jhchouuu commented Dec 19, 2025

Description

Changes

Uh oh!

CLAassistant commented Dec 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Jan 8, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

CLAassistant commented Dec 19, 2025 •

edited

Loading