Skip to content

xnnpack moe: support activation='silu' - #9930

Open
internetoftim wants to merge 1 commit into
google-ai-edge:mainfrom
internetoftim:moe-silu-activation
Open

xnnpack moe: support activation='silu'#9930
internetoftim wants to merge 1 commit into
google-ai-edge:mainfrom
internetoftim:moe-silu-activation

Conversation

@internetoftim

Copy link
Copy Markdown

What

Adds silu (x * sigmoid(x)) as a third accepted value for the moe custom op's flexbuffer activation option in the XNNPACK delegate kernel, alongside gelu and gelu_tanh. Four small hunks in tflite/delegates/xnnpack/moe_delegate_kernel.cc; purely additive — models not passing silu are 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

  • Verified against the flexbuffer parse path and activation dispatch; no test target exists for this kernel today.
  • End-to-end: a deepseek_v3 model extension for litert-torch's export_hf path (parity-tested vs transformers at fp32) drives this op; with this patch, a SiLU MoE toy export runs on the stock Python interpreter. That work lives in https://github.com/internetoftim/litert-torch/tree/kimi-moonlight-port and is being prepared for a separate litert-torch PR.

🤖 Generated with Claude Code

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>
@google-cla

google-cla Bot commented Sep 12, 2026

Copy link
Copy Markdown

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant