Skip to content

[MoE][DeepEPV2] Add PaddedStandard activation format - #54109

Open
bnellnm wants to merge 60 commits into
vllm-project:mainfrom
neuralmagic:padded-standard
Open

bnellnm wants to merge 60 commits into
vllm-project:mainfrom
neuralmagic:padded-standard

Conversation

@bnellnm

@bnellnm bnellnm commented Aug 27, 2026 •

Copy link
Copy Markdown
Collaborator

Motivation

DeepEPV2PrepareAndFinalize hands the experts an activation tensor whose row count is an upper bound rather than the exact token count, so some rows are inactive padding. The shape is still 2D (num_tokens, hidden_dim) — the padding is in the token dimension; nothing is padded along E. It arises because the number of tokens a rank receives after dispatch is data-dependent and not knowable on the host without a device sync, from two sources:

  • decode / cudagraph mode — the recv buffer is sized to the worst case (R * num_max_tokens_per_rank); only the [0, num_recv_tokens) prefix is written and the tail is left uninitialized.
  • prefill mode — tokens are grouped per-expert and each expert's block is padded up to expert_alignment; those gaps are zeroed.

Inactive rows are signalled by setting all of their topk_ids slots to -1, and ExpertTokensMetadata.psum_recv_per_rank carries the real/padding boundary on device (reading it host-side would break CUDA-graph capture).

Previously this was hidden behind the generic Standard activation format, so experts had no way to know their input might contain such rows. Kernels that can skip padding couldn't opt in, and kernels that can't handle it had no way to reject the layout. That is a correctness issue, not just a missed optimization: the decode-mode tail is uninitialized, so an expert that builds routing by scanning every row can read garbage that aliases a valid expert id and corrupt real tokens.

This PR introduces a distinct PaddedStandard activation format so that the padding contract is communicated explicitly from the prepare/finalize stage to the experts, and only experts that actually support padded input are selected.

Design

PaddedStandard is modeled as a superset of Standard: a padding-aware kernel treats unpadded input as the zero-padding case, so it accepts both formats, while a Standard-only kernel accepts only Standard.
BatchedExperts remains a disjoint format.

Declaring PaddedStandard is a contract: the kernel must either skip inactive rows or produce results for them that are safe to discard, and must honor the -1 markers rather than trusting row contents.

Rather than have each experts kernel enumerate a list of accepted formats, every kernel declares a single activation_format() and the superset relation is expressed by one directional helper on the enum:

class FusedMoEActivationFormat(Enum):                                                                                                  
    Standard = ("standard",)                                                                                                           
    BatchedExperts = ("batched_experts",)                                                                                              
    PaddedStandard = ("padded_standard",)                                                                                              
                                                                                                                                       
    def is_superset(self, pf_format: "FusedMoEActivationFormat") -> bool:                                                              
        """Whether this (experts) format is a superset of `pf_format`, i.e. an                                                         
        experts kernel declaring `self` can consume activations produced by a                                                          
        prepare/finalize step in `pf_format`."""                                                                                       
        if self == pf_format:                                                                                                          
            return True                                                                                                                
        return (                                                                                                                       
            self == FusedMoEActivationFormat.PaddedStandard                                                                            
            and pf_format == FusedMoEActivationFormat.Standard                                                                         
        )                                                                                                                              

The relation is directional and used as expert_format.is_superset(pf_format) everywhere: the experts format (self) must be a superset of the prepare/finalize format (the argument).

What changed

New activation format

  • Added FusedMoEActivationFormat.PaddedStandard and the is_superset() relation (modular_kernel.py).
  • DeepEPV2PrepareAndFinalize.activation_format now reports PaddedStandard instead of Standard (prepare_finalize/deepep_v2.py).

Experts declare a single activation format

  • FusedMoEExperts.activation_format() -> FusedMoEActivationFormat (kept single-valued; no list). Experts that support padded input declare PaddedStandard:
    • DeepGemmFP4Experts (Blackwell only)
    • HummingIndexedExperts
    • TrtLlmMxfp4Mxfp8ExpertsModular (mxfp4+mxfp8)
    • TrtLlmBf16ExpertsModular and TrtLlmFp8ExpertsModular — the TRTLLM-Gen
      kernels skip padding rows at tile granularity (only the modular variants;
      the monolithic classes stay Standard).
  • All other experts declare their single existing format (Standard or BatchedExperts).

