[TRTLLM-12507][feat] Per-expert lora support with Cutlass backend

[brb-nv]

Description

This MR adds per-expert lora support in Pytorch workflow with Cutlass backend. Currently supports eager-only. Cudagraph support to be added in a follow-up MR.

Test Coverage

$ pytest tests/unittest/_torch/lora/test_moe_layout.py -s -v
$ pytest tests/unittest/_torch/lora/test_moe_lora_validator.py -s -v
$ pytest tests/unittest/_torch/lora/test_moe_lora_extract.py -s -v
$ pytest tests/unittest/_torch/lora/test_moe_lora_op.py -s -v
$ pytest tests/unittest/_torch/modules/moe/test_cutlass_moe_op_smoke.py -s -v
$ pytest tests/unittest/_torch/lora/test_moe_lora_model_path.py -s -v

PR Checklist

Please review the following before submitting your PR:

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

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

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

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

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

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

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

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

GitHub Bot Help

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

Summary by CodeRabbit

Release Notes

• New Features

  • Added routed-expert LoRA support for Mixture-of-Experts layers, enabling low-rank adapter tuning on expert projections with CUTLASS backend support.

• Documentation

  • Added documentation for MoE LoRA including configuration options, supported backends, and runtime constraints.

• Tests

  • Added comprehensive test coverage for MoE LoRA validation, reference implementations, and CUDA operations.

[coderabbitai]

📝 Walkthrough

Walkthrough

Implemented routed-expert MoE LoRA support by extending the C++ runMoe API to accept per-request LoRA rank and weight tensors, integrating workspace management and kernel invocations; added Python custom op signatures, CutlassFusedMoE layer integration, validation constraints, reference implementations, and comprehensive unit tests covering layout generation, validation logic, and fused-kernel parity with PyTorch references.

Changes

Routed-expert MoE LoRA implementation

Layer / File(s)	Summary
C++ MoE LoRA infrastructure and helpers cpp/tensorrt_llm/thop/moeOp.cpp	MoeLoraRequestType enum, FusedMoeRunner state members for cuBLAS wrapper and LoraImpl cache, CUDA event management, and helper routines: lazy initialization, dtype mapping, cached LoraImpl construction, per-request-to-per-token rank/pointer expansion, LoRA parameter building, and workspace size computation.
Extended runMoe method: LoRA inputs and execution setup cpp/tensorrt_llm/thop/moeOp.cpp	Updated runMoe signature with optional per-request LoRA tensors and host metadata; runtime setup activates LoRA when fc1_lora_ranks present, performs conservative rejections (min-latency/alltoall/unsupported dtypes/quantization/CUDA graph), builds LoRA parameters, computes workspace size, and orchestrates all related state.
Kernel invocation with LoRA parameters cpp/tensorrt_llm/thop/moeOp.cpp	OSS and non-OSS Cutlass MoE GEMM callsites forward lora_active and populated lora_params with workspace pointer; destructor cleanup best-effort destroys CUDA event and clears LoraImpl cache.
Custom op and fake function signatures tensorrt_llm/_torch/custom_ops/torch_custom_ops.py	trtllm::fused_moe custom op and @torch.library.register_fake extended with optional LoRA tensors (fc1/fc2/gated ranks/weight pointers, host request/context) and lora_max_low_rank; min_latency_mode path raises RuntimeError if LoRA active.
LoRA layout helpers and validation constants tensorrt_llm/_torch/peft/lora/moe_layout.py, tensorrt_llm/_torch/peft/lora/validation.py	MOE_LORA_MODULES constant, make_per_expert_lora for seeded per-expert adapter generation, reference_moe_lora_delta for per-token routed delta, reference_swiglu_moe_lora as detailed fp32 SwiGLU reference; validation module enforces CUTLASS backend and unquantized base weights at construction time.
CutlassFusedMoE LoRA integration tensorrt_llm/_torch/modules/fused_moe/create_moe.py, tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py	CutlassFusedMoE detects MoE LoRA targets, extracts and validates LoRA tensors from lora_params into lora_kwargs dict, and forwards into fused_moe via **lora_kwargs; forward_impl/forward_chunk/run_moe extended with lora_params threading; resolve_moe_cls calls check_moe_lora_supported for early validation.
Model forward path integration tensorrt_llm/_torch/models/modeling_qwen_moe.py	QwenMoE.forward now forwards optional lora_params into self.experts call alongside routing inputs.
User-facing MoE LoRA documentation docs/source/features/lora.md	New "Routed-Expert MoE LoRA" section documenting supported backend (CUTLASS), dtype/quantization constraints, lora_target_modules configuration, per-expert adapter layout, CUDA-graph decode limitations, and specific rejection conditions.
Comprehensive unit and integration tests tests/unittest/_torch/lora/test_moe_layout.py, tests/unittest/_torch/lora/test_moe_lora_validator.py, tests/unittest/_torch/lora/test_moe_lora_op.py	Layout tests validate module list and per-expert independence; validator tests check target detection and CUTLASS/unquantized enforcement; fused op tests include shared builders, smoke tests (activation/rejection/incompleteness), and extensive bisection probes (no-LoRA, zero-LoRA, per-module, pairwise, aliasing) comparing against PyTorch reference.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Suggested labels

