iree-org / iree-org/wave

MoE Tracker

Open
#189 1 comment 0 reactions 0 assignees View on GitHub
tracker
Dominant language
Python
Stars
59
Forks
32
PR merge metrics
No merged PRs in 30d

Description

- [ ] Integrate MoE implementation from Wave as a SGLang test in [here](https://github.com/harsh-nod/sglang/blob/wave_moe/test/srt/test_wave_fused_moe.py)
- [ ] Implement Gather-GEMM for 1st MoE GEMM kernel (kernel that can do both gather and scatter in [here](https://github.com/sgl-project/sglang/blob/2d120f8b1873fe84a6861dfe5656539fa1950a8e/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py#L622-L791))
- [ ] Implement GEMM-Scatter for 2nd MoE GEMM Kernel (kernel that can do both gather and scatter in [here](https://github.com/sgl-project/sglang/blob/2d120f8b1873fe84a6861dfe5656539fa1950a8e/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py#L622-L791))
- [ ] Implement top-K/routing kernel on Wave as seen in [here](https://github.com/harsh-nod/sglang/blob/wave_moe/test/srt/test_wave_fused_moe.py#L140-L144)
- [ ] Implement moe align block size on wave as seen in [here](https://github.com/sgl-project/sglang/blob/2d120f8b1873fe84a6861dfe5656539fa1950a8e/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py#L549-L619)
- [ ] Implement moe sum kernel on wave as seen in [here](https://github.com/sgl-project/sglang/blob/2d120f8b1873fe84a6861dfe5656539fa1950a8e/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py#L1545-L1561)

Contributor guide

Open the contributing guide

Research direction

Start with test/srt/test_wave_fused_moe.py, then compare the referenced fused_moe.py sections for the GEMM, routing, block-size alignment, and sum kernels. Work through the checklist in order and use the linked test and implementation locations to define completion for each item.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.