Maybe_Jayden e3e9d70728 [Feat][MoE] SharedFusedMoE with native shared expert kernel and TP support (#778) ... ## Summary This PR implements `SharedExpertMLPKernel` (TileLang) and refactors `SharedFusedMoE` to support Tensor Parallelism for the shared expert path. Closes #777. --- ## Motivation The previous `SharedFusedMoE` accepted a user-provided `shared_experts_fn: Callable` (an external `nn.Module`) to compute the shared expert output. This design had two problems: 1. **No kernel ownership**: shared expert computation bypassed TileOPs' kernel layer entirely, losing performance visibility and benchmark coverage. 2. **No TP support**: no mechanism to shard shared expert weights across TP ranks, which is required for multi-GPU inference (e.g. Kimi K2 / DeepSeek-V3 with `tp_size=8`). --- ## Changes ### New files | File | Description | |------|-------------| | `tileops/kernels/moe/shared_expert_mlp.py` | `SharedExpertMLPKernel`: TileLang kernel implementing gate+up GEMM → SiLU-mul → down GEMM | | `tests/ops/test_moe_shared_fused_moe_distributed.py` | Multi-GPU TP tests via `torchrun`, comparing TileOPs vs vLLM `DeepseekV2MLP` | | `benchmarks/ops/bench_moe_shared_fused_moe.py` | Op-level benchmark with vLLM baseline | | `benchmarks/kernels/bench_moe_shared_expert_mlp.py` | Kernel-level benchmark | ### Modified files **`tileops/ops/moe/shared_fused_moe.py`** — interface refactor + TP support - Remove `shared_experts_fn: Callable` parameter - Add `shared_ffn_size: Optional[int]`, `tp_size: int = 1`, `tp_rank: int = 0` - Op accepts **complete** weights and shards them internally: - `shared_w_gate_up [2*F_s, H]` → ColumnParallel, split along `dim=0` - `shared_w_down [H, F_s]` → RowParallel, split along `dim=1` - Returns `(shared_out_partial, routed_out)`; when `tp_size > 1`, `shared_out_partial` is a partial sum — caller does `dist.all_reduce` **`tileops/kernels/moe/__init__.py`** — export `SharedExpertMLPKernel` **`tests/ops/test_moe_shared_fused_moe.py`** — rewritten with 3 tests (see Testing section) **`tests/ops/test_moe_fused_moe_distributed.py`** — fix stale parameter name `shared_experts_fn` → `shared_experts` --- ## Design ### TP split convention (matches vLLM `DeepseekV2MLP`) TileOPs mirrors vLLM's weight layout exactly. The caller passes complete weights; the op slices them with `narrow().contiguous()` per `tp_rank`. ``` gate_up_proj: MergedColumnParallelLinear → each rank holds [2*F_s/tp, H] down_proj: RowParallelLinear → each rank holds [H, F_s/tp] ``` ### Partial output contract ```python # tp_size > 1 usage shared_partial, routed_out = op(hidden, gating, w_gate_up, w_down, shared_w_gate_up=w_gu, # complete [2*F_s, H] shared_w_down=w_d) # complete [H, F_s] dist.all_reduce(shared_partial) # caller's responsibility final = shared_partial + routed_out ``` The routed expert path is **not affected** by TP sharding. --- ## Testing ### Single-GPU (`pytest -m smoke`) | Test | What it checks | |------|----------------| | `test_shared_fused_moe_basic` | `shared_out` vs float32 math reference (`atol=1e-2`); `routed_out` vs `FusedMoe` (`atol=1e-5`) | | `test_shared_fused_moe_none` | `shared_ffn_size=None` → `shared_out is None` | | `test_shared_fused_moe_tp` | Simulate `tp_size=2` on single GPU: accumulate partials, compare to per-shard float32 reference | All 3 passed ✅ ### Multi-GPU (`torchrun`) | Test | Config | Command | |------|--------|---------| | `smoke-tp2` | T=64, H=128, F_s=64, tp_size=2 | `torchrun --nproc_per_node=2 -m pytest ... -m smoke` | | `kimi-k2-tp2` | T=512, H=7168, F_s=18432, tp_size=2 | `torchrun --nproc_per_node=2 -m pytest ... -m full` | | `kimi-k2-tp8` | T=512, H=7168, F_s=18432, tp_size=8 | `torchrun --nproc_per_node=8 -m pytest ... -m full` | Compares TileOPs `SharedFusedMoE` partial + manual `dist.all_reduce` against vLLM `DeepseekV2MLP` (internal all-reduce). All passed ✅ Tests skip cleanly when not launched via `torchrun` (CI single-GPU environments). --- ## Benchmark **Hardware**: NVIDIA H200 | **Config**: Kimi K2 scale — E=384, K=8, H=7168, F=2048, F_s=18432, bf16, sigmoid routing | num_tokens | TileOPs (ms) | TileOPs (TFLOPS) | vLLM (ms) | vLLM (TFLOPS) | Ratio | |-----------|-------------|-----------------|----------|--------------|-------| | 1 | 0.95 | 1.58 | 0.50 | 3.01 | 0.53x | | 32 | 5.46 | 8.77 | 7.50 | 6.39 | **1.37x** | | 512 | 10.44 | 73.45 | 17.01 | 45.08 | **1.63x** | | 2048 | 19.99 | 153.39 | 24.06 | 127.48 | **1.20x** | | 4096 | 32.63 | 187.96 | 30.33 | 202.19 | 0.93x | TileOPs leads at medium batch sizes (T=32~2048). At T=1 vLLM is faster (memory-bound decode, lower cuBLAS overhead); at T=4096 vLLM's Triton kernel edges ahead in the compute-bound regime. --- ## Notes - `layout="padded"` is not supported in the new path (planned for removal). - `shared_ffn_size % tp_size == 0` is validated at construction time. - Routed expert TP (EP via `expert_map`) is a separate concern and unchanged. ## Follow-up - #828 — Add `pre_sharded` fast path to eliminate per-forward ~95MB weight copies in TP mode - #829 — Fix `enable_rasteration` → `enable_rasterization` typo across 40+ kernel configs Suggestions: `TPFixture` should inherit `FixtureBase` from `workloads.base`; consider separate GEMM tuning configs for gate_up vs down in `SharedExpertMLPKernel`. --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>		2026-04-07 17:42:09 +08:00
..
__init__.py	[Feat][MoE] Add MoePermuteAlignOp: token routing and padding for MoE grouped GEMM (#536)	2026-03-17 20:09:23 +08:00
bench_moe_shared_expert_mlp.py	[Feat][MoE] SharedFusedMoE with native shared expert kernel and TP support (#778)	2026-04-07 17:42:09 +08:00