api-compatible

Suggested reviewers

• venkywonka
• syuoni
• leslie-fang25
• amitz-nv
• 2ez4bz
• nv-guomingz

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

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

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically describes the main change: adding per-expert LoRA support with the Cutlass backend, which aligns with the comprehensive implementation described in the raw summary.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description check	✅ Passed	The PR description provides a brief summary of the feature (per-expert LoRA support), lists relevant test commands, and includes a completed PR checklist.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (4)

tests/unittest/_torch/lora/test_moe_layout.py (1)

18-20: ⚡ Quick win

Verify consistency with validation module's constant.

The test comment claims to check that the module names "match the validator's set," but it compares against a hardcoded set instead of importing MOE_LORA_MODULE_NAMES from validation.py. Import both constants and verify they match to prevent drift.

♻️ Proposed fix

+from tensorrt_llm._torch.peft.lora.validation import MOE_LORA_MODULE_NAMES
 from tensorrt_llm._torch.peft.lora.moe_layout import (
     MOE_LORA_MODULES,
     make_per_expert_lora,
     reference_moe_lora_delta,
 )


 def test_module_list_complete():
-    # Sanity check: the canonical module names match the validator's set.
-    assert set(MOE_LORA_MODULES) == {"moe_h_to_4h", "moe_4h_to_h", "moe_gate"}
+    """Verify that moe_layout.MOE_LORA_MODULES and validation.MOE_LORA_MODULE_NAMES
+    contain the same module names."""
+    assert set(MOE_LORA_MODULES) == MOE_LORA_MODULE_NAMES

🤖 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/lora/test_moe_layout.py` around lines 18 - 20, Replace
the hardcoded set assertion with a comparison against the validator's constant:
import MOE_LORA_MODULE_NAMES from the validation module and assert that
MOE_LORA_MODULES == MOE_LORA_MODULE_NAMES (using the existing symbol names
MOE_LORA_MODULES and MOE_LORA_MODULE_NAMES) so the test verifies the two
canonical lists stay consistent rather than comparing to a literal set.

tensorrt_llm/_torch/peft/lora/moe_layout.py (2)

165-165: 💤 Low value

Prefix unused unpacked variable with underscore.

The num_experts variable is unpacked from w3_w1.shape but never referenced in the function body. Prefix it with _ to signal it's intentionally unused.

♻️ Proposed fix

-    num_experts, two_inter, hidden_check = w3_w1.shape
+    _num_experts, two_inter, hidden_check = w3_w1.shape

🤖 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/peft/lora/moe_layout.py` at line 165, The tuple unpacking
for w3_w1 currently assigns num_experts but it's never used; change the unpack
to prefix the unused variable with an underscore (e.g., _num_experts, two_inter,
hidden_check = w3_w1.shape) so the intent is clear and linters won't flag
it—update the unpacking where w3_w1.shape is consumed in the moe_layout logic.

---

26-26: 💤 Low value

Consider moving device initialization inside the function.