TRTLLM mxfp4 expert split

  • TrtLlmMxfp4ExpertsModular was split into TrtLlmMxfp4ExpertsModularBase + two subclasses, because only the mxfp4+mxfp8 variant supports PaddedStandard. _supports_quant_scheme was moved down into the subclasses/monolithic class accordingly.

Format-compatibility checks go through is_superset

  • FusedMoEExperts.is_supported_config: cls.activation_format().is_superset(activation_format).
  • FusedMoEKernel._post_init_setup: asserts fused_experts.activation_format().is_superset(prepare_finalize.activation_format).
  • Test tooling modular_kernel_tools/common.py validity check uses the same relation instead of format equality.

Batched detection uses an is_batched property

  • BatchedExperts is disjoint, so instead of raw == BatchedExperts comparisons
    there is a FusedMoEActivationFormat.is_batched property that encapsulates the
    single legitimate identity check. All batched-detection sites (__init__
    validation, HummingExpertsBase.is_batched, cutlass use_batched_format,
    FusedMoEConfig.use_batched_activation_format, the oracle backend-remap
    branches, and mk_objects.py) use <format>.is_batched.

Centralized format selection

  • Added FusedMoEConfig.activation_format property that resolves the all2all format once: BatchedExperts for DeepEP-LL / NIXL-EP / batched_triton, PaddedStandard for DeepEP V2, else Standard.
  • Moved use_batched_activation_format from FusedMoEParallelConfig to FusedMoEConfig (now defined in terms of activation_format).
  • All per-dtype oracle selectors (fp8, int8, int_wna16, mxfp4, mxfp8, nvfp4, unquantized, w4a8, w4a8_int8, humming_utils) now use config.activation_format instead of each re-deriving the batched/standard branch.
  • Fixed a latent bug in oracle/fp8.py: the explicit-DeepGEMM branch tested == Standard (which would misroute a PaddedStandard config to BATCHED_DEEPGEMM); it now tests == BatchedExperts.

CompressedTensorsW4A4Mxfp4MoEMethod can reach DeepGemmFP4Experts

The method picked its FusedMoEExperts class from a hardcoded if-chain
(CutlassExpertsMxfp4._supports_current_device() → is_xpu() → Marlin) that never
consulted activation_format, so on SM100 it always chose CutlassExpertsMxfp4
(Standard) and DeepEPV2PrepareAndFinalize tripped the assert in
modular_kernel.py. --moe-backend=deep_gemm now joins b12x in deferring to
select_mxfp4_moe_backend, which is format-aware:

if moe.moe_backend == "b12x" or moe.moe_backend == "deep_gemm":                                                                        
    self.mxfp4_backend, experts_cls = select_mxfp4_moe_backend(moe)                                                                    

process_weights_after_loading gains the matching DEEPGEMM_MXFP4 branch.
DeepGemmFP4Experts has no process_weights_after_loading of its own, so without it
the DeepGEMM scale packing never ran and the weights fell through the else
catch-all into prepare_moe_fp4_layer_for_marlin.

Selection is unchanged for every other --moe-backend value, including auto: the
CUTLASS / XPU / Marlin chain is untouched. This is deliberately the narrow fix —
routing the remaining mxfp4 call sites (inc_mxfp4_moe.py, quark_moe.py) through
the oracle, and teaching the oracle about CutlassExpertsMxfp4, is a follow-up PR.

Testing

New tests

  • Padding test for deep_gemm in test_deepgemm.py
  • test_trtllm_bf16_moe_padding_robust — padding-robustness: poison the padding
    (NaN in the workspace arena / in the inactive rows of a hand-built PaddedStandard
    buffer whose topk_ids are -1) and assert the valid outputs stay finite and are
    bit-identical to a clean-padding run.
    These prove padding is excluded from results, not physical compute-skipping:
    MoE rows are independent in the token dimension and TRTLLM activation quant is
    per-row, so skipping and compute-then-discard give identical valid outputs. The
    compute-savings of tile/block-granularity skipping is a benchmark concern, not a
    pytest assertion.
  • Updated test_deepep_v2_moe.py to only use experts that support the new format.

