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

cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_device_path.h

@@ -33,13 +33,27 @@ 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).
+// the underlying cudaGraph(SplitK)GroupedGemm wrappers allocate workspace via
+// at::Tensor, which is not linkable from libmoe_gemm_src.a (that archive ends
+// up inside the TensorRT plugin shared object, which deliberately does 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).
+// Contract:
+//   mod:                 per-module device-resident scratch produced by
+//                        launchMoeLoraPointerExpand. The implementation repacks
+//                        it into a MoeLoraGemmGroupArrays, calls the problem
+//                        builder, and dispatches the in/out GEMMs.
+//   num_permuted_tokens: length of the permuted-row tables in mod.
+//   in_hidden_size:      K of the in-GEMM (the input row stride).
+//   max_lora_rank:       workspace ldd and worst-case rank cap.
+//   dtype_bytes:         sizeof(scalar) for the LoRA tensors.
+//   splitk_slices:       split-K factor for the in-GEMM.
+//   input_base:          base of the input matrix (M rows of in_hidden_size).
+//   output_base:         base of the module's output buffer; the implementation
+//                        accumulates into it, so callers must initialize it.
+//   data_type:           scalar dtype expected by the GEMM wrappers
+//                        (fp16/bf16/fp32).
+//   stream:              CUDA stream to launch onto.
 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);

cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_slot_expand.h

