NVIDIA · brb-nv · Jun 8, 2026 · Jun 4, 2026 · Jun 7, 2026
@@ -21,6 +21,7 @@
 #include "tensorrt_llm/common/cudaUtils.h"
 #include "tensorrt_llm/common/quantization.h"
 #include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h"
+#include "tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_device_path.h"
 #include <cstdint>
 #ifdef ENABLE_FP4
 #include <cuda_fp4.h>
@@ -68,6 +69,14 @@ struct LoraParams
 
     cudaEvent_t* memcpy_event_ptr;
 
+    // Device-side capture-safe LoRA path scratch. When device_path.enabled is
+    // true, the kernel uses launchMoeLoraPointerExpand, launchMoeLoraProblemBuilder,
+    // and cudaGraph(SplitK)GroupedGemm instead of the legacy host-pointer
+    // LoraImpl::run path. The pointers refer to persistent allocations owned by
+    // the calling FusedMoeRunner, so their addresses are stable across
+    // CUDA-graph captures and replays.
+    ::tensorrt_llm::kernels::cutlass_kernels::MoeLoraDevicePath device_path;
+
     LoraParams() = default;
 
     LoraParams(int num_reqs, int32_t const* fc1_lora_ranks, void const* const* fc1_lora_weight_ptrs,

@@ -0,0 +1,149 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "tensorrt_llm/common/config.h"
+
+#include <NvInferRuntime.h>
+
+#include <cstdint>
+#include <cuda_runtime.h>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::cutlass_kernels
+{
+
+// Forward declaration; the typedef below references it by name.
+struct MoeLoraDevicePathModule;
+
+// Function-pointer dispatch for the libtorch-dependent GEMM stage of the MoE
+// LoRA device path. The implementation lives in th_common (moeOp.cpp) because
+// the cudaGraph(SplitK)GroupedGemm wrappers allocate workspace via at::Tensor,
+// which cannot be linked from libmoe_gemm_src.a (that archive is also linked
+// into the TensorRT plugin, which must not depend on libtorch).
+//
+// It repacks mod into a MoeLoraGemmGroupArrays, runs the problem builder, and
+// dispatches the in/out GEMMs, accumulating into output_base (which the caller
+// must initialize). data_type is the scalar dtype (fp16/bf16/fp32).
+using MoeLoraDeviceRunFn = void (*)(MoeLoraDevicePathModule const& mod, int64_t num_permuted_tokens,
+    int64_t in_hidden_size, int64_t max_lora_rank, int64_t dtype_bytes, int64_t splitk_slices, void const* input_base,
+    void* output_base, nvinfer1::DataType data_type, cudaStream_t stream);
+
+// Per-module device-resident scratch for the MoE LoRA capture-safe path.
+// Pointers refer to device memory unless noted.
+//
+// The struct is typed with void* rather than the concrete
+// cutlass::gemm::GemmCoord* / int64_t* types so this header can be included
+// from moe_kernels.h without dragging in cutlass headers. The concrete types
+// are recovered at the call site (matching the contract documented in
+// moe_lora_problem_builder.h):
+//
+//   problem_sizes_*  -> cutlass::gemm::GemmCoord* (device, [P_max])
+//   a_ptrs_*/b/d     -> void**                    (device, [P_max])
+//   lda/ldb/ldd_*    -> int64_t*                  (device, [P_max])
+//   splitk_offsets   -> int64_t*                  (device, [P_max + 1])
+//   lowrank_ws_dev   -> void*                     (device, [P_max, max_lora_rank, dtype_bytes])
+//   host_max_*       -> cutlass::gemm::GemmCoord* (pinned host, [1])
+//
+// The split-K in-GEMM's partial-sum scratch is allocated internally by the
+// cuda_graph_split_k_grouped_gemm wrapper (sized from the host max-problem
+// hint); only the per-problem splitk_offsets are produced here.
+//
+// out_hidden_size is the trailing dimension of the module's output buffer; it
+// is inter_size for fc1/gated and hidden_size for fc2. The output base address
+// itself is passed directly to runMoeLoraDeviceModule at the call site.
+struct MoeLoraDevicePathModule
+{
+    // Per-source-token (rank, A_ptr, B_ptr) device mirrors, staged via a
+    // pinned-host to device async H2D in FusedMoeRunner::buildMoeLoraParams.
+    // These feed launchMoeLoraPointerExpand as ranks_src / ptrs_src.
+    int32_t const* ranks_src_dev = nullptr;
+    int64_t const* ptrs_src_dev = nullptr;
+
+    // Inner (A) and outer (B) dimensions for this module, fed to the
+    // pointer-expand kernel as dim_a / dim_b so it can compute the per-expert
+    // offset weight_index * dim * lora_rank. For fc1/gated this is
+    // (hidden_size, inter_size); for fc2 it is (inter_size, hidden_size).
+    int64_t dim_a = 0;
+    int64_t dim_b = 0;
+
+    // Per-permuted-row (rank, A_ptr + offset, B_ptr + offset).
+    int32_t* permuted_ranks_dev = nullptr;
+    int64_t* permuted_ptrs_dev = nullptr;
+
+    // cuda_graph_(split_k_)grouped_gemm-ready bundle.
+    void* problem_sizes_in_dev = nullptr;
+    void* problem_sizes_out_dev = nullptr;
+    void** a_ptrs_in_dev = nullptr;
+    void** b_ptrs_in_dev = nullptr;
+    void** d_ptrs_in_dev = nullptr;
+    void** b_ptrs_out_dev = nullptr;
+    void** d_ptrs_out_dev = nullptr;
+    int64_t* lda_in_dev = nullptr;
+    int64_t* ldb_in_dev = nullptr;
+    int64_t* ldd_in_dev = nullptr;
+    int64_t* ldb_out_dev = nullptr;
+    int64_t* ldd_out_dev = nullptr;
+    int64_t* splitk_offsets_dev = nullptr;
+
+    // Low-rank intermediate workspace shared between the in- and out-GEMM. The
+    // split-K partial-sum scratch is owned by the GEMM wrapper, not here.
+    void* lowrank_workspace_dev = nullptr;
+
+    // Host (pinned) per-call max problem size hints, required by the
+    // cuda_graph_*_grouped_gemm wrappers for kernel selection. The
+    // values are upper bounds (max_M, max_N, max_K) safe to fix at
+    // warmup time.
+    void* host_max_problem_in_pinned = nullptr;
+    void* host_max_problem_out_pinned = nullptr;
+
+    // Trailing dimension of the module's output buffer (inter_size for
+    // fc1/gated, hidden_size for fc2). The output base address is supplied
+    // directly to runMoeLoraDeviceModule at the call site.
+    int64_t out_hidden_size = 0;
+};
+
+// Top-level device-path bundle attached to LoraParams when the device LoRA
+// path is active. enabled == false means the FusedMoeRunner runs the legacy
+// host path.
+struct MoeLoraDevicePath
+{
+    bool enabled = false;
+
+    // Scalars common to all three modules. Fixed for the lifetime of the
+    // FusedMoeRunner once the scratch is allocated.
+    int64_t in_hidden_size = 0;
+    int64_t max_lora_rank = 0;
+    int64_t dtype_bytes = 0;
+    int64_t splitk_slices = 0;
+
+    bool has_gated = false;
+
+    // libtorch-bound GEMM dispatch entry point, populated by moeOp.cpp when the
+    // device path is enabled. nullptr means the device path is unavailable from
+    // this consumer (for example, the TensorRT plugin).
+    MoeLoraDeviceRunFn run = nullptr;
+
+    MoeLoraDevicePathModule fc1;
+    MoeLoraDevicePathModule fc2;
+    MoeLoraDevicePathModule gated;
+};
+
+} // namespace kernels::cutlass_kernels
+
+TRTLLM_NAMESPACE_END
@@ -0,0 +1,81 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "tensorrt_llm/common/config.h"
+
+#include <cstdint>
+#include <cuda_runtime.h>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::cutlass_kernels
+{
+
+// Device-side description of one LoRA module (fc1, fc2, or gated) for the
+// MoE per-token A/B pointer-table expansion. All pointers refer to device
+// memory.
+//
+// Inputs (per source token, indexed by the pre-permutation source row index):
+//   ranks_src: int32 [num_rows], per-source-token LoRA rank.
+//   ptrs_src:  int64 [num_rows * 2], pointer bits laid out as (A_ptr, B_ptr)
+//              per source token.
+//
+// Outputs (per permuted row, sized expanded_num_rows == num_rows * top_k):
+//   ranks_out: int32 [expanded_num_rows], per-permuted-row LoRA rank.
+//   ptrs_out:  int64 [expanded_num_rows * 2], per-permuted-row
+//              (A_ptr + offset, B_ptr + offset). The per-expert offset is
+//              weight_index * dim * rank * lora_dtype_bytes, so the consumer
+//              can reinterpret directly as the LoRA scalar type.
+//
+// dim_a and dim_b are the non-rank dimension of A and B respectively:
+// fc1/gated use (hidden_size, inter_size); fc2 uses (inter_size, hidden_size).
+struct MoeLoraExpandModule
+{
+    int32_t const* ranks_src = nullptr;
+    int64_t const* ptrs_src = nullptr;
+    int64_t dim_a = 0;
+    int64_t dim_b = 0;
+    int32_t* ranks_out = nullptr;
+    int64_t* ptrs_out = nullptr;
+};
+
+// Device-side replacement for the host-CPU pointer fan-out in
+// CutlassMoeFCRunner::setupLoraWorkspace. Reads per-source-token LoRA metadata
+// and permuted_rows, and writes per-permuted-row pointer tables directly into
+// device memory. It performs no host synchronization and no cudaMemcpyAsync
+// staging, so it is safe to launch from a captured CUDA graph.
+//
+// expert_first_token_offset has shape [num_experts_per_node + 1] (int64,
+// device-resident). The kernel uses it both to find the expert a permuted row
+// belongs to and to derive weight_index = local_expert_idx + start_expert for
+// the per-expert weight-buffer stride.
+//
+// lora_dtype_bytes is the size in bytes of the LoRA matrix scalar (e.g. 2 for
+// bf16/fp16). It scales the stride applied to the A/B pointers so consumers can
+// reinterpret the result directly as the appropriate scalar type.
+//
+// gated may be nullptr for non-gated activations; when non-null, the gated
+// module's outputs are produced in the same pass.
+void launchMoeLoraPointerExpand(int32_t const* permuted_rows, int64_t const* expert_first_token_offset,
+    int32_t num_experts_per_node, int32_t start_expert, int64_t num_rows, int64_t expanded_num_rows,
+    int64_t lora_dtype_bytes, MoeLoraExpandModule const& fc1, MoeLoraExpandModule const& fc2,
+    MoeLoraExpandModule const* gated, cudaStream_t stream);
+
+} // namespace kernels::cutlass_kernels
+
+TRTLLM_NAMESPACE_END
@@ -0,0 +1,118 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "tensorrt_llm/common/config.h"
+
+#include "cutlass/gemm_coord.h"
+
+#include <cstdint>
+#include <cuda_runtime.h>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::cutlass_kernels
+{
+
+// Caller-owned device-output bundle for one LoRA module. Each array is sized
+// for the maximum permuted-token count the FusedMoeRunner expects to see;
+// the builder fills the first num_permuted_tokens entries each call.
+//
+// Layout convention (mirrors attention LoRA in cuda_graph_grouped_gemm.h):
+//   In-GEMM:  D = A @ B with C aliased to D when there's no bias.
+//             A = input slice  [M=1, K=in_hidden_size]
+//             B = adapter A   [K=in_hidden_size, N=rank]
+//             D = lowrank slice [M=1, N=rank]
+//   Out-GEMM: D = A @ B with C aliased to D.
+//             A = lowrank slice [M=1, K=rank]    (= in-GEMM's D)
+//             B = adapter B    [K=rank, N=out_hidden_size]
+//             D = output slice [M=1, N=out_hidden_size]
+//
+// Because ptrC aliases ptrD in both GEMMs (no bias), only ptrD is exposed
+// per GEMM; the cuda_graph_grouped_gemm wrapper accepts the same address
+// for both. d_ptrs_in also serves as a_ptrs_out (the LoRA intermediate is
+// the input to the second GEMM); only one set of low-rank pointers is
+// produced for that reason.
+struct MoeLoraGemmGroupArrays
+{
+    // Per-problem (M, N, K) for the in-GEMM and out-GEMM respectively.
+    cutlass::gemm::GemmCoord* problem_sizes_in = nullptr;  // [P]
+    cutlass::gemm::GemmCoord* problem_sizes_out = nullptr; // [P]
+
+    // In-GEMM pointer arrays. ptr_c_in is implicit (== d_ptrs_in).
+    void** a_ptrs_in = nullptr; // [P]: input row pointer
+    void** b_ptrs_in = nullptr; // [P]: adapter A pointer (with per-expert offset)
+    void** d_ptrs_in = nullptr; // [P]: lowrank workspace row (also a_ptrs_out)
+
+    // Out-GEMM pointer arrays. ptr_c_out is implicit (== d_ptrs_out).
+    void** b_ptrs_out = nullptr; // [P]: adapter B pointer (with per-expert offset)
+    void** d_ptrs_out = nullptr; // [P]: output row pointer
+
+    // Leading dimensions. All row-major, fixed per problem given uniform
+    // input / lowrank-workspace / output strides.
+    int64_t* lda_in = nullptr;  // [P]: in_hidden_size
+    int64_t* ldb_in = nullptr;  // [P]: in_hidden_size (stride in adapter-A storage)
+    int64_t* ldd_in = nullptr;  // [P]: max_lora_rank (workspace stride)
+    int64_t* ldb_out = nullptr; // [P]: per-token rank (stride in adapter-B storage)
+    int64_t* ldd_out = nullptr; // [P]: out_hidden_size
+
+    // Per-problem exclusive prefix offset into the split-K scratch buffer
+    // used by the in-GEMM. Element [P] (one past the end) holds the total
+    // scratch size in fp32 elements, matching the layout that
+    // cuda_graph_split_k_grouped_gemm consumes.
+    int64_t* splitk_offsets = nullptr; // [P + 1]
+};
+
+// Device-side problem-and-pointer builder for one MoE LoRA module. It consumes
+// the per-permuted-row outputs of launchMoeLoraPointerExpand plus uniform
+// input, workspace, and output base addresses, and writes every device-resident
+// input the cuda_graph_(split_k_)grouped_gemm wrappers need.
+//
+// Inputs:
+//   ranks_dev: int32 [P], per-permuted-row LoRA rank.
+//   ptrs_dev:  int64 [P*2], per-permuted-row (A_ptr + offset, B_ptr + offset),
+//              already adjusted for the per-expert weight stride by the
+//              pointer-expand kernel.
+//
+// Base pointers (the per-token row offset is computed inside the kernel from
+// i * stride * dtype_bytes):
+//   input_base:        [P, in_hidden_size]
+//   lowrank_workspace: [P, max_lora_rank], reused as the in-GEMM output and
+//                      the out-GEMM input.
+//   output_base:       [P, out_hidden_size]
+//
+// Scalars:
+//   in_hidden_size:  K for the in-GEMM, also lda_in[i] and ldb_in[i].
+//   out_hidden_size: N for the out-GEMM, also ldd_out[i].
+//   max_lora_rank:   ldd_in[i], the workspace stride, fixed regardless of the
+//                    per-token rank so the GEMM lands at a known offset.
+//                    The out-GEMM's ldb_out[i] is the per-token rank (adapter B
+//                    is stored [out_hidden_size, rank]), not out_hidden_size.
+//   dtype_bytes:     scalar size in bytes (2 for bf16/fp16, 4 for fp32).
+//   splitk_slices:   split-K factor for the in-GEMM; drives the per-problem
+//                    split-K scratch stride.
+//
+// The split-K stride is a worst-case fixed value (max_lora_rank * splitk_slices
+// per problem) so the offsets can be computed from i alone without a prefix-sum.
+void launchMoeLoraProblemBuilder(int32_t const* ranks_dev, int64_t const* ptrs_dev, void const* input_base,
+    void* lowrank_workspace, void* output_base, int64_t num_permuted_tokens, int64_t in_hidden_size,
+    int64_t out_hidden_size, int64_t max_lora_rank, int64_t dtype_bytes, int64_t splitk_slices,
+    MoeLoraGemmGroupArrays const& out, cudaStream_t stream);
+
+} // namespace kernels::cutlass_kernels
+
+TRTLLM_NAMESPACE_END