Skip to content

Unoptimized MoE kernel #348

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Draft
panditsa wants to merge 66 commits into
base: main
Choose a base branch
from
Draft

Unoptimized MoE kernel #348

panditsa wants to merge 66 commits into from

Conversation

@panditsa
Copy link
Contributor

@panditsa panditsa commented Oct 8, 2025 •
edited
Loading

TODO:

  • Use tkw.topk based kernel
  • Cleanup

@panditsa panditsa force-pushed the sanket/fused_moe_kernel branch 3 times, most recently from 6a07f53 to c4d549b Compare October 20, 2025 01:13
@panditsa panditsa force-pushed the sanket/fused_moe_kernel branch from 5147aec to 5758b3d Compare November 12, 2025 00:56
panditsa and others added 26 commits November 11, 2025 16:59
@panditsa
sample kernel and test
530c765
@panditsa
shmem exp
6a4c77d
@panditsa
turn off min. shared allocs - temp
52c2fbc
@panditsa
updated moe block align kernel
b11616e
@panditsa
fixed test and kernel
cc5bb59
@panditsa
padded counts calculation
5ad81e7
@panditsa
update moe to use cumsum
5359f01
@panditsa
create inclusive and exclusive cumsum buffers
e25f528
@panditsa
cleanup
4ed4807
@panditsa
[HACK] node in custom.start, not in custom.implicit_captures
6c17f76
@panditsa
[Working] expert_ids calculated
54e1446
@panditsa
working sorted_token_ids
c268558
@panditsa
examples-simple
4192cae
@panditsa
exp. with scalar write
acc9d2c
@panditsa
cleanup
3a9fdb6
@panditsa
parametrize block_sizes
11b44ea
@panditsa
first draft
8fdd435
@panditsa
reference code, simplified from triton implementation
d7e8ab4
@panditsa
some simple gemm exps
a4660fe
@panditsa
indexed weight GEMM example
71ef5a1
@panditsa
remove unnecessary flags
5a47497
@panditsa
standalone gemm.py test that dynamically creates A
bfc80a0
@panditsa
reorder A for test
80d9eb1
@panditsa
WIP
72ff008
@panditsa
basic scatter working
4411a96
@panditsa
simple scatter with gemm working, gemm broken
00ca84a
@panditsa
use moe_gemm in moe.py and test
7af8d8f
@panditsa
silu_and_mul check
2492317
@panditsa
add second gemm and comments
f4e62fb
@panditsa
add a reduce_sum kernel
6ebc8a4
@panditsa
use the sum kernel
775cde5
@panditsa
in reference calculate topk ids
8256626
@panditsa
silu and mul fixes and test
08b6f31
@panditsa
working moe
f4551ad
@panditsa
update test.py
1908c1b
@panditsa
placeholder has no index, get_custom first
eb9babf
@panditsa
update all gemm examples
1ef98ac
@panditsa
fix block align
2461f2c
@panditsa
use wave moe_align_block_size kernel
32751b2
@panditsa
cleanup
7053605
@panditsa
wip
5ce7e7d
@panditsa
non-fast dims work? reduction
8f369fe
@panditsa
passing test
490539f
@panditsa
update the reduce kernel to add broadcast
c9c002d
@panditsa
working moe
15b0492
@panditsa
final cleanup
71b63b9
@panditsa
use wave topk
c1d31e9
@panditsa
WIP, large histogram
aa97c71
@panditsa
set the index only for specific dimension
10389ff
@panditsa
fused gemm example
75d1531
@panditsa
silu_and_mul update MILESTONE
f972cf8
@panditsa
tensor ops
dce2b15
@panditsa
examples large histogram
6cd3681
@panditsa
examples gemm-gemm
5758b3d
@panditsa
fix type propagation issues
3b05185
@panditsa
block scan changes still required
0441089
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.

1 participant

@panditsa