Results

Run on 8× B200 (SM100), FlashInfer 0.6.18, CUDA 13.0:

tests/kernels/moe/test_deepgemm.py                         20 passed                                                                   
tests/kernels/moe/test_trtllm_bf16_moe.py                   5 passed                                                                   
tests/kernels/moe/test_deepep_v2_moe.py (full matrix)      18 passed                                                                   
tests/kernels/moe/test_unquantized_backend_selection.py    25 passed                                                                   
tests/kernels/moe/test_moe_kernel_oracle.py                 1 passed                                                                   
tests/quantization/test_auto_round.py         97 passed, 5 failed                                                                      

The 5 test_auto_round.py failures are all test_auto_round_model end-to-end engine
starts and fail identically on an unmodified tree: this checkout's compiled _C
extension is stale relative to the Python wrappers (_C::gptq_marlin_repack()
signature mismatch; Failed to initialize MXFP4 GEMM: status=7 /
Failed to initialize the TMA descriptor 700 from cutlass_mxfp4_group_mm). Not
related to this change.

Tolerances in test_deepep_v2_moe.py are now scaled to the magnitude of the
reference output (_scaled_tolerances) rather than fixed. The output std grows
with the accumulation length, so a fixed atol was both too tight at the largest
shape — (222, 1024, 2048) fp8 failed with 25/454656 elements just over 6e-2,
the tail of a smooth error distribution rather than any real defect — and far too
loose at the smallest. Measured max_err/std is stable per dtype across every
MNK (bf16 ≤ 0.06, fp8 ≤ 0.33), so the fraction is set per dtype with ~2x headroom.
This tightens the small shapes substantially (bf16 (1, 256, 256): atol 1e-1 →
~2e-3) while removing the large-shape false failure.

test_deepep_v2_moe.py now pairs DeepEP V2 only with padding-aware TRTLLM experts
(flashinfer_trtllm for bf16, trtllm_fp8 for fp8); the flashinfer_cutlass
(Standard) path was dropped since it cannot consume the padded layout under the
is_superset gate.

Local-run caveat: on a single-node NVLink box DeepEP v2 needs EP_DISABLE_GIN=1,
because it otherwise requires NCCL GIN/RDMA for num_ranks > 1 and
ncclDevCommCreate fails. Multi-node genuinely needs GIN — do not set this
globally in CI.

Model evals

gsm8k, 1319 questions, 5-shot, DP=4 + --enable-expert-parallel --all2all-backend deepep_v2, no --enforce-eager (eager mode never populates
psum_recv_per_rank, so the padded decode layout under test would not appear).
tests/evals/gsm8k/configs/moe-refactor-dp-ep/config-b200-deepep-v2.txt:

model experts measured threshold
Qwen/Qwen3-30B-A3B TrtLlmBf16ExpertsModular 0.8893 0.88
Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8 TrtLlmFp8ExpertsModular 0.8954 0.88
nm-testing/Qwen3-30B-A3B-MXFP4A16 DeepGemmFP4Experts (psum) 0.8196 0.79

The mxfp4 row sits below the 0.88 W4A16 baseline because --moe-backend=deep_gemm
forces kFp8Dynamic128Sym activations, not because of padding: the same backend
without EP or DeepEP scores 0.8059, i.e. slightly lower than the padded EP run.

Backend selection was confirmed on the mxfp4 run rather than assumed — the server log
shows Using 'DEEPGEMM_MXFP4' Mxfp4 MoE backend / Using DeepGemmFP4Experts paired
with DeepEPV2PrepareAndFinalize. Before the deep_gemm routing fix this exact
configuration hit the activation-format assert in modular_kernel.py.

