[MoE] Generalize masked activation for padded layouts

[mgoin]

Summary

• add one shared count-aware MoE activation entry point for a flat valid prefix ([T, D]) and per-expert valid prefixes ([E, T, D])
• support the gated and non-gated activation set explicitly, with a fail-closed NotImplementedError when a future activation lacks a masked implementation
• route Humming, batched Triton, and batched Marlin through ApplyMoEActivationConfig instead of backend-specific SITU/clamped-SiLU branches
• select Humming GEMM layout automatically: indexed without EP, grouped_contiguous with EP, and grouped_masked for batched experts; environment overrides remain available but are validated choices
• pass DeepEP v2's cudagraph valid-prefix metadata into Humming so activation skips padded receive rows
• validate padded/invalid expert IDs before indexing expert_map in the Humming permute path

Why

DeepEP v1 and v2 expose different padded layouts. Legacy/batched execution has one valid prefix per expert, while DeepEP v2 has one valid prefix in a flat cudagraph receive buffer. Backend-specific activation branches had diverged and made SITU especially expensive in Kimi K3.

The shared op makes the layout contract explicit and specializes the flat versus batched branch at compile time. The final Kimi investigation also found a correctness issue before activation: DeepEP v2 uses -1 for padding, but Humming's permute preprocessing read expert_map[topk_id] before checking the ID. Invalid IDs now become the expert sentinel before any map access. The regression covers both -1 and a positive out-of-range ID.

Model validation

DeepSeek-V4 Flash

Tested deepseek-ai/DeepSeek-V4-Flash-0731 with automatic Humming selection and no VLLM_HUMMING_MOE_GEMM_TYPE override:

• DP=1, no DeepEP, Humming indexed, eager GSM8K 100: 0.96 accuracy, 0 invalid
• DP=2 / EP=2, DeepEP v2, Humming grouped_contiguous, eager GSM8K 100: 0.95 accuracy, 0 invalid

The close DP1/DP2 scores provide a model-level control for the new DeepEP v2 layout. A separate two-GPU correctness test covers the DSV4 topology (256 experts, top-k 6, K=4096, N=2048, MXFP4/BF16) including an idle DP rank.

Kimi K3

Served mgoin/Kimi-K3-pruned75 on four B300 GPUs with:

chg run -g=4 -- vllm serve mgoin/Kimi-K3-pruned75 \
  --tokenizer-mode kimi_k3 --trust-remote-code --dtype bfloat16 \
  --language-model-only \
  --data-parallel-size 4 --enable-expert-parallel \
  --all2all-backend deepep_v2 --moe-backend humming \
  --max-model-len 8192 --max-num-batched-tokens 512 \
  --max-num-seqs 16

• DP=4 / EP=4 selected DeepEPV2All2AllManager and Humming grouped_contiguous automatically
• no Humming GEMM environment override was set
• the default non-eager server captured 7 PIECEWISE + 5 FULL breakable CUDA graphs; this model's V2 runner reports compilation mode NONE, so this is CUDA-graph rather than TorchInductor validation
• a deterministic arithmetic request returned a coherent, correct answer (42)

Before the invalid-ID guard, the same run emitted repeated @ tokens; GSM8K 100 was 0% accurate with 100% invalid responses and every request exhausted 256 tokens. After the guard, chat-mode GSM8K 100 was 15% accurate / 33% invalid with 18,611 output tokens. Inspection showed the remaining invalid responses were mostly coherent but unfinished reasoning from this 75%-pruned checkpoint; a 10-question diagnostic at 1,024 tokens improved to 30% / 20% invalid. I do not treat the latter small sample as a model-quality result, but the disappearance of repeated-token corruption plus the focused native regression isolates the kernel fix.

Kimi required the 512-token batch cap on four B300s. A 2,048-token cap OOMed because DeepEP v2's worst-case receive/permute scratch must cover max_tokens * dp_size * topk rows.

Coverage and intentional exclusions

• flat masked activation: DeepEP v2 cudagraph padding
• batched masked activation: Marlin, batched Triton, and Humming grouped_masked GPU correctness tests
• standard unmasked activation: unchanged shared apply_moe_activation path
• DeepEP low-latency + Humming: intentionally rejected with a clear reason because low-latency dispatch does not support Humming's deferred input quantization contract

The full-model DSV4 and Kimi runs exercise DeepEP v2 grouped_contiguous. There is no fresh full-model run of BatchedExperts -> grouped_masked after the final rebase; that layout is covered at the kernel/backend-test level. Legacy DeepEP v1 was tested earlier during development, while this final validation round focused on v2 because it is the important new layout.

Tests

chg run -g=1 -- .venv/bin/python -m pytest \
  tests/kernels/core/test_activation.py -k masked_moe_activation -q
# 53 passed, 256 deselected

chg run -g=1 -- .venv/bin/python -m pytest \
  tests/kernels/moe/test_moe_permute_unpermute.py::test_moe_permute_ignores_invalid_expert_ids_with_scratch \
  -q -s
# 1 passed

chg run -g=2 -- env EP_DISABLE_GIN=1 .venv/bin/python -m pytest \
  tests/kernels/moe/test_deepep_v2_moe.py::test_deep_ep_v2_humming_dsv4_expert_topology \
  -q -s
# 1 passed

Also completed:

• focused DeepEP v2 Humming correctness/cudagraph suite
• full native extension build: 428/428 targets built and linked
• commit hooks: Ruff check/format, clang-format, typos, mypy, SPDX/config/import checks
• git diff origin/main...HEAD --check

DeepEP v2 model tests used an isolated user-space environment with NCCL 2.30.7 and NVSHMEM 3.3.24 plus EP_DISABLE_GIN=1; no sudo was required. The shared development .venv currently has NCCL 2.29.7, so its automatic DeepEP refresh cannot compile v2 GIN types. This is an environment-version limitation, not a vLLM native-extension failure.

Duplicate-work check

Open-PR searches for masked MoE activation Humming, SITU Humming, and apply_moe_activation batched found only this PR, so this does not duplicate another open change.

AI assistance

AI assistance was used to develop and test this change. The human submitter reviewed the changed lines, understands the implementation and validation, and is responsible for defending the change end to end.

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[mgoin]

/ci run

[github-actions]

✅ Triggered Buildkite CI #82656 for commit b7466fbc981a.

[mgoin]

/ci run

[github-actions]

✅ Triggered Buildkite CI #82709 for commit 8a1ec76bcfb5.

[mgoin]

/ci run

[github-actions]

✅ Triggered Buildkite CI #82773 for commit af1eb3a90874.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @mgoin.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mgoin]

/ci run

[github-actions]

✅ Triggered Buildkite CI #82928 for commit d954c9c36af6.

[mgoin]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83030 for commit 34ecf7caeb68.

[mgoin]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83815 for commit 24e145e7df48.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @mgoin.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mgoin]

/ci run

[github-actions]

✅ Triggered Buildkite CI #85147 for commit e436c6a23be8.