Conversation
github-actions
bot
added
in-progress
Collaborator
|
Can we move the dispatch and combine code into |
Member
Author
Not yet, we have to see how these patterns are used properly. For now, everything will be under example. Maybe another PR that extracts useful stuff and wraps it under ops and x. |
Contributor
There was a problem hiding this comment.
Pull request overview
Introduces an initial Iris-based expert-sharded MoE forward implementation (ported/simplified from Triton distributed MoE tests) along with a correctness test and a benchmark driver.
Changes:
- Adds a new distributed expert-sharded MoE example pipeline (routing/top-k, DP↔EP data movement via Iris, grouped matmul, reduce, and an optional fused matmul+combine kernel).
- Adds an example runner and a benchmark sweep script for MoE configurations.
- Adds a distributed pytest that compares the sharded pipeline against the single-device reference.
Reviewed changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated 11 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/examples/test_expert_sharded_moe.py | New distributed correctness test for the expert-sharded MoE pipeline. |
| examples/31_expert_sharded_moe/topk.py | Top-k routing + metadata construction for dispatch/combine indexing. |
| examples/31_expert_sharded_moe/reduce.py | Triton reduce kernel used for the final combine step (masked sum). |
| examples/31_expert_sharded_moe/ragged_metadata.py | Ragged metadata helpers for grouped expert computation. |
| examples/31_expert_sharded_moe/moe.py | Reference and Iris-based distributed MoE pipeline + Iris all-gather helper. |
| examples/31_expert_sharded_moe/grouped_matmul.py | Triton grouped/ragged GEMM for per-expert matmul. |
| examples/31_expert_sharded_moe/fused_exp_matmul_ep_to_dp.py | Fused expert matmul + EP→DP combine via Iris stores. |
| examples/31_expert_sharded_moe/expert_assignment.py | Expert-to-rank assignment/bitmask metadata. |
| examples/31_expert_sharded_moe/example_run.py | Standalone runner to validate sharded vs reference. |
| examples/31_expert_sharded_moe/dispatch.py | DP→EP dispatch implementation using Iris symmetric heap stores. |
| examples/31_expert_sharded_moe/combine.py | EP→DP combine implementation using Iris symmetric heap stores. |
| benchmark/examples/benchmark_moe.py | Benchmark/validation harness for the distributed MoE example. |
mawad-amd
approved these changes
Feb 26, 2026
Collaborator
mawad-amd
left a comment
mawad-amd
left a comment
There was a problem hiding this comment.
Do we need do add a license here or something? Please merge whenever you are ready -- CI is down at the moment and it will take a bit to bring back up.
Member
Author
I need to fix the fused implementation. Will merge after. |
Member
Author
|
I will run all tests locally before i merged, nw! @mawad-amd |
This was
linked to
issues
Feb 26, 2026
Member
Author
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Expanding the MoE (distributed) work done in original triton_kernels (credit to them: https://github.com/triton-lang/triton/blob/main/python/triton_kernels/tests/test_distributed.py) to use Iris' symmetric memory abstraction and working on exposing fine-grained overlap.
Initial Benchmarking
Submission Checklist