[None][feat] DSv4 prep: MoE routing and backend support#15402
Conversation
|
/bot run --disable-fail-fast |
|
/bot kill |
41dcd0e to
148479f
Compare
|
PR_Github #54519 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughAdds DeepSeek-V4 MoE support: a new ChangesDeepSeek-V4 MoE: Gate Kernel, Routing, SwiGLU Clamp, input_ids Threading
Sequence Diagram(s)sequenceDiagram
participant Model as Model forward
participant ConfigMoE as ConfigurableMoE.forward_impl
participant Scheduler as ExternalCommMoEScheduler.forward
participant Backend as TRTLLMGenFusedMoE.forward_impl
participant Routing as DeepSeekV4MoeRoutingMethod.apply
participant GateOp as trtllm::gate_forward (CUDA)
participant RunMoE as run_fp8_block_scale_moe
Model->>ConfigMoE: x, router_logits, input_ids=input_ids
ConfigMoE->>Scheduler: forward(x, router_logits, input_ids=input_ids)
Scheduler->>Backend: forward_impl(x, router_logits, input_ids=input_ids)
Backend->>Routing: apply(router_logits, input_ids) [separated routing]
Routing->>GateOp: gate_forward(scores_in, bias, input_ids, tid2eid, out_weights, out_indices, topk, route_scale, is_hash)
GateOp-->>Routing: out_weights, out_indices written
Routing-->>Backend: token_selected_experts, token_final_scales
Backend->>RunMoE: run_fp8_block_scale_moe(..., gemm1_clamp_limit=swiglu_limit_scalar)
RunMoE-->>Backend: MoE output
Backend-->>Scheduler: output tensor
Scheduler-->>ConfigMoE: output tensor
ConfigMoE-->>Model: output tensor
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Suggested reviewers
✨ Finishing Touches🧪 Generate unit tests (beta)
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 12
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py (1)
1137-1142:⚠️ Potential issue | 🟠 Major | ⚡ Quick winKeep
input_idsoptional at theforward_chunkboundary.This is typed as optional but has no default, so existing callers of
forward_chunk(x, router_logits, output_dtype=...)now fail before reaching the DSv4-aware path.Proposed fix
def forward_chunk( self, x: Union[torch.Tensor, Fp4QuantizedTensor], router_logits: torch.Tensor, - input_ids: Optional[torch.IntTensor], + input_ids: Optional[torch.IntTensor] = None, output_dtype: Optional[torch.dtype] = None,🤖 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/modules/fused_moe/fused_moe_cutlass.py` around lines 1137 - 1142, The forward_chunk method has an input_ids parameter typed as Optional[torch.IntTensor] but lacks a default value, causing existing callers that skip this parameter to fail. Add a default value of None to the input_ids parameter definition to make it truly optional and backward compatible with callers that use forward_chunk(x, router_logits, output_dtype=...) without providing input_ids.cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h (1)
2-2:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winUpdate the header year for this modified C++ header.
This file is modified in this PR, but the copyright range still ends at 2025.
Proposed fix
- * Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved.As per coding guidelines, “All TensorRT-LLM source files should contain an NVIDIA copyright header including the year of latest meaningful modification.”
🤖 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.h` at line 2, Update the copyright year range in the header comment at the top of runner.h to reflect the current year, since this file is being modified in this PR. Change the copyright line from ending at 2025 to the current year to comply with TensorRT-LLM coding guidelines that require the copyright header to include the year of the latest meaningful modification.Source: Coding guidelines
tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py (1)
441-443:⚠️ Potential issue | 🟠 Major | ⚡ Quick winReport routing-method separated routing here too.
DeepSeekV4 sets
routing_method.requires_separated_routing = True, andforward_impl()honors it, but_supports_load_balancer()still returnsFalsefor quantized/non-DP cases. Any caller using this flag to decide whether top-k routing is external will misclassify DeepSeekV4.Proposed fix
""" + if self.routing_method.requires_separated_routing: + return True if self._requires_separated_routing(): return True return self.use_dp and self.parallel_size > 1🤖 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/modules/fused_moe/fused_moe_trtllm_gen.py` around lines 441 - 443, The `_supports_load_balancer()` method needs to consistently report when separated routing is required. Currently it returns True when `_requires_separated_routing()` is true in the first condition, but ensure this logic properly propagates to all code paths so that callers using this method to determine whether top-k routing is external (such as for DeepSeekV4 which sets `routing_method.requires_separated_routing = True`) correctly classify cases with separated routing enabled. Verify that the separated routing requirement is properly accounted for in the return statement logic, making sure the method reports True whenever separated routing is in effect regardless of the DP or parallel_size configuration.tensorrt_llm/_torch/modules/fused_moe/create_moe.py (1)
428-445:⚠️ Potential issue | 🔴 CriticalTritonFusedMoE should not accept
swiglu_limit_scalar, but validation allows it.
TritonFusedMoEis included in the validation check at line 302, which permits bothswiglu_limitandswiglu_limit_scalarto be provided. However, the constructor signature (at line 1466) only acceptsswiglu_limit,swiglu_alpha, andswiglu_beta— there is noswiglu_limit_scalarparameter. The constructor call at line 428-444 does not passswiglu_limit_scalar, whereasDeepGemmFusedMoE(line 414-421) receives both parameters.The implementation processes only per-expert tensors (
swiglu_limit), not scalar values. Either removeTritonFusedMoEfrom the validation list at line 302, or add explicit rejection logic ifswiglu_limit_scalaris provided whenmoe_cls == TritonFusedMoE.🤖 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/modules/fused_moe/create_moe.py` around lines 428 - 445, There is a mismatch between validation logic and the TritonFusedMoE constructor signature. The validation logic permits both swiglu_limit and swiglu_limit_scalar for TritonFusedMoE, but the TritonFusedMoE constructor only accepts swiglu_limit, swiglu_alpha, and swiglu_beta—not swiglu_limit_scalar. Either remove TritonFusedMoE from the validation block that permits swiglu_limit_scalar to be provided together with swiglu_limit, or add an explicit assertion or error check in the TritonFusedMoE branch that rejects swiglu_limit_scalar if it is provided, similar to how the assertion already rejects apply_router_weight_on_input.
🧹 Nitpick comments (3)
tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py (1)
624-626: ⚡ Quick winAdd
strict=Trueto thezip()call for defensive length-check.While
x_list,router_logits_list, andinput_ids_listare all derived from the samechunk_size_listand the empty-chunk substitution loop (lines 607–620) maintains synchronization, making the length match explicit improves maintainability and guards against future refactoring mistakes.♻️ Proposed fix
- for idx_chunk, (x_chunk, router_logits_chunk, input_ids_chunk) in enumerate( - zip(x_list, router_logits_list, input_ids_list) - ): + for idx_chunk, (x_chunk, router_logits_chunk, input_ids_chunk) in enumerate( + zip(x_list, router_logits_list, input_ids_list, strict=True) + ):🤖 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/modules/fused_moe/moe_scheduler.py` around lines 624 - 626, Add the `strict=True` parameter to the zip() call that iterates over x_list, router_logits_list, and input_ids_list to enforce that all three lists have equal length and protect against future refactoring mistakes that could cause length mismatches.cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu (1)
313-313: ⚡ Quick winConsider centralizing
kTOPKconstant.
kTOPK = 6is duplicated incustomMoeRoutingKernels.cu(line 313) andmoeGateOp.cpp(line 29). If this value changes, both files must be updated.♻️ Suggested refactor
Define
kTOPKin the header and reference it from both files:// In customMoeRoutingKernels.h +static constexpr int kGateForwardTopK = 6; + void gate_forward(void* scores_in, ...); // In customMoeRoutingKernels.cu -static constexpr int kTOPK = 6; +static constexpr int kTOPK = kGateForwardTopK; // In moeGateOp.cpp -static constexpr int kTOPK = 6; +static constexpr int kTOPK = tensorrt_llm::kernels::kGateForwardTopK;Also applies to: 437-481
🤖 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/customMoeRoutingKernels.cu` at line 313, The kTOPK constant is duplicated in customMoeRoutingKernels.cu (line 313) and moeGateOp.cpp (line 29), creating maintenance issues if the value needs to change. Define kTOPK as a static constexpr int in a shared header file that both customMoeRoutingKernels.cu and moeGateOp.cpp can include, then remove the duplicate static constexpr int kTOPK = 6 definitions from both implementation files and replace them with includes of the header. This ensures a single source of truth for the constant value.tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py (1)
1019-1023: Movegemm1_clamp_limitafteruse_dpto preserve backward compatibility and avoid future positional-argument binding issues.The current parameter order places
gemm1_clamp_limitbeforeoutput, which could cause silent misinterpretation if future callers use positional arguments beyondtopk_ids. All current call sites are safe (they use keyword arguments or stop attopk_ids), but reordering prevents potential breakage if the calling pattern evolves.Apply the same parameter order to both the main function definition and the
register_fakeimplementation.Proposed adjustment
Move
gemm1_clamp_limitto the end of the parameter list:act_type: int = 0, - gemm1_clamp_limit: Optional[float] = None, output: Optional[torch.Tensor] = None, tune_max_num_tokens: int = 8192, - use_dp: bool = False) -> torch.Tensor: + use_dp: bool = False, + gemm1_clamp_limit: Optional[float] = None) -> torch.Tensor:Apply to both lines 1019–1023 and 1120–1124.
🤖 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 1019 - 1023, The function signature has gemm1_clamp_limit positioned before output, which compromises backward compatibility with positional arguments. Move the gemm1_clamp_limit parameter to the end of the parameter list (after use_dp) in both the main function definition and the register_fake implementation to ensure positional argument safety and prevent potential silent misinterpretation if calling patterns evolve in the future.
🤖 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/customMoeRoutingKernels.cu`:
- Around line 333-334: The constexpr variable kExpertsPerThread at line 333
divides nExperts by WARP_SIZE without validating that nExperts is divisible by
WARP_SIZE. While current supported expert counts (256, 384) happen to work, add
a static_assert immediately before the kExpertsPerThread declaration to enforce
the compile-time constraint that nExperts must be evenly divisible by WARP_SIZE,
preventing future bugs if new expert counts are added that violate this
assumption.
- Around line 372-379: The kernel reads token_id from input_ids[global_warp_id]
and directly indexes into tid2eid without validating that token_id is within
bounds (token_id < vocab_size). Add vocab_size as a kernel parameter and insert
bounds checking before the line expert_ids = tid2eid + token_id * topK to ensure
token_id does not exceed the valid range. If validation fails, handle it
appropriately (skip processing or set a default value). Alternatively, if
runtime bounds checking is not feasible, add clear documentation in the kernel
and Python wrapper stating that input_ids must contain valid token indices
within [0, vocab_size), and implement corresponding validation in the Python
wrapper before kernel invocation.
In
`@cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_problem_builder.cu`:
- Line 72: The out-GEMM problem size calculation in the GemmCoord instantiation
now unconditionally sets N=out_hidden_size, but the test reference code expects
N to be 0 when rank=0 via a conditional (rank > 0) ? out_hidden_size : 0. Since
the test vector includes rank=0 cases and compares against this reference at the
verification point, the test will fail. Either revert the kernel code to use the
same conditional logic to match the test reference, or update the test reference
to use the new unconditional behavior and verify that CUTLASS correctly handles
the rank=0 GEMM case.
In `@tensorrt_llm/_torch/modules/fused_moe/create_moe.py`:
- Around line 300-304: The assertion validation logic checks both the
swiglu_limit and swiglu_limit_scalar parameters in the condition, but the error
message only references swiglu_limit. Update the error message string in the
assertion to mention both parameters so users providing either swiglu_limit or
swiglu_limit_scalar receive a clear and complete error message when using an
unsupported backend with the moe_cls validation.
In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py`:
- Around line 1504-1505: Add `strict=True` parameter to both `zip()` calls to
enforce that the three lists being zipped have identical length. In
`tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py` lines 1504-1505,
modify the `zip(x_list, router_logits_list, input_ids_list)` call to include
`strict=True`. In `tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py`
lines 1285-1286, apply the same change to the corresponding `zip()` call that
combines three chunk-aligned lists. This prevents silent truncation when lists
are expected to have matching lengths by design.
In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py`:
- Around line 469-478: The validation assertion currently allows per-expert
swiglu_limit (tensor variant) when has_deepseek_fp8_block_scales is true, but
the FP8 runner at line 772 only forwards swiglu_limit_scalar (the scalar
variant), causing the per-expert tensor to be silently ignored. Remove
has_deepseek_fp8_block_scales from the assertion condition to reject per-expert
swiglu_limit on the FP8 path, ensuring only swiglu_limit_scalar is used with the
deepseek_fp8_block_scales configuration.
In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py`:
- Around line 789-790: Add strict=True parameter to the zip() call that iterates
over x_list, router_logits_list, and input_ids_list to ensure all three lists
have the same length and prevent silent truncation of mismatched items. This
will make the code raise an error immediately if the lists diverge in length
rather than silently dropping trailing chunks.
In `@tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py`:
- Around line 202-207: The function signature accepts activation_clamp through
**kwargs but no longer has it as an explicit parameter, causing existing code
that passes activation_clamp= to silently drop that argument with no warning.
Add logic after the function signature to check if activation_clamp exists in
kwargs, then either map it to swiglu_limit_scalar (if swiglu_limit_scalar is
None) or raise an error if both parameters are provided simultaneously. Remove
activation_clamp from kwargs after handling it to prevent it from being passed
downstream where it is not expected. This preserves backward compatibility while
preventing silent failures for callers using the deprecated parameter name.
- Around line 395-449: The exception handling is too broad in two places. First,
at the MPI local-rank discovery block (where comm.Split_type and Get_rank are
called), change the `except Exception:` clause to specifically catch
`mpi4py.MPI.Exception` to avoid masking unexpected failures. Second, at the
dist.init_process_group call, the bare `except RuntimeError:` assumes all
RuntimeErrors are port-race collisions and automatically retries. Instead, catch
RuntimeError as exc, check if the error message contains "address already in
use" or "eaddrinuse" (case-insensitive), and only proceed with the retry logic
if this condition is true; otherwise re-raise the exception to surface actual
NCCL or environment configuration errors.
In `@tensorrt_llm/_torch/modules/fused_moe/quantization.py`:
- Around line 5353-5357: The loop that empties the Parameters (w3_w1_weight,
w3_w1_weight_scale, w2_weight_scale) does not record metadata before reseating
them. Since the class inherits pre_reload_weights() which only rebuilds
parameters listed in module.rebuild_tensor_metadata, these buffers cannot be
properly restored during reload. Before emptying these parameters in the loop,
capture their original metadata by recording their shapes and dtypes in
module.rebuild_tensor_metadata, or use an existing replacement helper method to
maintain the reload contract. This ensures that when load_weights() is called
later, it can properly restore these tensors before copying data into them.
In `@tensorrt_llm/_torch/modules/fused_moe/routing.py`:
- Around line 553-558: The hashed routing path in the gate_forward call passes
self.callable_e_score_correction_bias() directly to
torch.ops.trtllm.gate_forward, but this callable can return None when there is
no bias (as documented for hashed DeepSeekV4). The gate_forward operation
expects a torch.Tensor for the bias parameter, not None. Before calling
gate_forward in the is_hashed block, normalize the bias by checking if
callable_e_score_correction_bias() returns None and providing an appropriate
zero tensor or default tensor value that matches the expected shape and device,
ensuring the dispatcher receives a valid tensor instead of None.
In `@tensorrt_llm/quantization/utils/fp8_utils.py`:
- Around line 497-499: The swiglu_limit activation semantics are inconsistent
across three files, treating zero and negative values destructively. In
tensorrt_llm/quantization/utils/fp8_utils.py lines 497-499, change the
has_swiglu_limit check from testing "swiglu_limit is not None" to "swiglu_limit
is not None and swiglu_limit > 0.0" (optionally also require math.isfinite for
additional robustness). In cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp lines
172-183, modify the logic that sets args.has_gemm1_clamp_limit_value to only
enable it when the limit is positive; otherwise explicitly clear both the flag
and value. In tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py lines
613-617, change the NotImplementedError condition to only trigger for positive
clamp limits, treating None or non-positive values as a no-op clamp instead of
an error.
---
Outside diff comments:
In `@cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h`:
- Line 2: Update the copyright year range in the header comment at the top of
runner.h to reflect the current year, since this file is being modified in this
PR. Change the copyright line from ending at 2025 to the current year to comply
with TensorRT-LLM coding guidelines that require the copyright header to include
the year of the latest meaningful modification.
In `@tensorrt_llm/_torch/modules/fused_moe/create_moe.py`:
- Around line 428-445: There is a mismatch between validation logic and the
TritonFusedMoE constructor signature. The validation logic permits both
swiglu_limit and swiglu_limit_scalar for TritonFusedMoE, but the TritonFusedMoE
constructor only accepts swiglu_limit, swiglu_alpha, and swiglu_beta—not
swiglu_limit_scalar. Either remove TritonFusedMoE from the validation block that
permits swiglu_limit_scalar to be provided together with swiglu_limit, or add an
explicit assertion or error check in the TritonFusedMoE branch that rejects
swiglu_limit_scalar if it is provided, similar to how the assertion already
rejects apply_router_weight_on_input.
In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py`:
- Around line 1137-1142: The forward_chunk method has an input_ids parameter
typed as Optional[torch.IntTensor] but lacks a default value, causing existing
callers that skip this parameter to fail. Add a default value of None to the
input_ids parameter definition to make it truly optional and backward compatible
with callers that use forward_chunk(x, router_logits, output_dtype=...) without
providing input_ids.
In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py`:
- Around line 441-443: The `_supports_load_balancer()` method needs to
consistently report when separated routing is required. Currently it returns
True when `_requires_separated_routing()` is true in the first condition, but
ensure this logic properly propagates to all code paths so that callers using
this method to determine whether top-k routing is external (such as for
DeepSeekV4 which sets `routing_method.requires_separated_routing = True`)
correctly classify cases with separated routing enabled. Verify that the
separated routing requirement is properly accounted for in the return statement
logic, making sure the method reports True whenever separated routing is in
effect regardless of the DP or parallel_size configuration.
---
Nitpick comments:
In `@cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu`:
- Line 313: The kTOPK constant is duplicated in customMoeRoutingKernels.cu (line
313) and moeGateOp.cpp (line 29), creating maintenance issues if the value needs
to change. Define kTOPK as a static constexpr int in a shared header file that
both customMoeRoutingKernels.cu and moeGateOp.cpp can include, then remove the
duplicate static constexpr int kTOPK = 6 definitions from both implementation
files and replace them with includes of the header. This ensures a single source
of truth for the constant value.
In `@tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py`:
- Around line 1019-1023: The function signature has gemm1_clamp_limit positioned
before output, which compromises backward compatibility with positional
arguments. Move the gemm1_clamp_limit parameter to the end of the parameter list
(after use_dp) in both the main function definition and the register_fake
implementation to ensure positional argument safety and prevent potential silent
misinterpretation if calling patterns evolve in the future.
In `@tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py`:
- Around line 624-626: Add the `strict=True` parameter to the zip() call that
iterates over x_list, router_logits_list, and input_ids_list to enforce that all
three lists have equal length and protect against future refactoring mistakes
that could cause length mismatches.
🪄 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: 42c24abb-ea2e-42f4-b47d-e3167db8586f
📒 Files selected for processing (49)
3rdparty/fetch_content.jsoncpp/tensorrt_llm/kernels/customMoeRoutingKernels.cucpp/tensorrt_llm/kernels/customMoeRoutingKernels.hcpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_heuristic.cppcpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_problem_builder.cucpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cucpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cucpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.hcpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernelTopK.cuhcpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cucpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.hcpp/tensorrt_llm/thop/CMakeLists.txtcpp/tensorrt_llm/thop/dsv3RouterGemmOp.cppcpp/tensorrt_llm/thop/fp8BlockScaleMoe.cppcpp/tensorrt_llm/thop/moeGateOp.cpptensorrt_llm/_torch/custom_ops/cpp_custom_ops.pytensorrt_llm/_torch/custom_ops/torch_custom_ops.pytensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.pytensorrt_llm/_torch/distributed/moe_alltoall.pytensorrt_llm/_torch/models/modeling_gemma4.pytensorrt_llm/_torch/modules/fused_moe/__init__.pytensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.pytensorrt_llm/_torch/modules/fused_moe/configurable_moe.pytensorrt_llm/_torch/modules/fused_moe/create_moe.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.pytensorrt_llm/_torch/modules/fused_moe/interface.pytensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.pytensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.pytensorrt_llm/_torch/modules/fused_moe/moe_op_backend.pytensorrt_llm/_torch/modules/fused_moe/moe_scheduler.pytensorrt_llm/_torch/modules/fused_moe/ops/moe_op_deepgemm.pytensorrt_llm/_torch/modules/fused_moe/quantization.pytensorrt_llm/_torch/modules/fused_moe/routing.pytensorrt_llm/_torch/modules/gated_mlp.pytensorrt_llm/_torch/modules/linear.pytensorrt_llm/_torch/modules/swiglu.pytensorrt_llm/quantization/utils/fp8_utils.pytensorrt_llm/tools/layer_wise_benchmarks/runner.pytests/unittest/_torch/modules/moe/test_moe_module.pytests/unittest/_torch/modules/test_moe_load_balancer.pytests/unittest/_torch/modules/test_moe_routing.pytests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.pytests/unittest/_torch/thop/serial/test_moe_gate.py
💤 Files with no reviewable changes (1)
- tensorrt_llm/_torch/modules/linear.py
|
PR_Github #54520 [ kill ] triggered by Bot. Commit: |
|
PR_Github #54519 [ run ] completed with state |
|
PR_Github #54520 [ kill ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #54525 [ run ] triggered by Bot. Commit: |
|
PR_Github #54525 [ run ] completed with state
|
|
Suggestion: collapse swiglu_limit_scalar into swiglu_limit; don't add it to backends that consume the per-expert tensor. swiglu_limit (per-expert tensor) and swiglu_limit_scalar (uniform float) are the same concept (the SwiGLU clamp limit) in two representations. The split leaks an internal kernel-ABI difference into the public module/backend constructors, and the scalar is plumbed into backends that never read it. Who actually consumes the scalar — only the FP8 separate-activation path: fused_moe_deepgemm.py:726 → silu_and_mul_masked_post_quant_fwd(swiglu_limit=self.swiglu_limit_scalar) Why one public param is enough: the model always provides a uniform clamp (DSv4 config.swiglu_limit = 10.0). The per-expert tensor is derived internally for NVFP4 only, because the clamp is folded with per-expert fc31_alpha (quantization.py:3954 swiglu_limit.data.div_(fc31_alpha)) — that division is what makes it genuinely per-expert, not the public API. And the tensor→scalar reduction already exists: mega_moe_cute_dsl.py:_resolve_gate_up_clamp takes the per-layer swiglu_limit tensor, asserts uniformity, and reduces to one float. Proposed shape: Keep a single public swiglu_limit: Union[float, torch.Tensor] (or just float, since both known consumers — DSv4 and GPT-OSS — use a uniform limit). Drop swiglu_limit_scalar from MoE.init and from every backend constructor. **For the MOE API part, I will do the refactor in TRTLLM-13501 -- [TRTLLM-13501](https://jirasw.nvidia.com/browse/TRTLLM-13501)** |
xxi-nv
left a comment
xxi-nv
left a comment
There was a problem hiding this comment.
To unblock the merge, For the MOE API part, I will do the refactor in
TRTLLM-13501 -- TRTLLM-13501
148479f to
42aeb73
Compare
|
PR_Github #55294 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #55311 [ run ] triggered by Bot. Commit: |
|
PR_Github #55311 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #55319 [ run ] triggered by Bot. Commit: |
|
PR_Github #55319 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #55321 [ run ] triggered by Bot. Commit: |
|
PR_Github #55321 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #55325 [ run ] triggered by Bot. Commit: |
|
PR_Github #55325 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
yuanjingx87
left a comment
yuanjingx87
left a comment
There was a problem hiding this comment.
approve on oss-compliance perspective.
|
PR_Github #55347 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #55351 [ run ] triggered by Bot. Commit: |
|
PR_Github #55347 [ run ] completed with state |
|
PR_Github #55351 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #55448 [ run ] triggered by Bot. Commit: |
|
PR_Github #55448 [ run ] completed with state |
… import Between 1.3.0rc19 and rc20, an import-time deep_gemm.set_pdl() call (added in NVIDIA#15402) instantiated DeepGEMM's DeviceRuntime (cuBLASLt handle + 32 MiB workspace tensor), creating a CUDA context (~0.5-1.2 GiB with loaded modules, arch-dependent) in every process importing tensorrt_llm: the trtllm-bench parent (which never launches a kernel) and every MPI worker on its default device before torch.cuda.set_device. Those contexts are resident when the KV cache pool is sized from free GPU memory, shrinking the pool (RTX 6000D: 48.33 -> 47.39 GiB, max concurrent requests 513 -> 503) and causing ~10% throughput regressions on memory-tight GPUs (this bug and nvbugs 6405760, 6418453, 6419078) via a long low-batch scheduling tail. The root cause was fixed on main by NVIDIA#15632, which defers PDL configuration to PyTorchModelEngine init. Add a regression test asserting that importing tensorrt_llm neither configures DeepGEMM PDL nor creates a CUDA context, so the import path stays free of device side effects. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com> Signed-off-by: xxi <xxi@nvidia.com> Co-authored-by: Fanrong Li <lfr-0531@users.noreply.github.com> Co-authored-by: xxi <xxi@nvidia.com> Co-authored-by: xingfei xi <95731198+xxi-nv@users.noreply.github.com>
6 participants
Summary by CodeRabbit
Release Notes
New Features
Improvements
Tests
Description
This is part of the DSv4 split from umbrella PR #14751. It isolates shared MoE/routing backend support so it can be reviewed independently from the DSv4 model and sparse attention wiring.
Included here:
gate_forward,dsv3_router_gemm_op, andRoutingKernelTopK.cuh.swiglu_limitplumbing, and DeepSeekV4 routing interface changes.CUTLASS + NVFP4on SM100-family fallback.This PR has no dependency on earlier DSv4 split PRs after current
main. PR-9 will depend on this PR for the final DSv4 model/API/tokenizer wiring.Intentionally excluded: DSv4 model/config/tokenizer files, sparse attention backend/cache-manager files, compressor/mHC implementation, IndexerTopK implementation, and attention fusion custom ops.
Risk: this touches shared MoE backend code used by non-DSv4 consumers. The default-path behavior is covered by existing MoE module/routing tests below.
Test Coverage
Build/install:
CCACHE_DIR=/home/scratch.fanrongl_coreai/ccache_trtllm python3 ./scripts/build_wheel.py --trt_root /usr/local/tensorrt --benchmarks --use_ccache --cuda_architectures "90-real;100-real" --configure_cmake./.venv-3.12/bin/pip install --force-reinstall --no-deps build/tensorrt_llm-1.3.0rc19-cp312-cp312-linux_x86_64.whlSmoke/static checks:
tensorrt_llm.__file__resolves to this PR worktree.fused_moe,moe_alltoall.py,torch.ops.trtllm.gate_forward,torch.ops.trtllm.dsv3_router_gemm_op,torch.ops.trtllm.fp8_block_scale_moe_runner, andtensorrt_llm.deep_gemm_cpp_tllm.DeepSeekV4MoeRoutingMethod.apply(..., input_ids=...)accepts optionalinput_ids;Gemma4MoeRoutingMethod.applyis restored toapply(router_logits)and PR-8 no longer changesmodeling_gemma4.py.git diff --check github/main..HEAD./.venv-3.12/bin/pre-commit run --files $(git diff --name-only HEAD)before the final rebase.GPU tests after checking
nvidia-smi; used idleCUDA_VISIBLE_DEVICES=1:timeout 25m ./.venv-3.12/bin/python -m pytest tests/unittest/_torch/modules/moe/test_moe_module.py tests/unittest/_torch/thop/serial/test_moe_gate.py-> 270 passed, 2116 skipped.timeout 20m ./.venv-3.12/bin/python -m pytest tests/unittest/_torch/modules/test_moe_routing.py tests/unittest/_torch/modules/test_moe_load_balancer.py-> 527 passed, 87 skipped.timeout 20m ./.venv-3.12/bin/python -m pytest tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py-> 48 passed.After rebasing onto latest
main(0ec325040087af5f98d040fe1dcf59273c19d209), new main only touchedtests/integration/defs/perf/utils.pyand did not overlap PR-8 files. Rebase verification:modeling_gemma4.pyis absent fromgithub/main..HEAD.git diff --check github/main..HEADpassed.timeout 20m ./.venv-3.12/bin/python -m pytest tests/unittest/_torch/modules/test_moe_routing.py tests/unittest/_torch/modules/test_moe_load_balancer.py tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py-> 575 passed, 87 skipped.CUDA_VISIBLE_DEVICES=1:timeout 45m ./.venv-3.12/bin/python -m pytest -q tests/unittest/_torch/modules/test_moe_routing.py tests/unittest/_torch/modules/test_moe_load_balancer.py tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py-> 575 passed, 87 skipped.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.