Repository navigation
Conversation
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
WalkthroughThe change adds ChangesMoE activation format handling
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
Merge Risk: 🟠 High · up to 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)
✅ Passed checks (4 passed)
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. Comment |
Signed-off-by: Bill Nell <bnell@redhat.com>
Signed-off-by: Bill Nell <bnell@redhat.com>
|
/ci run |
|
❌ This PR is 76 commits behind upstream |
|
/ci run |
|
✅ Triggered Buildkite CI #99279 for commit |
|
/ci run |
|
✅ Triggered Buildkite CI #99405 for commit |
|
/ci retry |
|
✅ Queued 1 failed job(s) for retry in Buildkite CI #99405. |
|
/ci run |
|
❌ This PR is 1 commit behind upstream |
|
/ci run |
|
✅ Triggered Buildkite CI #99617 for commit |
|
/ci retry |
|
✅ The previous CI build is still running: https://buildkite.com/vllm/ci/builds/99617 |
|
/ci run |
|
❌ This PR is 19 commits behind upstream |
Motivation
DeepEPV2PrepareAndFinalizehands 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 alongE. 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:R * num_max_tokens_per_rank); only the[0, num_recv_tokens)prefix is written and the tail is left uninitialized.expert_alignment; those gaps are zeroed.Inactive rows are signalled by setting all of their
topk_idsslots to-1, andExpertTokensMetadata.psum_recv_per_rankcarries the real/padding boundary on device (reading it host-side would break CUDA-graph capture).Previously this was hidden behind the generic
Standardactivation 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
PaddedStandardactivation 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
PaddedStandardis modeled as a superset ofStandard: a padding-aware kernel treats unpadded input as the zero-padding case, so it accepts both formats, while aStandard-only kernel accepts onlyStandard.BatchedExpertsremains a disjoint format.Declaring
PaddedStandardis a contract: the kernel must either skip inactive rows or produce results for them that are safe to discard, and must honor the-1markers 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: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
FusedMoEActivationFormat.PaddedStandardand theis_superset()relation (modular_kernel.py).DeepEPV2PrepareAndFinalize.activation_formatnow reportsPaddedStandardinstead ofStandard(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 declarePaddedStandard:DeepGemmFP4Experts(Blackwell only)HummingIndexedExpertsTrtLlmMxfp4Mxfp8ExpertsModular(mxfp4+mxfp8)TrtLlmBf16ExpertsModularandTrtLlmFp8ExpertsModular— the TRTLLM-Genkernels skip padding rows at tile granularity (only the modular variants;
the monolithic classes stay
Standard).StandardorBatchedExperts).TRTLLM mxfp4 expert split
TrtLlmMxfp4ExpertsModularwas split intoTrtLlmMxfp4ExpertsModularBase+ two subclasses, because only the mxfp4+mxfp8 variant supportsPaddedStandard._supports_quant_schemewas moved down into the subclasses/monolithic class accordingly.Format-compatibility checks go through
is_supersetFusedMoEExperts.is_supported_config:cls.activation_format().is_superset(activation_format).FusedMoEKernel._post_init_setup: assertsfused_experts.activation_format().is_superset(prepare_finalize.activation_format).modular_kernel_tools/common.pyvalidity check uses the same relation instead of format equality.Batched detection uses an
is_batchedpropertyBatchedExpertsis disjoint, so instead of raw== BatchedExpertscomparisonsthere is a
FusedMoEActivationFormat.is_batchedproperty that encapsulates thesingle legitimate identity check. All batched-detection sites (
__init__validation,
HummingExpertsBase.is_batched, cutlassuse_batched_format,FusedMoEConfig.use_batched_activation_format, the oracle backend-remapbranches, and
mk_objects.py) use<format>.is_batched.Centralized format selection
FusedMoEConfig.activation_formatproperty that resolves the all2all format once:BatchedExpertsfor DeepEP-LL / NIXL-EP /batched_triton,PaddedStandardfor DeepEP V2, elseStandard.use_batched_activation_formatfromFusedMoEParallelConfigtoFusedMoEConfig(now defined in terms ofactivation_format).fp8,int8,int_wna16,mxfp4,mxfp8,nvfp4,unquantized,w4a8,w4a8_int8,humming_utils) now useconfig.activation_formatinstead of each re-deriving the batched/standard branch.oracle/fp8.py: the explicit-DeepGEMM branch tested== Standard(which would misroute aPaddedStandardconfig toBATCHED_DEEPGEMM); it now tests== BatchedExperts.CompressedTensorsW4A4Mxfp4MoEMethodcan reachDeepGemmFP4ExpertsThe method picked its
FusedMoEExpertsclass from a hardcoded if-chain(
CutlassExpertsMxfp4._supports_current_device()→is_xpu()→ Marlin) that neverconsulted
activation_format, so on SM100 it always choseCutlassExpertsMxfp4(
Standard) andDeepEPV2PrepareAndFinalizetripped the assert inmodular_kernel.py.--moe-backend=deep_gemmnow joinsb12xin deferring toselect_mxfp4_moe_backend, which is format-aware:process_weights_after_loadinggains the matchingDEEPGEMM_MXFP4branch.DeepGemmFP4Expertshas noprocess_weights_after_loadingof its own, so without itthe DeepGEMM scale packing never ran and the weights fell through the
elsecatch-all into
prepare_moe_fp4_layer_for_marlin.Selection is unchanged for every other
--moe-backendvalue, includingauto: theCUTLASS / XPU / Marlin chain is untouched. This is deliberately the narrow fix —
routing the remaining mxfp4 call sites (
inc_mxfp4_moe.py,quark_moe.py) throughthe oracle, and teaching the oracle about
CutlassExpertsMxfp4, is a follow-up PR.Testing
New tests
test_deepgemm.pytest_trtllm_bf16_moe_padding_robust— padding-robustness: poison the padding(NaN in the workspace arena / in the inactive rows of a hand-built
PaddedStandardbuffer whose
topk_idsare-1) and assert the valid outputs stay finite and arebit-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.
test_deepep_v2_moe.pyto only use experts that support the new format.Results
Run on 8× B200 (SM100), FlashInfer 0.6.18, CUDA 13.0:
The 5
test_auto_round.pyfailures are alltest_auto_round_modelend-to-end enginestarts and fail identically on an unmodified tree: this checkout's compiled
_Cextension 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 700fromcutlass_mxfp4_group_mm). Notrelated to this change.
Tolerances in
test_deepep_v2_moe.pyare now scaled to the magnitude of thereference output (
_scaled_tolerances) rather than fixed. The output std growswith the accumulation length, so a fixed
atolwas both too tight at the largestshape —
(222, 1024, 2048)fp8 failed with 25/454656 elements just over6e-2,the tail of a smooth error distribution rather than any real defect — and far too
loose at the smallest. Measured
max_err/stdis stable per dtype across everyMNK (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):atol1e-1 →~2e-3) while removing the large-shape false failure.
test_deepep_v2_moe.pynow pairs DeepEP V2 only with padding-aware TRTLLM experts(
flashinfer_trtllmfor bf16,trtllm_fp8for fp8); theflashinfer_cutlass(
Standard) path was dropped since it cannot consume the padded layout under theis_supersetgate.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 > 1andncclDevCommCreatefails. Multi-node genuinely needs GIN — do not set thisglobally in CI.
Model evals
gsm8k, 1319 questions, 5-shot, DP=4 +
--enable-expert-parallel --all2all-backend deepep_v2, no--enforce-eager(eager mode never populatespsum_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:Qwen/Qwen3-30B-A3BTrtLlmBf16ExpertsModularQwen/Qwen3-Coder-30B-A3B-Instruct-FP8TrtLlmFp8ExpertsModularnm-testing/Qwen3-30B-A3B-MXFP4A16DeepGemmFP4Experts(psum)The mxfp4 row sits below the 0.88 W4A16 baseline because
--moe-backend=deep_gemmforces
kFp8Dynamic128Symactivations, not because of padding: the same backendwithout 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 DeepGemmFP4Expertspairedwith
DeepEPV2PrepareAndFinalize. Before thedeep_gemmrouting fix this exactconfiguration 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:1526allocates the finalize output astorch.empty_like(hidden_states)at the unpadded 2880, sooutput.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:415slicesself.experts(...)[:, : self.hidden_size]). Separately,tests/evals/gpt_oss/test_gpqa_correctness.py:128hardcodes--enforce-eager, whichleaves
psum_recv_per_rankunpopulated, so this config could not exercise decodepadding even if it started.
Notes