Moe #1

Merged

nirmie merged 68 commits intoMoE-cleanfrom

Mar 4, 2026

Commits on Jan 23, 2026

sample kernel and test

panditsa
authored and
nirmie
committed
shmem exp

panditsa
authored and
nirmie
committed
turn off min. shared allocs - temp

panditsa
authored and
nirmie
committed
updated moe block align kernel

panditsa
authored and
nirmie
committed
fixed test and kernel

panditsa
authored and
nirmie
committed
padded counts calculation

panditsa
authored and
nirmie
committed
update moe to use cumsum

panditsa
authored and
nirmie
committed
create inclusive and exclusive cumsum buffers

panditsa
authored and
nirmie
committed
cleanup

panditsa
authored and
nirmie
committed
[HACK] node in custom.start, not in custom.implicit_captures

panditsa
authored and
nirmie
committed
[Working] expert_ids calculated

panditsa
authored and
nirmie
committed
working sorted_token_ids

panditsa
authored and
nirmie
committed
examples-simple

panditsa
authored and
nirmie
committed
exp. with scalar write

panditsa
authored and
nirmie
committed
cleanup

panditsa
authored and
nirmie
committed
parametrize block_sizes

panditsa
authored and
nirmie
committed
first draft

panditsa
authored and
nirmie
committed
reference code, simplified from triton implementation

panditsa
authored and
nirmie
committed
some simple gemm exps

panditsa
authored and
nirmie
committed
indexed weight GEMM example

panditsa
authored and
nirmie
committed
remove unnecessary flags

panditsa
authored and
nirmie
committed
standalone gemm.py test that dynamically creates A

panditsa
authored and
nirmie
committed
reorder A for test

panditsa
authored and
nirmie
committed
WIP

panditsa
authored and
nirmie
committed
basic scatter working

panditsa
authored and
nirmie
committed
simple scatter with gemm working, gemm broken

panditsa
authored and
nirmie
committed
removed one condition

panditsa
authored and
nirmie
committed
scatter A gemm working

panditsa
authored and
nirmie
committed
fixed expert based scatter gemm test

panditsa
authored and
nirmie
committed
reorder test based on complexity

panditsa
authored and
nirmie
committed
hackfixme

panditsa
authored and
nirmie
committed
scatter gemm with padding value

panditsa
authored and
nirmie
committed
remove scatter_gemm.mlir file

panditsa
authored and
nirmie
committed
working scatter-gather gemm for one expert

panditsa
authored and
nirmie
committed
hackfixme

panditsa
authored and
nirmie
committed
moe gemm example

panditsa
authored and
nirmie
committed
use moe_gemm in moe.py and test

panditsa
authored and
nirmie
committed
silu_and_mul check

panditsa
authored and
nirmie
committed
add second gemm and comments

panditsa
authored and
nirmie
committed
add a reduce_sum kernel

panditsa
authored and
nirmie
committed
use the sum kernel

panditsa
authored and
nirmie
committed
in reference calculate topk ids

panditsa
authored and
nirmie
committed
silu and mul fixes and test

panditsa
authored and
nirmie
committed
working moe

panditsa
authored and
nirmie
committed
update test.py

panditsa
authored and
nirmie
committed
placeholder has no index, get_custom first

panditsa
authored and
nirmie
committed
update all gemm examples

panditsa
authored and
nirmie
committed
fix block align

panditsa
authored and
nirmie
committed
use wave moe_align_block_size kernel

panditsa
authored and
nirmie
committed
cleanup

panditsa
authored and
nirmie
committed
wip

panditsa
authored and
nirmie
committed
non-fast dims work? reduction

panditsa
authored and
nirmie
committed
passing test

panditsa
authored and
nirmie
committed
update the reduce kernel to add broadcast

panditsa
authored and
nirmie
committed
working moe

panditsa
authored and
nirmie
committed
final cleanup

panditsa
authored and
nirmie
committed
use wave topk

panditsa
authored and
nirmie
committed
WIP, large histogram

panditsa
authored and
nirmie
committed
set the index only for specific dimension

panditsa
authored and
nirmie
committed
fused gemm example

panditsa
authored and
nirmie
committed
silu_and_mul update MILESTONE

panditsa
authored and
nirmie
committed
tensor ops

panditsa
authored and
nirmie
committed
examples large histogram

panditsa
authored and
nirmie
committed
examples gemm-gemm

panditsa
authored and
nirmie
committed
fix type propagation issues

panditsa
authored and
nirmie
committed
block scan changes still required

panditsa
authored and
nirmie
committed

Commits on Mar 3, 2026

MoE working with PyTorch code chunks for gather, scatter, and routing/expert IDs
nirmie
committed

Commits on Mar 4, 2026

Merge branch 'MoE-clean' into MoE
nirmie
authored