xnnpack moe: support activation='silu' - #9930
Open
internetoftim wants to merge 1 commit into
Open
Conversation
Adds SiLU (x * sigmoid(x)) as a third accepted value for the 'moe' custom op's flexbuffer 'activation' option, alongside gelu and gelu_tanh. Additive only: models that do not pass 'silu' are unaffected. Motivation: DeepSeek-V3-family MoE checkpoints (e.g. Moonshot's Moonlight-16B-A3B, the Kimi-VL text tower) use SiLU expert MLPs; today the kernel rejects them at prepare time, forcing the sequential dense fallback which computes all experts. Their non-renormalized sigmoid-scaled top weights already work as-is, since this kernel uses top_weights exactly as passed - activation was the only semantic gap. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Adds
silu(x * sigmoid(x)) as a third accepted value for themoecustom op's flexbufferactivationoption in the XNNPACK delegate kernel, alongsidegeluandgelu_tanh. Four small hunks intflite/delegates/xnnpack/moe_delegate_kernel.cc; purely additive — models not passingsiluare byte-for-byte unaffected.Why
DeepSeek-V3-architecture MoE checkpoints (e.g. Moonshot's Moonlight-16B-A3B and the Kimi-VL text tower) use SiLU expert MLPs. Today the kernel rejects
activation='silu'at prepare time, which forces such models onto the sequential dense-fallback path (all experts computed per token) — roughly an order of magnitude more weight traffic per decode step for a 64-expert/top-6 model.Everything else about DeepSeek-style routing already works with this kernel: its sigmoid-scaled, non-renormalized top weights are used exactly as passed (verified against a litert-torch export of a DeepseekV3ForCausalLM toy model — with a flexbuffer-patched activation, logits matched a GELU-substituted eager reference at 1.7e-5, confirming activation is the only semantic gap on the CPU path).
The SiLU helper follows the file's existing pattern of host-side scalar activation helpers (
Gelu,GeluTanh), including the same TODO trajectory of eventually lowering to an xnn unary op.The ML-Drift/GPU parser (
ml_drift_delegate/delegate/composite/moe_experts_parser.cc) still accepts gelu only; this PR deliberately does not touch it since I can't validate the GPU expert body's activation from outside. Happy to extend if a maintainer confirms the GPU path can take SiLU.Testing
🤖 Generated with Claude Code