While torch.device("cpu") is immutable and won't cause the classic mutable-default bug, the coding guidelines prefer avoiding function calls in default arguments. Move the default initialization inside the function body for consistency.

♻️ Proposed refactor

     *,
     dtype: torch.dtype = torch.bfloat16,
-    device: torch.device = torch.device("cpu"),
+    device: torch.device | None = None,
     seed: Optional[int] = None,
 ) -> Dict[str, torch.Tensor]:
     """Generate a per-expert (A, B) LoRA tensor pair for an MoE module.
 
     Returns shapes `A: [E, rank, in_dim]` and `B: [E, out_dim, rank]`, with an
     independent low-rank pair per expert, suitable for `torch.stack(...)` in
     `tensorrt_llm/lora_manager.py`.
 
     Args:
         num_experts: number of experts in this MoE layer.
         rank: LoRA rank for this module.
         in_dim: input hidden size (e.g. `hidden_size` for moe_h_to_4h).
         out_dim: output hidden size (e.g. `intermediate_size` for moe_h_to_4h).
         dtype: tensor dtype.
-        device: tensor device.
+        device: tensor device (defaults to CPU).
         seed: optional torch RNG seed for deterministic generation.
 
     Returns:
         Dict with keys "A" (shape [E, rank, in_dim]) and "B" (shape [E, out_dim, rank]).
     """
+    if device is None:
+        device = torch.device("cpu")
     if seed is not 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/peft/lora/moe_layout.py` at line 26, The parameter
currently uses a call in its default value: device: torch.device =
torch.device("cpu"); change the signature to accept device:
Optional[torch.device] = None (or device=None) and inside the function body set
device = device if device is not None else torch.device("cpu"); update any type
hints/imports as needed and keep the variable name 'device' so usages in the
function (e.g., in moe_layout-related logic) remain unchanged.

tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py (1)

383-474: 💤 Low value

Consider defensive access for num_seqs key.

Line 452 accesses lora_params["num_seqs"] directly. If the LoRA infrastructure ever passes lora_params without this key, the resulting KeyError would be less informative than a targeted error message. This is a minor robustness concern.

🛡️ Optional defensive fix

-        num_seqs = lora_params["num_seqs"]
+        num_seqs = lora_params.get("num_seqs")
+        if num_seqs is None:
+            raise ValueError(
+                "lora_params must contain 'num_seqs' when MoE LoRA tensors are active"
+            )

🤖 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 383
- 474, The _extract_moe_lora_tensors method uses lora_params["num_seqs"]
directly which can raise an unhelpful KeyError; update _extract_moe_lora_tensors
to defensively check for the presence and validity of "num_seqs" (e.g., use
lora_params.get("num_seqs") and verify it's an int > 0) and raise a clear
ValueError if missing/invalid, so subsequent slices like
kernel_ranks["fc1"][:num_seqs] and lora_params["host_request_types"][:num_seqs]
are safe and the error message names the missing/invalid "num_seqs" key and the
layer (reference self.layer_idx and function name _extract_moe_lora_tensors).

🤖 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 `@docs/source/features/lora.md`:
- Around line 193-199: The example snippet is missing the torch import required
for torch.bfloat16; add an import statement for torch at the top of the example
so the call to make_per_expert_lora (and the variable fc1_adapter) can use
torch.bfloat16 without error. Ensure the snippet includes "import torch" before
the call to make_per_expert_lora.

---

Nitpick comments:
In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py`:
- Around line 383-474: The _extract_moe_lora_tensors method uses
lora_params["num_seqs"] directly which can raise an unhelpful KeyError; update
_extract_moe_lora_tensors to defensively check for the presence and validity of
"num_seqs" (e.g., use lora_params.get("num_seqs") and verify it's an int > 0)
and raise a clear ValueError if missing/invalid, so subsequent slices like
kernel_ranks["fc1"][:num_seqs] and lora_params["host_request_types"][:num_seqs]
are safe and the error message names the missing/invalid "num_seqs" key and the
layer (reference self.layer_idx and function name _extract_moe_lora_tensors).

In `@tensorrt_llm/_torch/peft/lora/moe_layout.py`:
- Line 165: The tuple unpacking for w3_w1 currently assigns num_experts but it's
never used; change the unpack to prefix the unused variable with an underscore
(e.g., _num_experts, two_inter, hidden_check = w3_w1.shape) so the intent is
clear and linters won't flag it—update the unpacking where w3_w1.shape is
consumed in the moe_layout logic.
- Line 26: The parameter currently uses a call in its default value: device:
torch.device = torch.device("cpu"); change the signature to accept device:
Optional[torch.device] = None (or device=None) and inside the function body set
device = device if device is not None else torch.device("cpu"); update any type
hints/imports as needed and keep the variable name 'device' so usages in the
function (e.g., in moe_layout-related logic) remain unchanged.

In `@tests/unittest/_torch/lora/test_moe_layout.py`:
- Around line 18-20: Replace the hardcoded set assertion with a comparison
against the validator's constant: import MOE_LORA_MODULE_NAMES from the
validation module and assert that MOE_LORA_MODULES == MOE_LORA_MODULE_NAMES
(using the existing symbol names MOE_LORA_MODULES and MOE_LORA_MODULE_NAMES) so
the test verifies the two canonical lists stay consistent rather than comparing
to a literal set.

🪄 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: c0ad5f96-e7fe-4325-b6b7-11e93f61d4ae

📥 Commits

Reviewing files that changed from the base of the PR and between 54259ed and 00cae91.

📒 Files selected for processing (11)

• cpp/tensorrt_llm/thop/moeOp.cpp
• docs/source/features/lora.md
• tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
• tensorrt_llm/_torch/models/modeling_qwen_moe.py
• tensorrt_llm/_torch/modules/fused_moe/create_moe.py
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py
• tensorrt_llm/_torch/peft/lora/moe_layout.py
• tensorrt_llm/_torch/peft/lora/validation.py
• tests/unittest/_torch/lora/test_moe_layout.py
• tests/unittest/_torch/lora/test_moe_lora_op.py
• tests/unittest/_torch/lora/test_moe_lora_validator.py

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #51391 [ run ] triggered by Bot. Commit: 71eec3c Link to invocation

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #51401 [ run ] triggered by Bot. Commit: 1e5cfe2 Link to invocation

[tensorrt-cicd]

PR_Github #51391 [ run ] completed with state ABORTED. Commit: 71eec3c

Link to invocation

[tensorrt-cicd]

PR_Github #51401 [ run ] completed with state FAILURE. Commit: 1e5cfe2
/LLM/main/L0_MergeRequest_PR pipeline #40812 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

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

CI Agent Failure Analysis

Link to invocation

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #51456 [ run ] triggered by Bot. Commit: 1e5cfe2 Link to invocation

[QiJune]

LGTM

[tensorrt-cicd]

PR_Github #51456 [ run ] completed with state SUCCESS. Commit: 1e5cfe2
/LLM/main/L0_MergeRequest_PR pipeline #40864 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

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

CI Agent Failure Analysis

Link to invocation

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #51648 [ run ] triggered by Bot. Commit: 1e5cfe2 Link to invocation

[tensorrt-cicd]

PR_Github #51648 [ run ] completed with state SUCCESS. Commit: 1e5cfe2
/LLM/main/L0_MergeRequest_PR pipeline #41031 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

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

CI Agent Failure Analysis

Link to invocation

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #51681 [ run ] triggered by Bot. Commit: 1e5cfe2 Link to invocation

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #51699 [ run ] triggered by Bot. Commit: 1e5cfe2 Link to invocation

[tensorrt-cicd]

PR_Github #51699 [ run ] completed with state SUCCESS. Commit: 1e5cfe2
/LLM/main/L0_MergeRequest_PR pipeline #41076 completed with status: 'SUCCESS'

CI Report

Link to invocation

[tensorrt-cicd]

PR_Github #51681 [ run ] completed with state SUCCESS. Commit: 1e5cfe2
/LLM/main/L0_MergeRequest_PR pipeline #41061 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

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

CI Agent Failure Analysis

Link to invocation