Skip to content

Wave RDNA4 MoE kernel#1090

Open
nirmie wants to merge 83 commits intoiree-org:mainfrom
nirmie:MoE-clean
Open

Wave RDNA4 MoE kernel#1090
nirmie wants to merge 83 commits intoiree-org:mainfrom
nirmie:MoE-clean

Conversation

@nirmie
Copy link
Contributor

@nirmie nirmie commented Mar 9, 2026

Wave Moe Kernel off of @panditsa initial work

  • 32 threads/wavefront (RDNA4)
  • multiple microkernels (working on fusing it, different hardware constraints between block_align and actual gemms cause issues)
  • Can combine gather/scatter kernels with GEMM (will do soon)

@nirmie nirmie force-pushed the MoE-clean branch 3 times, most recently from 96ff51a to fc631d5 Compare March 12, 2026 23:12
panditsa and others added 27 commits March 16, 2026 15:58
@panditsa
sample kernel and test
a6ef302 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
shmem exp
17620b4 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
turn off min. shared allocs - temp
f91cbfb 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
updated moe block align kernel
5f4f805 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
fixed test and kernel
1809e71 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
padded counts calculation
8e21bc7 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
update moe to use cumsum
f96aaf8 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
create inclusive and exclusive cumsum buffers
3d12b26 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
cleanup
d9414fa 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
[HACK] node in custom.start, not in custom.implicit_captures
5eb4380 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
[Working] expert_ids calculated
29742e9 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
working sorted_token_ids
32298fc 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
examples-simple
09eeb5f 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
exp. with scalar write
b7b72d8 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
cleanup
dd1d78d 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
parametrize block_sizes
befb397 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
first draft
c715099 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
reference code, simplified from triton implementation
c9313d4 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
some simple gemm exps
675ed94 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
indexed weight GEMM example
934b377 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
remove unnecessary flags
0b95add 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
standalone gemm.py test that dynamically creates A
f1ffc55 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
reorder A for test
2d1b1ff 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
WIP
a72129e 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
basic scatter working
5695e48 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
simple scatter with gemm working, gemm broken
a4d5112 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
removed one condition
5736f3d 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
panditsa and others added 24 commits March 16, 2026 15:58
@panditsa
use wave topk
3aac56b 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
WIP, large histogram
9d878ba 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
set the index only for specific dimension
9c942e3 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
fused gemm example
226a429 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
silu_and_mul update MILESTONE
673c6a0 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
tensor ops
7920d01 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
examples large histogram
9265bc1 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
examples gemm-gemm
1e1b64e 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
fix type propagation issues
53ec21e 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
block scan changes still required
82f9eac 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
add side-effects to write and atomic ops
84718c6 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
undo hacks on decompose scan ops
2e999b2 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
fix reduce sum to work on larger shapes
ad5f015 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
remove unused examples
532112d 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
fix the get users for iterate start
72ab532 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
delete duplicate writes
1f17e2d 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
fixing spec decode
5486d3f 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
fix nested iterates
5bb3c30 
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@nirmie
MoE working with PyTorch code chunks for gather, scatter, and routing…
a580485 
…/expert IDs

Signed-off-by: Nirmal Senthilkumar <nirmalsent@gmail.com>
@nirmie
added wave based gather/scatter for MoE
7e4c09c 
Signed-off-by: Nirmal Senthilkumar <nirmalsent@gmail.com>
@nirmie
prefix sum calculations for experts works, still need to implement ex…
148ca24 
…pert IDs fill using Wave

Signed-off-by: Nirmal Senthilkumar <nirmalsent@gmail.com>
@nirmie
all parts of MoE are in Wave code but are in microkernels
2414a48 
Signed-off-by: Nirmal Senthilkumar <nirmalsent@gmail.com>
@nirmie
Merged new moe kernels into old moe file, all test cases pass with mi…
403fca9 
…crokernels architecture

Signed-off-by: Nirmal Senthilkumar <nirmalsent@gmail.com>
@nirmie
Added skip for non-RDNA4 architecture for pytest
ca84af8 
Signed-off-by: Nirmal Senthilkumar <nirmalsent@gmail.com>
@nirmie
Clean up fused MoE test and reduce_sum kernel
722eb99 
Signed-off-by: Nirmal Senthilkumar <nirmalsent@gmail.com>
nirmie and others added 2 commits March 23, 2026 20:58
@nirmie @claude
Fuse scatter1+SiLU+gather2 into single block-space kernel
4b210d3 
Since scatter and gather use the same sorted_ids, the token-space
round-trip cancels. SiLU can be applied directly in block/slot space,
eliminating 3 kernel launches and 2 intermediate buffers.

---------

Signed-off-by: Nirmal Senthilkumar <nirmalsent@gmail.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
@nirmie @claude
Add GEMM BLOCK_N tuning and unit tests for MoE kernels
3caaac0 
BLOCK_N=min(n,128) gives ~1.7x speedup on GEMM1 (RDNA4) by increasing
N-tile reuse per wave. Unit tests cover each kernel individually from
smallest to largest, with full pipeline integration tests last.

---------

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Nirmal Senthilkumar <nirmalsent@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@nirmie @panditsa