Not covered: tests/evals/gpt_oss/configs/gpt-oss-20b-sm100-fi-mxfp4-mxfp8-trtllm-deepep-v2.yaml.
gpt-oss's hidden size 2880 is rounded up to 3072 by
DeepEPV2PrepareAndFinalize.maybe_roundup_layer_hidden_size (2880 × bf16 = 5760 B →
6144 B transfer atom), but modular_kernel.py:1526 allocates the finalize output as
torch.empty_like(hidden_states) at the unpadded 2880, so output.copy_(combined_x)
gets a 2880↔3072 shape conflict. That allocation does not honor the padded-output
contract the model side already assumes (gpt_oss.py:415 slices
self.experts(...)[:, : self.hidden_size]). Separately,
tests/evals/gpt_oss/test_gpqa_correctness.py:128 hardcodes --enforce-eager, which
leaves psum_recv_per_rank unpopulated, so this config could not exercise decode
padding even if it started.

Notes

  • AI assistance was used in preparing this change (tests + comments).
  • Not a duplicate: introduces a new activation-format category rather than modifying existing dispatch behavior; no open PR adds a padded standard MoE format.

Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
@mergify mergify Bot added the quantization label Aug 27, 2026
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
@mergify mergify Bot added gpt-oss Related to GPT-OSS models nvidia rocm Related to AMD ROCm labels Aug 27, 2026
@mergify mergify Bot added the intel-gpu Related to Intel GPU label Aug 27, 2026
@github-project-automation github-project-automation Bot moved this to Todo in AMD Aug 27, 2026
@mergify mergify Bot added the cpu Related to CPU backends label Aug 27, 2026
@mergify

mergify Bot commented Sep 2, 2026

Copy link
Copy Markdown
Contributor

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

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

@mergify mergify Bot added the needs-rebase label Sep 2, 2026
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
@coderabbitai

coderabbitai Bot commented Sep 2, 2026 •

Copy link
Copy Markdown
Contributor

Review Change Stack

Walkthrough

The change adds PaddedStandard activation-format support, centralizes format selection in FusedMoEConfig, updates compatibility checks to use format supersets, and replaces direct BatchedExperts comparisons across MoE backends and expert implementations.

Changes

MoE activation format handling

