[kernels][fused_moe] Add another fused EP MoE kernels - #3040
[kernels][fused_moe] Add another fused EP MoE kernels#3040rupengliu-meta wants to merge 7 commits into
Conversation
Summary: Bring the fused expert-parallel MoE TPU kernels into the repo as a self-contained, dependency-free package. The main kernel runs the entire EP MoE (gather, GMM1, activation, GMM2) with the ICI reduce-scatter fused into a single Pallas call, replacing the usual post-kernel all-reduce and keeping the collective overlapped with compute. A second, smaller kernel is included as a worked example of fusing the upstream all-gather into the grouped matmul (per-round push schedule) and is not performance-tuned. The code is fully standalone: device-specific tuned block-size tables and all internal references were stripped for external release, the two kernels share a single trimmed grouped-matmul base, and a README documents usage and future optimization directions. Test Plan: TODO - numerical verification requires a multi-device TPU mesh, which was not available. On CPU, verified the package byte-compiles and that both entry points (fused_moe_func_rs and the AG+GMM1 example) import cleanly with no unresolved references. Signed-off-by: rupengliu-meta <rupengliu@meta.com>
5a18fbb to
e71e5ab
Compare
rupengliu-meta
changed the title
|
@rupengliu-meta Thanks a lot for the contribution! Do you have a branch to integrate the kernel with the model like Qwen 3.5 or Qwen coder so we can compare the performance E2E? Thank you! |
|
@helloworld1 @kyuyeunk I will try my best to find the bandwidth in the following two weeks or so to do the integration. Also I already added an example integration file. We have an internal benchmarking dashboard with inhouse model, as mentioned, the current fused moe kernel outperforms the existing moe implementation by 15-50% for moe layer and ~10-25% e2e, the gain is more with larger batch size |
|
I did a hack integration 08b0347 and benchmark the qwen 3.5 397B model on 8K/1K concurrency 256 case. MoE layer latency Decode: |
what is the block sizes you are using, the block sizes need to trigger the expert weight caching, which avoid weights loading during repeated gmm_tile, feel free to share the kernel trace with your benchmark, prefill should see the bigger gain. |
|
ok, from your PR, it seems you didn't do any block size tuning? |
|
@rupengliu-meta Could you give us some insight / heuristic for a good tuning? |
yes, please use below heuristic: same-expert weight caching. It's not a flag; it turns on automatically when the buffers can hold all of one expert's weight tiles at once: num_w1_bufs >= (n1/tile_n1) * (k1/tile_k1) Core heuristic (split by M): Small M (each expert touched ~once): caching is useless. Use full-N tiles, single buffer (…,1,1) to minimize per-step overhead. FP8: at decode, you should see hbm bandwidth ultilization reaches 80%+, mostly bandwidth bound. At prefill, it should be purely compute bound, you shouldn't see any dma wait and a2a collective latency etc The kernel supports up to 16k max batched token count right now |
QiliangCui
added
ready
| dst_ref=w1_scale_buf_ref.at[buf_id], | ||
| sem=w1_sem_ref.at[buf_id], | ||
| ).start() | ||
| if has_bias: |
There was a problem hiding this comment.
The weight DMA at 519–549 and scale DMA at 550–584 branch on act_fn and num_n1 > 1 to load gate half from _gate_offset = n_id*tile_n1 and up half from _up_offset = out_n1 + n_id*tile_n1. The bias DMA here is unconditional and loads a contiguous 2*tile_n1 slab from pl.ds(n_id * w1_dma_n, w1_dma_n). Compute at 673–687 then slices the buffer as [gate | up].
Failure: with act_fn='silu', has_bias=True, and intermediate_size large enough that num_n1 > 1 (e.g. intermediate_size=1024, tile_n1=256, num_n1=4), w1_bias_buf_ref[:, tile_n1:2*tile_n1] contains the next gate slice instead of the up bias (which lives at w1_bias[E, :, out_n1 + n_id*tile_n1 : ...]). GMM1 adds gate biases where up biases belong — silent wrong-answer for any MoE-with-MLP-bias model (GPT-OSS-style).
Note: w2_bias DMA at 615–616 is fine — GMM2 has no fused activation and no gate/up split.
Fix: mirror the weight/scale DMA branching: when act_fn and num_n1 > 1, issue two copies for the gate and up halves into w1_bias_buf_ref.at[:, :tile_n1] and w1_bias_buf_ref.at[:, tile_n1:].
| else: | ||
| out_f32 = out_3d.astype(jnp.float32) | ||
| pnw = pnw_raw | ||
| out_f32 = out_f32 * pnw[None, None, :] |
There was a problem hiding this comment.
Lines 323–325 compute:
out_f32 = out_f32 * pnw
var = mean(out_f32 ** 2)
out_3d = out_f32 * rsqrt(var + 1e-8)
i.e. y = (x*w) * rsqrt(mean((x*w)^2) + eps). Standard RMSNorm (cf. tpu_inference/layers/jax/layers.py:96–98) is var = mean(x^2), normed = x * rsqrt(var + eps), then normed *= scale — normalize FIRST, scale LAST.
Failure: for x=[1,1], w=[0,1] (so pnw = w + 1 = [1,2]):
- reference:
rms(x)=1,normed=[1,1], scaled= [1.0, 2.0] - kernel:
out_f32=[1,2],var=(1+4)/2=2.5,rsqrt≈0.6325, output≈[0.6325, 1.265](>36% divergence)
For uniform w=1 the output is half the reference — the weight is normalized away.
Fix: compute variance from out_f32 before applying pnw, then multiply the normalized result by pnw.
| hidden_in_spec = data_p_spec if sp_enabled else P() | ||
| moe_out_spec = (P(combine_partition_axes(MLP_DATA, expert_axis)) | ||
| if sp_enabled else P()) | ||
| fp8_post_gather = ((not sp_enabled) and w1_scale is not None |
There was a problem hiding this comment.
The reassignment overrides the caller's True if sp_enabled OR w1_scale is None OR w2_scale is None. The flag has double duty: it drives the kernel's fp8_direct_write (which does require quantized weights)
AND _all_gather_token_hidden's activation quant (which does not depend on weight quant). So on SP-off + bf16 the caller opting into comm compression gets silently downgraded to bf16 with no log.
Fix: either split into two flags (kernel_fp8_direct_write vs comm_fp8_activation), or log at INFO when a caller-supplied True is downgraded.
|
no need to merge the code, feel free to optimize it internally, we can share the optimized code offline. much appreciated! (please optimize the fp4 code path, thanks) |
Signed-off-by: rupengliu-meta <rupengliu@meta.com>
rupengliu-meta
added
the
ready
|
I did the micro benchmark and found the kernel in general performs better for higher number of tokens to a certain point. I tried different combo of tiling combos and heuristic without finding a config that can work well across mulitple workloads. I think selectively using this kernel for based on the token size and workload can provide E2E performance gain.
|
helloworld1
left a comment
helloworld1
left a comment
There was a problem hiding this comment.
The kernel is good starting point as mega fused kernel. Though multiple optimization / improvement can be made here. The block size tuning logic can be improved. The util gmm_v2_gather_scatter has many code sharable with fused_moe_gmm. But for now it is good to keep everything encapsulated in the fused_moe directory.
I would change to "fused_moe_v2" to avoid confusion with existing fused_moe kernel.
this area is slower is expected due to the tiling restriction of the moe kernel. And this domain is not meaningful for us at all (with speculative decoding) |
5 participants
Summary:
Upload a new kernel by @rupengliu-meta @bangshengtang @zongweiz @yaochengji . Bring the fused expert-parallel MoE TPU kernels into the repo as a self-contained, dependency-free package. The main kernel runs the entire EP MoE (gather, GMM1, activation, GMM2) with the ICI A2A fused into a single Pallas call, replacing the usual post-kernel all-reduce and keeping the collective overlapped with compute. A second, smaller kernel is included as a worked example of fusing the upstream all-gather into the grouped matmul (per-round push schedule) and is not performance-tuned.
The code is fully standalone: device-specific tuned block-size tables and all internal references were stripped for external release, the two kernels share a single trimmed grouped-matmul base, and a README documents usage and future optimization directions.
Perf gain:
Currently we see 15-55% perf gain for entire MoE layer compared with existing MoE implementation; Around 10%-30% e2e gain on different batch_sizes/prefill_decode mode with this kernel
Description
Start with a short description of what the PR does and how this is a change from
the past.
The rest of the description includes relevant details and context, examples:
If the change fixes a Github issue, please include a link, e.g.,:
FIXES: #123456
Tests
Please describe how you tested this change, and include any instructions and/or
commands to reproduce.
Checklist
Before submitting this PR, please make sure: