[TRTLLM-12507][feat] Cudagraph support for per-expert lora in Cutlass backend - Part 2

[brb-nv]

Description

This is Part 2 of routed-expert MoE LoRA on the CUTLASS backend, building on the device path merged in Part 1 (#14923). It makes the device path fully CUDA-graph capturable with per-request adapters under a single persistent graph:

• Adds an on-device slot -> token expansion kernel (moe_lora_slot_expand.cu/.h) that runs every replay, materializing the per-source-token (rank, A_ptr, B_ptr) tables from persistent device slot tables + token_to_slot.
• moeOp.cpp: slot-indexed input schema, persistent device slot-table buffers sized via reserve_lora_host_buffers, captured async H2D from stable pinned buffers, and auto-enabling the device path for the slot-indexed schema.
• Python wiring (cuda_graph_lora_params.py, cuda_graph_lora_manager.py, fused_moe_cutlass.py, layer.py): produces the slot-indexed fused_moe kwargs at decode time and packs the per-module pointer tables behind a fixed cache.

Test Coverage

$ ./cpp/build/tests/unit_tests/kernels/moeLoraSlotExpandTest
$ pytest tests/unittest/_torch/lora/test_moe_lora_op.py -s -v
$ pytest tests/unittest/_torch/lora/test_moe_lora_device_path.py -s -v
$ pytest tests/unittest/_torch/lora/test_moe_lora_cuda_graph_params.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 slot-indexed LoRA routing for Mixture of Experts layers, enabling efficient multi-adapter support in CUDA graphs.
  • Added buffer pre-reservation mechanism for stable CUDA graph replay.
  • Extended MoE LoRA operator to support both per-request and slot-indexed input schemas.

• Documentation

  • Updated routed-expert MoE LoRA configuration guidance.

• Tests

  • Added comprehensive unit and integration tests for slot-indexed LoRA routing and CUDA graph capture scenarios.

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #53103 [ run ] triggered by Bot. Commit: 5069057 Link to invocation

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR adds slot-indexed tensor routing for routed-expert MoE LoRA within CUDA graphs. Per-token adapter selection is now driven by slot indices and slot-table lookups (instead of per-request expansion), enabling capture-safe replay with updated slot assignments. The implementation spans a new CUDA slot-expansion kernel, Python LoRA infrastructure for slot-table management, custom op extensions, C++ runner integration with buffer pre-reservation, and comprehensive tests.

Changes

MoE LoRA Slot-Indexed CUDA-Graph Expansion

Layer / File(s)	Summary
Slot Expansion Kernel (CUDA) cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_slot_expand.h, cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_slot_expand.cu, cpp/tests/unit_tests/kernels/CMakeLists.txt, cpp/tests/unit_tests/kernels/moeLoraSlotExpandTest.cu	New header/implementation define MoeLoraSlotExpandModule and launchMoeLoraSlotExpand, which expand per-token slot indices into per-token LoRA ranks and pointer triplets for fc1/fc2/gated modules. Invalid slots produce zero ranks and null pointers. Unit tests validate multi-slot behavior, gated activation, out-of-range handling, and no-op edge cases with CPU reference comparison.
Python LoRA CUDA-Graph Slot-Table Infrastructure tensorrt_llm/_torch/peft/lora/layer.py, tensorrt_llm/_torch/peft/lora/cuda_graph_lora_params.py, tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py, tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py	New MoeLoraLayer marker class enables LoRA discovery in routed-expert MoE. CudaGraphLoraParams adds pinned token_to_slot_host buffer and get_moe_slot_inputs(layer_idx, module_id) to retrieve per-module slot ranks and packed pointer tensors. CudaGraphLoraManager pre-reserves MoE LoRA scratch during init. CutlassFusedMoE initializes LoRA marker, exposes reserve_moe_lora_cuda_graph_workspace, and adds _extract_moe_lora_tensors_cuda_graph helper to build slot-indexed op kwargs.
Custom Op Slot-Indexed Tensor Interface tensorrt_llm/_torch/custom_ops/torch_custom_ops.py	fused_moe custom op and fake schema extended with optional slot-indexed LoRA tensors (fc1_slot_lora_ranks/weight_ptrs, fc2_slot_lora_ranks/weight_ptrs, gated_slot_lora_ranks/weight_ptrs, token_to_slot). Runtime validation treats slot-indexed LoRA as active, rejects LoRA-rejection min-latency path, and enforces mutual exclusivity with per-request schema. Non-min-latency path forwards new tensors to run_moe runner call.
MoE Runner Slot-Indexed Schema & Device-Path Selection cpp/tensorrt_llm/thop/moeOp.cpp	FusedMoeRunner::runMoe accepts slot-indexed LoRA tensors and validates mutual exclusivity with per-request inputs. Device-path automatic selection when slot-indexed inputs present. New reserveLoraHostBuffers API pre-allocates pinned-host and device-path scratch with worst-case warmup sizing (exported as reserve_lora_host_buffers to Python). New persistent device tensors for token_to_slot and per-slot rank/pointer tables with capacity tracking via ensureLoraSlotTableBuffers. buildMoeLoraParams refactored to dual-mode: per-request host-driven expansion vs. slot-indexed async H2D + launchMoeLoraSlotExpand device-path expansion.
Kernel & Integration Tests tests/unittest/_torch/lora/test_moe_lora_cuda_graph_params.py, tests/unittest/_torch/lora/test_moe_lora_device_path.py, tests/unittest/_torch/lora/test_moe_lora_op.py	test_moe_lora_cuda_graph_params.py validates get_moe_slot_inputs pointer packing, buffer caching/address stability, and _extract_moe_lora_tensors_cuda_graph op-consumable kwarg wiring. test_moe_lora_device_path.py expands coverage with slot-indexed CUDA-graph tests: eager/replay equivalence, multi-adapter routing, buffer pre-reservation stability across captures, and in-place slot pointer reassignment reflecting on replay. test_moe_lora_op.py adds smoke tests confirming slot-indexed/per-request equivalence and rejection of mixed schema inputs.
Documentation Updates cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_device_path.h, docs/source/features/lora.md	MoeLoraDeviceRunFn contract rewritten into structured parameter list. DoRA pointer handling clarified as unused. Routed-expert LoRA adapter modules no longer require moe_gate and moe_4h_to_h to be present together.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#14801: Both PRs extend the same routed-expert MoE LoRA plumbing in tensorrt_llm/thop/moeOp.cpp/torch_custom_ops.py/fused_moe_cutlass.py (adding LoRA-related runMoe/fused_moe inputs and kernel kwargs), with the main PR further specializing that flow to the new slot-indexed CUDA-graph expansion path.
• NVIDIA/TensorRT-LLM#14923: The main PR builds on the retrieved PR's existing CUDA-graph-safe "device-path" LoRA infrastructure (e.g., LoraParams/moeOp.cpp/moe_kernels wiring) by adding slot-indexed MoE-LoRA slot table expansion (launchMoeLoraSlotExpand) and routing through that same device-path execution path.

Suggested reviewers

• byshiue
• yizhang-nv
• QiJune
• nv-guomingz
• govind-ramnarayan

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 64.29% 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
Description check	✅ Passed	The description covers key implementation details (slot expansion kernel, persistent buffers, Python wiring), test coverage commands, and verification of PR checklist items. It explains both what was changed and why.
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.
Title check	✅ Passed	The PR title clearly and specifically describes the main feature: adding CUDA graph support for per-expert LoRA in the Cutlass backend, directly matching the changeset's core objective.

_{✏️ 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: 6

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/modules/fused_moe/fused_moe_cutlass.py (1)

485-496: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Teach _moe_lora_active() about slot-indexed CUDA-graph inputs.

In CUDA-graph mode lora_params only carries cuda_graph_params, so this helper always returns False. That bypasses both safety checks in forward_impl(), including the multi-chunk rejection at Lines 1428-1437. A chunked CUDA-graph MoE LoRA call will therefore reuse the full-batch slot tables for every chunk instead of failing fast.

Suggested 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
+            return any(
+                cuda_graph_params.get_moe_slot_inputs(self.layer_idx, int(module_type))
+                is not None
+                for module_type in (
+                    LoraModuleType.MOE_H_TO_4H,
+                    LoraModuleType.MOE_4H_TO_H,
+                    LoraModuleType.MOE_GATE,
+                )
+            )
         layer_params = lora_params.get(self.layer_idx, {})
         if not layer_params:
             return False
         return any(
             int(LoraModuleType.from_string(name)) in layer_params

🤖 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 485
- 496, _moe_lora_active currently ignores cuda-graph slot-indexed inputs because
in CUDA-graph mode lora_params carries its data under the "cuda_graph_params"
envelope; update _moe_lora_active to unwrap cuda_graph_params when present
(e.g., check if "cuda_graph_params" in lora_params and then use that dict) and
then inspect the per-layer entry for the presence of any MoE LoRA module names
from self._MOE_LORA_MODULE_NAMES by converting names via
LoraModuleType.from_string and checking membership in the layer's keys; ensure
the logic handles both the normal and slot-indexed shapes (nested dicts) so the
function returns True when any routed-expert MoE LoRA tensor exists for
self.layer_idx.

cpp/tensorrt_llm/thop/moeOp.cpp (2)

2-2: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Update the copyright year for this 2026 modification.

This file was meaningfully modified in 2026, but the header still ends at 2025. Please bump the NVIDIA copyright year to keep the header compliant.

🤖 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/moeOp.cpp` at line 2, Update the copyright header year
from "2022-2025" to "2022-2026" in the file's top-of-file header comment (the
copyright line present in moeOp.cpp) so the header reflects the 2026
modification; ensure formatting and punctuation remain unchanged aside from the
year range.

Source: Coding guidelines

---

1412-1468: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

The post-capture reallocation guard is not actually latched.

mLoraCaptureObserved is only set inside checkLoraReallocSafeDuringCapture(), which means a graph captured with already-sized buffers leaves the flag false. After that, a later reserve_lora_host_buffers() or larger eager request can still reallocate the pinned/device buffers whose addresses were baked into the captured graph. Please mark capture as soon as any LoRA call runs under isCapturing(stream) and route all host/slot/device growth through the same guard.

🤖 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/moeOp.cpp` around lines 1412 - 1468, The realloc guard
flag mLoraCaptureObserved is only set inside
checkLoraReallocSafeDuringCapture(), so growth paths can still reallocate after
a graph has been captured; fix by detecting capture earlier and routing all
buffer-growth through the same guard: in reserveLoraHostBuffers (and any public
growth entrypoints that call ensureLoraExpandBuffers, ensureLoraSlotTableBuffers
or ensureLoraDeviceScratch) acquire the current CUDA stream (or add a
cudaStream_t parameter), call checkLoraReallocSafeDuringCapture(stream,
requested, current) before any ensure* call to both set mLoraCaptureObserved
when capturing and to reject unsafe growth, and ensure the same check is invoked
inside
ensureLoraExpandBuffers/ensureLoraSlotTableBuffers/ensureLoraDeviceScratch if
they can be called independently.

🤖 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/thop/moeOp.cpp`:
- Around line 1808-1852: The code performs cudaMemcpyAsync from caller CPU
tensors (token_to_slot, fc1_slot_lora_ranks, fc1_slot_lora_weight_ptrs,
fc2_slot_lora_ranks, fc2_slot_lora_weight_ptrs) which must be pinned for safe
CUDA-graph capture; add checks (e.g. TORCH_CHECK(tensor.is_pinned(), "... must
be pinned host tensors for captured async H2D")) for each of those tensors
before the h2d_slot lambda/copies (or wrap them in a helper that validates
is_pinned()) and fail early with a clear error message identifying the offending
tensor name so callers know to provide pinned memory.
- Around line 1817-1859: The code copies fc2 and gated slot tables to device
without validating their shapes causing possible OOB reads; before the h2d_slot
calls add the same TORCH_CHECKs used for fc1 (validate fc2_slot_lora_weight_ptrs
has dim()==2 and size == [num_slots,3] and fc2_slot_lora_ranks
size(0)==num_slots) and, when has_gated is true, validate
gated_slot_lora_weight_ptrs and gated_slot_lora_ranks similarly (use
CHECK_CPU_INPUT where appropriate), then proceed to call h2d_slot and the gated
copies; reference fc2_slot_lora_weight_ptrs, fc2_slot_lora_ranks,
gated_slot_lora_weight_ptrs, gated_slot_lora_ranks, num_slots, and h2d_slot.

In `@docs/source/features/lora.md`:
- Line 149: Update the LoRA docs to reflect the actual MoE adapter requirements
and CUDA-graph behavior: change the statement that “moe_gate and moe_4h_to_h are
required” to require both moe_h_to_4h and moe_4h_to_h (with moe_gate optional)
consistent with fused_moe_cutlass.py and
tests/unittest/_torch/lora/test_moe_lora_extract.py, and revise the CUDA-graph
section to note that CUDA-graph capture/support exists for routed-expert MoE
LoRA (referencing use_cuda_graph_mode, CudaGraphLoraManager and
tests/unittest/_torch/lora/test_moe_lora_device_path.py) rather than saying
CUDA-graph is entirely rejected for MoE LoRA.

In `@tensorrt_llm/_torch/custom_ops/torch_custom_ops.py`:
- Around line 315-326: The current checks only look at
fc1_lora_ranks/fc1_slot_lora_ranks; update the logic to detect any per-request
LoRA tensor present and any slot-indexed LoRA tensor present instead of keying
only on fc1_. Concretely, replace the lora_active bool and the mutual-exclusion
check with predicates that OR together all per-request LoRA symbols (e.g.,
fc1_lora_ranks, fc2_lora_ranks, ... any fc*_lora_* per-request tensors) and
separately OR together all slot-indexed symbols (e.g., fc1_slot_lora_ranks,
fc2_slot_lora_ranks, token_to_slot, ... any fc*_slot_lora_* / slot-related
tensors), then: (1) if min_latency_mode and either predicate is true raise the
same RuntimeError, and (2) if both predicates are true raise the
mutual-exclusion RuntimeError; use the existing error messages and reference the
same symbols (fc1_lora_ranks, fc1_slot_lora_ranks, token_to_slot) in the code to
locate where to change.

In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py`:
- Around line 634-642: The code currently returns None when either
slot_ranks["fc1"] or slot_ranks["fc2"] is missing during CUDA-graph capture,
silently dropping partial MoE LoRA configs; change the logic after the loop over
slot_to_kernel to detect partial configurations and raise an explicit error
instead of returning None: use the same consistency check as the eager path (if
exactly one of slot_ranks["fc1"] or slot_ranks["fc2"] is None) then raise a
ValueError describing the partial MoE LoRA config for self.layer_idx, otherwise
(both None) allow the existing None return; update the block that references
slot_ranks/slot_ptrs and cuda_graph_params.get_moe_slot_inputs to implement this
check.

In `@tensorrt_llm/_torch/peft/lora/cuda_graph_lora_params.py`:
- Around line 225-233: Guard the runtime tokens_per_seq against the preallocated
host capacity: compute the host capacity (e.g., max_tokens =
self.token_to_slot_host.numel()), then clamp the effective per-sequence token
count (use effective_tokens_per_seq = min(tokens_per_seq, max_tokens //
batch_size_or_slot_count)) before building token_slots (repeat_interleave on
slot_ids_t) and adjust num_tokens accordingly so
self.token_to_slot_host[:num_tokens].copy_(token_slots) never tries to copy more
elements than the preallocated buffer; apply the same clamping logic
consistently where sorted_ids and the MoE workspace reservation are assumed
sized from max_batch_size * max_tokens_per_seq and optionally raise a clear
error if the requested runtime tokens exceed the preallocated budget.

---

Outside diff comments:
In `@cpp/tensorrt_llm/thop/moeOp.cpp`:
- Line 2: Update the copyright header year from "2022-2025" to "2022-2026" in
the file's top-of-file header comment (the copyright line present in moeOp.cpp)
so the header reflects the 2026 modification; ensure formatting and punctuation
remain unchanged aside from the year range.
- Around line 1412-1468: The realloc guard flag mLoraCaptureObserved is only set
inside checkLoraReallocSafeDuringCapture(), so growth paths can still reallocate
after a graph has been captured; fix by detecting capture earlier and routing
all buffer-growth through the same guard: in reserveLoraHostBuffers (and any
public growth entrypoints that call ensureLoraExpandBuffers,
ensureLoraSlotTableBuffers or ensureLoraDeviceScratch) acquire the current CUDA
stream (or add a cudaStream_t parameter), call
checkLoraReallocSafeDuringCapture(stream, requested, current) before any ensure*
call to both set mLoraCaptureObserved when capturing and to reject unsafe
growth, and ensure the same check is invoked inside
ensureLoraExpandBuffers/ensureLoraSlotTableBuffers/ensureLoraDeviceScratch if
they can be called independently.

In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py`:
- Around line 485-496: _moe_lora_active currently ignores cuda-graph
slot-indexed inputs because in CUDA-graph mode lora_params carries its data
under the "cuda_graph_params" envelope; update _moe_lora_active to unwrap
cuda_graph_params when present (e.g., check if "cuda_graph_params" in
lora_params and then use that dict) and then inspect the per-layer entry for the
presence of any MoE LoRA module names from self._MOE_LORA_MODULE_NAMES by
converting names via LoraModuleType.from_string and checking membership in the
layer's keys; ensure the logic handles both the normal and slot-indexed shapes
(nested dicts) so the function returns True when any routed-expert MoE LoRA
tensor exists for self.layer_idx.

🪄 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: 3f4ef076-e048-4ea7-96f7-94c3253b5665

📥 Commits

Reviewing files that changed from the base of the PR and between 358505c and 9a47f95.

📒 Files selected for processing (15)

• cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_device_path.h
• cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_slot_expand.h
• cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_slot_expand.cu
• cpp/tensorrt_llm/thop/moeOp.cpp
• cpp/tests/unit_tests/kernels/CMakeLists.txt
• cpp/tests/unit_tests/kernels/moeLoraSlotExpandTest.cu
• docs/source/features/lora.md
• tensorrt_llm/_torch/custom_ops/torch_custom_ops.py
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py
• tensorrt_llm/_torch/peft/lora/cuda_graph_lora_manager.py
• tensorrt_llm/_torch/peft/lora/cuda_graph_lora_params.py
• tensorrt_llm/_torch/peft/lora/layer.py
• tests/unittest/_torch/lora/test_moe_lora_cuda_graph_params.py
• tests/unittest/_torch/lora/test_moe_lora_device_path.py
• tests/unittest/_torch/lora/test_moe_lora_op.py

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #53115 [ run ] triggered by Bot. Commit: fa6de78 Link to invocation

[tensorrt-cicd]

PR_Github #53103 [ run ] completed with state ABORTED. Commit: 5069057

Link to invocation

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #53158 [ run ] triggered by Bot. Commit: 5aa968c Link to invocation

[tensorrt-cicd]

PR_Github #53115 [ run ] completed with state ABORTED. Commit: fa6de78

Link to invocation

[tensorrt-cicd]

PR_Github #53158 [ run ] completed with state FAILURE. Commit: 5aa968c
/LLM/main/L0_MergeRequest_PR pipeline #42361 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 #53342 [ run ] triggered by Bot. Commit: 5aa968c Link to invocation

[tensorrt-cicd]

PR_Github #53342 [ run ] completed with state SUCCESS. Commit: 5aa968c
/LLM/main/L0_MergeRequest_PR pipeline #42522 completed with status: 'SUCCESS'

CI Report

Link to invocation