[TRTLLM-12507][feat] Cudagraph support for routed-expert MoE LoRA with Cutlass backend - Part 1 - #14923
Merged
brb-nv merged 2 commits intoJun 8, 2026
Conversation
…h Cutlass backend - Part 1 Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Contributor
📝 WalkthroughWalkthroughThis PR adds a CUDA-graph-capture-safe device-path for MoE LoRA execution by moving pointer expansion and GEMM problem building from the CPU to GPU. New kernel interfaces replace host-side fan-out with on-device processing, enabling full CUDA graph capture for LoRA pipelines via environment variable opt-in. ChangesMoE LoRA Device-Path Implementation
Sequence Diagram(s)sequenceDiagram
participant Host
participant PinnedHost as Pinned Host<br/>Buffers
participant GPU as GPU Device
participant GemmKernel as Grouped<br/>GEMM Kernel
Host->>PinnedHost: expandPerRequestLoraTo<br/>(ranks, ptrs)
Host->>GPU: async H2D copy<br/>(pinned → device)
Host->>GPU: launchMoeLoraPointerExpand<br/>(permuted_rows, expert_offsets)
GPU->>GPU: expand ranks/ptrs<br/>per expert
Host->>GPU: launchMoeLoraProblemBuilder<br/>(expanded ranks/ptrs)
GPU->>GPU: build GEMM<br/>descriptors
Host->>GemmKernel: cudaGraphGroupedGemm<br/>(GEMM arrays)
GemmKernel->>GPU: compute LoRA<br/>fc1, fc2, gated
🎯 4 (Complex) | ⏱️ ~60 minutes 🚥 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 9
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
cpp/tensorrt_llm/thop/moeOp.cpp (2)
1515-1539:⚠️ Potential issue | 🔴 Critical | ⚡ Quick winValidate every incoming LoRA rank against
lora_max_low_rank.
lora_max_low_rankis what sizeslowrank_workspaceand the max-problem hints, but the per-request rank tensors are never checked against it. A caller can still pass a larger rank here, which makes the device path build GEMM problems wider than the allocated scratch.Suggested fix
TORCH_CHECK(lora_max_low_rank > 0, "MoE LoRA requires lora_max_low_rank > 0; got ", lora_max_low_rank); @@ int64_t const num_seqs = fc1_lora_ranks->size(0); bool const has_gated = is_gated_activation && gated_lora_ranks.has_value(); + + auto validate_rank_tensor = [&](char const* name, torch::Tensor const& ranks_tensor) + { + CHECK_CPU_INPUT(ranks_tensor, at::ScalarType::Int) + auto const* rank_data = ranks_tensor.data_ptr<int32_t>(); + for (int64_t i = 0; i < ranks_tensor.size(0); ++i) + { + TORCH_CHECK(rank_data[i] >= 0 && rank_data[i] <= lora_max_low_rank, name, "[", i, "]=", + rank_data[i], " is outside [0, ", lora_max_low_rank, "]."); + } + }; + validate_rank_tensor("fc1_lora_ranks", *fc1_lora_ranks); + validate_rank_tensor("fc2_lora_ranks", *fc2_lora_ranks); + if (has_gated) + { + validate_rank_tensor("gated_lora_ranks", *gated_lora_ranks); + }
593-598:⚠️ Potential issue | 🟠 Major | ⚡ Quick winSkip the legacy LoRA workspace on the device path.
When
device_path.enabledis true, this still computes and attaches the cuBLASlora_workspace, butmoeLoraDeviceRunImpluses the persistent device scratch instead. That duplicates LoRA scratch per stream and can turn largetop_k/rank requests into avoidable OOMs.Suggested fix
- if (lora_params_opt.has_value()) + if (lora_params_opt.has_value() && !lora_params_opt->device_path.enabled) { auto const lora_dtype = loraTypeFromActDtype(mActivationDtype); lora_workspace_size = computeLoraWorkspaceSize(lora_params_opt->fc1_lora_impl, lora_params_opt->fc2_lora_impl, num_rows, experts_per_token, lora_params_opt->num_reqs, lora_dtype); } @@ - if (lora_active) + if (lora_active && !lora_params.device_path.enabled) { lora_params.workspace = workspace_info.lora_workspace; }Also applies to: 632-635
🤖 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 593 - 598, The code currently computes and attaches the legacy cuBLAS lora_workspace even when device_path.enabled is true; update the logic around the lora_params_opt block (the call to computeLoraWorkspaceSize that sets lora_workspace_size) to skip computing/attaching the legacy lora_workspace when device_path.enabled is true (since moeLoraDeviceRunImpl uses the persistent device scratch), and apply the same conditional change to the similar block around lines handling lora_workspace at the other site (the 632-635 occurrence); specifically, only call computeLoraWorkspaceSize and allocate/set lora_workspace_size if lora_params_opt.has_value() AND device_path.enabled is false.
🧹 Nitpick comments (2)
cpp/tests/unit_tests/kernels/moeLoraPointerExpandTest.cu (1)
225-323: ⚡ Quick winCoverage is insufficient for the two non-happy-path branches.
These cases only cover the shared-memory lookup path with fully populated expert ranges. Please add one test in
cpp/tests/unit_tests/kernels/moeLoraPointerExpandTest.cuwhereexpanded_num_rows > expert_first_token_offset.back()to assert ghost rows stay zeroed, and one test withnum_experts_per_node = 1025to force the global-memory scan incpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_pointer_expand.cu.As per coding guidelines,
tests/**: "Act as a QA engineer reviewing test changes and coverage ... state whether coverage is sufficient, insufficient, or needs follow-up outside the PR."🤖 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/tests/unit_tests/kernels/moeLoraPointerExpandTest.cu` around lines 225 - 323, Add two new TEST_F cases to MoeLoraPointerExpandTest: (1) a "GhostRowsRemainZero" test that sets expanded_num_rows greater than expert_first_token_offset.back() (e.g., expanded_num_rows = expert_first_token_offset.back() + 2) and uses runAndCompare(permuted_rows, expert_first_token_offset, ..., expanded_num_rows, lora_dtype_bytes, fc1, fc2, /*gated=*/nullptr) and asserts the extra/ghost rows are zeroed (use the same style as existing tests and the runAndCompare helper to validate output); (2) a "ForceGlobalScanNumExperts1025" test that sets num_experts_per_node = 1025 to exercise the global-memory scan path in moe_lora_pointer_expand.cu and otherwise mirror one of the existing simple cases (populate permuted_rows, expert_first_token_offset, fc1/fc2, expanded_num_rows accordingly). Ensure both tests use the MoeLoraPointerExpandTest fixture and reference runAndCompare so they integrate with existing verification logic.cpp/tests/unit_tests/kernels/moeLoraProblemBuilderTest.cu (1)
301-329: ⚡ Quick winCoverage is insufficient for
splitk_offsets == nullptr.Every live case allocates
out.splitk_offsets, so thelaunch_countbranch and the kernel's null-offset path never run. Please add a case incpp/tests/unit_tests/kernels/moeLoraProblemBuilderTest.cuthat leavessplitk_offsetsnull and verifies the remaining arrays still match the reference.As per coding guidelines,
tests/**: "Act as a QA engineer reviewing test changes and coverage ... state whether coverage is sufficient, insufficient, or needs follow-up outside the PR."🤖 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/tests/unit_tests/kernels/moeLoraProblemBuilderTest.cu` around lines 301 - 329, The test never covers the kernel path where splitk_offsets == nullptr because every case allocates out.splitk_offsets; add a new test case in moeLoraProblemBuilderTest (use the TEST_F BoundaryCases block or a new test) that creates a MoeLoraGemmGroupArrays instance with out.splitk_offsets left null, call launchMoeLoraProblemBuilder with that arrays struct (and appropriate input_base/lowrank_workspace/output_base and parameters) and synchronize the stream, then verify the produced arrays against the reference (reuse runAndCompare or implement an equivalent comparison) to exercise the launch_count/null-offset branch in the kernel.
🤖 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/cutlass_kernels/include/moe_lora_problem_builder.h`:
- Around line 65-71: The comment for the leading-dimension pointer ldb_out in
moe_lora_problem_builder.h is incorrect; update it to match the function
contract (ldb_out[i] == out_hidden_size) so callers understand the true
stride/semantics. Locate the declaration of ldb_out (int64_t* ldb_out = nullptr)
and change its inline comment to state that it represents the per-problem/output
hidden-size stride (out_hidden_size) rather than "per-token rank"; ensure the
wording aligns with ldb_in/ldd_out comments and the documented contract in the
header.
In `@cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu`:
- Around line 61-68: The copyright header in moe_kernels.cu (the top-of-file
header above the includes such as the lines referencing
moe_lora_pointer_expand.h and moe_util_kernels.h) still ends at 2025; update the
year range to include 2026 in the NVIDIA copyright block at the top of the file
so the header reflects the file modification year.
- Around line 3883-3900: The device-path branch can leave stale data in LoRA
output workspaces when dp.run skips rank-0 rows; before calling
runMoeLoraDeviceModule for dp.fc1/dp.gated (and similarly for fc2 later),
zero-fill the device buffers (lora_fc1_result, lora_fc2_result, lora_gated_out)
with an async memset so skipped slots become deterministic. Compute the byte
size from expanded_num_rows, the per-row output width (hidden_size or fc output
width) and element size (dp.dtype_bytes or moeLoraNvInferType<byte size>) and
call cudaMemsetAsync or the project’s async zero routine on each buffer on the
same stream prior to invoking runMoeLoraDeviceModule; keep the zeroing in the
same device-path branch where lora_params.device_path and dp.run are used.
- Around line 3728-3769: runMoe() still performs host staging of routing buffers
and records memcpy_event_ptr even when lora_params.device_path.enabled is true,
which breaks the intended device-only (CUDA-graph-safe) path; wrap the host D2H
copy and the setting of memcpy_event_ptr (the code that writes into
host_lora_workspace_ and records memcpy_event_ptr) in a guard like if
(!lora_params.device_path.enabled) so that when device_path.enabled is true the
host staging and event recording are skipped; update references in runMoe() to
avoid creating host dependencies only when device_path is disabled.
In
`@cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_pointer_expand.cu`:
- Around line 106-112: In the moe_lora_pointer_expand kernel, the early-return
when expert_idx >= num_experts_per_node leaves stale values in ranks_out and
ptrs_out; before returning, explicitly zero the output slot(s) for this thread:
set the corresponding ranks_out element(s) to 0 and clear the associated
ptrs_out pointer/offset fields so padding ("ghost") rows become no-op on reused
device buffers. Locate the branch that currently does "if (expert_idx >=
num_experts_per_node) { return; }" and insert zeroing of the per-row ranks_out
and ptrs_out entries for the current row/thread (use the same indexing variables
used elsewhere in the kernel) before returning.
In
`@cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_problem_builder.cu`:
- Around line 59-95: The builder uses per-token runtime rank (ranks[i]) to set
problem_sizes_in/out but lays out workspace and ldd_in using max_lora_rank,
which can cause out-of-bounds if rank > max_lora_rank; add an explicit runtime
check inside the loop (where ranks, problem_sizes_in/problem_sizes_out,
d_ptrs_in, ldd_in are set) that validates rank <= max_lora_rank and
returns/errors (e.g., throw or return a clear error code) if violated, update
any surrounding function contract/comments to document this invariant, and add a
regression test that constructs a token with rank > max_lora_rank to assert the
builder fails instead of corrupting memory.
In `@cpp/tensorrt_llm/thop/moeOp.cpp`:
- Around line 1319-1329: The current checkLoraReallocSafeDuringCapture only
blocks growth while the current call is inside a CUDA capture, but doesn't
prevent later reallocations that would invalidate previously captured graphs;
add persistent tracking and enforce prohibition of growth after any capture was
observed. Introduce a member flag (e.g., bool has_seen_capture_ or
lora_capture_observed_) on the MoE/LoRA operator object, set it when
tensorrt_llm::common::isCapturing(stream) returns true inside
checkLoraReallocSafeDuringCapture, and change the logic so that any attempted
growth (when requested > current) is rejected if either the current stream is
capturing OR has_seen_capture_ is true; update the error message in
checkLoraReallocSafeDuringCapture to reflect that growth is forbidden because a
capture has previously occurred. Ensure the flag is thread-safe or only modified
under the same synchronization already used for buffer management.
- Around line 1290-1300: Reject invalid/mojibake request-type values and
negative context lengths before computing repeat and touching expand_*_data:
validate that the parsed req_type (from req_types[req_id]) is one of the
expected MoeLoraRequestType enum values (e.g. kGENERATION or the context branch)
and that ctx_lens[req_id] is non-negative, then compute repeat from that
validated value and TORCH_CHECK both repeat >= 0 and produced + repeat <=
num_tokens with an explicit error message referencing req_id,
host_request_types/host_context_lengths, produced and num_tokens so malformed
host inputs fail cleanly instead of allowing negative repeats or out-of-bounds
writes; update the logic around MoeLoraRequestType, req_types, ctx_lens,
produced, and expand_*_data accordingly.
In `@cpp/tests/unit_tests/kernels/CMakeLists.txt`:
- Around line 98-101: The test target moeLoraSlotExpandTest is missing from the
conditional block that registers MoE tests; add a line to register it alongside
the other tests so the target moeLoraSlotExpandTest (source file
moeLoraSlotExpandTest.cu) is built when USING_OSS_CUTLASS_MOE_GEMM is true—i.e.,
inside the same if(USING_OSS_CUTLASS_MOE_GEMM) ... endif() add an
add_gtest(moeLoraSlotExpandTest moeLoraSlotExpandTest.cu) entry to ensure CI
runs that kernel test.
---
Outside diff comments:
In `@cpp/tensorrt_llm/thop/moeOp.cpp`:
- Around line 593-598: The code currently computes and attaches the legacy
cuBLAS lora_workspace even when device_path.enabled is true; update the logic
around the lora_params_opt block (the call to computeLoraWorkspaceSize that sets
lora_workspace_size) to skip computing/attaching the legacy lora_workspace when
device_path.enabled is true (since moeLoraDeviceRunImpl uses the persistent
device scratch), and apply the same conditional change to the similar block
around lines handling lora_workspace at the other site (the 632-635 occurrence);
specifically, only call computeLoraWorkspaceSize and allocate/set
lora_workspace_size if lora_params_opt.has_value() AND device_path.enabled is
false.
---
Nitpick comments:
In `@cpp/tests/unit_tests/kernels/moeLoraPointerExpandTest.cu`:
- Around line 225-323: Add two new TEST_F cases to MoeLoraPointerExpandTest: (1)
a "GhostRowsRemainZero" test that sets expanded_num_rows greater than
expert_first_token_offset.back() (e.g., expanded_num_rows =
expert_first_token_offset.back() + 2) and uses runAndCompare(permuted_rows,
expert_first_token_offset, ..., expanded_num_rows, lora_dtype_bytes, fc1, fc2,
/*gated=*/nullptr) and asserts the extra/ghost rows are zeroed (use the same
style as existing tests and the runAndCompare helper to validate output); (2) a
"ForceGlobalScanNumExperts1025" test that sets num_experts_per_node = 1025 to
exercise the global-memory scan path in moe_lora_pointer_expand.cu and otherwise
mirror one of the existing simple cases (populate permuted_rows,
expert_first_token_offset, fc1/fc2, expanded_num_rows accordingly). Ensure both
tests use the MoeLoraPointerExpandTest fixture and reference runAndCompare so
they integrate with existing verification logic.
In `@cpp/tests/unit_tests/kernels/moeLoraProblemBuilderTest.cu`:
- Around line 301-329: The test never covers the kernel path where
splitk_offsets == nullptr because every case allocates out.splitk_offsets; add a
new test case in moeLoraProblemBuilderTest (use the TEST_F BoundaryCases block
or a new test) that creates a MoeLoraGemmGroupArrays instance with
out.splitk_offsets left null, call launchMoeLoraProblemBuilder with that arrays
struct (and appropriate input_base/lowrank_workspace/output_base and parameters)
and synchronize the stream, then verify the produced arrays against the
reference (reuse runAndCompare or implement an equivalent comparison) to
exercise the launch_count/null-offset branch in the kernel.
🪄 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: 99b8176f-020c-451a-a0ef-222ba05ac20f
📒 Files selected for processing (13)
cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.hcpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_device_path.hcpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_pointer_expand.hcpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_problem_builder.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/kernels/cutlass_kernels/moe_gemm/moe_lora_problem_builder.cucpp/tensorrt_llm/thop/moeOp.cppcpp/tests/unit_tests/kernels/CMakeLists.txtcpp/tests/unit_tests/kernels/moeLoraPointerExpandTest.cucpp/tests/unit_tests/kernels/moeLoraProblemBuilderTest.cutests/unittest/_torch/lora/test_moe_lora_device_path.pytests/unittest/_torch/lora/test_moe_lora_op.py
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #51930 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #51930 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #52124 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #52124 [ run ] completed with state
|
leslie-fang25
approved these changes
Jun 5, 2026
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #52417 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #52417 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #52456 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #52456 [ run ] completed with state |
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
brb-nv
enabled auto-merge (squash)
June 7, 2026 19:39
Collaborator
|
PR_Github #52601 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #52601 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #52624 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #52624 [ run ] completed with state |
2ez4bz
pushed a commit
to 2ez4bz/TensorRT-LLM
that referenced
this pull request
Jun 8, 2026
…h Cutlass backend - Part 1 (NVIDIA#14923) Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
brb-nv
added a commit
to brb-nv/TensorRT-LLM
that referenced
this pull request
Jun 9, 2026
…h Cutlass backend Squashed combination of the full NVIDIA#14881 work (15 commits) prior to rebase onto main. Part 1 (NVIDIA#14923) has already been merged into main. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Merged
1 task
This was referenced Jun 22, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This is to enable cudagraph support for routed-expert MoE LoRA with Cutlass backend. This is first part of a larger change: routed-expert MoE LoRA with Cutlass backend. Currently, the device path is guarded by an env variable
TLLM_MOE_LORA_USE_DEVICE_PATH- I plan to make this the default path in a future MR after thorough testing.Goal is to enable a fully on-device path for routed-expert MoE LoRA in the Cutlass backend instead of the legacy host-pointer LoRA path (this had a cudaEventSynchronize which made it not amenable to cudagraph). Key functionality:
launchMoeLoraPointerExpandbuilds per-permuted-row (rank, A, B) pointer tables on device from the per-source-token on host.launchMoeLoraProblemBuilderbuilds the per-token grouped-GEMM problem sizes / pointer arrays etc on device.Implementation detail:
moeLoraDeviceRunImpl) therefore lives inth_commonas the wrappers allocate workspace viaat::Tensor.moe_kernels.cureaches the dispatch GEMM through theMoeLoraDeviceRunFnfunction pointer.If there are suggestions of how we can do this better, kindly let me know.
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.