[None][chore] Simplify MoE lora - #15330
Conversation
0fecc97 to
653d04f
Compare
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthroughThe PR replaces the legacy MoE LoRA "device-path / cuBLAS" execution with a unified capture-safe "grouped-GEMM" core. C++ header types are renamed ( ChangesMoE LoRA Device-Path → Grouped-GEMM Migration
Sequence Diagram(s)sequenceDiagram
participant PY as fused_moe (Python)
participant OP as FusedMoeRunner (C++ TorchBind)
participant BUILD as buildMoeLoraParams
participant SCRATCH as ensureLoraDeviceScratch
participant KERNEL as moe_kernels.cu
PY->>OP: run_moe(..., lora_ranks, lora_weight_ptrs, token_to_slot)
OP->>OP: capturability check (reject per-request under CUDA graph)
OP->>BUILD: lora tensors + slot-indexed data
BUILD->>SCRATCH: allocate LoraGroupedGemmBuffers (kGroupedGemmSplitKSlices)
SCRATCH-->>BUILD: fc1/fc2/gated device buffers
BUILD->>BUILD: populateLoraGroupedGemmModule(fc1, fc2, gated)
BUILD->>BUILD: grouped_gemm.enabled=true, run=moeLoraGroupedGemmRunImpl
BUILD-->>OP: LoraParams{grouped_gemm}
OP->>KERNEL: launchMoeKernel(LoraParams)
KERNEL->>KERNEL: setupLoraWorkspace → grouped-GEMM early-return branch
KERNEL->>KERNEL: loraFC1 → zero + runMoeLoraGroupedGemmModule(fc1[, gated])
KERNEL->>KERNEL: loraFC2 → zero + runMoeLoraGroupedGemmModule(fc2)
KERNEL-->>OP: output
OP-->>PY: output tensor
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu (1)
3735-3767:⚠️ Potential issue | 🟠 Major | ⚡ Quick winValidate grouped-GEMM metadata before pointer expansion.
This branch now trusts
dp.dtype_bytesand each module’s dimensions/buffers to compute byte offsets inlaunchMoeLoraPointerExpand. If the caller passes stale dimensions, a mismatched dtype size, or a partially populated module, the result is wrong expert weight addresses or a device illegal access before the later GEMM wrapper guard can help.🛡️ Proposed validation before constructing expand modules
auto const& dp = lora_params.grouped_gemm; + int64_t const expectedDtypeBytes = static_cast<int64_t>(sizeof(ScaleBiasType)); + TLLM_CHECK_WITH_INFO(dp.dtype_bytes == expectedDtypeBytes, + "Grouped-GEMM LoRA dtype_bytes does not match the runner scalar type."); + + auto validateLoraModule = [](MoeLoraGroupedGemmModule const& mod, char const* name, + int64_t expectedDimA, int64_t expectedDimB) + { + TLLM_CHECK_WITH_INFO(mod.ranks_src_dev != nullptr && mod.ptrs_src_dev != nullptr + && mod.permuted_ranks_dev != nullptr && mod.permuted_ptrs_dev != nullptr, + "Grouped-GEMM LoRA %s module is missing pointer-expand buffers.", name); + TLLM_CHECK_WITH_INFO(mod.dim_a == expectedDimA && mod.dim_b == expectedDimB, + "Grouped-GEMM LoRA %s module dimensions do not match the MoE runner dimensions.", name); + }; + validateLoraModule(dp.fc1, "fc1", hidden_size, inter_size); + validateLoraModule(dp.fc2, "fc2", inter_size, hidden_size); + if (is_gated_activation) + { + validateLoraModule(dp.gated, "gated", hidden_size, inter_size); + } + // Translate per-module grouped-GEMM metadata into the MoeLoraExpandModule🤖 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/cutlass_kernels/moe_gemm/moe_kernels.cu` around lines 3735 - 3767, Add validation of the grouped-GEMM metadata (dp) before constructing the MoeLoraExpandModule objects fc1_mod, fc2_mod, and gated_mod. Validate that dp.dtype_bytes is valid and that each module's dimensions (dim_a, dim_b) and device buffers (ranks_src_dev, ptrs_src_dev, permuted_ranks_dev, permuted_ptrs_dev) are properly populated and non-null for fc1 and fc2. For gated_mod, perform the same validation only when is_gated_activation is true. This validation should occur before constructing any of the expand modules and before calling launchMoeLoraPointerExpand to prevent incorrect expert weight address computation or device illegal memory access.tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py (1)
516-523:⚠️ Potential issue | 🟠 Major | ⚡ Quick winDetect CUDA-graph MoE LoRA params in
_moe_lora_active.
_extract_moe_lora_tensorshandlesuse_cuda_graph_modewithout a per-layer eager dict, but_moe_lora_activereturnsFalsebefore checkingcuda_graph_params. That lets slot-indexed MoE LoRA bypass the multi-chunk rejection and stray-param guard.Proposed fix
def _moe_lora_active(self, lora_params: Optional[Dict]) -> bool: """Return True when lora_params carries routed-expert MoE LoRA tensors for this layer, meaning run_moe would fuse a LoRA delta. """ if not lora_params or self.layer_idx is None: return False + if lora_params.get("use_cuda_graph_mode", False): + cuda_graph_params = lora_params.get("cuda_graph_params") + if cuda_graph_params is None: + return False + layer_module2key = getattr(cuda_graph_params, "layer_module2key", {}) + return any( + (self.layer_idx, int(LoraModuleType.from_string(name))) in layer_module2key + for name in MOE_LORA_MODULE_NAMES + ) layer_params = lora_params.get(self.layer_idx, {}) if not layer_params: return False🤖 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 516 - 523, The `_moe_lora_active` method currently only checks `lora_params` to determine if MoE LoRA is active, but it should also check `cuda_graph_params` since `_extract_moe_lora_tensors` handles use_cuda_graph_mode without a per-layer eager dict. Update the method to check both `lora_params` and `cuda_graph_params` dictionaries, returning True if any MoE LoRA module types are found in either parameter dictionary for the current layer index.
🤖 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 `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py`:
- Around line 413-416: The targets set constructed from lora_config's
lora_target_modules (at line 413) uses raw strings without normalization, but
the has_moe_lora_targets check at line 398 normalizes strings to lowercase for
comparison. This inconsistency can cause mixed-case configs to behave
incorrectly. Normalize the targets by converting each string in
lora_target_modules to lowercase when constructing the targets set, ensuring
consistency with the normalization performed in has_moe_lora_targets.
In `@tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_cutlass.py`:
- Around line 207-226: The non-min-latency code path for run_moe is only
appending 10 arguments but the C++ binding expects 17 trailing arguments in
total. Update the run_moe_args list extension in the if not min_latency_mode
block to include all 17 required positional arguments by adding 7 more None
values and 0 values to match the full schema expected by the TorchBind binding,
even though all LoRA-related arguments should be None or 0 since LoRA is not
applied in compute_moe.
In `@tests/unittest/llmapi/test_llm_pytorch.py`:
- Around line 1200-1202: Update the docstring comment that describes the test's
expectation about eager and CUDA-graph mode outputs to clarify that they are not
expected to be bit-identical, aligning it with the actual test assertion logic
that only verifies general correctness rather than exact token-for-token
matching. The comment currently contradicts what the test actually validates, so
revise the description to accurately reflect that the outputs may differ
slightly due to floating-point precision while still being functionally correct.
---
Outside diff comments:
In `@cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu`:
- Around line 3735-3767: Add validation of the grouped-GEMM metadata (dp) before
constructing the MoeLoraExpandModule objects fc1_mod, fc2_mod, and gated_mod.
Validate that dp.dtype_bytes is valid and that each module's dimensions (dim_a,
dim_b) and device buffers (ranks_src_dev, ptrs_src_dev, permuted_ranks_dev,
permuted_ptrs_dev) are properly populated and non-null for fc1 and fc2. For
gated_mod, perform the same validation only when is_gated_activation is true.
This validation should occur before constructing any of the expand modules and
before calling launchMoeLoraPointerExpand to prevent incorrect expert weight
address computation or device illegal memory access.
In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py`:
- Around line 516-523: The `_moe_lora_active` method currently only checks
`lora_params` to determine if MoE LoRA is active, but it should also check
`cuda_graph_params` since `_extract_moe_lora_tensors` handles
use_cuda_graph_mode without a per-layer eager dict. Update the method to check
both `lora_params` and `cuda_graph_params` dictionaries, returning True if any
MoE LoRA module types are found in either parameter dictionary for the current
layer index.
🪄 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: a89651ca-07af-482e-8ac9-9fbeaeb18325
📒 Files selected for processing (21)
cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.hcpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_grouped_gemm.hcpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cucpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_pointer_expand.cucpp/tensorrt_llm/thop/moeOp.cppcpp/tests/unit_tests/kernels/moeLoraProblemBuilderTest.cudocs/source/features/lora.mdtensorrt_llm/_torch/custom_ops/torch_custom_ops.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.pytensorrt_llm/_torch/modules/fused_moe/moe_scheduler.pytensorrt_llm/_torch/modules/fused_moe/ops/moe_op_cutlass.pytensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.pytensorrt_llm/_torch/peft/lora/cuda_graph_lora_params.pytensorrt_llm/_torch/peft/lora/layer.pytensorrt_llm/_torch/peft/lora/moe_layout.pytensorrt_llm/_torch/peft/lora/validation.pytests/unittest/_torch/lora/test_moe_lora_cuda_graph_params.pytests/unittest/_torch/lora/test_moe_lora_extract.pytests/unittest/_torch/lora/test_moe_lora_grouped_gemm.pytests/unittest/_torch/lora/test_moe_lora_op.pytests/unittest/llmapi/test_llm_pytorch.py
|
PR_Github #55077 [ run ] triggered by Bot. Commit: |
|
PR_Github #55077 [ run ] completed with state
|
Shixiaowei02
left a comment
There was a problem hiding this comment.
LGTM for the document changes.
653d04f to
f91501f
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #56241 [ run ] triggered by Bot. Commit: |
|
PR_Github #56241 [ run ] completed with state
|
fb7c4bc to
0733e42
Compare
|
PR_Github #56416 [ run ] triggered by Bot. Commit: |
|
PR_Github #56416 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #56681 [ run ] triggered by Bot. Commit: |
|
PR_Github #56681 [ run ] completed with state
|
0733e42 to
d70bb01
Compare
d70bb01 to
71860d2
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #56731 [ run ] triggered by Bot. Commit: |
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
71860d2 to
94b1a10
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #56758 [ run ] triggered by Bot. Commit: |
|
PR_Github #56731 [ run ] completed with state |
|
PR_Github #56758 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #56982 [ run ] triggered by Bot. Commit: |
|
PR_Github #56982 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #57025 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #57027 [ run ] triggered by Bot. Commit: |
|
PR_Github #57025 [ run ] completed with state |
|
PR_Github #57027 [ run ] completed with state |
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Description
This MR collapses the routed-expert MoE LoRA code from three execution variants down to a single grouped-GEMM core fed by two input front-ends (eager per-request and CUDA-graph slot-indexed) - with no change to runtime behavior.
Currently, there are 3 execution methods for routed-expert MoE LoRA:
TLLM_MOE_LORA_USE_DEVICE_PATH.This MR removed the first one. Goal is to minimize codepaths for maintainability and share as much code as possible for eager and cudagraph paths.

This is just a refactor guarded by existing tests.
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.