[None][feat] Share Expert Fusion with cherry-pick #11143

[leslie-fang25]

Summary by CodeRabbit

• New Features

  • Added fused shared-expert support for DeepSeek-V3 FP8 block-scale MoE execution.
  • Shared experts can now be merged into routed expert processing, improving integration and reducing separate execution.
  • Added configuration support for selecting fused or unfused shared-expert execution.
  • Added microbenchmark options for shared-expert counts and execution mode.

• Bug Fixes

  • Improved expert routing, workspace sizing, weight validation, and output aggregation when shared experts are enabled.

• Tests

  • Added coverage for fused shared-expert accuracy, layout validation, and tactic reuse.

Description

Support shared expert fusion in trtllm gen MoE backend. Currently limits support to FP8 block-wise quantization and TP parallelization with this PR.

Test Coverage

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[leslie-fang25]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #59388 [ run ] triggered by Bot. Commit: dd9c892 Link to invocation

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds fused shared-expert support for DeepSeek FP8 block-scale MoE routing, execution, weight layout, model integration, benchmarks, and tests. Routing and GEMM dimensions now include appended shared experts, whose outputs use fixed weight 1.0F.

Changes

Fused shared-expert MoE support

Layer / File(s)	Summary
Routing contracts and DeepSeek writeback cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/*, cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.*	Routing metadata, TopK buffers, and DeepSeek writeback now include fused shared experts.
Fused-aware FP8 execution pipeline cpp/tensorrt_llm/thop/*, cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu	Workspace, tensor validation, configuration selection, and GEMM wiring use total expert counts; existing call sites pass zero for the new parameter.
Backend propagation and fused weight layout tensorrt_llm/_torch/modules/fused_moe/*, tensorrt_llm/_torch/custom_ops/*	Backend APIs propagate fused-expert counts and allocate/copy shared expert weights into trailing slots.
DeepSeek model integration tensorrt_llm/_torch/models/modeling_deepseekv3.py	DeepSeek layers fuse shared experts and skip separate shared-expert execution when fusion is enabled.
Benchmark and test coverage tests/microbenchmarks/bench_moe/*, tests/unittest/_torch/modules/moe/test_moe_backend.py, tests/integration/test_lists/*	Benchmark modes and targeted tests cover fused routing accuracy, tactics, and weight layouts.

Estimated code review effort: 5 (Critical) | ~90 minutes

Sequence Diagram(s)

sequenceDiagram
  participant DeepseekV3ForCausalLM
  participant ConfigurableMoE
  participant TRTLLMOpBackend
  participant FP8BlockScaleMoeRunner
  participant routingMainKernel
  DeepseekV3ForCausalLM->>ConfigurableMoE: fuse shared experts
  ConfigurableMoE->>TRTLLMOpBackend: run fused MoE
  TRTLLMOpBackend->>FP8BlockScaleMoeRunner: pass fused expert count
  FP8BlockScaleMoeRunner->>routingMainKernel: route expanded TopK
  routingMainKernel-->>FP8BlockScaleMoeRunner: routed and shared expert indices

Loading

Possibly related PRs

• NVIDIA/TensorRT-LLM#11143: Implements closely related fused/shared-expert routing and runner changes.

Suggested reviewers: christinaz, lfr-0531, yizhang-nv

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 29.69% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title matches the main change: adding shared expert fusion support, with only minor wording noise.
Description check	✅ Passed	The description follows the template and states the feature, but the Test Coverage section is empty.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 9

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (3)

tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py (1)

857-862: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Add num_fused_shared_experts to unique_id().

get_cache_key() uses str(runner.unique_id()), and num_fused_shared_experts now changes get_valid_tactics(). Leaving it out can collide fused and unfused configs and reuse a tactic that should have been filtered out.

🔧 Suggested fix

     def unique_id(self):
         return (self.top_k, self.intermediate_size, self.local_num_experts,
-                self.act_type)
+                self.act_type, self.num_fused_shared_experts)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py` around lines 857 -
862, Update the MoE runner’s unique_id() to include num_fused_shared_experts
alongside the existing tactic-validity fields, ensuring get_cache_key()
distinguishes fused and unfused configurations while preserving the current
tuple-based cache key behavior.

cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu (1)

128-231: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

SigmoidRenorm and MiniMax2 silently drop fused shared experts instead of rejecting them.

Both branches build a routingCustom::Data with no shared-expert-fusion fields, so a non-zero numFusedSharedExpert is quietly ignored here — unlike the Llama4 (line 290) and Renormalize/Default (line 336) branches, which were given an explicit TLLM_CHECK_WITH_INFO(numFusedSharedExpert == 0, ...) guard. Without the same guard, a caller requesting fusion with these two routing methods gets wrong output (missing shared experts) with no diagnostic, instead of a clear failure.

🐛 Proposed fix

     else if (routingMethodType == RoutingMethodType::SigmoidRenorm)
     {
+        TLLM_CHECK_WITH_INFO(
+            numFusedSharedExpert == 0, "SigmoidRenorm routing method does not support fusing shared expert");
         // SigmoidRenorm: sigmoid(logit) → topK → renormalize.

     else if (routingMethodType == RoutingMethodType::MiniMax2)
     {
+        TLLM_CHECK_WITH_INFO(
+            numFusedSharedExpert == 0, "MiniMax2 routing method does not support fusing shared expert");
         // MiniMaxM2: sigmoid(logit) + bias → topK → renormalize un-biased sigmoid scores.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu` around
lines 128 - 231, Add an explicit TLLM_CHECK_WITH_INFO(numFusedSharedExpert == 0,
...) guard to both the SigmoidRenorm and MiniMax2 branches before constructing
routingCustom::Data, matching the existing Llama4 and Renormalize/Default
handling. Reject non-zero fused shared experts rather than silently proceeding
without fusion.

cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingDeepSeek.cu (1)

193-200: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

mPtrExpertCounts zero-init must include fused shared experts.

RoutingDeepSeek.cu:193-199 only clears 2 * params.mNumExperts from the pre-expansion launch, but routingIndicesHistogramKernel / routingIndicesOffsetsKernel run after data.mNumExperts is expanded and use that expanded range. The fused shared-expert tail can keep stale counts.

🐛 Proposed fix

-        int32_t expertCountsNum = 2 * params.mNumExperts;
+        int32_t expertCountsNum = 2 * (params.mNumExperts + params.mNumFusedSharedExperts);

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingDeepSeek.cu`
around lines 193 - 200, Update the mPtrExpertCounts initialization in the
routing setup to clear the full post-expansion expert-count range, including the
fused shared-expert tail used by routingIndicesHistogramKernel and
routingIndicesOffsetsKernel. Base the initArr length on the expanded expert
count rather than only 2 * params.mNumExperts, while preserving the existing
guarded initialization and launch-stride behavior.

🧹 Nitpick comments (2)

tests/unittest/_torch/modules/moe/test_moe_backend.py (1)

1231-1370: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy lift

Add benchmark-path regression coverage.

These tests sufficiently cover backend accuracy and fused layout, but not the newly changed benchmark wiring. Add tests/microbenchmarks/bench_moe/test_shared_experts.py covering unsupported backend/quantization rejection, fused-slot initialization, and multi-GPU reduction semantics.

As per path instructions, state coverage sufficiency and suggest a concrete test filename for TensorRT-LLM test changes.

Also applies to: 1377-1489

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/modules/moe/test_moe_backend.py` around lines 1231 -
1370, Extend benchmark-path regression coverage by adding tests in
tests/microbenchmarks/bench_moe/test_shared_experts.py for unsupported backend
or quantization rejection, fused-slot initialization, and multi-GPU reduction
semantics. Keep the existing TRTLLM backend accuracy and fused-layout coverage
in tests/unittest/_torch/modules/moe/test_moe_backend.py, which is sufficient
for those behaviors; use this benchmark test file for the changed benchmark
wiring.

Source: Path instructions

cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp (1)

401-405: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Duplicated TLLM_MOE_FUSED_MIN_TILEN env-var parsing.

fusedMinTileN (here) and fusedMinTileNFallback (L473-477) independently read the same env var with the same default via getenv/atoi. Extract this into a shared free function/constant to avoid drift if the env var name or default ever changes.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp` around lines 401 - 405, Extract
the shared TLLM_MOE_FUSED_MIN_TILEN getenv/atoi logic and default value into a
single free function or constant, then update both fusedMinTileN and
fusedMinTileNFallback to reuse it. Remove the duplicated parsing while
preserving the current environment-variable name and fallback value.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernel.h`:
- Around line 152-154: Initialize mNumFusedSharedExperts and
mTotalExpertsPerToken in KernelParamsBase with explicit zero defaults, matching
the existing scalar member initialization pattern and ensuring safe no-fusion
behavior before setBaseParams() runs.

In `@cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp`:
- Around line 143-149: In the enclosing function, add an upfront TORCH_CHECK
that num_fused_shared_experts.value_or(0) is non-negative, then replace the
repeated > 0 ternaries in the num_total_experts, total_experts_per_token, and
num_total_local_experts calculations with num_fused_shared_experts.value_or(0).
Apply the same simplification at the duplicated call sites in getValidConfigs
and run, including the args.num_fused_shared_experts assignment, while
preserving consistent totals across all paths.
- Around line 215-216: Update the external routing-buffer validation in the FP8
MoE routing setup to account for total_experts_per_token, which includes
num_fused_shared_experts, rather than validating only against top_k. Ensure
topk_weights and topk_ids are rejected unless they contain enough entries for
every token consumed by the fused shared-experts path, preventing downstream
out-of-bounds access.

In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py`:
- Around line 339-354: Update the shared-expert fusion gate in TRTLLMGenFusedMoE
to also require model_config.mapping.moe_ep_size == 1, alongside the existing
backend, data-parallel, and quantization checks. Keep fusion disabled for EP
configurations while preserving the current behavior for supported non-EP
configurations.

In `@tests/microbenchmarks/bench_moe/build.py`:
- Around line 287-309: The shared-expert MLP is currently created only for the
unfused mode, leaving fused shared-expert weight slots uninitialized. Update the
conditional around GatedMLP construction to build and initialize shared_mlp
whenever model.n_shared_experts > 0, preserve the _UnfusedSharedMoE wrapping for
unfused mode, and call moe.fuse_shared_expert(shared_mlp) for fused mode.
- Around line 294-300: Update the MoE construction around the routed branch and
shared_mlp so both TTP branches use consistent reduction semantics before their
outputs are summed. Prefer enabling reduction for shared_mlp to match the routed
MoE’s reduce_results=True, or keep both outputs local and add one reduction
after summation; ensure the baseline timing includes the required multi-GPU
communication.
- Around line 121-160: Add complete type annotations to every new function in
the affected sites: in tests/microbenchmarks/bench_moe/build.py lines 121-160,
annotate __init__ with a None return, __getattr__ parameters and return, and
forward parameters and return; in
tests/unittest/_torch/modules/moe/test_moe_backend.py lines 1177-1218, annotate
routing-helper parameters and returns plus _write_fused_shared_expert_slots;
lines 1231-1370, annotate expert_info, the test return, and nested run_moe; and
lines 1377-1489, add the test’s -> None return annotation. Preserve existing
behavior and use appropriate concrete or existing project typing symbols.
- Around line 181-205: Update the shared-expert validation near the existing
parallel-layout guard to reject or skip fused candidates unless they use the
TRTLLM backend with FP8_BLOCK_SCALES quantization. Apply this only when
model.n_shared_experts > 0 and shared_expert_mode is not "unfused"; preserve
unfused behavior and ensure unsupported candidates cannot proceed while
recording n_shared_experts > 0.

In `@tests/microbenchmarks/bench_moe/specs.py`:
- Around line 111-116: Update ModelSpec’s shared_expert_mode annotation to
Literal["fused", "unfused"] and extend its validation to reject n_shared_experts
values below zero. Ensure invalid modes and negative shared-expert counts fail
explicitly rather than being silently interpreted or disabled, while preserving
valid configurations.

---

Outside diff comments:
In
`@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingDeepSeek.cu`:
- Around line 193-200: Update the mPtrExpertCounts initialization in the routing
setup to clear the full post-expansion expert-count range, including the fused
shared-expert tail used by routingIndicesHistogramKernel and
routingIndicesOffsetsKernel. Base the initArr length on the expanded expert
count rather than only 2 * params.mNumExperts, while preserving the existing
guarded initialization and launch-stride behavior.

In `@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu`:
- Around line 128-231: Add an explicit TLLM_CHECK_WITH_INFO(numFusedSharedExpert
== 0, ...) guard to both the SigmoidRenorm and MiniMax2 branches before
constructing routingCustom::Data, matching the existing Llama4 and
Renormalize/Default handling. Reject non-zero fused shared experts rather than
silently proceeding without fusion.

In `@tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py`:
- Around line 857-862: Update the MoE runner’s unique_id() to include
num_fused_shared_experts alongside the existing tactic-validity fields, ensuring
get_cache_key() distinguishes fused and unfused configurations while preserving
the current tuple-based cache key behavior.

---

Nitpick comments:
In `@cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp`:
- Around line 401-405: Extract the shared TLLM_MOE_FUSED_MIN_TILEN getenv/atoi
logic and default value into a single free function or constant, then update
both fusedMinTileN and fusedMinTileNFallback to reuse it. Remove the duplicated
parsing while preserving the current environment-variable name and fallback
value.

In `@tests/unittest/_torch/modules/moe/test_moe_backend.py`:
- Around line 1231-1370: Extend benchmark-path regression coverage by adding
tests in tests/microbenchmarks/bench_moe/test_shared_experts.py for unsupported
backend or quantization rejection, fused-slot initialization, and multi-GPU
reduction semantics. Keep the existing TRTLLM backend accuracy and fused-layout
coverage in tests/unittest/_torch/modules/moe/test_moe_backend.py, which is
sufficient for those behaviors; use this benchmark test file for the changed
benchmark wiring.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 1cd7a453-bbf7-4d7f-90cf-f91d0b03cb57

📥 Commits

Reviewing files that changed from the base of the PR and between 5977a19 and dd9c892.

📒 Files selected for processing (23)

• cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingDeepSeek.cu
• cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernel.h
• cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu
• cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h
• cpp/tensorrt_llm/thop/cuteDslMoeUtilsOp.cpp
• cpp/tensorrt_llm/thop/fp4BlockScaleMoe.cpp
• cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp
• cpp/tensorrt_llm/thop/fp8PerTensorScaleMoe.cpp
• cpp/tensorrt_llm/thop/mxFp4BlockScaleMoe.cpp
• tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py
• tensorrt_llm/_torch/models/modeling_deepseekv3.py
• tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py
• tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py
• tensorrt_llm/_torch/modules/fused_moe/quantization.py
• tests/integration/test_lists/test-db/l0_b200.yml
• tests/integration/test_lists/test-db/l0_b300.yml
• tests/microbenchmarks/bench_moe/build.py
• tests/microbenchmarks/bench_moe/case_runner.py
• tests/microbenchmarks/bench_moe/cli.py
• tests/microbenchmarks/bench_moe/mapping.py
• tests/microbenchmarks/bench_moe/specs.py
• tests/unittest/_torch/modules/moe/test_moe_backend.py

[tensorrt-cicd]

PR_Github #59388 [ run ] completed with state SUCCESS. Commit: dd9c892
/LLM/main/L0_MergeRequest_PR pipeline #47862 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[leslie-fang25]

/bot run --disable-fail-fast

[BowenFu]

LGTM — Share-Expert Fusion is opt-in via TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION (default "0"). With it unset, num_fused_shared_expert=0 everywhere: weight shapes are unchanged, the fused-shared-expert kernel branches never fire, and DeepseekV3MoE.forward keeps the existing shared_experts path — MoE compute for existing models is byte-identical. Adds unit + b200/b300 tests.

[leslie-fang25]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #60135 [ run ] triggered by Bot. Commit: ff3227e Link to invocation

[tensorrt-cicd]

PR_Github #60135 [ run ] completed with state SUCCESS. Commit: ff3227e
/LLM/main/L0_MergeRequest_PR pipeline #48516 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[xxi-nv]

Thanks for the cherry-pick.

[sunnyqgg]

LGTM

[leslie-fang25]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61526 [ run ] triggered by Bot. Commit: 5016db4 Link to invocation

[tensorrt-cicd]

PR_Github #61526 [ run ] completed with state SUCCESS. Commit: 5016db4
/LLM/main/L0_MergeRequest_PR pipeline #49741 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[leslie-fang25]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61663 [ run ] triggered by Bot. Commit: 5016db4 Link to invocation

[tensorrt-cicd]

PR_Github #61663 [ run ] completed with state SUCCESS. Commit: 5016db4
/LLM/main/L0_MergeRequest_PR pipeline #49868 completed with status: 'SUCCESS'

CI Report

Link to invocation