@@ -0,0 +1,72 @@
+/*
+ * 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
+// slot -> per-source-token expansion. All pointers refer to device memory.
+//
+// Inputs (indexed by adapter slot):
+//   slot_ranks: int32 [num_slots], per-slot LoRA rank (0 == slot inactive).
+//   slot_ptrs:  int64 [num_slots * 3], pointer bits laid out as
+//               (A_ptr, B_ptr, dora_ptr) per slot; dora is ignored.
+//
+// Outputs (indexed by source token, sized [num_tokens]):
+//   ranks_out: int32 [num_tokens], per-source-token LoRA rank.
+//   ptrs_out:  int64 [num_tokens * 2], per-source-token (A_ptr, B_ptr).
+//
+// These outputs match the (ranks_src, ptrs_src) inputs that
+// launchMoeLoraPointerExpand consumes, so its results feed directly into the
+// pointer-expand stage.
+struct MoeLoraSlotExpandModule
+{
+    int32_t const* slot_ranks = nullptr;
+    int64_t const* slot_ptrs = nullptr;
+    int32_t* ranks_out = nullptr;
+    int64_t* ptrs_out = nullptr;
+};
+
+// Device-side slot -> per-source-token expansion for routed-expert MoE LoRA.
+// For each source token t in [0, num_tokens), reads slot = token_to_slot[t]
+// and copies the slot's (rank, A_ptr, B_ptr) into the per-token output tables
+// for every module. Runs entirely on the stream with no host synchronization,
+// so it is safe to record into a CUDA graph: replaying it re-reads the
+// (in-place updated) device slot tables and token_to_slot, picking up new
+// adapter assignments without re-capture.
+//
+// token_to_slot: int32 [num_tokens] (device). Entries outside [0, num_slots)
+// are treated as inactive (rank 0, null pointers) rather than indexing out of
+// bounds.
+//
+// gated may be nullptr for non-gated activations; when non-null, the gated
+// module's outputs are produced in the same pass.
+void launchMoeLoraSlotExpand(int32_t const* token_to_slot, int64_t num_tokens, int64_t num_slots,
+    MoeLoraSlotExpandModule const& fc1, MoeLoraSlotExpandModule const& fc2, MoeLoraSlotExpandModule const* gated,
+    cudaStream_t stream);
+
+} // namespace kernels::cutlass_kernels
+
+TRTLLM_NAMESPACE_END

cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_slot_expand.cu

@@ -0,0 +1,113 @@
+/*
+ * 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.
+ */
+
+#include "tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_slot_expand.h"
+
+#include "tensorrt_llm/common/assert.h"
+#include "tensorrt_llm/common/cudaUtils.h"
+
+#include <cstdint>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels::cutlass_kernels
+{
+
+namespace
+{
+
+// Threads-per-block. The kernel is a small per-token gather (one slot lookup,
+// a handful of stores per module), so block size only affects occupancy. 256
+// is a safe default for the Hopper / Blackwell targets shared with the rest of
+// the MoE path.
+constexpr int kBlockSize = 256;
+
+// Per-module gather. Inlined into the main kernel so a single thread handles
+// fc1, fc2, and (optionally) gated for its token after one slot lookup.
+__device__ inline void slotExpandOneModule(MoeLoraSlotExpandModule const& mod, int64_t t, int32_t slot, bool valid_slot)
+{
+    if (valid_slot)
+    {
+        mod.ranks_out[t] = mod.slot_ranks[slot];
+        mod.ptrs_out[2 * t + 0] = mod.slot_ptrs[3 * slot + 0];
+        mod.ptrs_out[2 * t + 1] = mod.slot_ptrs[3 * slot + 1];
+    }
+    else
+    {
+        // Out-of-range slot: emit an inactive (rank 0, null pointer) entry so the
+        // downstream pointer-expand / grouped-GEMM treats this token as a no-op
+        // instead of indexing out of bounds.
+        mod.ranks_out[t] = 0;
+        mod.ptrs_out[2 * t + 0] = 0;
+        mod.ptrs_out[2 * t + 1] = 0;
+    }
+}
+
+// One thread per source token. Reads its adapter slot from token_to_slot and
+// scatters the slot's (rank, A_ptr, B_ptr) into the per-token output tables for
+// each module.
+__global__ void moeLoraSlotExpandKernel(int32_t const* __restrict__ token_to_slot, int64_t num_tokens,
+    int64_t num_slots, MoeLoraSlotExpandModule fc1, MoeLoraSlotExpandModule fc2, MoeLoraSlotExpandModule gated,
+    bool has_gated)
+{
+    int64_t const t = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+    if (t >= num_tokens)
+    {
+        return;
+    }
+
+    int32_t const slot = token_to_slot[t];
+    bool const valid_slot = (slot >= 0) && (static_cast<int64_t>(slot) < num_slots);
+
+    slotExpandOneModule(fc1, t, slot, valid_slot);
+    slotExpandOneModule(fc2, t, slot, valid_slot);
+    if (has_gated)
+    {
+        slotExpandOneModule(gated, t, slot, valid_slot);
+    }
+}
+
+} // namespace
+
+void launchMoeLoraSlotExpand(int32_t const* token_to_slot, int64_t num_tokens, int64_t num_slots,
+    MoeLoraSlotExpandModule const& fc1, MoeLoraSlotExpandModule const& fc2, MoeLoraSlotExpandModule const* gated,
+    cudaStream_t stream)
+{
+    if (num_tokens <= 0)
+    {
+        return;
+    }
+    TLLM_CHECK_WITH_INFO(token_to_slot != nullptr, "token_to_slot must be non-null");
+    TLLM_CHECK_WITH_INFO(num_slots > 0, "num_slots must be positive");
+    TLLM_CHECK_WITH_INFO(
+        fc1.slot_ranks != nullptr && fc1.slot_ptrs != nullptr && fc1.ranks_out != nullptr && fc1.ptrs_out != nullptr,
+        "fc1 slot-expand module is missing a buffer");
+    TLLM_CHECK_WITH_INFO(
+        fc2.slot_ranks != nullptr && fc2.slot_ptrs != nullptr && fc2.ranks_out != nullptr && fc2.ptrs_out != nullptr,
+        "fc2 slot-expand module is missing a buffer");
+
+    bool const has_gated = gated != nullptr;
+    MoeLoraSlotExpandModule const gated_arg = has_gated ? *gated : MoeLoraSlotExpandModule{};
+
+    int64_t const grid = (num_tokens + kBlockSize - 1) / kBlockSize;
+    moeLoraSlotExpandKernel<<<static_cast<unsigned int>(grid), kBlockSize, 0, stream>>>(
+        token_to_slot, num_tokens, num_slots, fc1, fc2, gated_arg, has_gated);
+    sync_check_cuda_error(stream);
+}
+
+} // namespace kernels::cutlass_kernels
+
+TRTLLM_NAMESPACE_END