Layer / File(s) Summary
Activation format contract and configuration
vllm/model_executor/layers/fused_moe/modular_kernel.py, vllm/model_executor/layers/fused_moe/config.py, vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py, tests/kernels/moe/modular_kernel_tools/common.py
FusedMoEActivationFormat now includes PaddedStandard, is_batched, and is_superset. FusedMoEConfig.activation_format selects the format for each backend. Kernel and test validation use superset compatibility.
Backend and kernel format propagation
vllm/model_executor/layers/fused_moe/oracle/*, vllm/model_executor/layers/fused_moe/eep_reconfigure.py, vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py, vllm/model_executor/layers/quantization/utils/humming_utils.py
Backend selection and kernel construction read activation_format directly. Batched execution checks use activation_format.is_batched.
Expert format implementations and registrations
vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py, vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py, vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py, tests/kernels/moe/modular_kernel_tools/mk_objects.py
Selected expert implementations now declare PaddedStandard. TensorRT-LLM modular experts split quantization support by concrete subclass. DeepEP V2 test registration uses the new format.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant FusedMoEConfig
  participant BackendSelector
  participant PrepareFinalize
  participant FusedMoEExperts
  participant FusedMoEKernel
  FusedMoEConfig->>BackendSelector: provide activation_format
  BackendSelector->>PrepareFinalize: select prepare/finalize implementation
  PrepareFinalize->>FusedMoEKernel: provide activation format
  FusedMoEExperts->>FusedMoEKernel: provide expert format
  FusedMoEKernel->>FusedMoEExperts: validate is_superset compatibility
Loading

Merge Risk: 🟠 High · up to d784c

Several DeepEP V2, EEP, LoRA, and MXFP8 configurations can fail during backend selection or kernel initialization. These compatibility regressions should be fixed before merge.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 26.47% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 68 functions across 20 files. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Title check ✅ Passed The title clearly and concisely identifies the main change: adding the PaddedStandard activation format for the DeepEP V2 MoE path.
Description check ✅ Passed The description directly explains the PaddedStandard format, its compatibility rules, affected experts and backends, implementation changes, and validation results.

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@mergify mergify Bot removed the needs-rebase label Sep 2, 2026
@bnellnm
bnellnm marked this pull request as ready for review September 3, 2026 00:23
@bnellnm
bnellnm requested a review from a team October 30, 2026 15:26
@mergify mergify Bot added the ci/build label Oct 30, 2026
@bnellnm

bnellnm commented Nov 1, 2026

Copy link
Copy Markdown
Collaborator Author

/ci run

@github-actions

github-actions Bot commented Nov 1, 2026

Copy link
Copy Markdown

❌ This PR is 76 commits behind upstream main. The bot could not merge upstream main into this branch (API returned 403: user doesn't have permission to update head repository). No new CI build was started. Either merge or rebase onto the latest main yourself and rerun /ci run, or use /ci run --allow-stale to test this branch as is.

@bnellnm

bnellnm commented Nov 2, 2026

Copy link
Copy Markdown
Collaborator Author

/ci run

@github-actions

github-actions Bot commented Nov 2, 2026

Copy link
Copy Markdown

✅ Triggered Buildkite CI #99279 for commit 08be661e5116.

Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
@bnellnm

bnellnm commented Nov 3, 2026

Copy link
Copy Markdown
Collaborator Author

/ci run

@github-actions

github-actions Bot commented Nov 3, 2026

Copy link
Copy Markdown

✅ Triggered Buildkite CI #99405 for commit 05921ca9b148.

Signed-off-by: Bill Nell <bnell@redhat.com>
@bnellnm

bnellnm commented Nov 3, 2026

Copy link
Copy Markdown
Collaborator Author

/ci retry

@github-actions

github-actions Bot commented Nov 3, 2026

Copy link
Copy Markdown

✅ Queued 1 failed job(s) for retry in Buildkite CI #99405.

@bnellnm

bnellnm commented Nov 3, 2026

Copy link
Copy Markdown
Collaborator Author

/ci run

@github-actions

github-actions Bot commented Nov 3, 2026

Copy link
Copy Markdown

❌ This PR is 1 commit behind upstream main. The bot could not merge upstream main into this branch (API returned 403: user doesn't have permission to update head repository). No new CI build was started. Either merge or rebase onto the latest main yourself and rerun /ci run, or use /ci run --allow-stale to test this branch as is.

@bnellnm

bnellnm commented Nov 3, 2026

Copy link
Copy Markdown
Collaborator Author

/ci run

@github-actions

github-actions Bot commented Nov 3, 2026

Copy link
Copy Markdown

✅ Triggered Buildkite CI #99617 for commit a2b27b64efcc.

@robertgshaw2-redhat
robertgshaw2-redhat enabled auto-merge (squash) November 4, 2026 17:47
@robertgshaw2-redhat

Copy link
Copy Markdown
Collaborator

/ci retry

@github-actions

github-actions Bot commented Nov 4, 2026

Copy link
Copy Markdown

✅ The previous CI build is still running: https://buildkite.com/vllm/ci/builds/99617

@mgoin

mgoin commented Nov 4, 2026

Copy link
Copy Markdown
Member

/ci run

@github-actions

github-actions Bot commented Nov 4, 2026

Copy link
Copy Markdown

❌ This PR is 19 commits behind upstream main. The bot could not merge upstream main into this branch (API returned 403: user doesn't have permission to update head repository). No new CI build was started. Either merge or rebase onto the latest main yourself and rerun /ci run, or use /ci run --allow-stale to test this branch as is.

This branch has not been deployed

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

Labels

ci/build cpu Related to CPU backends deepseek Related to DeepSeek models DSv4.1 Related to DeepSeek-V4.1 models gpt-oss Related to GPT-OSS models intel-gpu Related to Intel GPU nixl NIXL KV connector / multi-tiered KV offloading (KVCR) nvidia quantization ready ONLY add when PR is ready to merge/full CI is needed rocm Related to AMD ROCm

Projects

Status: Todo
Status: Ready
Status: Ready

Development

Successfully merging this pull request may close these issues.

5 participants