[None][feat] Share Expert Fusion with cherry-pick #11143 - #15297
Conversation
b4f8941 to
8763371
Compare
1fdd087 to
dd9c892
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #59388 [ run ] triggered by Bot. Commit: |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
WalkthroughAdds fused shared-expert support for DeepSeek FP8 block-scale MoE routing, execution, weight layout, model integration, benchmarks, and tests. Routing and GEMM dimensions include appended shared experts with weight ChangesFused shared-expert MoE support
Estimated code review effort: 5 (Critical) | ~90 minutes Sequence Diagram(s)sequenceDiagram
participant DeepseekV3MoE
participant ConfigurableMoE
participant TRTLLMOpBackend
participant FP8BlockScaleMoERunner
participant routingMainKernel
DeepseekV3MoE->>ConfigurableMoE: fuse shared experts
ConfigurableMoE->>TRTLLMOpBackend: run fused MoE
TRTLLMOpBackend->>FP8BlockScaleMoERunner: pass fused expert count
FP8BlockScaleMoERunner->>routingMainKernel: route expanded TopK
routingMainKernel-->>FP8BlockScaleMoERunner: return routed and shared expert indices
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
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 winAdd
num_fused_shared_expertstounique_id().
get_cache_key()usesstr(runner.unique_id()), andnum_fused_shared_expertsnow changesget_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
SigmoidRenormandMiniMax2silently drop fused shared experts instead of rejecting them.Both branches build a
routingCustom::Datawith no shared-expert-fusion fields, so a non-zeronumFusedSharedExpertis quietly ignored here — unlike theLlama4(line 290) andRenormalize/Default(line 336) branches, which were given an explicitTLLM_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
mPtrExpertCountszero-init must include fused shared experts.
RoutingDeepSeek.cu:193-199only clears2 * params.mNumExpertsfrom the pre-expansion launch, butroutingIndicesHistogramKernel/routingIndicesOffsetsKernelrun afterdata.mNumExpertsis 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 liftAdd 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.pycovering 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 winDuplicated
TLLM_MOE_FUSED_MIN_TILENenv-var parsing.
fusedMinTileN(here) andfusedMinTileNFallback(L473-477) independently read the same env var with the same default viagetenv/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
📒 Files selected for processing (23)
cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingDeepSeek.cucpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernel.hcpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cucpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.hcpp/tensorrt_llm/thop/cuteDslMoeUtilsOp.cppcpp/tensorrt_llm/thop/fp4BlockScaleMoe.cppcpp/tensorrt_llm/thop/fp8BlockScaleMoe.cppcpp/tensorrt_llm/thop/fp8PerTensorScaleMoe.cppcpp/tensorrt_llm/thop/mxFp4BlockScaleMoe.cpptensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.pytensorrt_llm/_torch/models/modeling_deepseekv3.pytensorrt_llm/_torch/modules/fused_moe/configurable_moe.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.pytensorrt_llm/_torch/modules/fused_moe/moe_op_backend.pytensorrt_llm/_torch/modules/fused_moe/quantization.pytests/integration/test_lists/test-db/l0_b200.ymltests/integration/test_lists/test-db/l0_b300.ymltests/microbenchmarks/bench_moe/build.pytests/microbenchmarks/bench_moe/case_runner.pytests/microbenchmarks/bench_moe/cli.pytests/microbenchmarks/bench_moe/mapping.pytests/microbenchmarks/bench_moe/specs.pytests/unittest/_torch/modules/moe/test_moe_backend.py
dd9c892 to
0d1bde2
Compare
|
PR_Github #59388 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63687 [ run ] triggered by Bot. Commit: |
|
PR_Github #63687 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63733 [ run ] triggered by Bot. Commit: |
|
PR_Github #63733 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63763 [ run ] triggered by Bot. Commit: |
|
PR_Github #63763 [ run ] completed with state |
unittest/auto_deploy/singlegpu/custom_ops/moe/test_trtllm_moe.py and unittest/auto_deploy/singlegpu/transformations/library/test_moe_fusion.py fail in DGX_B200-AutoDeploy-1 for any PR rebased onto current main: identical failures on pipelines 51974 (this PR) and 51977 (PR NVIDIA#17225, zero file overlap). Suspect commit 89bba4c (NVIDIA#15297), which modifies the trtllm_moe custom op and Blackwell blockScaleMoe kernels. Waived pending an NVBug; the entries will be updated with the bug link once filed. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…-scale MoE The TRTLLM-Gen small-tile (tileN 8/16) dynB batched-GEMM cubins flakily hit an illegal memory access in the gemm2 K-loop. PR NVIDIA#15297 already added a WAR for this exact fault -- restricting tactics to tileN >= 32 -- but scoped it to the fused shared-expert path via num_fused_shared_experts > 0. Shared-expert fusion is opt-in (TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION, default off), so DeepSeek-R1 FP8 TP=8 runs unfused and never reached the exclusion. It faults identically: the defect is in the shared small-tile cubins, not caused by expert fusion. Apply the exclusion for every caller, at both selection sites (getValidConfigs and the tileN == -1 fallback in run()). Measured on DeepSeek-R1 FP8 EP=1 (B200, SM100f): the warmup shapes that fault (1/2/8 tokens) were the only ones able to select tileN 8/16, while the 12288 token shape gets tileN 64/128 and always passed. After the change every shape from 1 to 12288 tokens offers only tileN >= 32, and all five Phase-2 warmup shapes complete where the second previously crashed. The tiles stay in mSupportedTileN: the ctor builds one runner per tile and each asserts a non-empty passing-config list, so the exclusion must happen at tactic-selection time rather than by dropping the tile. This is safe because every FP8 block-scale MoE shape retains a tileN >= 32 tactic -- verified across DeepSeek-R1 (EP 1/4/8), Qwen3-235B and Qwen3-30B -- so the tactic list is never emptied. The waiver is kept: an independent illegal memory access remains in the Phase-2 autotuner warmup, which this change does not address. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…-scale MoE The TRTLLM-Gen small-tile (tileN 8/16) dynB batched-GEMM cubins flakily hit an illegal memory access in the gemm2 K-loop. PR NVIDIA#15297 already added a WAR for this exact fault -- restricting tactics to tileN >= 32 -- but scoped it to the fused shared-expert path via num_fused_shared_experts > 0. Shared-expert fusion is opt-in (TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION, default off), so DeepSeek-R1 FP8 TP=8 runs unfused and never reached the exclusion. It faults identically: the defect is in the shared small-tile cubins, not caused by expert fusion. Apply the exclusion for every caller, at both selection sites (getValidConfigs and the tileN == -1 fallback in run()). Measured on DeepSeek-R1 FP8 EP=1 (B200, SM100f): the warmup shapes that fault (1/2/8 tokens) were the only ones able to select tileN 8/16, while the 12288 token shape gets tileN 64/128 and always passed. After the change every shape from 1 to 12288 tokens offers only tileN >= 32, and all five Phase-2 warmup shapes complete where the second previously crashed. The tiles stay in mSupportedTileN: the ctor builds one runner per tile and each asserts a non-empty passing-config list, so the exclusion must happen at tactic-selection time rather than by dropping the tile. This is safe because every FP8 block-scale MoE shape retains a tileN >= 32 tactic -- verified across DeepSeek-R1 (EP 1/4/8), Qwen3-235B and Qwen3-30B -- so the tactic list is never emptied. The waiver is kept: an independent illegal memory access remains in the Phase-2 autotuner warmup, which this change does not address. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…-scale MoE The TRTLLM-Gen small-tile (tileN 8/16) dynB batched-GEMM cubins flakily hit an illegal memory access in the gemm2 K-loop. PR NVIDIA#15297 already added a WAR for this exact fault -- restricting tactics to tileN >= 32 -- but scoped it to the fused shared-expert path via num_fused_shared_experts > 0. Shared-expert fusion is opt-in (TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION, default off), so DeepSeek-R1 FP8 TP=8 runs unfused and never reached the exclusion. It faults identically: the defect is in the shared small-tile cubins, not caused by expert fusion. Apply the exclusion for every caller, at both selection sites (getValidConfigs and the tileN == -1 fallback in run()). Measured on DeepSeek-R1 FP8 EP=1 (B200, SM100f): the warmup shapes that fault (1/2/8 tokens) were the only ones able to select tileN 8/16, while the 12288 token shape gets tileN 64/128 and always passed. After the change every shape from 1 to 12288 tokens offers only tileN >= 32, and all five Phase-2 warmup shapes complete where the second previously crashed. The tiles stay in mSupportedTileN: the ctor builds one runner per tile and each asserts a non-empty passing-config list, so the exclusion must happen at tactic-selection time rather than by dropping the tile. This is safe because every FP8 block-scale MoE shape retains a tileN >= 32 tactic -- verified across DeepSeek-R1 (EP 1/4/8), Qwen3-235B and Qwen3-30B -- so the tactic list is never emptied. The waiver is kept: an independent illegal memory access remains in the Phase-2 autotuner warmup, which this change does not address. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…-scale MoE The TRTLLM-Gen small-tile (tileN 8/16) dynB batched-GEMM cubins flakily hit an illegal memory access in the gemm2 K-loop. PR NVIDIA#15297 already added a WAR for this exact fault -- restricting tactics to tileN >= 32 -- but scoped it to the fused shared-expert path via num_fused_shared_experts > 0. Shared-expert fusion is opt-in (TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION, default off), so DeepSeek-R1 FP8 TP=8 runs unfused and never reached the exclusion. It faults identically: the defect is in the shared small-tile cubins, not caused by expert fusion. Apply the exclusion for every caller, at both selection sites (getValidConfigs and the tileN == -1 fallback in run()). Measured on DeepSeek-R1 FP8 EP=1 (B200, SM100f): the warmup shapes that fault (1/2/8 tokens) were the only ones able to select tileN 8/16, while the 12288 token shape gets tileN 64/128 and always passed. After the change every shape from 1 to 12288 tokens offers only tileN >= 32, and all five Phase-2 warmup shapes complete where the second previously crashed. The tiles stay in mSupportedTileN: the ctor builds one runner per tile and each asserts a non-empty passing-config list, so the exclusion must happen at tactic-selection time rather than by dropping the tile. This is safe because every FP8 block-scale MoE shape retains a tileN >= 32 tactic -- verified across DeepSeek-R1 (EP 1/4/8), Qwen3-235B and Qwen3-30B -- so the tactic list is never emptied. The waiver is kept: an independent illegal memory access remains in the Phase-2 autotuner warmup, which this change does not address. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…-scale MoE The TRTLLM-Gen small-tile (tileN 8/16) dynB batched-GEMM cubins flakily hit an illegal memory access in the gemm2 K-loop. PR NVIDIA#15297 already added a WAR for this exact fault -- restricting tactics to tileN >= 32 -- but scoped it to the fused shared-expert path via num_fused_shared_experts > 0. Shared-expert fusion is opt-in (TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION, default off), so DeepSeek-R1 FP8 TP=8 runs unfused and never reached the exclusion. It faults identically: the defect is in the shared small-tile cubins, not caused by expert fusion. Apply the exclusion for every caller, at both selection sites (getValidConfigs and the tileN == -1 fallback in run()). Measured on DeepSeek-R1 FP8 EP=1 (B200, SM100f): the warmup shapes that fault (1/2/8 tokens) were the only ones able to select tileN 8/16, while the 12288 token shape gets tileN 64/128 and always passed. After the change every shape from 1 to 12288 tokens offers only tileN >= 32, and all five Phase-2 warmup shapes complete where the second previously crashed. The tiles stay in mSupportedTileN: the ctor builds one runner per tile and each asserts a non-empty passing-config list, so the exclusion must happen at tactic-selection time rather than by dropping the tile. This is safe because every FP8 block-scale MoE shape retains a tileN >= 32 tactic -- verified across DeepSeek-R1 (EP 1/4/8), Qwen3-235B and Qwen3-30B -- so the tactic list is never emptied. The waiver is kept: an independent illegal memory access remains in the Phase-2 autotuner warmup, which this change does not address. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…-scale MoE The TRTLLM-Gen small-tile (tileN 8/16) dynB batched-GEMM cubins flakily hit an illegal memory access in the gemm2 K-loop. PR NVIDIA#15297 already added a WAR for this exact fault -- restricting tactics to tileN >= 32 -- but scoped it to the fused shared-expert path via num_fused_shared_experts > 0. Shared-expert fusion is opt-in (TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION, default off), so DeepSeek-R1 FP8 TP=8 runs unfused and never reached the exclusion. It faults identically: the defect is in the shared small-tile cubins, not caused by expert fusion. Apply the exclusion for every caller, at both selection sites (getValidConfigs and the tileN == -1 fallback in run()). Measured on DeepSeek-R1 FP8 EP=1 (B200, SM100f): the warmup shapes that fault (1/2/8 tokens) were the only ones able to select tileN 8/16, while the 12288 token shape gets tileN 64/128 and always passed. After the change every shape from 1 to 12288 tokens offers only tileN >= 32, and all five Phase-2 warmup shapes complete where the second previously crashed. The tiles stay in mSupportedTileN: the ctor builds one runner per tile and each asserts a non-empty passing-config list, so the exclusion must happen at tactic-selection time rather than by dropping the tile. This is safe because every FP8 block-scale MoE shape retains a tileN >= 32 tactic -- verified across DeepSeek-R1 (EP 1/4/8), Qwen3-235B and Qwen3-30B -- so the tactic list is never emptied. The waiver is kept: an independent illegal memory access remains in the Phase-2 autotuner warmup, which this change does not address. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…-scale MoE The TRTLLM-Gen small-tile (tileN 8/16) dynB batched-GEMM cubins flakily hit an illegal memory access in the gemm2 K-loop. PR NVIDIA#15297 already added a WAR for this exact fault -- restricting tactics to tileN >= 32 -- but scoped it to the fused shared-expert path via num_fused_shared_experts > 0. Shared-expert fusion is opt-in (TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION, default off), so DeepSeek-R1 FP8 TP=8 runs unfused and never reached the exclusion. It faults identically: the defect is in the shared small-tile cubins, not caused by expert fusion. Apply the exclusion for every caller, at both selection sites (getValidConfigs and the tileN == -1 fallback in run()). Measured on DeepSeek-R1 FP8 EP=1 (B200, SM100f): the warmup shapes that fault (1/2/8 tokens) were the only ones able to select tileN 8/16, while the 12288 token shape gets tileN 64/128 and always passed. After the change every shape from 1 to 12288 tokens offers only tileN >= 32, and all five Phase-2 warmup shapes complete where the second previously crashed. The tiles stay in mSupportedTileN: the ctor builds one runner per tile and each asserts a non-empty passing-config list, so the exclusion must happen at tactic-selection time rather than by dropping the tile. This is safe because every FP8 block-scale MoE shape retains a tileN >= 32 tactic -- verified across DeepSeek-R1 (EP 1/4/8), Qwen3-235B and Qwen3-30B -- so the tactic list is never emptied. The waiver is kept: an independent illegal memory access remains in the Phase-2 autotuner warmup, which this change does not address. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…-scale MoE The TRTLLM-Gen small-tile (tileN 8/16) dynB batched-GEMM cubins flakily hit an illegal memory access in the gemm2 K-loop. PR NVIDIA#15297 already added a WAR for this exact fault -- restricting tactics to tileN >= 32 -- but scoped it to the fused shared-expert path via num_fused_shared_experts > 0. Shared-expert fusion is opt-in (TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION, default off), so DeepSeek-R1 FP8 TP=8 runs unfused and never reached the exclusion. It faults identically: the defect is in the shared small-tile cubins, not caused by expert fusion. Apply the exclusion for every caller, at both selection sites (getValidConfigs and the tileN == -1 fallback in run()). Measured on DeepSeek-R1 FP8 EP=1 (B200, SM100f): the warmup shapes that fault (1/2/8 tokens) were the only ones able to select tileN 8/16, while the 12288 token shape gets tileN 64/128 and always passed. After the change every shape from 1 to 12288 tokens offers only tileN >= 32, and all five Phase-2 warmup shapes complete where the second previously crashed. The tiles stay in mSupportedTileN: the ctor builds one runner per tile and each asserts a non-empty passing-config list, so the exclusion must happen at tactic-selection time rather than by dropping the tile. This is safe because every FP8 block-scale MoE shape retains a tileN >= 32 tactic -- verified across DeepSeek-R1 (EP 1/4/8), Qwen3-235B and Qwen3-30B -- so the tactic list is never emptied. The waiver is kept: an independent illegal memory access remains in the Phase-2 autotuner warmup, which this change does not address. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…-scale MoE The TRTLLM-Gen small-tile (tileN 8/16) dynB batched-GEMM cubins flakily hit an illegal memory access in the gemm2 K-loop. PR NVIDIA#15297 already added a WAR for this exact fault -- restricting tactics to tileN >= 32 -- but scoped it to the fused shared-expert path via num_fused_shared_experts > 0. Shared-expert fusion is opt-in (TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION, default off), so DeepSeek-R1 FP8 TP=8 runs unfused and never reached the exclusion. It faults identically: the defect is in the shared small-tile cubins, not caused by expert fusion. Apply the exclusion for every caller, at both selection sites (getValidConfigs and the tileN == -1 fallback in run()). Measured on DeepSeek-R1 FP8 EP=1 (B200, SM100f): the warmup shapes that fault (1/2/8 tokens) were the only ones able to select tileN 8/16, while the 12288 token shape gets tileN 64/128 and always passed. After the change every shape from 1 to 12288 tokens offers only tileN >= 32, and all five Phase-2 warmup shapes complete where the second previously crashed. The tiles stay in mSupportedTileN: the ctor builds one runner per tile and each asserts a non-empty passing-config list, so the exclusion must happen at tactic-selection time rather than by dropping the tile. This is safe because every FP8 block-scale MoE shape retains a tileN >= 32 tactic -- verified across DeepSeek-R1 (EP 1/4/8), Qwen3-235B and Qwen3-30B -- so the tactic list is never emptied. The waiver is kept: an independent illegal memory access remains in the Phase-2 autotuner warmup, which this change does not address. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…-scale MoE The TRTLLM-Gen small-tile (tileN 8/16) dynB batched-GEMM cubins flakily hit an illegal memory access in the gemm2 K-loop. PR NVIDIA#15297 already added a WAR for this exact fault -- restricting tactics to tileN >= 32 -- but scoped it to the fused shared-expert path via num_fused_shared_experts > 0. Shared-expert fusion is opt-in (TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION, default off), so DeepSeek-R1 FP8 TP=8 runs unfused and never reached the exclusion. It faults identically: the defect is in the shared small-tile cubins, not caused by expert fusion. Apply the exclusion for every caller, at both selection sites (getValidConfigs and the tileN == -1 fallback in run()). Measured on DeepSeek-R1 FP8 EP=1 (B200, SM100f): the warmup shapes that fault (1/2/8 tokens) were the only ones able to select tileN 8/16, while the 12288 token shape gets tileN 64/128 and always passed. After the change every shape from 1 to 12288 tokens offers only tileN >= 32, and all five Phase-2 warmup shapes complete where the second previously crashed. The tiles stay in mSupportedTileN: the ctor builds one runner per tile and each asserts a non-empty passing-config list, so the exclusion must happen at tactic-selection time rather than by dropping the tile. This is safe because every FP8 block-scale MoE shape retains a tileN >= 32 tactic -- verified across DeepSeek-R1 (EP 1/4/8), Qwen3-235B and Qwen3-30B -- so the tactic list is never emptied. The waiver is kept: an independent illegal memory access remains in the Phase-2 autotuner warmup, which this change does not address. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…-scale MoE The TRTLLM-Gen small-tile (tileN 8/16) dynB batched-GEMM cubins flakily hit an illegal memory access in the gemm2 K-loop. PR NVIDIA#15297 already added a WAR for this exact fault -- restricting tactics to tileN >= 32 -- but scoped it to the fused shared-expert path via num_fused_shared_experts > 0. Shared-expert fusion is opt-in (TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION, default off), so DeepSeek-R1 FP8 TP=8 runs unfused and never reached the exclusion. It faults identically: the defect is in the shared small-tile cubins, not caused by expert fusion. Apply the exclusion for every caller, at both selection sites (getValidConfigs and the tileN == -1 fallback in run()). Measured on DeepSeek-R1 FP8 EP=1 (B200, SM100f): the warmup shapes that fault (1/2/8 tokens) were the only ones able to select tileN 8/16, while the 12288 token shape gets tileN 64/128 and always passed. After the change every shape from 1 to 12288 tokens offers only tileN >= 32, and all five Phase-2 warmup shapes complete where the second previously crashed. The tiles stay in mSupportedTileN: the ctor builds one runner per tile and each asserts a non-empty passing-config list, so the exclusion must happen at tactic-selection time rather than by dropping the tile. This is safe because every FP8 block-scale MoE shape retains a tileN >= 32 tactic -- verified across DeepSeek-R1 (EP 1/4/8), Qwen3-235B and Qwen3-30B -- so the tactic list is never emptied. The waiver is kept: an independent illegal memory access remains in the Phase-2 autotuner warmup, which this change does not address. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…-scale MoE The TRTLLM-Gen small-tile (tileN 8/16) dynB batched-GEMM cubins flakily hit an illegal memory access in the gemm2 K-loop. PR NVIDIA#15297 already added a WAR for this exact fault -- restricting tactics to tileN >= 32 -- but scoped it to the fused shared-expert path via num_fused_shared_experts > 0. Shared-expert fusion is opt-in (TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION, default off), so DeepSeek-R1 FP8 TP=8 runs unfused and never reached the exclusion. It faults identically: the defect is in the shared small-tile cubins, not caused by expert fusion. Apply the exclusion for every caller, at both selection sites (getValidConfigs and the tileN == -1 fallback in run()). Measured on DeepSeek-R1 FP8 EP=1 (B200, SM100f): the warmup shapes that fault (1/2/8 tokens) were the only ones able to select tileN 8/16, while the 12288 token shape gets tileN 64/128 and always passed. After the change every shape from 1 to 12288 tokens offers only tileN >= 32, and all five Phase-2 warmup shapes complete where the second previously crashed. The tiles stay in mSupportedTileN: the ctor builds one runner per tile and each asserts a non-empty passing-config list, so the exclusion must happen at tactic-selection time rather than by dropping the tile. This is safe because every FP8 block-scale MoE shape retains a tileN >= 32 tactic -- verified across DeepSeek-R1 (EP 1/4/8), Qwen3-235B and Qwen3-30B -- so the tactic list is never emptied. The waiver is kept: an independent illegal memory access remains in the Phase-2 autotuner warmup, which this change does not address. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…-scale MoE The TRTLLM-Gen small-tile (tileN 8/16) dynB batched-GEMM cubins flakily hit an illegal memory access in the gemm2 K-loop. PR NVIDIA#15297 already added a WAR for this exact fault -- restricting tactics to tileN >= 32 -- but scoped it to the fused shared-expert path via num_fused_shared_experts > 0. Shared-expert fusion is opt-in (TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION, default off), so DeepSeek-R1 FP8 TP=8 runs unfused and never reached the exclusion. It faults identically: the defect is in the shared small-tile cubins, not caused by expert fusion. Apply the exclusion for every caller, at both selection sites (getValidConfigs and the tileN == -1 fallback in run()). Measured on DeepSeek-R1 FP8 EP=1 (B200, SM100f): the warmup shapes that fault (1/2/8 tokens) were the only ones able to select tileN 8/16, while the 12288 token shape gets tileN 64/128 and always passed. After the change every shape from 1 to 12288 tokens offers only tileN >= 32, and all five Phase-2 warmup shapes complete where the second previously crashed. The tiles stay in mSupportedTileN: the ctor builds one runner per tile and each asserts a non-empty passing-config list, so the exclusion must happen at tactic-selection time rather than by dropping the tile. This is safe because every FP8 block-scale MoE shape retains a tileN >= 32 tactic -- verified across DeepSeek-R1 (EP 1/4/8), Qwen3-235B and Qwen3-30B -- so the tactic list is never emptied. The waiver is kept: an independent illegal memory access remains in the Phase-2 autotuner warmup, which this change does not address. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…-scale MoE The TRTLLM-Gen small-tile (tileN 8/16) dynB batched-GEMM cubins flakily hit an illegal memory access in the gemm2 K-loop. PR NVIDIA#15297 already added a WAR for this exact fault -- restricting tactics to tileN >= 32 -- but scoped it to the fused shared-expert path via num_fused_shared_experts > 0. Shared-expert fusion is opt-in (TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION, default off), so DeepSeek-R1 FP8 TP=8 runs unfused and never reached the exclusion. It faults identically: the defect is in the shared small-tile cubins, not caused by expert fusion. Apply the exclusion for every caller, at both selection sites (getValidConfigs and the tileN == -1 fallback in run()). Measured on DeepSeek-R1 FP8 EP=1 (B200, SM100f): the warmup shapes that fault (1/2/8 tokens) were the only ones able to select tileN 8/16, while the 12288 token shape gets tileN 64/128 and always passed. After the change every shape from 1 to 12288 tokens offers only tileN >= 32, and all five Phase-2 warmup shapes complete where the second previously crashed. The tiles stay in mSupportedTileN: the ctor builds one runner per tile and each asserts a non-empty passing-config list, so the exclusion must happen at tactic-selection time rather than by dropping the tile. This is safe because every FP8 block-scale MoE shape retains a tileN >= 32 tactic -- verified across DeepSeek-R1 (EP 1/4/8), Qwen3-235B and Qwen3-30B -- so the tactic list is never emptied. The waiver is kept: an independent illegal memory access remains in the Phase-2 autotuner warmup, which this change does not address. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…-scale MoE The TRTLLM-Gen small-tile (tileN 8/16) dynB batched-GEMM cubins flakily hit an illegal memory access in the gemm2 K-loop. PR NVIDIA#15297 already added a WAR for this exact fault -- restricting tactics to tileN >= 32 -- but scoped it to the fused shared-expert path via num_fused_shared_experts > 0. Shared-expert fusion is opt-in (TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION, default off), so DeepSeek-R1 FP8 TP=8 runs unfused and never reached the exclusion. It faults identically: the defect is in the shared small-tile cubins, not caused by expert fusion. Apply the exclusion for every caller, at both selection sites (getValidConfigs and the tileN == -1 fallback in run()). Measured on DeepSeek-R1 FP8 EP=1 (B200, SM100f): the warmup shapes that fault (1/2/8 tokens) were the only ones able to select tileN 8/16, while the 12288 token shape gets tileN 64/128 and always passed. After the change every shape from 1 to 12288 tokens offers only tileN >= 32, and all five Phase-2 warmup shapes complete where the second previously crashed. The tiles stay in mSupportedTileN: the ctor builds one runner per tile and each asserts a non-empty passing-config list, so the exclusion must happen at tactic-selection time rather than by dropping the tile. This is safe because every FP8 block-scale MoE shape retains a tileN >= 32 tactic -- verified across DeepSeek-R1 (EP 1/4/8), Qwen3-235B and Qwen3-30B -- so the tactic list is never emptied. The waiver is kept: an independent illegal memory access remains in the Phase-2 autotuner warmup, which this change does not address. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…-scale MoE The TRTLLM-Gen small-tile (tileN 8/16) dynB batched-GEMM cubins flakily hit an illegal memory access in the gemm2 K-loop. PR NVIDIA#15297 already added a WAR for this exact fault -- restricting tactics to tileN >= 32 -- but scoped it to the fused shared-expert path via num_fused_shared_experts > 0. Shared-expert fusion is opt-in (TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION, default off), so DeepSeek-R1 FP8 TP=8 runs unfused and never reached the exclusion. It faults identically: the defect is in the shared small-tile cubins, not caused by expert fusion. Apply the exclusion for every caller, at both selection sites (getValidConfigs and the tileN == -1 fallback in run()). Measured on DeepSeek-R1 FP8 EP=1 (B200, SM100f): the warmup shapes that fault (1/2/8 tokens) were the only ones able to select tileN 8/16, while the 12288 token shape gets tileN 64/128 and always passed. After the change every shape from 1 to 12288 tokens offers only tileN >= 32, and all five Phase-2 warmup shapes complete where the second previously crashed. The tiles stay in mSupportedTileN: the ctor builds one runner per tile and each asserts a non-empty passing-config list, so the exclusion must happen at tactic-selection time rather than by dropping the tile. This is safe because every FP8 block-scale MoE shape retains a tileN >= 32 tactic -- verified across DeepSeek-R1 (EP 1/4/8), Qwen3-235B and Qwen3-30B -- so the tactic list is never emptied. The waiver is kept: an independent illegal memory access remains in the Phase-2 autotuner warmup, which this change does not address. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
Summary by CodeRabbit (updated)
l0_b200andl0_b300CI test lists.Dev Engineer Review
Code changes
WarpSizerange.num_fused_shared_expertsthrough C++, Torch, custom-op, and backend layers.TLLM_MOE_FUSED_MIN_TILEN.Configuration and test-list review
l0_b200.ymlandl0_b300.ymladd the two matching test functions.TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION.acb5fe4. A failed helper run also occurred for commit56c493e.QA Engineer Review
Test-list changes
Modified files:
tests/integration/test_lists/test-db/l0_b200.ymltests/integration/test_lists/test-db/l0_b300.ymlAdded entries:
test_trtllm_fp8_block_scales_fused_shared_expertstest_trtllm_fp8_block_scales_fuse_shared_expert_layoutBoth tests are registered in both CI test lists.
Test-code changes
Added to
tests/unittest/_torch/modules/moe/test_moe_backend.py:test_trtllm_fp8_block_scales_fused_shared_expertstest_trtllm_fp8_block_scales_fuse_shared_expert_layoutThe tests cover accuracy, tactic capture and replay, autotuner candidates, and fused weight and scale layout.
The microbenchmark changes add fused and unfused shared-expert execution paths. They are not registered in CI or QA test lists.
Verdict: needs follow-up. CBTS coverage data is unavailable, and several reported CI pipelines failed.
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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.