[None][feat] Enable MEGAMOE_DEEPGEMM backend for DeepSeek V4#14129
Merged
Tabrizian merged 1 commit intoMay 14, 2026
Merged
Conversation
V4-Flash routed experts use the MXFP4 byte layout that MegaMoEDeepGemm's fused fp8_fp4_mega_moe kernel was designed for, but several wiring gaps kept users from selecting the backend end-to-end: - llmapi: MoeConfig.backend Literal omitted "MEGAMOE_DEEPGEMM", so YAML overrides via --extra_llm_api_options failed Pydantic validation before any backend code ran. - create_moe: the model unconditionally forwards swiglu_limit_scalar (V4-Flash uses uniform SwiGLU clamp), but the allow-list assert and per-backend dispatch never threaded it through MegaMoEDeepGemm. Add MegaMoEDeepGemm to the assert (with explicit per-expert-tensor rejection, since DG only supports the scalar form) and map swiglu_limit_scalar -> activation_clamp in the dispatch branch. - mega_moe_deepgemm: _resolve_ep_pg required torch.distributed to be initialized, which is the Ray / DeviceMesh contract. TRT-LLM's default executor uses raw mpi4py and never calls init_process_group. Add an MPI fallback that, when dist is uninitialized, derives rank / world_size from MPI.COMM_WORLD, picks an intra-node local rank via Split_type(COMM_TYPE_SHARED) with OMPI_COMM_WORLD_LOCAL_RANK env-var fallback, uses setdefault for the four rendezvous vars so Slurm / Ray / torchrun setups stay intact, bcasts MASTER_ADDR / MASTER_PORT, and wraps init_process_group with one retry on a fresh port to absorb the close() -> NCCL bind() race on busy hosts. - can_implement: accept SM100-family (SM100 + SM103). DG's fp8_fp4_mega_moe cubin is sm_100f SASS, compatible with both B200 and B300. - quantization: in W4A8MXFP4MXFP8MegaMoEDeepGemmMethod._transform_main_weights, free the three raw checkpoint Parameters (w3_w1_weight, w3_w1_weight_scale, w2_weight_scale) after the DG transform builds the DG-form _t_l1 / _t_l2 views. DG allocates fresh storage for L1 (interleaved) and both SF tensors, so those are pure duplicates; L2 weight is passed through unchanged and stays alive via _t_l2[0]. data_ptr asserts guard the invariant so a future DG that view-passes L1 surfaces here at model-build time instead of dangling at forward. Frees ~560 MB/layer => ~24 GiB/rank on V4-Flash (43 MoE layers, EP=4), enough headroom for the 8.69 GiB combine_token_buffer SymmBuffer introduced by DG nv_dev_67fc648. - tests: parametrize TestDeepSeekV4Flash::test_nvfp4_4gpus_static_eplb with the new MEGAMOE_DEEPGEMM backend and register it in l0_dgx_b300_ds.yml. The TP=4 EP=4 attention-DP layout is the one validated below; this enforces the end-to-end wire-up so a future config-plumbing regression cannot slip through silently. Validated on DeepSeek-V4-Flash, B200 TP=4 EP=4 attention-DP: - gsm8k 1319 samples = 95.49 % (TRTLLM ref 95.11 %). - trtllm-bench ISL=1024 OSL=1024 500 reqs: +15.3 % throughput and -12.7 % latency vs TRTLLM MoE backend on the same model and config. Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
Barry-Delaney
requested review from
HuiGao-NV and
schetlur-nv
and removed request for
a team
May 14, 2026 08:29
Barry-Delaney
changed the title
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #48338 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #48338 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #48384 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #48384 [ run ] completed with state
|
Collaborator
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #48403 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #48403 [ run ] completed with state |
lfr-0531
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 10, 2026
…14129) Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
lfr-0531
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 11, 2026
…14129) Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
lfr-0531
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 12, 2026
…14129) Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com>
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
4 participants
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.
Enables MegaMoE backend on DeepSeek-V4 and added related CI tests.