From 64caf61cef2bb0a6f339ea4039025568587a82f0 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Wed, 3 Jun 2026 17:24:28 -0700 Subject: [PATCH 1/2] [TRTLLM-12507][feat] Cudagraph support for routed-expert MoE LoRA with Cutlass backend - Part 1 Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../cutlass_kernels/include/moe_kernels.h | 9 + .../include/moe_lora_device_path.h | 149 +++++ .../include/moe_lora_pointer_expand.h | 81 +++ .../include/moe_lora_problem_builder.h | 116 ++++ .../cutlass_kernels/moe_gemm/moe_kernels.cu | 150 ++++- .../moe_gemm/moe_lora_pointer_expand.cu | 161 +++++ .../moe_gemm/moe_lora_problem_builder.cu | 147 +++++ cpp/tensorrt_llm/thop/moeOp.cpp | 569 ++++++++++++++++-- cpp/tests/unit_tests/kernels/CMakeLists.txt | 5 + .../kernels/moeLoraPointerExpandTest.cu | 325 ++++++++++ .../kernels/moeLoraProblemBuilderTest.cu | 331 ++++++++++ .../_torch/lora/test_moe_lora_device_path.py | 184 ++++++ .../unittest/_torch/lora/test_moe_lora_op.py | 48 ++ 13 files changed, 2220 insertions(+), 55 deletions(-) create mode 100644 cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_device_path.h create mode 100644 cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_pointer_expand.h create mode 100644 cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_problem_builder.h create mode 100644 cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_pointer_expand.cu create mode 100644 cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_problem_builder.cu create mode 100644 cpp/tests/unit_tests/kernels/moeLoraPointerExpandTest.cu create mode 100644 cpp/tests/unit_tests/kernels/moeLoraProblemBuilderTest.cu create mode 100644 tests/unittest/_torch/lora/test_moe_lora_device_path.py diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h index 4e76d2be5614..216877a4ffc7 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h @@ -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 #ifdef ENABLE_FP4 #include @@ -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, diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_device_path.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_device_path.h new file mode 100644 index 000000000000..2d2c133255c4 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_device_path.h @@ -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 + +#include +#include + +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 diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_pointer_expand.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_pointer_expand.h new file mode 100644 index 000000000000..c2f509c4ed9c --- /dev/null +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_pointer_expand.h @@ -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 +#include + +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 diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_problem_builder.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_problem_builder.h new file mode 100644 index 000000000000..a1a43c427a9e --- /dev/null +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_problem_builder.h @@ -0,0 +1,116 @@ +/* + * 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 +#include + +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 ldb_out[i] and 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. +// 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 diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu index a0b08a3df775..1bca18a54ea9 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu @@ -58,7 +58,14 @@ #include "tensorrt_llm/kernels/preQuantScaleKernel.h" #include "tensorrt_llm/kernels/quantization.cuh" +#include "tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_pointer_expand.h" #include "tensorrt_llm/kernels/cutlass_kernels/include/moe_util_kernels.h" +// NOTE: the device-path GEMM dispatch (cudaGraph(SplitK)GroupedGemm, +// launchMoeLoraProblemBuilder) is not called here. Those wrappers pull in +// libtorch via at::Tensor, and this file is archived into libmoe_gemm_src.a, +// which the TensorRT plugin also links and must keep libtorch-free. The +// dispatch is reached through the LoraParams::device_path.run function pointer, +// populated in moeOp.cpp. #ifndef CUDART_VERSION #error CUDART_VERSION Undefined! @@ -3657,6 +3664,49 @@ void CutlassMoeFCRunner +constexpr nvinfer1::DataType moeLoraNvInferType() +{ + if constexpr (std::is_same_v) + { + return nvinfer1::DataType::kHALF; + } +#if defined(ENABLE_BF16) + else if constexpr (std::is_same_v) + { + return nvinfer1::DataType::kBF16; + } +#endif + else if constexpr (std::is_same_v) + { + return nvinfer1::DataType::kFLOAT; + } + else + { + static_assert(sizeof(ScaleBiasType) == 0, "MoE LoRA device path supports fp16/bf16/fp32 only."); + } +} + template bool CutlassMoeFCRunner::setupLoraWorkspace( int64_t expanded_num_rows, int64_t num_rows, int64_t inter_size, int64_t hidden_size, int start_expert, @@ -3675,6 +3725,49 @@ bool CutlassMoeFCRunner(permuted_data_); } - void* lora_workspace = lora_params.workspace; - void* tmp_lora_fc_result = static_cast(lora_fc1_result); - int64_t num_valid_tokens = host_expert_first_token_offset[num_experts_per_node]; - int64_t num_reqs_lora = std::min(num_valid_tokens, static_cast(num_reqs * num_experts_per_node)); + // Device-path branch, running entirely on the stream. setupLoraWorkspace + // has already populated the per-permuted-row ranks and pointers for fc1 and + // gated via launchMoeLoraPointerExpand. + if (lora_params.device_path.enabled) + { + auto const& dp = lora_params.device_path; + nvinfer1::DataType const data_type = moeLoraNvInferType(); - ::tensorrt_llm::kernels::Lora_run(fc1_lora_impl.get(), num_valid_tokens, num_reqs_lora, input, - host_permuted_fc1_lora_ranks.data(), host_permuted_fc1_weight_ptrs.data(), 0, &tmp_lora_fc_result, - lora_workspace, stream); + runMoeLoraDeviceModule(dp.fc1, expanded_num_rows, /*in_hidden_size=*/hidden_size, dp.max_lora_rank, + dp.dtype_bytes, dp.splitk_slices, /*input_base=*/static_cast(input), + /*output_base=*/static_cast(lora_fc1_result), dp.run, data_type, stream); - if (is_gated_activation) + if (is_gated_activation) + { + runMoeLoraDeviceModule(dp.gated, expanded_num_rows, /*in_hidden_size=*/hidden_size, dp.max_lora_rank, + dp.dtype_bytes, dp.splitk_slices, /*input_base=*/static_cast(input), + /*output_base=*/static_cast(lora_gated_out), dp.run, data_type, stream); + } + } + else { - void* tmp_lora_gated_result = static_cast(lora_gated_out); + void* lora_workspace = lora_params.workspace; + void* tmp_lora_fc_result = static_cast(lora_fc1_result); + int64_t num_valid_tokens = host_expert_first_token_offset[num_experts_per_node]; + int64_t num_reqs_lora = std::min(num_valid_tokens, static_cast(num_reqs * num_experts_per_node)); + ::tensorrt_llm::kernels::Lora_run(fc1_lora_impl.get(), num_valid_tokens, num_reqs_lora, input, - host_permuted_gated_lora_ranks.data(), host_permuted_gated_weight_ptrs.data(), 0, &tmp_lora_gated_result, + host_permuted_fc1_lora_ranks.data(), host_permuted_fc1_weight_ptrs.data(), 0, &tmp_lora_fc_result, lora_workspace, stream); + + if (is_gated_activation) + { + void* tmp_lora_gated_result = static_cast(lora_gated_out); + ::tensorrt_llm::kernels::Lora_run(fc1_lora_impl.get(), num_valid_tokens, num_reqs_lora, input, + host_permuted_gated_lora_ranks.data(), host_permuted_gated_weight_ptrs.data(), 0, + &tmp_lora_gated_result, lora_workspace, stream); + } } // add bias and reorder @@ -3849,6 +3964,21 @@ void CutlassMoeFCRunner(fc1_result_); } + // Device-path branch, mirroring loraFC1's branch. It consumes the + // per-permuted-row ranks and pointers that setupLoraWorkspace produced via + // launchMoeLoraPointerExpand. num_tokens here is expanded_num_rows from + // runMoe (top_k * num_rows). + if (lora_params.device_path.enabled) + { + auto const& dp = lora_params.device_path; + nvinfer1::DataType const data_type = moeLoraNvInferType(); + runMoeLoraDeviceModule(dp.fc2, num_tokens, /*in_hidden_size=*/inter_size, dp.max_lora_rank, dp.dtype_bytes, + dp.splitk_slices, /*input_base=*/static_cast(input), + /*output_base=*/static_cast(lora_fc2_result_), dp.run, data_type, stream); + sync_check_cuda_error(stream); + return; + } + void* lora_workspace = lora_params.workspace; int64_t num_valid_tokens = host_expert_first_token_offset[num_experts_per_node]; void* tmp_lora_fc_result = static_cast(lora_fc2_result_); diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_pointer_expand.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_pointer_expand.cu new file mode 100644 index 000000000000..a12e2cadb286 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_pointer_expand.cu @@ -0,0 +1,161 @@ +/* + * 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_pointer_expand.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::cutlass_kernels +{ + +namespace +{ + +// Threads-per-block. The kernel is bandwidth-bound, so block size mainly +// affects occupancy. 256 is a good default for Hopper/Blackwell. +constexpr int kBlockSize = 256; + +// Cap on num_experts_per_node staged in shared memory for the expert lookup. +// Above this the kernel falls back to a global-memory scan (still correct). +// Set well above realistic values (typical MoE uses 8-64 experts per node). +constexpr int kMaxExpertsInSmem = 1024; + +// Per-module expansion. Inlined into the main kernel so we only pay one +// permuted_rows[i] and expert lookup per output row. +__device__ inline void expandOneModule( + MoeLoraExpandModule const& mod, int64_t i, int32_t source_index, int64_t weight_index, int64_t lora_dtype_bytes) +{ + int32_t const rank = mod.ranks_src[source_index]; + + // Per-expert byte offsets: weight_index * dim * rank * sizeof(scalar). + int64_t const a_stride = weight_index * mod.dim_a * rank * lora_dtype_bytes; + int64_t const b_stride = weight_index * mod.dim_b * rank * lora_dtype_bytes; + + int64_t const a_src = mod.ptrs_src[2 * source_index + 0]; + int64_t const b_src = mod.ptrs_src[2 * source_index + 1]; + + // Pointer arithmetic in raw bytes (uintptr_t-equivalent). Consumers + // reinterpret to the LoRA scalar type with no further offset, matching + // the existing host-loop semantics in setupLoraWorkspace. + mod.ptrs_out[2 * i + 0] = a_src + a_stride; + mod.ptrs_out[2 * i + 1] = b_src + b_stride; + mod.ranks_out[i] = rank; +} + +// One thread per permuted row: find its expert via search over +// expert_first_token_offset (staged in shared memory), compute +// source_index = permuted_rows[i] % num_rows, and expand fc1, fc2, and +// (optionally) gated. Rank 0 is a per-token no-op in the grouped GEMM, so no +// global "any-token-has-lora" reduction is needed. +__global__ void moeLoraPointerExpandKernel(int32_t const* __restrict__ permuted_rows, + int64_t const* __restrict__ 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 fc1, + MoeLoraExpandModule fc2, MoeLoraExpandModule gated, bool has_gated) +{ + // Stage expert_first_token_offset in shared memory once per block; every + // thread reads it during the expert lookup below. + extern __shared__ int64_t smem_first_token_offset[]; + bool const use_smem = num_experts_per_node + 1 <= kMaxExpertsInSmem; + if (use_smem) + { + for (int e = threadIdx.x; e < num_experts_per_node + 1; e += blockDim.x) + { + smem_first_token_offset[e] = expert_first_token_offset[e]; + } + __syncthreads(); + } + + int64_t const i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= expanded_num_rows) + { + return; + } + + // Find expert_idx s.t. first_offset[expert_idx] <= i < first_offset[expert_idx + 1]. + // Linear scan; num_experts_per_node is small enough (~8-64 typical) that + // a binary search adds branch divergence with no meaningful speedup. + int64_t const* offsets = use_smem ? smem_first_token_offset : expert_first_token_offset; + int32_t expert_idx = 0; + for (int32_t e = 0; e < num_experts_per_node; ++e) + { + if (offsets[e + 1] > i) + { + expert_idx = e; + break; + } + expert_idx = e + 1; + } + // Tokens past the last valid offset (padding "ghost" rows) get + // expert_idx == num_experts_per_node; drop them so weight_index cannot run + // off the expert table. + if (expert_idx >= num_experts_per_node) + { + return; + } + + int64_t const weight_index = static_cast(expert_idx) + start_expert; + int32_t const source_index = static_cast(permuted_rows[i] % num_rows); + + expandOneModule(fc1, i, source_index, weight_index, lora_dtype_bytes); + expandOneModule(fc2, i, source_index, weight_index, lora_dtype_bytes); + if (has_gated) + { + expandOneModule(gated, i, source_index, weight_index, lora_dtype_bytes); + } +} + +} // namespace + +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) +{ + if (expanded_num_rows <= 0) + { + return; + } + TLLM_CHECK_WITH_INFO(permuted_rows != nullptr, "permuted_rows must be non-null"); + TLLM_CHECK_WITH_INFO(expert_first_token_offset != nullptr, "expert_first_token_offset must be non-null"); + TLLM_CHECK_WITH_INFO(num_experts_per_node > 0, "num_experts_per_node must be positive"); + TLLM_CHECK_WITH_INFO(num_rows > 0, "num_rows must be positive"); + TLLM_CHECK_WITH_INFO(lora_dtype_bytes > 0, "lora_dtype_bytes must be positive"); + + bool const has_gated = gated != nullptr; + MoeLoraExpandModule const gated_arg = has_gated ? *gated : MoeLoraExpandModule{}; + + int64_t const grid = (expanded_num_rows + kBlockSize - 1) / kBlockSize; + // Reserve shared memory only when the expert table fits. Above the cap the + // kernel falls back to global-memory reads (still correct, no staging), so + // we pass 0 bytes to avoid allocating shared memory we will not touch. + int const smem_entries = num_experts_per_node + 1; + size_t const smem_bytes + = (smem_entries <= kMaxExpertsInSmem) ? static_cast(smem_entries) * sizeof(int64_t) : 0; + + moeLoraPointerExpandKernel<<(grid), kBlockSize, smem_bytes, stream>>>(permuted_rows, + expert_first_token_offset, num_experts_per_node, start_expert, num_rows, expanded_num_rows, lora_dtype_bytes, + fc1, fc2, gated_arg, has_gated); + sync_check_cuda_error(stream); +} + +} // namespace kernels::cutlass_kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_problem_builder.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_problem_builder.cu new file mode 100644 index 000000000000..cd4dd30991cb --- /dev/null +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_problem_builder.cu @@ -0,0 +1,147 @@ +/* + * 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_problem_builder.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::cutlass_kernels +{ + +namespace +{ + +// Threads-per-block. The kernel is bandwidth-bound, so block size mainly +// controls occupancy. 256 is a good default for Hopper/Blackwell. +constexpr int kBlockSize = 256; + +// One thread per permuted row writes all output arrays. Each store stream is +// contiguous, so accesses coalesce; there is no inter-thread communication. +__global__ void moeLoraProblemBuilderKernel(int32_t const* __restrict__ ranks, int64_t const* __restrict__ ptrs, + int64_t input_base, int64_t lowrank_workspace, int64_t 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, + cutlass::gemm::GemmCoord* __restrict__ problem_sizes_in, cutlass::gemm::GemmCoord* __restrict__ problem_sizes_out, + void** __restrict__ a_ptrs_in, void** __restrict__ b_ptrs_in, void** __restrict__ d_ptrs_in, + void** __restrict__ b_ptrs_out, void** __restrict__ d_ptrs_out, int64_t* __restrict__ lda_in, + int64_t* __restrict__ ldb_in, int64_t* __restrict__ ldd_in, int64_t* __restrict__ ldb_out, + int64_t* __restrict__ ldd_out, int64_t* __restrict__ splitk_offsets) +{ + int64_t const i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= num_permuted_tokens) + { + // The +1 splitk_offsets sentinel (one past num_permuted_tokens) is + // written by thread 0 of the last block; everyone else returns. + if (i == num_permuted_tokens && splitk_offsets != nullptr) + { + splitk_offsets[num_permuted_tokens] = num_permuted_tokens * max_lora_rank * splitk_slices; + } + return; + } + + int32_t const rank = ranks[i]; + int64_t const a_ptr_bits = ptrs[2 * i + 0]; + int64_t const b_ptr_bits = ptrs[2 * i + 1]; + + // Problem sizes: each permuted token gets its own (M=1) GEMM. This matches + // worst-case scheduling with no run-length aggregation; a future + // optimization can aggregate consecutive identical-adapter tokens. + problem_sizes_in[i] = cutlass::gemm::GemmCoord(1, rank, static_cast(in_hidden_size)); + problem_sizes_out[i] = cutlass::gemm::GemmCoord(1, static_cast(out_hidden_size), rank); + + // Pointer rows. dtype_bytes scales the per-row stride so the same + // builder serves bf16/fp16/fp32 adapters without templating. + int64_t const in_row_stride = in_hidden_size * dtype_bytes; + int64_t const work_row_stride = max_lora_rank * dtype_bytes; + int64_t const out_row_stride = out_hidden_size * dtype_bytes; + + a_ptrs_in[i] = reinterpret_cast(input_base + i * in_row_stride); + b_ptrs_in[i] = reinterpret_cast(a_ptr_bits); + d_ptrs_in[i] = reinterpret_cast(lowrank_workspace + i * work_row_stride); + b_ptrs_out[i] = reinterpret_cast(b_ptr_bits); + d_ptrs_out[i] = reinterpret_cast(output_base + i * out_row_stride); + + // Leading dimensions. For the in-/out- GEMMs, lda/ldd correspond to the + // input row-stride / workspace row-stride / output row-stride; ldb is + // the per-problem stride in the LoRA adapter's storage and matches + // the cuda_graph_grouped_gemm convention used by attention LoRA + // (loraOp.cpp): + // in-GEMM: adapter A stored as [rank, in_hidden_size] + // -> ldb_in = in_hidden_size + // out-GEMM: adapter B stored as [out_hidden_size, rank] + // -> ldb_out = rank (per-token, since per-token rank + // can differ in slot-indexed multi-LoRA mode) + lda_in[i] = in_hidden_size; + ldb_in[i] = in_hidden_size; + ldd_in[i] = max_lora_rank; + ldb_out[i] = rank; + ldd_out[i] = out_hidden_size; + + // Split-K scratch offsets. Worst-case fixed stride (independent of + // per-token rank) so each thread computes its own offset locally; no + // cross-thread scan needed. + if (splitk_offsets != nullptr) + { + splitk_offsets[i] = i * max_lora_rank * splitk_slices; + } +} + +} // namespace + +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) +{ + if (num_permuted_tokens <= 0) + { + return; + } + TLLM_CHECK_WITH_INFO(ranks_dev != nullptr, "ranks_dev must be non-null"); + TLLM_CHECK_WITH_INFO(ptrs_dev != nullptr, "ptrs_dev must be non-null"); + TLLM_CHECK_WITH_INFO(out.problem_sizes_in != nullptr, "problem_sizes_in must be non-null"); + TLLM_CHECK_WITH_INFO(out.problem_sizes_out != nullptr, "problem_sizes_out must be non-null"); + TLLM_CHECK_WITH_INFO(out.a_ptrs_in && out.b_ptrs_in && out.d_ptrs_in && out.b_ptrs_out && out.d_ptrs_out, + "All ptr_*_in/out arrays must be non-null"); + TLLM_CHECK_WITH_INFO( + out.lda_in && out.ldb_in && out.ldd_in && out.ldb_out && out.ldd_out, "All ld* arrays must be non-null"); + TLLM_CHECK_WITH_INFO(dtype_bytes > 0, "dtype_bytes must be positive"); + TLLM_CHECK_WITH_INFO(max_lora_rank > 0, "max_lora_rank must be positive"); + TLLM_CHECK_WITH_INFO(in_hidden_size > 0 && out_hidden_size > 0, "hidden sizes must be positive"); + TLLM_CHECK_WITH_INFO(splitk_slices > 0, "splitk_slices must be positive"); + + // Launch one extra thread so the splitk_offsets[num_permuted_tokens] + // sentinel can be filled by exactly one thread (cleaner than a + // dedicated tail launch). + int64_t const launch_count = num_permuted_tokens + (out.splitk_offsets != nullptr ? 1 : 0); + int64_t const grid = (launch_count + kBlockSize - 1) / kBlockSize; + + moeLoraProblemBuilderKernel<<(grid), kBlockSize, 0, stream>>>(ranks_dev, ptrs_dev, + reinterpret_cast(input_base), reinterpret_cast(lowrank_workspace), + reinterpret_cast(output_base), num_permuted_tokens, in_hidden_size, out_hidden_size, max_lora_rank, + dtype_bytes, splitk_slices, out.problem_sizes_in, out.problem_sizes_out, out.a_ptrs_in, out.b_ptrs_in, + out.d_ptrs_in, out.b_ptrs_out, out.d_ptrs_out, out.lda_in, out.ldb_in, out.ldd_in, out.ldb_out, out.ldd_out, + out.splitk_offsets); + sync_check_cuda_error(stream); +} + +} // namespace kernels::cutlass_kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/moeOp.cpp b/cpp/tensorrt_llm/thop/moeOp.cpp index 454dbe07dcb9..824b1909ddd4 100644 --- a/cpp/tensorrt_llm/thop/moeOp.cpp +++ b/cpp/tensorrt_llm/thop/moeOp.cpp @@ -21,12 +21,17 @@ #include "moe_kernels.h" #endif // Always include the public header for moe_gemm_kernels.h +#include "cutlass/gemm_coord.h" #include "tensorrt_llm/kernels/cutlass_kernels/include/moe_gemm_kernels.h" +#include "tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_device_path.h" +#include "tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_problem_builder.h" #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/cublasMMWrapper.h" +#include "tensorrt_llm/common/dataType.h" #include "tensorrt_llm/common/opUtils.h" #include "tensorrt_llm/common/workspace.h" +#include "tensorrt_llm/kernels/cuda_graph_grouped_gemm.h" #include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h" #include "tensorrt_llm/kernels/cutlass_kernels/include/cutlass_kernel_selector.h" #include "tensorrt_llm/kernels/lora/lora.h" @@ -68,6 +73,70 @@ enum class MoeLoraRequestType : int32_t kGENERATION = 1 }; +// --------------------------------------------------------------------------- +// libtorch-bound implementation of MoeLoraDeviceRunFn. +// +// The per-module GEMM dispatch for the device LoRA path: builds the per-token +// problem descriptors on device via launchMoeLoraProblemBuilder, then +// dispatches cudaGraph(SplitK)GroupedGemm. The latter allocates workspace via +// at::Tensor, so this lives in th_common (which links libtorch); moe_kernels.cu +// reaches it through LoraParams::device_path.run, keeping libmoe_gemm_src.a +// (and the TensorRT plugin) libtorch-free. +// --------------------------------------------------------------------------- +inline void moeLoraDeviceRunImpl(::tensorrt_llm::kernels::cutlass_kernels::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) +{ + TLLM_CHECK_WITH_INFO(mod.permuted_ranks_dev != nullptr, + "Device-path LoRA module is missing permuted ranks buffer (forgot to populate device_path?)."); + + // Repack the device-resident scratch into the bundle the problem-builder + // consumes. The typed casts recover the concrete pointer types that + // MoeLoraDevicePathModule stores as void* for header decoupling. + ::tensorrt_llm::kernels::cutlass_kernels::MoeLoraGemmGroupArrays arrays{}; + arrays.problem_sizes_in = static_cast(mod.problem_sizes_in_dev); + arrays.problem_sizes_out = static_cast(mod.problem_sizes_out_dev); + arrays.a_ptrs_in = mod.a_ptrs_in_dev; + arrays.b_ptrs_in = mod.b_ptrs_in_dev; + arrays.d_ptrs_in = mod.d_ptrs_in_dev; + arrays.b_ptrs_out = mod.b_ptrs_out_dev; + arrays.d_ptrs_out = mod.d_ptrs_out_dev; + arrays.lda_in = mod.lda_in_dev; + arrays.ldb_in = mod.ldb_in_dev; + arrays.ldd_in = mod.ldd_in_dev; + arrays.ldb_out = mod.ldb_out_dev; + arrays.ldd_out = mod.ldd_out_dev; + arrays.splitk_offsets = mod.splitk_offsets_dev; + + ::tensorrt_llm::kernels::cutlass_kernels::launchMoeLoraProblemBuilder(mod.permuted_ranks_dev, mod.permuted_ptrs_dev, + input_base, mod.lowrank_workspace_dev, output_base, num_permuted_tokens, in_hidden_size, mod.out_hidden_size, + max_lora_rank, dtype_bytes, splitk_slices, arrays, stream); + sync_check_cuda_error(stream); + + // The cuda_graph_(split_k_)grouped_gemm wrappers accept ldc == ldd when C + // aliases D (the no-bias case). The problem-builder produces a single + // ldd_in / ldd_out per stage, reused for ldcGpu below. + auto* host_max_in = static_cast(mod.host_max_problem_in_pinned); + auto* host_max_out = static_cast(mod.host_max_problem_out_pinned); + + // kMinKN mirrors the value attention LoRA uses for kernel selection. The + // wrappers fall back to the smaller-tile family when min(K, N) < kMinKN. + constexpr int kMinKN = 16; + + ::tensorrt_llm::kernels::cudaGraphSplitKGroupedGemm(arrays.problem_sizes_in, static_cast(num_permuted_tokens), + arrays.a_ptrs_in, arrays.b_ptrs_in, arrays.d_ptrs_in, arrays.d_ptrs_in, arrays.lda_in, arrays.ldb_in, + arrays.ldd_in, arrays.ldd_in, + /*isLoraIn=*/true, data_type, static_cast(splitk_slices), kMinKN, host_max_in, arrays.splitk_offsets, + stream); + sync_check_cuda_error(stream); + + ::tensorrt_llm::kernels::cudaGraphGroupedGemm(arrays.problem_sizes_out, static_cast(num_permuted_tokens), + arrays.d_ptrs_in /*== a_ptrs_out*/, arrays.b_ptrs_out, arrays.d_ptrs_out, arrays.d_ptrs_out, arrays.ldd_in, + arrays.ldb_out, arrays.ldd_out, arrays.ldd_out, + /*isLoraIn=*/false, data_type, kMinKN, host_max_out, stream); + sync_check_cuda_error(stream); +} + class FusedMoeRunner : public torch::CustomClassHolder { public: @@ -249,6 +318,16 @@ class FusedMoeRunner : public torch::CustomClassHolder mGemm1Profiles = mKernelRunner->getTactics(MoeGemmId::GEMM_1); mGemm2Profiles = mKernelRunner->getTactics(MoeGemmId::GEMM_2); cuInit(0); + + // Device-LoRA-path opt-in for the per-request schema. Any non-empty + // value other than "0"/"OFF"/"off" enables the capture-safe on-device + // LoRA path (pointer-expand + problem-builder + grouped-GEMM) instead of + // the legacy host-pointer path, matching LORA_USE_UNIFIED_GEMM. + if (char const* envv = std::getenv("TLLM_MOE_LORA_USE_DEVICE_PATH")) + { + std::string val(envv); + mUseDeviceLoraPath = !val.empty() && val != "0" && val != "OFF" && val != "off"; + } } ~FusedMoeRunner() @@ -297,7 +376,7 @@ class FusedMoeRunner : public torch::CustomClassHolder bool use_dynamic_fc2_scale = false, // Routed-expert LoRA inputs (all optional; presence of fc1_lora_ranks activates LoRA). // Each *_ranks : CPU int32 [num_seqs] - // Each *_weights : CPU int64 [num_seqs, 3] -- (A_ptr, B_ptr, DoRA_ptr-unused) + // Each *_weights : CPU int64 [num_seqs, 3], holding (A_ptr, B_ptr, DoRA_ptr); DoRA unused. torch::optional const& fc1_lora_ranks = torch::nullopt, torch::optional const& fc1_lora_weight_ptrs = torch::nullopt, torch::optional const& fc2_lora_ranks = torch::nullopt, @@ -492,20 +571,24 @@ class FusedMoeRunner : public torch::CustomClassHolder "MoE LoRA only supports fp16 and bf16 activation dtypes."); TORCH_CHECK(mWeightDtype == c10::ScalarType::Half || mWeightDtype == c10::ScalarType::BFloat16, "MoE LoRA only supports unquantized fp16/bf16 expert weights."); - // CUDA-graph capture is incompatible with the kernel's LoRA path, - // which performs a host-side `cudaEventSynchronize` and CPU-side - // per-token pointer expansion inside `setupLoraWorkspace`. The - // event-synchronize cannot be recorded into a graph, so reject here - // with a clear message to avoid a segfault during capture. - TORCH_CHECK(!tensorrt_llm::common::isCapturing(stream), - "MoE LoRA is not supported under CUDA graph capture. The fused-MoE kernel's " - "LoRA path performs a host-side cudaEventSynchronize after a D2H pointer-expansion copy, " - "which is not capturable. Run the LoRA path eagerly, or disable MoE LoRA when capturing."); + // CUDA-graph capture is only safe on the device LoRA path. The legacy + // host path performs a host-side cudaEventSynchronize and per-token + // pointer expansion in setupLoraWorkspace, plus host-side run-length + // encoding in LoraImpl::run, none of which is capturable. The device + // path (launchMoeLoraPointerExpand and runMoeLoraDeviceModule in + // moe_kernels.cu) runs entirely on the stream and is opted into via + // TLLM_MOE_LORA_USE_DEVICE_PATH. + TORCH_CHECK(mUseDeviceLoraPath || !tensorrt_llm::common::isCapturing(stream), + "MoE LoRA + CUDA graph capture requires the device LoRA path. The per-request schema runs " + "the legacy host path by default, which performs a host-side cudaEventSynchronize after a " + "D2H pointer-expansion copy and is not capturable. Set TLLM_MOE_LORA_USE_DEVICE_PATH=1, run " + "LoRA eagerly, or disable MoE LoRA when capturing."); } // Build LoraParams up-front so we can compute the required cuBLAS workspace before allocation. auto lora_params_opt = buildMoeLoraParams(fc1_lora_ranks, fc1_lora_weight_ptrs, fc2_lora_ranks, fc2_lora_weight_ptrs, gated_lora_ranks, gated_lora_weight_ptrs, host_request_types, host_context_lengths, - /*num_tokens=*/num_rows, hidden_size, inter_size, mActivationDtype, lora_max_low_rank, is_gated_act); + /*num_tokens=*/num_rows, hidden_size, inter_size, mActivationDtype, lora_max_low_rank, is_gated_act, stream, + static_cast(experts_per_token)); size_t lora_workspace_size = 0; if (lora_params_opt.has_value()) { @@ -909,14 +992,95 @@ class FusedMoeRunner : public torch::CustomClassHolder // Sync event used by setupLoraWorkspace (kernel waits on this before reading // host-side permuted_rows arrays). Created lazily. cudaEvent_t mLoraMemcpyEvent = nullptr; - // Scratch storage for the per-token expanded LoRA pointer/rank arrays. - // Reused across calls; .clear() drops content but retains capacity. - std::vector mLoraExpandFC1WeightPtrs; - std::vector mLoraExpandFC2WeightPtrs; - std::vector mLoraExpandGatedWeightPtrs; - std::vector mLoraExpandFC1Ranks; - std::vector mLoraExpandFC2Ranks; - std::vector mLoraExpandGatedRanks; + + // Pinned-host and persistent-device buffers for the capture-safe MoE LoRA + // path. The pinned-host tensors hold the per-token expanded LoRA tables + // (ranks and weight-pointer pairs) so the in-op async H2D into the device + // mirrors is graph-capturable; an async H2D from pageable host memory + // silently becomes synchronous and breaks capture. Both tensors are sized + // at mLoraHostBufCapacity (max_num_tokens) and reused across calls so the + // source and destination addresses are stable across capture and replay. + // Only the first num_tokens entries are valid each call. + at::Tensor mLoraExpandFC1RanksPinned; // [max_num_tokens] int32 + at::Tensor mLoraExpandFC1WeightPtrsPinned; // [max_num_tokens * 2] int64 (A, B) + at::Tensor mLoraExpandFC2RanksPinned; // [max_num_tokens] int32 + at::Tensor mLoraExpandFC2WeightPtrsPinned; // [max_num_tokens * 2] int64 + at::Tensor mLoraExpandGatedRanksPinned; // [max_num_tokens] int32 + at::Tensor mLoraExpandGatedWeightPtrsPinned; // [max_num_tokens * 2] int64 + at::Tensor mLoraExpandFC1RanksDevice; + at::Tensor mLoraExpandFC1WeightPtrsDevice; + at::Tensor mLoraExpandFC2RanksDevice; + at::Tensor mLoraExpandFC2WeightPtrsDevice; + at::Tensor mLoraExpandGatedRanksDevice; + at::Tensor mLoraExpandGatedWeightPtrsDevice; + // Tracks how many entries were populated this call so the H2D copies only + // the live portion. Per module; gated may be inactive for non-gated layers. + int64_t mLoraExpandFC1Size = 0; + int64_t mLoraExpandFC2Size = 0; + int64_t mLoraExpandGatedSize = 0; + // Highest max_num_tokens we have allocated storage for. Grown lazily by + // buildMoeLoraParams; resizing reallocates and changes the buffer addresses. + int64_t mLoraHostBufCapacity = 0; + + // Persistent device-resident scratch backing the capture-safe MoE LoRA + // path. One LoraDevicePathBuffers per module (fc1, fc2, gated). All + // at::Tensor members are allocated by ensureLoraDeviceScratch and reused + // across calls so the addresses baked into a captured graph remain valid + // for replay. Pointers from these tensors are packed into + // LoraParams::device_path by buildMoeLoraParams when the device path is taken. + struct LoraDevicePathBuffers + { + // Per-permuted-row (rank, A_ptr + offset, B_ptr + offset). + at::Tensor permuted_ranks; // int32 [P_max] + at::Tensor permuted_ptrs; // int64 [2 * P_max] + + // Grouped-GEMM bundle. Concrete types restored at the LoraParams boundary. + at::Tensor problem_sizes_in; // int8 [P_max * sizeof(GemmCoord)] + at::Tensor problem_sizes_out; // int8 [P_max * sizeof(GemmCoord)] + at::Tensor a_ptrs_in; // int64 [P_max] + at::Tensor b_ptrs_in; // int64 [P_max] + at::Tensor d_ptrs_in; // int64 [P_max] + at::Tensor b_ptrs_out; // int64 [P_max] + at::Tensor d_ptrs_out; // int64 [P_max] + at::Tensor lda_in; // int64 [P_max] + at::Tensor ldb_in; // int64 [P_max] + at::Tensor ldd_in; // int64 [P_max] + at::Tensor ldb_out; // int64 [P_max] + at::Tensor ldd_out; // int64 [P_max] + at::Tensor splitk_offsets; // int64 [P_max + 1] + + // GEMM data-flow buffers. The split-K in-GEMM's partial-sum scratch is + // allocated internally by cuda_graph_split_k_grouped_gemm, so only the + // low-rank intermediate is owned here. + at::Tensor lowrank_workspace; // dtype [P_max * max_lora_rank] + + // Pinned-host single GemmCoord upper bounds; required by the + // cuda_graph_*_grouped_gemm wrappers for kernel selection. + at::Tensor host_max_problem_in; // int8 pinned [sizeof(GemmCoord)] + at::Tensor host_max_problem_out; // int8 pinned [sizeof(GemmCoord)] + }; + + LoraDevicePathBuffers mFc1DeviceBuf; + LoraDevicePathBuffers mFc2DeviceBuf; + LoraDevicePathBuffers mGatedDeviceBuf; + + // Tracks the shape parameters baked into the current scratch + // allocation. (Re)allocation is required if any of these grows or if + // the dtype changes. + int64_t mLoraDeviceScratchCapacity = 0; // P_max = max(num_tokens * top_k) + int64_t mLoraDeviceScratchMaxLoraRank = 0; + int64_t mLoraDeviceScratchDtypeBytes = 0; + int64_t mLoraDeviceScratchSplitKSlices = 0; + bool mLoraDeviceScratchHasGated = false; + + // Set from the TLLM_MOE_LORA_USE_DEVICE_PATH environment variable at + // construction time. Selects the capture-safe device LoRA path. + bool mUseDeviceLoraPath = false; + + // Split-K slice count for the device-path low-rank in-GEMM. Mirrors the + // value LoraImpl uses internally so the device-path split-K scratch is sized + // identically. + static constexpr int64_t kDevicePathSplitKSlices = 16; void freeProfileWorkspace() { @@ -1086,9 +1250,14 @@ class FusedMoeRunner : public torch::CustomClassHolder // num_tokens: total tokens flowing through this op (used as a consistency check) // // Outputs the two `expand_*` vectors with shapes [num_tokens] / [num_tokens * 2]. + // Writes the [num_tokens] expanded LoRA tables into the caller-owned + // pinned-host buffers expand_ranks_data ([num_tokens] int32) and + // expand_ptrs_data ([num_tokens * 2] int64; each pair is (A, B) as + // raw pointer bits stored in int64). The buffers must already be + // allocated to at least num_tokens / num_tokens * 2 elements. void expandPerRequestLoraTo(torch::Tensor const& ranks, torch::Tensor const& weight_ptrs, torch::Tensor const& host_request_types, torch::Tensor const& host_context_lengths, int64_t num_tokens, - std::vector& expand_ranks, std::vector& expand_ptrs) + int32_t* expand_ranks_data, int64_t* expand_ptrs_data) { CHECK_CPU_INPUT(ranks, at::ScalarType::Int) CHECK_CPU_INPUT(weight_ptrs, at::ScalarType::Long) @@ -1109,28 +1278,31 @@ class FusedMoeRunner : public torch::CustomClassHolder auto const* req_types = static_cast(host_request_types.data_ptr()); auto const* ctx_lens = static_cast(host_context_lengths.data_ptr()); - expand_ranks.clear(); - expand_ptrs.clear(); - expand_ranks.reserve(num_tokens); - expand_ptrs.reserve(num_tokens * 2); - int64_t produced = 0; for (int64_t req_id = 0; req_id < num_seqs; ++req_id) { int32_t const rank = rank_data[req_id]; - void const* const a_ptr = reinterpret_cast(ptr_data[req_id * 3 + 0]); - void const* const b_ptr = reinterpret_cast(ptr_data[req_id * 3 + 1]); + int64_t const a_ptr = ptr_data[req_id * 3 + 0]; + int64_t const b_ptr = ptr_data[req_id * 3 + 1]; // ptr_data[req_id * 3 + 2] is the optional DoRA magnitude vector pointer; ignored here // (MoE+DoRA is rejected at load time, see tensorrt_llm/lora_manager.py). auto const req_type = static_cast(req_types[req_id]); int64_t const repeat = (req_type == MoeLoraRequestType::kGENERATION) ? int64_t{1} : static_cast(ctx_lens[req_id]); + // Guard the destination writes BEFORE producing them. expand_*_data + // point at fixed-capacity pinned buffers sized for num_tokens, so a + // malformed host_context_lengths (summing past num_tokens) must be a + // clean error rather than an out-of-bounds write into pinned memory. + TORCH_CHECK(produced + repeat <= num_tokens, "MoE LoRA per-request expansion overran the ", num_tokens, + "-token buffer at request ", req_id, " (produced ", produced, " + ", repeat, + "). Check host_request_types / host_context_lengths against the op's token count."); for (int64_t i = 0; i < repeat; ++i) { - expand_ranks.push_back(rank); - expand_ptrs.push_back(a_ptr); - expand_ptrs.push_back(b_ptr); + int64_t const t = produced + i; + expand_ranks_data[t] = rank; + expand_ptrs_data[2 * t + 0] = a_ptr; + expand_ptrs_data[2 * t + 1] = b_ptr; } produced += repeat; } @@ -1138,10 +1310,192 @@ class FusedMoeRunner : public torch::CustomClassHolder " tokens but op input has ", num_tokens, " tokens."); } + // Guard against reallocating MoE-LoRA scratch while a CUDA graph is being + // captured on `stream`. Reallocation would hand out fresh device/pinned + // addresses, silently invalidating the copies and kernels already recorded + // into earlier captured graphs (which keep replaying against the old + // addresses). Convert that silent corruption into a loud, actionable error. + // No-op when not capturing or when stream is null (warmup pre-sizing). + void checkLoraReallocSafeDuringCapture(cudaStream_t stream, int64_t requested, int64_t current) const + { + if (stream == nullptr || !tensorrt_llm::common::isCapturing(stream)) + { + return; + } + TORCH_CHECK(false, "MoE LoRA scratch (current capacity ", current, ") is too small for ", requested, + " entries during CUDA graph capture. Growing it would invalidate addresses baked into " + "already-captured graphs. Run the device LoRA path eagerly through the worst-case shape before " + "capture so the scratch is pre-sized."); + } + + // Internal helper: (re)allocate the six pinned-host + six device tensor + // pairs to hold capacity expanded tokens. Called by buildMoeLoraParams + // (lazy on first call at a given size). The (re)allocation drops the + // previous storage; callers must make sure any in-flight CUDA graph that + // references the old addresses has either been destroyed or never replays + // again. + void ensureLoraExpandBuffers(int64_t capacity) + { + auto const pinned_int_opts = at::TensorOptions().dtype(at::kInt).pinned_memory(true); + auto const pinned_long_opts = at::TensorOptions().dtype(at::kLong).pinned_memory(true); + auto const dev_int_opts = at::TensorOptions().dtype(at::kInt).device(at::kCUDA); + auto const dev_long_opts = at::TensorOptions().dtype(at::kLong).device(at::kCUDA); + + mLoraExpandFC1RanksPinned = at::empty({capacity}, pinned_int_opts); + mLoraExpandFC2RanksPinned = at::empty({capacity}, pinned_int_opts); + mLoraExpandGatedRanksPinned = at::empty({capacity}, pinned_int_opts); + mLoraExpandFC1WeightPtrsPinned = at::empty({capacity * 2}, pinned_long_opts); + mLoraExpandFC2WeightPtrsPinned = at::empty({capacity * 2}, pinned_long_opts); + mLoraExpandGatedWeightPtrsPinned = at::empty({capacity * 2}, pinned_long_opts); + + mLoraExpandFC1RanksDevice = at::empty({capacity}, dev_int_opts); + mLoraExpandFC2RanksDevice = at::empty({capacity}, dev_int_opts); + mLoraExpandGatedRanksDevice = at::empty({capacity}, dev_int_opts); + mLoraExpandFC1WeightPtrsDevice = at::empty({capacity * 2}, dev_long_opts); + mLoraExpandFC2WeightPtrsDevice = at::empty({capacity * 2}, dev_long_opts); + mLoraExpandGatedWeightPtrsDevice = at::empty({capacity * 2}, dev_long_opts); + } + + // Allocate the per-module device-path scratch for the capture-safe LoRA + // path. The buffers are sized in permuted tokens (P = num_tokens * top_k) + // and the per-token LoRA rank upper bound max_lora_rank; both feed the + // pointer-expand, problem-builder, and cuda_graph_*_grouped_gemm kernels. + // + // The function is idempotent at or below the current capacity and + // reallocates only when one of (capacity, max_lora_rank, dtype_bytes, + // splitk_slices, has_gated) grows. Reallocation drops the previous storage, + // so callers must ensure any in-flight CUDA graph referencing the old + // addresses has been destroyed or will not replay. + // + // The host-side max-problem-size pins hold one GemmCoord each; the value is + // a worst-case upper bound, independent of per-call data. + void ensureLoraDeviceScratch(int64_t capacity, int64_t max_lora_rank, int64_t dtype_bytes, int64_t splitk_slices, + bool has_gated, cudaStream_t stream = nullptr) + { + TORCH_CHECK(capacity > 0, "device-path capacity must be positive; got ", capacity); + TORCH_CHECK(max_lora_rank > 0, "device-path max_lora_rank must be positive; got ", max_lora_rank); + TORCH_CHECK(dtype_bytes > 0, "device-path dtype_bytes must be positive; got ", dtype_bytes); + TORCH_CHECK(splitk_slices > 0, "device-path splitk_slices must be positive; got ", splitk_slices); + + bool const need_resize = capacity > mLoraDeviceScratchCapacity || max_lora_rank > mLoraDeviceScratchMaxLoraRank + || dtype_bytes != mLoraDeviceScratchDtypeBytes || splitk_slices != mLoraDeviceScratchSplitKSlices + || (has_gated && !mLoraDeviceScratchHasGated); + if (!need_resize) + { + return; + } + // Refuse to grow device scratch mid-capture (see helper for rationale). + checkLoraReallocSafeDuringCapture(stream, capacity, mLoraDeviceScratchCapacity); + + // Grow each field to the requested upper bound and remember the + // dtype/rank/splitk combo so subsequent calls can early-exit. + int64_t const new_capacity = std::max(capacity, mLoraDeviceScratchCapacity); + int64_t const new_max_lora_rank = std::max(max_lora_rank, mLoraDeviceScratchMaxLoraRank); + bool const new_has_gated = mLoraDeviceScratchHasGated || has_gated; + + // c10::ScalarType for the lowrank workspace. The kernel treats the + // buffer opaquely (per-byte stride is dtype_bytes), so we pick a + // dtype with matching element size to keep at::Tensor accounting + // sensible; consumers cast via .data_ptr(). + c10::ScalarType const dtype_scalar = (dtype_bytes == 2) ? at::kBFloat16 + : (dtype_bytes == 4) ? at::kFloat + : at::kByte; + // Callers should pass bf16/fp16 (2 bytes). Other sizes still work at the + // byte level, but this assertion catches accidental misuse. + TORCH_CHECK(dtype_bytes == 1 || dtype_bytes == 2 || dtype_bytes == 4, + "device-path lowrank workspace dtype_bytes must be 1/2/4; got ", dtype_bytes); + + auto const dev_int8_opts = at::TensorOptions().dtype(at::kByte).device(at::kCUDA); + auto const dev_int32_opts = at::TensorOptions().dtype(at::kInt).device(at::kCUDA); + auto const dev_int64_opts = at::TensorOptions().dtype(at::kLong).device(at::kCUDA); + auto const dev_dtype_opts = at::TensorOptions().dtype(dtype_scalar).device(at::kCUDA); + auto const pinned_int8_opts = at::TensorOptions().dtype(at::kByte).pinned_memory(true); + + // sizeof(cutlass::gemm::GemmCoord) == sizeof(int) * 3 in practice; + // we ask for the exact byte count at allocation time so the bound + // tracks any cutlass struct-layout change. + int64_t const gemm_coord_bytes = static_cast(sizeof(cutlass::gemm::GemmCoord)); + + auto alloc_one = [&](LoraDevicePathBuffers& mod) + { + mod.permuted_ranks = at::empty({new_capacity}, dev_int32_opts); + mod.permuted_ptrs = at::empty({new_capacity * 2}, dev_int64_opts); + + mod.problem_sizes_in = at::empty({new_capacity * gemm_coord_bytes}, dev_int8_opts); + mod.problem_sizes_out = at::empty({new_capacity * gemm_coord_bytes}, dev_int8_opts); + + mod.a_ptrs_in = at::empty({new_capacity}, dev_int64_opts); + mod.b_ptrs_in = at::empty({new_capacity}, dev_int64_opts); + mod.d_ptrs_in = at::empty({new_capacity}, dev_int64_opts); + mod.b_ptrs_out = at::empty({new_capacity}, dev_int64_opts); + mod.d_ptrs_out = at::empty({new_capacity}, dev_int64_opts); + + mod.lda_in = at::empty({new_capacity}, dev_int64_opts); + mod.ldb_in = at::empty({new_capacity}, dev_int64_opts); + mod.ldd_in = at::empty({new_capacity}, dev_int64_opts); + mod.ldb_out = at::empty({new_capacity}, dev_int64_opts); + mod.ldd_out = at::empty({new_capacity}, dev_int64_opts); + mod.splitk_offsets = at::empty({new_capacity + 1}, dev_int64_opts); + + mod.lowrank_workspace = at::empty({new_capacity * new_max_lora_rank}, dev_dtype_opts); + + mod.host_max_problem_in = at::empty({gemm_coord_bytes}, pinned_int8_opts); + mod.host_max_problem_out = at::empty({gemm_coord_bytes}, pinned_int8_opts); + }; + + alloc_one(mFc1DeviceBuf); + alloc_one(mFc2DeviceBuf); + if (new_has_gated) + { + alloc_one(mGatedDeviceBuf); + } + + mLoraDeviceScratchCapacity = new_capacity; + mLoraDeviceScratchMaxLoraRank = new_max_lora_rank; + mLoraDeviceScratchDtypeBytes = dtype_bytes; + mLoraDeviceScratchSplitKSlices = splitk_slices; + mLoraDeviceScratchHasGated = new_has_gated; + } + + // Pack the per-module at::Tensor scratch into the typed pointer bundle + // attached to LoraParams. The buffers are owned by FusedMoeRunner, so the + // resulting pointers stay valid as long as the runner outlives the + // LoraParams use. dim_a/dim_b, ranks_src_dev, and out_hidden_size are filled + // in by buildMoeLoraParams; the output base is passed directly to + // runMoeLoraDeviceModule at the call site. + void populateLoraDevicePathModule( + LoraDevicePathBuffers& mod, ::tensorrt_llm::kernels::cutlass_kernels::MoeLoraDevicePathModule& out) const + { + out.permuted_ranks_dev = mod.permuted_ranks.data_ptr(); + out.permuted_ptrs_dev = mod.permuted_ptrs.data_ptr(); + + out.problem_sizes_in_dev = mod.problem_sizes_in.data_ptr(); + out.problem_sizes_out_dev = mod.problem_sizes_out.data_ptr(); + out.a_ptrs_in_dev = reinterpret_cast(mod.a_ptrs_in.data_ptr()); + out.b_ptrs_in_dev = reinterpret_cast(mod.b_ptrs_in.data_ptr()); + out.d_ptrs_in_dev = reinterpret_cast(mod.d_ptrs_in.data_ptr()); + out.b_ptrs_out_dev = reinterpret_cast(mod.b_ptrs_out.data_ptr()); + out.d_ptrs_out_dev = reinterpret_cast(mod.d_ptrs_out.data_ptr()); + out.lda_in_dev = mod.lda_in.data_ptr(); + out.ldb_in_dev = mod.ldb_in.data_ptr(); + out.ldd_in_dev = mod.ldd_in.data_ptr(); + out.ldb_out_dev = mod.ldb_out.data_ptr(); + out.ldd_out_dev = mod.ldd_out.data_ptr(); + out.splitk_offsets_dev = mod.splitk_offsets.data_ptr(); + + out.lowrank_workspace_dev = mod.lowrank_workspace.data_ptr(); + out.host_max_problem_in_pinned = mod.host_max_problem_in.data_ptr(); + out.host_max_problem_out_pinned = mod.host_max_problem_out.data_ptr(); + + // out_hidden_size is set by buildMoeLoraParams; default it here. + out.out_hidden_size = 0; + } + // Build a populated LoraParams from the optional CPU tensors. Caller is // responsible for setting `lora_params.workspace` (the cuBLAS scratch). // Returns std::nullopt when LoRA is inactive (no fc1 ranks tensor). - // Mutates the mLoraExpand* member vectors. + // Mutates the mLoraExpand* pinned tensors and queues an async H2D into + // the device mirrors on stream. std::optional<::tensorrt_llm::kernels::LoraParams> buildMoeLoraParams( torch::optional const& fc1_lora_ranks, torch::optional const& fc1_lora_weight_ptrs, @@ -1151,7 +1505,8 @@ class FusedMoeRunner : public torch::CustomClassHolder torch::optional const& gated_lora_weight_ptrs, torch::optional const& host_request_types, torch::optional const& host_context_lengths, int64_t num_tokens, int64_t hidden_size, - int64_t inter_size, c10::ScalarType act_dtype, int64_t lora_max_low_rank, bool is_gated_activation) + int64_t inter_size, c10::ScalarType act_dtype, int64_t lora_max_low_rank, bool is_gated_activation, + cudaStream_t stream, int experts_per_token) { if (!fc1_lora_ranks.has_value()) { @@ -1164,8 +1519,8 @@ class FusedMoeRunner : public torch::CustomClassHolder TORCH_CHECK(host_request_types.has_value() && host_context_lengths.has_value(), "MoE LoRA requires host_request_types and host_context_lengths CPU tensors."); // For gated activations (e.g. SwiGLU) the kernel's setupLoraWorkspace - // unconditionally dereferences `lora_params.gated_lora_ranks` / - // `gated_lora_weight_ptrs`, so the caller MUST provide them. + // unconditionally dereferences gated_lora_ranks and gated_lora_weight_ptrs, + // so the caller must provide them. if (is_gated_activation) { TORCH_CHECK(gated_lora_ranks.has_value() && gated_lora_weight_ptrs.has_value(), @@ -1182,36 +1537,160 @@ class FusedMoeRunner : public torch::CustomClassHolder int64_t const num_seqs = fc1_lora_ranks->size(0); bool const has_gated = is_gated_activation && gated_lora_ranks.has_value(); + // Ensure pinned/device buffers can hold num_tokens entries. + // Idempotent at-or-below current capacity. + if (num_tokens > mLoraHostBufCapacity) + { + checkLoraReallocSafeDuringCapture(stream, num_tokens, mLoraHostBufCapacity); + ensureLoraExpandBuffers(num_tokens); + mLoraHostBufCapacity = num_tokens; + } + expandPerRequestLoraTo(*fc1_lora_ranks, *fc1_lora_weight_ptrs, *host_request_types, *host_context_lengths, - num_tokens, mLoraExpandFC1Ranks, mLoraExpandFC1WeightPtrs); + num_tokens, mLoraExpandFC1RanksPinned.data_ptr(), + mLoraExpandFC1WeightPtrsPinned.data_ptr()); expandPerRequestLoraTo(*fc2_lora_ranks, *fc2_lora_weight_ptrs, *host_request_types, *host_context_lengths, - num_tokens, mLoraExpandFC2Ranks, mLoraExpandFC2WeightPtrs); + num_tokens, mLoraExpandFC2RanksPinned.data_ptr(), + mLoraExpandFC2WeightPtrsPinned.data_ptr()); + mLoraExpandFC1Size = num_tokens; + mLoraExpandFC2Size = num_tokens; if (has_gated) { expandPerRequestLoraTo(*gated_lora_ranks, *gated_lora_weight_ptrs, *host_request_types, - *host_context_lengths, num_tokens, mLoraExpandGatedRanks, mLoraExpandGatedWeightPtrs); + *host_context_lengths, num_tokens, mLoraExpandGatedRanksPinned.data_ptr(), + mLoraExpandGatedWeightPtrsPinned.data_ptr()); + mLoraExpandGatedSize = num_tokens; } else { - mLoraExpandGatedRanks.clear(); - mLoraExpandGatedWeightPtrs.clear(); + mLoraExpandGatedSize = 0; } + // Queue an async H2D into the persistent device mirrors. The copy + // source is pinned, so the async copy is truly async and capturable, and + // the destination is a persistent device buffer with a stable address + // across captures. The device path consumes these mirrors via + // launchMoeLoraPointerExpand; the legacy host path ignores them and + // reads the pinned host pointers through LoraParams below. + auto issue_h2d = [&](at::Tensor const& src, at::Tensor& dst, int64_t numel) + { + if (numel == 0) + { + return; + } + TLLM_CUDA_CHECK(cudaMemcpyAsync(dst.data_ptr(), src.data_ptr(), + static_cast(numel) * src.element_size(), cudaMemcpyHostToDevice, stream)); + }; + issue_h2d(mLoraExpandFC1RanksPinned, mLoraExpandFC1RanksDevice, mLoraExpandFC1Size); + issue_h2d(mLoraExpandFC1WeightPtrsPinned, mLoraExpandFC1WeightPtrsDevice, mLoraExpandFC1Size * 2); + issue_h2d(mLoraExpandFC2RanksPinned, mLoraExpandFC2RanksDevice, mLoraExpandFC2Size); + issue_h2d(mLoraExpandFC2WeightPtrsPinned, mLoraExpandFC2WeightPtrsDevice, mLoraExpandFC2Size * 2); + issue_h2d(mLoraExpandGatedRanksPinned, mLoraExpandGatedRanksDevice, mLoraExpandGatedSize); + issue_h2d(mLoraExpandGatedWeightPtrsPinned, mLoraExpandGatedWeightPtrsDevice, mLoraExpandGatedSize * 2); + auto impls = getOrCreateLoraImpls(hidden_size, inter_size, act_dtype, static_cast(lora_max_low_rank)); + // The host-side LoRA path (LoraImpl::run) reads the per-token ranks and + // pointers through these raw host pointers, which point at the pinned + // host tensors populated above. ::tensorrt_llm::kernels::LoraParams lora_params{ static_cast(num_seqs), - mLoraExpandFC1Ranks.data(), - mLoraExpandFC1WeightPtrs.data(), - mLoraExpandFC2Ranks.data(), - mLoraExpandFC2WeightPtrs.data(), + mLoraExpandFC1RanksPinned.data_ptr(), + reinterpret_cast(mLoraExpandFC1WeightPtrsPinned.data_ptr()), + mLoraExpandFC2RanksPinned.data_ptr(), + reinterpret_cast(mLoraExpandFC2WeightPtrsPinned.data_ptr()), impls.first, impls.second, /*workspace=*/nullptr, // caller fills in &mLoraMemcpyEvent, - has_gated ? mLoraExpandGatedRanks.data() : nullptr, - has_gated ? mLoraExpandGatedWeightPtrs.data() : nullptr, + has_gated ? mLoraExpandGatedRanksPinned.data_ptr() : nullptr, + has_gated ? reinterpret_cast(mLoraExpandGatedWeightPtrsPinned.data_ptr()) : nullptr, }; + + // Device-LoRA-path scratch. Allocate the per-module device-resident + // buffers and pack their pointers into lora_params.device_path. The + // device path is taken when the env-var opts in (TLLM_MOE_LORA_USE_DEVICE_PATH). + bool const use_device_path = mUseDeviceLoraPath; + if (use_device_path) + { + int64_t const dtype_bytes = static_cast(common::getDTypeSize(loraTypeFromActDtype(act_dtype))); + int64_t const capacity = num_tokens * static_cast(experts_per_token); + // Pass stream so a mid-capture resize (which would invalidate + // previously captured graphs) is rejected with a clear error + // rather than silently corrupting replay. + ensureLoraDeviceScratch(capacity, lora_max_low_rank, dtype_bytes, kDevicePathSplitKSlices, + /*has_gated=*/has_gated, stream); + + auto& dp = lora_params.device_path; + dp.enabled = true; + dp.in_hidden_size = hidden_size; + dp.max_lora_rank = lora_max_low_rank; + dp.dtype_bytes = dtype_bytes; + dp.splitk_slices = kDevicePathSplitKSlices; + dp.has_gated = has_gated; + // Populate the libtorch-bound GEMM dispatch entry point so + // runMoeLoraDeviceModule in moe_kernels.cu can call through + // it without dragging libtorch into libmoe_gemm_src.a. + dp.run = &moeLoraDeviceRunImpl; + populateLoraDevicePathModule(mFc1DeviceBuf, dp.fc1); + populateLoraDevicePathModule(mFc2DeviceBuf, dp.fc2); + if (has_gated) + { + populateLoraDevicePathModule(mGatedDeviceBuf, dp.gated); + } + + // Per-module dim_a/dim_b describe the LoRA adapter shape the + // pointer-expand kernel offsets into; per-module out_hidden_size + // describes the LoRA delta sink the problem-builder kernel writes + // into. The runner passes the output base (lora_fc1_result_ / + // lora_fc2_result_ / lora_gated_out) directly to + // runMoeLoraDeviceModule at the loraFC1/loraFC2 call sites so the + // GEMMs land where the downstream bias/reorder kernels expect. + // + // For fc1 (and gated): adapter A is [hidden, rank], B is [rank, inter]. + // For fc2: adapter A is [inter, rank], B is [rank, hidden]. + dp.fc1.dim_a = hidden_size; + dp.fc1.dim_b = inter_size; + dp.fc1.ranks_src_dev = mLoraExpandFC1RanksDevice.data_ptr(); + dp.fc1.ptrs_src_dev = mLoraExpandFC1WeightPtrsDevice.data_ptr(); + dp.fc1.out_hidden_size = inter_size; + + dp.fc2.dim_a = inter_size; + dp.fc2.dim_b = hidden_size; + dp.fc2.ranks_src_dev = mLoraExpandFC2RanksDevice.data_ptr(); + dp.fc2.ptrs_src_dev = mLoraExpandFC2WeightPtrsDevice.data_ptr(); + dp.fc2.out_hidden_size = hidden_size; + + if (has_gated) + { + dp.gated.dim_a = hidden_size; + dp.gated.dim_b = inter_size; + dp.gated.ranks_src_dev = mLoraExpandGatedRanksDevice.data_ptr(); + dp.gated.ptrs_src_dev = mLoraExpandGatedWeightPtrsDevice.data_ptr(); + dp.gated.out_hidden_size = inter_size; + } + + // Pinned-host max-problem-size hints used by cuda_graph_*_grouped_gemm + // for kernel selection. Values are upper bounds safe to fix at + // warmup time (M=1 since each problem is one row; N/K depend on + // module direction and max_lora_rank). + auto fill_max_problem = [](void* host_ptr, int m, int n, int k) + { + auto* coord = static_cast(host_ptr); + *coord = cutlass::gemm::GemmCoord(m, n, k); + }; + // In-GEMM: M=1, N=max_lora_rank, K=in_dim. Out-GEMM: M=1, N=out_dim, K=max_lora_rank. + fill_max_problem(dp.fc1.host_max_problem_in_pinned, 1, lora_max_low_rank, hidden_size); + fill_max_problem(dp.fc1.host_max_problem_out_pinned, 1, inter_size, lora_max_low_rank); + fill_max_problem(dp.fc2.host_max_problem_in_pinned, 1, lora_max_low_rank, inter_size); + fill_max_problem(dp.fc2.host_max_problem_out_pinned, 1, hidden_size, lora_max_low_rank); + if (has_gated) + { + fill_max_problem(dp.gated.host_max_problem_in_pinned, 1, lora_max_low_rank, hidden_size); + fill_max_problem(dp.gated.host_max_problem_out_pinned, 1, inter_size, lora_max_low_rank); + } + } + return lora_params; } diff --git a/cpp/tests/unit_tests/kernels/CMakeLists.txt b/cpp/tests/unit_tests/kernels/CMakeLists.txt index 95d33e421050..fbef4ed88c2a 100644 --- a/cpp/tests/unit_tests/kernels/CMakeLists.txt +++ b/cpp/tests/unit_tests/kernels/CMakeLists.txt @@ -95,6 +95,11 @@ target_link_libraries(routingKernelsTest PRIVATE Python3::Python) add_gtest(moeLoadBalanceKernelTest moeLoadBalanceKernelTest.cpp) +if(USING_OSS_CUTLASS_MOE_GEMM) + add_gtest(moeLoraPointerExpandTest moeLoraPointerExpandTest.cu) + add_gtest(moeLoraProblemBuilderTest moeLoraProblemBuilderTest.cu) +endif() + add_gtest(eaglePackDataTest eaglePackDataTest.cpp) add_gtest(sparseKvCacheTest sparseKvCacheTest.cu) add_gtest(prepareCustomMaskTest prepareCustomMaskTest.cpp) diff --git a/cpp/tests/unit_tests/kernels/moeLoraPointerExpandTest.cu b/cpp/tests/unit_tests/kernels/moeLoraPointerExpandTest.cu new file mode 100644 index 000000000000..4ece53ffcb6a --- /dev/null +++ b/cpp/tests/unit_tests/kernels/moeLoraPointerExpandTest.cu @@ -0,0 +1,325 @@ +/* + * 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/common/cudaUtils.h" +#include "tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_pointer_expand.h" + +#include +#include +#include + +namespace +{ + +using ::tensorrt_llm::kernels::cutlass_kernels::launchMoeLoraPointerExpand; +using ::tensorrt_llm::kernels::cutlass_kernels::MoeLoraExpandModule; + +// Host-side reference reproducing the per-permuted-row pointer arithmetic +// from CutlassMoeFCRunner::setupLoraWorkspace. Same control flow as the +// device kernel; used only as ground truth for parity checks. +struct RefModule +{ + std::vector ranks_src; + std::vector ptrs_src; + int64_t dim_a; + int64_t dim_b; + std::vector ranks_out; + std::vector ptrs_out; +}; + +void cpuExpand(std::vector const& permuted_rows, std::vector 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, RefModule& fc1, RefModule& fc2, RefModule* gated) +{ + auto expand_one = [&](RefModule& mod, int64_t i, int32_t source_index, int64_t weight_index) + { + int32_t const rank = mod.ranks_src[source_index]; + int64_t const a_stride = weight_index * mod.dim_a * rank * lora_dtype_bytes; + int64_t const b_stride = weight_index * mod.dim_b * rank * lora_dtype_bytes; + mod.ptrs_out[2 * i + 0] = mod.ptrs_src[2 * source_index + 0] + a_stride; + mod.ptrs_out[2 * i + 1] = mod.ptrs_src[2 * source_index + 1] + b_stride; + mod.ranks_out[i] = rank; + }; + + fc1.ranks_out.assign(expanded_num_rows, 0); + fc1.ptrs_out.assign(expanded_num_rows * 2, 0); + fc2.ranks_out.assign(expanded_num_rows, 0); + fc2.ptrs_out.assign(expanded_num_rows * 2, 0); + if (gated) + { + gated->ranks_out.assign(expanded_num_rows, 0); + gated->ptrs_out.assign(expanded_num_rows * 2, 0); + } + + for (int32_t expert_idx = 0; expert_idx < num_experts_per_node; ++expert_idx) + { + int64_t const weight_index = static_cast(expert_idx) + start_expert; + for (int64_t i = expert_first_token_offset[expert_idx]; i < expert_first_token_offset[expert_idx + 1]; ++i) + { + int32_t const source_index = static_cast(permuted_rows[i] % num_rows); + expand_one(fc1, i, source_index, weight_index); + expand_one(fc2, i, source_index, weight_index); + if (gated) + { + expand_one(*gated, i, source_index, weight_index); + } + } + } +} + +template +T* deviceUpload(std::vector const& host) +{ + T* dev = nullptr; + auto const bytes = host.size() * sizeof(T); + if (bytes > 0) + { + TLLM_CUDA_CHECK(cudaMalloc(&dev, bytes)); + TLLM_CUDA_CHECK(cudaMemcpy(dev, host.data(), bytes, cudaMemcpyHostToDevice)); + } + return dev; +} + +template +T* deviceAllocZero(size_t count) +{ + T* dev = nullptr; + auto const bytes = count * sizeof(T); + TLLM_CUDA_CHECK(cudaMalloc(&dev, bytes)); + TLLM_CUDA_CHECK(cudaMemset(dev, 0, bytes)); + return dev; +} + +template +void deviceDownload(T* dev, std::vector& host) +{ + if (host.empty()) + { + return; + } + TLLM_CUDA_CHECK(cudaMemcpy(host.data(), dev, host.size() * sizeof(T), cudaMemcpyDeviceToHost)); +} + +class MoeLoraPointerExpandTest : public ::testing::Test +{ +protected: + void SetUp() override + { + TLLM_CUDA_CHECK(cudaStreamCreate(&mStream)); + } + + void TearDown() override + { + for (auto* p : mAllocations) + { + (void) cudaFree(p); + } + (void) cudaStreamDestroy(mStream); + } + + template + T* upload(std::vector const& host) + { + T* p = deviceUpload(host); + if (p != nullptr) + { + mAllocations.push_back(p); + } + return p; + } + + template + T* allocZero(size_t count) + { + T* p = deviceAllocZero(count); + mAllocations.push_back(p); + return p; + } + + // Run the kernel against ref and assert the device outputs match. + void runAndCompare(std::vector const& permuted_rows, std::vector 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, RefModule& fc1_ref, RefModule& fc2_ref, RefModule* gated_ref) + { + cpuExpand(permuted_rows, expert_first_token_offset, num_experts_per_node, start_expert, num_rows, + expanded_num_rows, lora_dtype_bytes, fc1_ref, fc2_ref, gated_ref); + + auto* permuted_rows_dev = upload(permuted_rows); + auto* offsets_dev = upload(expert_first_token_offset); + + auto build_module = [&](RefModule const& r) + { + MoeLoraExpandModule m; + m.ranks_src = upload(r.ranks_src); + m.ptrs_src = upload(r.ptrs_src); + m.dim_a = r.dim_a; + m.dim_b = r.dim_b; + m.ranks_out = allocZero(expanded_num_rows); + m.ptrs_out = allocZero(expanded_num_rows * 2); + return m; + }; + + MoeLoraExpandModule fc1_dev = build_module(fc1_ref); + MoeLoraExpandModule fc2_dev = build_module(fc2_ref); + MoeLoraExpandModule gated_dev{}; + MoeLoraExpandModule const* gated_dev_ptr = nullptr; + if (gated_ref != nullptr) + { + gated_dev = build_module(*gated_ref); + gated_dev_ptr = &gated_dev; + } + + launchMoeLoraPointerExpand(permuted_rows_dev, offsets_dev, num_experts_per_node, start_expert, num_rows, + expanded_num_rows, lora_dtype_bytes, fc1_dev, fc2_dev, gated_dev_ptr, mStream); + TLLM_CUDA_CHECK(cudaStreamSynchronize(mStream)); + + // Compare per-module. + auto compare = [&](RefModule const& ref_mod, MoeLoraExpandModule const& dev_mod, char const* name) + { + std::vector host_ranks(expanded_num_rows, 0); + std::vector host_ptrs(expanded_num_rows * 2, 0); + deviceDownload(dev_mod.ranks_out, host_ranks); + deviceDownload(dev_mod.ptrs_out, host_ptrs); + for (int64_t i = 0; i < expanded_num_rows; ++i) + { + EXPECT_EQ(host_ranks[i], ref_mod.ranks_out[i]) << name << " rank mismatch at i=" << i; + EXPECT_EQ(host_ptrs[2 * i + 0], ref_mod.ptrs_out[2 * i + 0]) << name << " A ptr mismatch at i=" << i; + EXPECT_EQ(host_ptrs[2 * i + 1], ref_mod.ptrs_out[2 * i + 1]) << name << " B ptr mismatch at i=" << i; + } + }; + + compare(fc1_ref, fc1_dev, "fc1"); + compare(fc2_ref, fc2_dev, "fc2"); + if (gated_ref != nullptr) + { + compare(*gated_ref, gated_dev, "gated"); + } + } + + cudaStream_t mStream{}; + std::vector mAllocations; +}; + +// Helper: build a "fake but distinct" pointer for source token s of module +// tag. Encoding the (tag, s, side) lets the test cheaply verify the +// kernel reads the right slot of ptrs_src. The high bits guarantee +// (ptr + per-expert-byte-offset) doesn't alias another (tag, s, side). +int64_t fakePtr(int tag, int32_t s, int side) +{ + return (static_cast(tag) << 56) | (static_cast(side) << 48) | (static_cast(s + 1) << 32); +} + +// Smallest non-trivial case: 4 source tokens, 3 experts, top_k=2 so the +// permuted batch has 8 rows. Per-expert, no gated. +TEST_F(MoeLoraPointerExpandTest, PerExpertNoGated) +{ + int32_t const num_experts_per_node = 3; + int32_t const start_expert = 0; + int64_t const num_rows = 4; + int64_t const expanded_num_rows = 8; // top_k=2 + + // (expert_id assignment is irrelevant to the kernel; we only need + // expert_first_token_offset for the lookup and permuted_rows for the + // source-index reverse.) + std::vector permuted_rows = {0, 4, 1, 5, 2, 6, 3, 7}; + std::vector expert_first_token_offset = {0, 3, 5, 8}; + + RefModule fc1{}; + fc1.dim_a = 16; // "hidden_size" + fc1.dim_b = 32; // "inter_size" + fc1.ranks_src = {2, 0, 4, 1}; + fc1.ptrs_src.resize(num_rows * 2); + for (int32_t s = 0; s < num_rows; ++s) + { + fc1.ptrs_src[2 * s + 0] = fakePtr(/*tag=*/1, s, /*side=*/0); + fc1.ptrs_src[2 * s + 1] = fakePtr(/*tag=*/1, s, /*side=*/1); + } + + RefModule fc2{}; + fc2.dim_a = 32; + fc2.dim_b = 16; + fc2.ranks_src = {1, 2, 0, 3}; + fc2.ptrs_src.resize(num_rows * 2); + for (int32_t s = 0; s < num_rows; ++s) + { + fc2.ptrs_src[2 * s + 0] = fakePtr(/*tag=*/2, s, /*side=*/0); + fc2.ptrs_src[2 * s + 1] = fakePtr(/*tag=*/2, s, /*side=*/1); + } + + runAndCompare(permuted_rows, expert_first_token_offset, num_experts_per_node, start_expert, num_rows, + expanded_num_rows, /*lora_dtype_bytes=*/2, fc1, fc2, /*gated=*/nullptr); +} + +// Gated activation: three modules, exercises the gated arg path. +TEST_F(MoeLoraPointerExpandTest, GatedActivation) +{ + int32_t const num_experts_per_node = 4; + int32_t const start_expert = 2; // exercises start_expert != 0 + int64_t const num_rows = 5; + int64_t const expanded_num_rows = 10; + + std::vector permuted_rows = {0, 5, 1, 6, 2, 7, 3, 8, 4, 9}; + std::vector expert_first_token_offset = {0, 2, 5, 7, 10}; + + auto build_basic = [&](int tag, int64_t dim_a, int64_t dim_b) + { + RefModule m{}; + m.dim_a = dim_a; + m.dim_b = dim_b; + m.ranks_src = {3, 0, 1, 4, 2}; + m.ptrs_src.resize(num_rows * 2); + for (int32_t s = 0; s < num_rows; ++s) + { + m.ptrs_src[2 * s + 0] = fakePtr(tag, s, 0); + m.ptrs_src[2 * s + 1] = fakePtr(tag, s, 1); + } + return m; + }; + + RefModule fc1 = build_basic(/*tag=*/1, /*hidden=*/8, /*inter=*/24); + RefModule fc2 = build_basic(/*tag=*/2, /*inter=*/24, /*hidden=*/8); + RefModule gated = build_basic(/*tag=*/3, /*hidden=*/8, /*inter=*/24); + runAndCompare(permuted_rows, expert_first_token_offset, num_experts_per_node, start_expert, num_rows, + expanded_num_rows, /*lora_dtype_bytes=*/2, fc1, fc2, &gated); +} + +// Non-trivial lora_dtype_bytes (e.g. fp32 = 4) to verify the stride scaling +// flows through the offset arithmetic. +TEST_F(MoeLoraPointerExpandTest, Fp32StrideBytes) +{ + int32_t const num_experts_per_node = 2; + int32_t const start_expert = 0; + int64_t const num_rows = 2; + int64_t const expanded_num_rows = 4; + + std::vector permuted_rows = {0, 1, 0, 1}; + std::vector expert_first_token_offset = {0, 2, 4}; + + RefModule fc1{}; + fc1.dim_a = 4; + fc1.dim_b = 8; + fc1.ranks_src = {2, 3}; + fc1.ptrs_src = {fakePtr(1, 0, 0), fakePtr(1, 0, 1), fakePtr(1, 1, 0), fakePtr(1, 1, 1)}; + + RefModule fc2 = fc1; + fc2.dim_a = 8; + fc2.dim_b = 4; + + runAndCompare(permuted_rows, expert_first_token_offset, num_experts_per_node, start_expert, num_rows, + expanded_num_rows, /*lora_dtype_bytes=*/4, fc1, fc2, /*gated=*/nullptr); +} + +} // namespace diff --git a/cpp/tests/unit_tests/kernels/moeLoraProblemBuilderTest.cu b/cpp/tests/unit_tests/kernels/moeLoraProblemBuilderTest.cu new file mode 100644 index 000000000000..9b47b4ee8328 --- /dev/null +++ b/cpp/tests/unit_tests/kernels/moeLoraProblemBuilderTest.cu @@ -0,0 +1,331 @@ +/* + * 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/common/cudaUtils.h" +#include "tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_problem_builder.h" + +#include "cutlass/gemm_coord.h" + +#include +#include +#include + +namespace +{ + +using ::tensorrt_llm::kernels::cutlass_kernels::launchMoeLoraProblemBuilder; +using ::tensorrt_llm::kernels::cutlass_kernels::MoeLoraGemmGroupArrays; + +// Host-side reference reproducing the builder's per-row logic. Same +// formulas as the kernel; used only as parity ground truth. +struct RefOutputs +{ + std::vector problem_sizes_in; + std::vector problem_sizes_out; + std::vector a_ptrs_in; // store as int64 bits for simple compare + std::vector b_ptrs_in; + std::vector d_ptrs_in; + std::vector b_ptrs_out; + std::vector d_ptrs_out; + std::vector lda_in; + std::vector ldb_in; + std::vector ldd_in; + std::vector ldb_out; + std::vector ldd_out; + std::vector splitk_offsets; +}; + +RefOutputs cpuReference(std::vector const& ranks, std::vector const& ptrs, int64_t input_base, + int64_t lowrank_workspace, int64_t output_base, int64_t in_hidden_size, int64_t out_hidden_size, + int64_t max_lora_rank, int64_t dtype_bytes, int64_t splitk_slices) +{ + int64_t const P = static_cast(ranks.size()); + RefOutputs r; + r.problem_sizes_in.resize(P); + r.problem_sizes_out.resize(P); + r.a_ptrs_in.resize(P); + r.b_ptrs_in.resize(P); + r.d_ptrs_in.resize(P); + r.b_ptrs_out.resize(P); + r.d_ptrs_out.resize(P); + r.lda_in.resize(P); + r.ldb_in.resize(P); + r.ldd_in.resize(P); + r.ldb_out.resize(P); + r.ldd_out.resize(P); + r.splitk_offsets.resize(P + 1); + for (int64_t i = 0; i < P; ++i) + { + int32_t const rank = ranks[i]; + r.problem_sizes_in[i] = cutlass::gemm::GemmCoord(1, rank, static_cast(in_hidden_size)); + r.problem_sizes_out[i] = cutlass::gemm::GemmCoord(1, static_cast(out_hidden_size), rank); + + int64_t const in_row_stride = in_hidden_size * dtype_bytes; + int64_t const work_row_stride = max_lora_rank * dtype_bytes; + int64_t const out_row_stride = out_hidden_size * dtype_bytes; + + r.a_ptrs_in[i] = input_base + i * in_row_stride; + r.b_ptrs_in[i] = ptrs[2 * i + 0]; + r.d_ptrs_in[i] = lowrank_workspace + i * work_row_stride; + r.b_ptrs_out[i] = ptrs[2 * i + 1]; + r.d_ptrs_out[i] = output_base + i * out_row_stride; + + r.lda_in[i] = in_hidden_size; + r.ldb_in[i] = in_hidden_size; + r.ldd_in[i] = max_lora_rank; + r.ldb_out[i] = rank; + r.ldd_out[i] = out_hidden_size; + + r.splitk_offsets[i] = i * max_lora_rank * splitk_slices; + } + r.splitk_offsets[P] = P * max_lora_rank * splitk_slices; + return r; +} + +template +T* deviceUpload(std::vector const& host) +{ + if (host.empty()) + { + return nullptr; + } + T* dev = nullptr; + TLLM_CUDA_CHECK(cudaMalloc(&dev, host.size() * sizeof(T))); + TLLM_CUDA_CHECK(cudaMemcpy(dev, host.data(), host.size() * sizeof(T), cudaMemcpyHostToDevice)); + return dev; +} + +template +T* deviceAllocZero(size_t count) +{ + T* dev = nullptr; + TLLM_CUDA_CHECK(cudaMalloc(&dev, count * sizeof(T))); + TLLM_CUDA_CHECK(cudaMemset(dev, 0, count * sizeof(T))); + return dev; +} + +template +void deviceDownload(T const* dev, std::vector& host) +{ + if (host.empty()) + { + return; + } + TLLM_CUDA_CHECK(cudaMemcpy(host.data(), dev, host.size() * sizeof(T), cudaMemcpyDeviceToHost)); +} + +class MoeLoraProblemBuilderTest : public ::testing::Test +{ +protected: + void SetUp() override + { + TLLM_CUDA_CHECK(cudaStreamCreate(&mStream)); + } + + void TearDown() override + { + for (auto* p : mAllocations) + { + (void) cudaFree(p); + } + (void) cudaStreamDestroy(mStream); + } + + template + T* upload(std::vector const& h) + { + T* p = deviceUpload(h); + if (p) + { + mAllocations.push_back(p); + } + return p; + } + + template + T* allocZero(size_t n) + { + T* p = deviceAllocZero(n); + mAllocations.push_back(p); + return p; + } + + void runAndCompare(std::vector const& ranks, std::vector const& ptrs, int64_t input_base, + int64_t lowrank_workspace, int64_t output_base, int64_t in_hidden_size, int64_t out_hidden_size, + int64_t max_lora_rank, int64_t dtype_bytes, int64_t splitk_slices) + { + auto const P = static_cast(ranks.size()); + RefOutputs ref = cpuReference(ranks, ptrs, input_base, lowrank_workspace, output_base, in_hidden_size, + out_hidden_size, max_lora_rank, dtype_bytes, splitk_slices); + + int32_t* ranks_dev = upload(ranks); + int64_t* ptrs_dev = upload(ptrs); + + MoeLoraGemmGroupArrays out; + out.problem_sizes_in + = reinterpret_cast(allocZero(P * sizeof(cutlass::gemm::GemmCoord))); + out.problem_sizes_out + = reinterpret_cast(allocZero(P * sizeof(cutlass::gemm::GemmCoord))); + out.a_ptrs_in = reinterpret_cast(allocZero(P)); + out.b_ptrs_in = reinterpret_cast(allocZero(P)); + out.d_ptrs_in = reinterpret_cast(allocZero(P)); + out.b_ptrs_out = reinterpret_cast(allocZero(P)); + out.d_ptrs_out = reinterpret_cast(allocZero(P)); + out.lda_in = allocZero(P); + out.ldb_in = allocZero(P); + out.ldd_in = allocZero(P); + out.ldb_out = allocZero(P); + out.ldd_out = allocZero(P); + out.splitk_offsets = allocZero(P + 1); + + launchMoeLoraProblemBuilder(ranks_dev, ptrs_dev, reinterpret_cast(input_base), + reinterpret_cast(lowrank_workspace), reinterpret_cast(output_base), P, in_hidden_size, + out_hidden_size, max_lora_rank, dtype_bytes, splitk_slices, out, mStream); + TLLM_CUDA_CHECK(cudaStreamSynchronize(mStream)); + + // Compare device outputs to host reference. + auto check_int64 = [&](char const* name, int64_t* dev, std::vector const& ref_vec) + { + std::vector host(ref_vec.size(), 0); + deviceDownload(dev, host); + for (size_t i = 0; i < ref_vec.size(); ++i) + { + EXPECT_EQ(host[i], ref_vec[i]) << name << " mismatch at i=" << i; + } + }; + auto check_ptr_array = [&](char const* name, void** dev, std::vector const& ref_vec) + { check_int64(name, reinterpret_cast(dev), ref_vec); }; + auto check_problem_sizes + = [&](char const* name, cutlass::gemm::GemmCoord* dev, std::vector const& ref_vec) + { + std::vector host(ref_vec.size()); + TLLM_CUDA_CHECK(cudaMemcpy( + host.data(), dev, ref_vec.size() * sizeof(cutlass::gemm::GemmCoord), cudaMemcpyDeviceToHost)); + for (size_t i = 0; i < ref_vec.size(); ++i) + { + EXPECT_EQ(host[i].m(), ref_vec[i].m()) << name << " M mismatch at i=" << i; + EXPECT_EQ(host[i].n(), ref_vec[i].n()) << name << " N mismatch at i=" << i; + EXPECT_EQ(host[i].k(), ref_vec[i].k()) << name << " K mismatch at i=" << i; + } + }; + + check_problem_sizes("problem_sizes_in", out.problem_sizes_in, ref.problem_sizes_in); + check_problem_sizes("problem_sizes_out", out.problem_sizes_out, ref.problem_sizes_out); + check_ptr_array("a_ptrs_in", out.a_ptrs_in, ref.a_ptrs_in); + check_ptr_array("b_ptrs_in", out.b_ptrs_in, ref.b_ptrs_in); + check_ptr_array("d_ptrs_in", out.d_ptrs_in, ref.d_ptrs_in); + check_ptr_array("b_ptrs_out", out.b_ptrs_out, ref.b_ptrs_out); + check_ptr_array("d_ptrs_out", out.d_ptrs_out, ref.d_ptrs_out); + check_int64("lda_in", out.lda_in, ref.lda_in); + check_int64("ldb_in", out.ldb_in, ref.ldb_in); + check_int64("ldd_in", out.ldd_in, ref.ldd_in); + check_int64("ldb_out", out.ldb_out, ref.ldb_out); + check_int64("ldd_out", out.ldd_out, ref.ldd_out); + check_int64("splitk_offsets", out.splitk_offsets, ref.splitk_offsets); + } + + cudaStream_t mStream{}; + std::vector mAllocations; +}; + +// "Pretend" adapter pointers. The kernel treats these as opaque bits, so +// we use easily-distinguishable patterns to catch indexing mistakes. +int64_t fakeAdapter(int tag, int32_t i, int side) +{ + return (static_cast(tag) << 56) | (static_cast(side) << 48) | (static_cast(i + 1) << 32); +} + +TEST_F(MoeLoraProblemBuilderTest, Bf16Smoke) +{ + int64_t const in_hidden_size = 16; + int64_t const out_hidden_size = 32; + int64_t const max_lora_rank = 8; + int64_t const dtype_bytes = 2; + int64_t const splitk_slices = 4; + + int64_t const input_base = static_cast(0x1'0000'0000ull); + int64_t const lowrank_workspace = static_cast(0x2'0000'0000ull); + int64_t const output_base = static_cast(0x3'0000'0000ull); + + std::vector ranks = {2, 0, 4, 1, 8, 3}; + std::vector ptrs; + for (int32_t i = 0; i < static_cast(ranks.size()); ++i) + { + ptrs.push_back(fakeAdapter(/*tag=*/1, i, /*side=*/0)); + ptrs.push_back(fakeAdapter(/*tag=*/1, i, /*side=*/1)); + } + + runAndCompare(ranks, ptrs, input_base, lowrank_workspace, output_base, in_hidden_size, out_hidden_size, + max_lora_rank, dtype_bytes, splitk_slices); +} + +TEST_F(MoeLoraProblemBuilderTest, Fp32StrideBytes) +{ + int64_t const in_hidden_size = 12; + int64_t const out_hidden_size = 24; + int64_t const max_lora_rank = 16; + int64_t const dtype_bytes = 4; + int64_t const splitk_slices = 8; + + int64_t const input_base = static_cast(0x4'0000'0000ull); + int64_t const lowrank_workspace = static_cast(0x5'0000'0000ull); + int64_t const output_base = static_cast(0x6'0000'0000ull); + + std::vector ranks = {1, 16, 8}; + std::vector ptrs; + for (int32_t i = 0; i < static_cast(ranks.size()); ++i) + { + ptrs.push_back(fakeAdapter(/*tag=*/2, i, /*side=*/0)); + ptrs.push_back(fakeAdapter(/*tag=*/2, i, /*side=*/1)); + } + + runAndCompare(ranks, ptrs, input_base, lowrank_workspace, output_base, in_hidden_size, out_hidden_size, + max_lora_rank, dtype_bytes, splitk_slices); +} + +// Cover an empty call (no-op) and a single-token call (smallest live case) +// to lock down the corner cases the larger tests don't exercise. +TEST_F(MoeLoraProblemBuilderTest, BoundaryCases) +{ + int64_t const in_hidden_size = 8; + int64_t const out_hidden_size = 8; + int64_t const max_lora_rank = 4; + int64_t const dtype_bytes = 2; + int64_t const splitk_slices = 2; + int64_t const input_base = static_cast(0x7'0000'0000ull); + int64_t const lowrank_workspace = static_cast(0x8'0000'0000ull); + int64_t const output_base = static_cast(0x9'0000'0000ull); + + // Empty call: P = 0, no allocations needed; launch should be a no-op. + { + MoeLoraGemmGroupArrays empty{}; + launchMoeLoraProblemBuilder(nullptr, nullptr, reinterpret_cast(input_base), + reinterpret_cast(lowrank_workspace), reinterpret_cast(output_base), + /*num_permuted_tokens=*/0, in_hidden_size, out_hidden_size, max_lora_rank, dtype_bytes, splitk_slices, + empty, mStream); + TLLM_CUDA_CHECK(cudaStreamSynchronize(mStream)); + } + + // Single-token call: P = 1, exercises the +1 sentinel write at index 1. + { + std::vector ranks = {3}; + std::vector ptrs = {fakeAdapter(3, 0, 0), fakeAdapter(3, 0, 1)}; + runAndCompare(ranks, ptrs, input_base, lowrank_workspace, output_base, in_hidden_size, out_hidden_size, + max_lora_rank, dtype_bytes, splitk_slices); + } +} + +} // namespace diff --git a/tests/unittest/_torch/lora/test_moe_lora_device_path.py b/tests/unittest/_torch/lora/test_moe_lora_device_path.py new file mode 100644 index 000000000000..de0d72f5177d --- /dev/null +++ b/tests/unittest/_torch/lora/test_moe_lora_device_path.py @@ -0,0 +1,184 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the *device path* of routed-expert MoE LoRA in +`torch.ops.trtllm.fused_moe`. + +The device path (opted into for the per-request schema via +`TLLM_MOE_LORA_USE_DEVICE_PATH=1`) performs the per-token pointer expansion, +problem building, and grouped GEMMs entirely on the CUDA stream via the new +on-device kernels, instead of the legacy host-pointer LoRA path. This test +checks device-path eager correctness vs. both the legacy host path and an fp32 +PyTorch reference, exercising the pointer-expand / problem-builder / +grouped-GEMM kernels. + +It requires a CUDA GPU and the built `trtllm::fused_moe` op. +""" + +import pytest +import torch + +from tensorrt_llm._torch.peft.lora.moe_layout import make_per_expert_lora, reference_swiglu_moe_lora + +_TRTLLM_AVAILABLE = hasattr(torch.ops, "trtllm") and hasattr(torch.ops.trtllm, "fused_moe") + +requires_cuda_and_op = pytest.mark.skipif( + not torch.cuda.is_available() or not _TRTLLM_AVAILABLE, + reason="Requires CUDA and built TensorRT-LLM C++ extension (torch.ops.trtllm.fused_moe).", +) + + +@pytest.fixture(autouse=True) +def _isolate_moe_runner_cache(): + """Give every test a fresh cached FusedMoeRunner and release device scratch + afterward. + + The device path is selected per-runner at construction from + TLLM_MOE_LORA_USE_DEVICE_PATH, and the runner is cached at module level by + MoERunner. Clearing the cache before each test forces a fresh runner that + re-reads the env var; clearing + empty_cache afterward releases the + per-runner device scratch so it cannot alias later allocations. + """ + from tensorrt_llm._torch.custom_ops.torch_custom_ops import MoERunner + + MoERunner.runner_dict.clear() + yield + MoERunner.runner_dict.clear() + if torch.cuda.is_available(): + torch.cuda.synchronize() + torch.cuda.empty_cache() + + +# Adapters drawn from N(0, 1) blow up the SwiGLU intermediate at these shapes; +# scale them down so the legitimate output stays O(1)-O(10) and the bf16 noise +# stays well under the tolerance (see the rationale in test_moe_lora_op.py). +_LORA_SCALE = 0.25 +_RTOL = 5e-2 +_ATOL = 1.0 + + +def _build_base_inputs( + num_tokens, hidden_size, inter_size, num_experts, top_k, dtype, device, seed=0 +): + torch.manual_seed(seed) + x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) + w3_w1 = torch.randn(num_experts, 2 * inter_size, hidden_size, dtype=dtype, device=device) * 0.02 + w2 = torch.randn(num_experts, hidden_size, inter_size, dtype=dtype, device=device) * 0.02 + logits = torch.randn(num_tokens, num_experts, dtype=torch.float32, device=device) + topk_scores, topk_ids = torch.topk(logits, k=top_k, dim=-1) + topk_scores = torch.softmax(topk_scores, dim=-1) + return x, w3_w1, w2, topk_ids.to(torch.int32), topk_scores.to(torch.float32) + + +def _make_adapter_set(num_experts, rank, hidden_size, inter_size, dtype, device, base_seed): + """Three scaled per-expert adapters (fc1/gate-side, gated/up-side, fc2).""" + + def _scaled(*args, seed): + a = make_per_expert_lora(*args, dtype=dtype, device=device, seed=seed) + a["A"].mul_(_LORA_SCALE) + a["B"].mul_(_LORA_SCALE) + return a + + fc1 = _scaled(num_experts, rank, hidden_size, inter_size, seed=base_seed + 0) + gated = _scaled(num_experts, rank, hidden_size, inter_size, seed=base_seed + 1) + fc2 = _scaled(num_experts, rank, inter_size, hidden_size, seed=base_seed + 2) + return {"fc1": fc1, "gated": gated, "fc2": fc2} + + +def _per_request_kwargs(num_tokens, adapters, rank): + """Single-request per-request schema covering all tokens with one adapter.""" + fc1, gated, fc2 = adapters["fc1"], adapters["gated"], adapters["fc2"] + return dict( + fc1_lora_ranks=torch.tensor([rank], dtype=torch.int32, device="cpu"), + fc1_lora_weight_ptrs=torch.tensor( + [[fc1["A"].data_ptr(), fc1["B"].data_ptr(), 0]], dtype=torch.int64, device="cpu" + ), + fc2_lora_ranks=torch.tensor([rank], dtype=torch.int32, device="cpu"), + fc2_lora_weight_ptrs=torch.tensor( + [[fc2["A"].data_ptr(), fc2["B"].data_ptr(), 0]], dtype=torch.int64, device="cpu" + ), + gated_lora_ranks=torch.tensor([rank], dtype=torch.int32, device="cpu"), + gated_lora_weight_ptrs=torch.tensor( + [[gated["A"].data_ptr(), gated["B"].data_ptr(), 0]], dtype=torch.int64, device="cpu" + ), + host_request_types=torch.zeros(1, dtype=torch.int32, device="cpu"), + host_context_lengths=torch.tensor([num_tokens], dtype=torch.int32, device="cpu"), + lora_max_low_rank=rank, + ) + + +def _call_fused_moe(x, w3_w1, w2, topk_ids, topk_scores, output_dtype, lora_kwargs): + common = dict( + input=x, + token_selected_experts=topk_ids, + token_final_scales=topk_scores, + fc1_expert_weights=w3_w1, + fc1_expert_biases=None, + fc2_expert_weights=w2, + fc2_expert_biases=None, + output_dtype=output_dtype, + quant_scales=[], + ) + common.update(lora_kwargs) + return torch.ops.trtllm.fused_moe(**common)[0] + + +def _reference(x, w3_w1, w2, topk_ids, topk_scores, adapters): + return reference_swiglu_moe_lora( + x, + w3_w1, + w2, + topk_ids, + topk_scores, + fc1_a=adapters["fc1"]["A"], + fc1_b=adapters["fc1"]["B"], + gated_a=adapters["gated"]["A"], + gated_b=adapters["gated"]["B"], + fc2_a=adapters["fc2"]["A"], + fc2_b=adapters["fc2"]["B"], + ) + + +@requires_cuda_and_op +def test_device_path_eager_matches_host_and_reference(monkeypatch): + """Per-request schema on the device path (env-var opt-in) must match both + the legacy host path and the fp32 PyTorch reference. Exercises the on-device + pointer-expand / problem-builder / grouped-GEMM kernels in eager mode. + """ + from tensorrt_llm._torch.custom_ops.torch_custom_ops import MoERunner + + device = torch.device("cuda") + dtype = torch.bfloat16 + num_tokens, hidden_size, inter_size = 16, 128, 256 + num_experts, top_k, rank = 4, 2, 8 + + x, w3_w1, w2, topk_ids, topk_scores = _build_base_inputs( + num_tokens, hidden_size, inter_size, num_experts, top_k, dtype, device + ) + adapters = _make_adapter_set( + num_experts, rank, hidden_size, inter_size, dtype, device, base_seed=300 + ) + lora_kwargs = _per_request_kwargs(num_tokens, adapters, rank) + + # Host path (device path env explicitly disabled), fresh runner. + monkeypatch.setenv("TLLM_MOE_LORA_USE_DEVICE_PATH", "0") + MoERunner.runner_dict.clear() + try: + out_host = _call_fused_moe(x, w3_w1, w2, topk_ids, topk_scores, dtype, dict(lora_kwargs)) + finally: + MoERunner.runner_dict.clear() + + # Device path (env opt-in), fresh runner. + monkeypatch.setenv("TLLM_MOE_LORA_USE_DEVICE_PATH", "1") + MoERunner.runner_dict.clear() + try: + out_device = _call_fused_moe(x, w3_w1, w2, topk_ids, topk_scores, dtype, dict(lora_kwargs)) + finally: + MoERunner.runner_dict.clear() + + out_ref = _reference(x, w3_w1, w2, topk_ids, topk_scores, adapters) + + assert torch.isfinite(out_device).all() + torch.testing.assert_close(out_device, out_ref, rtol=_RTOL, atol=_ATOL) + # Host vs device path are different reduction orders but should agree + # within the same bf16 tolerance. + torch.testing.assert_close(out_device, out_host, rtol=_RTOL, atol=_ATOL) diff --git a/tests/unittest/_torch/lora/test_moe_lora_op.py b/tests/unittest/_torch/lora/test_moe_lora_op.py index 5a609b5ca4b4..4a270145a3d2 100644 --- a/tests/unittest/_torch/lora/test_moe_lora_op.py +++ b/tests/unittest/_torch/lora/test_moe_lora_op.py @@ -162,6 +162,54 @@ def test_moe_per_expert_lora_changes_output(): assert diff > 1e-3, f"LoRA had no observable effect (mean abs diff={diff})" +@requires_cuda_and_op +def test_moe_lora_rejects_overlong_context_lengths(): + """A per-request expansion whose host_context_lengths sum past the op's token + count must raise cleanly instead of overrunning the fixed-capacity pinned + expansion buffer. + + The per-token (rank, A, B) tables are written into buffers sized for + num_tokens; a context request claiming more tokens than the op actually has + would, without the bounds guard in expandPerRequestLoraTo, scribble past the + end of pinned memory. Here a single context request declares 2 * num_tokens, + so the expansion must fail fast. + """ + device = torch.device("cuda") + dtype = torch.bfloat16 + num_tokens, hidden_size, inter_size = 8, 128, 256 + num_experts, top_k = 4, 2 + rank = 8 + + x, w3_w1, w2, topk_ids, topk_scores = _build_base_inputs( + num_tokens, hidden_size, inter_size, num_experts, top_k, dtype, device + ) + fc1_adapter = make_per_expert_lora( + num_experts, rank, hidden_size, inter_size, dtype=dtype, device=device, seed=10 + ) + fc2_adapter = make_per_expert_lora( + num_experts, rank, inter_size, hidden_size, dtype=dtype, device=device, seed=11 + ) + lora_kwargs = _build_lora_request_buffers( + num_tokens, + fc1_adapter["A"], + fc1_adapter["B"], + fc2_adapter["A"], + fc2_adapter["B"], + rank=rank, + ) + # Single context request (host_request_types == 0) whose declared context + # length exceeds the op's token count, so the expansion overruns by design. + lora_kwargs["host_request_types"] = torch.zeros(1, dtype=torch.int32, device="cpu") + lora_kwargs["host_context_lengths"] = torch.tensor( + [2 * num_tokens], dtype=torch.int32, device="cpu" + ) + + with pytest.raises((RuntimeError, ValueError)): + _call_fused_moe( + x, w3_w1, w2, topk_ids, topk_scores, output_dtype=dtype, lora_kwargs=lora_kwargs + ) + + @requires_cuda_and_op def test_moe_lora_rejected_in_min_latency_mode(): device = torch.device("cuda") From 5e311de657e59442ee20f7ead45b54cd43e2f970 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Sun, 7 Jun 2026 12:38:06 -0700 Subject: [PATCH 2/2] address comments from coderabbit Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../include/moe_lora_problem_builder.h | 4 +- .../cutlass_kernels/moe_gemm/moe_kernels.cu | 23 +++- .../moe_gemm/moe_lora_pointer_expand.cu | 19 +++- .../moe_gemm/moe_lora_problem_builder.cu | 5 + cpp/tensorrt_llm/thop/moeOp.cpp | 77 ++++++++++--- .../kernels/moeLoraPointerExpandTest.cu | 103 +++++++++++++++++- .../kernels/moeLoraProblemBuilderTest.cu | 39 ++++++- 7 files changed, 244 insertions(+), 26 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_problem_builder.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_problem_builder.h index a1a43c427a9e..a71e89b87864 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_problem_builder.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_problem_builder.h @@ -97,9 +97,11 @@ struct MoeLoraGemmGroupArrays // // 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 ldb_out[i] and ldd_out[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. diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu index 1bca18a54ea9..a99f42003e47 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2020-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. @@ -3888,6 +3888,14 @@ auto CutlassMoeFCRunner(); + // The device-path GEMM skips rank-0 rows, but the bias/reorder paths + // read lora_fc1_result_ for every valid row. Zero the buffer first so + // skipped rows are a deterministic no-op. It is contiguous and holds + // both the gated and fc1 halves when gated, so one memset covers both. + size_t const fc1_result_bytes = static_cast(expanded_num_rows) * static_cast(inter_size) + * (is_gated_activation ? 2u : 1u) * sizeof(ScaleBiasType); + TLLM_CUDA_CHECK(cudaMemsetAsync(lora_fc1_result_, 0, fc1_result_bytes, stream)); + runMoeLoraDeviceModule(dp.fc1, expanded_num_rows, /*in_hidden_size=*/hidden_size, dp.max_lora_rank, dp.dtype_bytes, dp.splitk_slices, /*input_base=*/static_cast(input), /*output_base=*/static_cast(lora_fc1_result), dp.run, data_type, stream); @@ -3972,6 +3980,13 @@ void CutlassMoeFCRunner(); + + // As in loraFC1, zero the output so rank-0 rows the GEMM skips do not + // feed stale data into the downstream add. + size_t const fc2_result_bytes + = static_cast(num_tokens) * static_cast(hidden_size) * sizeof(ScaleBiasType); + TLLM_CUDA_CHECK(cudaMemsetAsync(lora_fc2_result_, 0, fc2_result_bytes, stream)); + runMoeLoraDeviceModule(dp.fc2, num_tokens, /*in_hidden_size=*/inter_size, dp.max_lora_rank, dp.dtype_bytes, dp.splitk_slices, /*input_base=*/static_cast(input), /*output_base=*/static_cast(lora_fc2_result_), dp.run, data_type, stream); @@ -4216,7 +4231,11 @@ void CutlassMoeFCRunner& host_permuted_rows = host_lora_workspace_.host_permuted_rows; std::vector& host_expert_first_token_offset = host_lora_workspace_.host_expert_first_token_offset; diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_pointer_expand.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_pointer_expand.cu index a12e2cadb286..d130d917ad33 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_pointer_expand.cu +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_pointer_expand.cu @@ -60,6 +60,16 @@ __device__ inline void expandOneModule( mod.ranks_out[i] = rank; } +// Reset one module's output slot to a rank-0 no-op. The device-path scratch is +// persistent and reused, so ghost rows must be explicitly zeroed; otherwise +// stale ranks or pointers survive into the next grouped GEMM. +__device__ inline void zeroOneModule(MoeLoraExpandModule const& mod, int64_t i) +{ + mod.ranks_out[i] = 0; + mod.ptrs_out[2 * i + 0] = 0; + mod.ptrs_out[2 * i + 1] = 0; +} + // One thread per permuted row: find its expert via search over // expert_first_token_offset (staged in shared memory), compute // source_index = permuted_rows[i] % num_rows, and expand fc1, fc2, and @@ -105,9 +115,16 @@ __global__ void moeLoraPointerExpandKernel(int32_t const* __restrict__ permuted_ } // Tokens past the last valid offset (padding "ghost" rows) get // expert_idx == num_experts_per_node; drop them so weight_index cannot run - // off the expert table. + // off the expert table. Zero their output slots first so reused scratch + // becomes a deterministic rank-0 no-op. if (expert_idx >= num_experts_per_node) { + zeroOneModule(fc1, i); + zeroOneModule(fc2, i); + if (has_gated) + { + zeroOneModule(gated, i); + } return; } diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_problem_builder.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_problem_builder.cu index cd4dd30991cb..f1b6da4ca189 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_problem_builder.cu +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_lora_problem_builder.cu @@ -19,6 +19,7 @@ #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" +#include #include TRTLLM_NAMESPACE_BEGIN @@ -57,6 +58,10 @@ __global__ void moeLoraProblemBuilderKernel(int32_t const* __restrict__ ranks, i } int32_t const rank = ranks[i]; + // The workspace row and ldd_in[i] use max_lora_rank, so a larger rank makes + // the in-GEMM write past its slice. Callers validate ranks host-side (see + // moeOp.cpp); this assert is a debug-build backstop. + assert(rank <= max_lora_rank); int64_t const a_ptr_bits = ptrs[2 * i + 0]; int64_t const b_ptr_bits = ptrs[2 * i + 1]; diff --git a/cpp/tensorrt_llm/thop/moeOp.cpp b/cpp/tensorrt_llm/thop/moeOp.cpp index 824b1909ddd4..b69e2987771f 100644 --- a/cpp/tensorrt_llm/thop/moeOp.cpp +++ b/cpp/tensorrt_llm/thop/moeOp.cpp @@ -590,7 +590,11 @@ class FusedMoeRunner : public torch::CustomClassHolder /*num_tokens=*/num_rows, hidden_size, inter_size, mActivationDtype, lora_max_low_rank, is_gated_act, stream, static_cast(experts_per_token)); size_t lora_workspace_size = 0; - if (lora_params_opt.has_value()) + // The device path uses persistent device scratch and never touches the + // legacy cuBLAS lora_workspace, so skip computing/allocating it there to + // avoid duplicating LoRA scratch per stream (and the resulting OOM risk + // at large top_k/rank). + 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, @@ -629,7 +633,7 @@ class FusedMoeRunner : public torch::CustomClassHolder // LoraParams is either the populated one we just built or a default-constructed empty one (use_lora=false). ::tensorrt_llm::kernels::LoraParams lora_params = lora_params_opt.value_or(::tensorrt_llm::kernels::LoraParams{}); - if (lora_active) + if (lora_active && !lora_params.device_path.enabled) { lora_params.workspace = workspace_info.lora_workspace; } @@ -1022,6 +1026,12 @@ class FusedMoeRunner : public torch::CustomClassHolder // buildMoeLoraParams; resizing reallocates and changes the buffer addresses. int64_t mLoraHostBufCapacity = 0; + // Set once a CUDA-graph capture has been observed on the LoRA path. After + // that, growing the persistent scratch is forbidden even outside capture, + // since a captured graph keeps replaying against the freed addresses. + // Mutable so the const capture-safety check can record it. + mutable bool mLoraCaptureObserved = false; + // Persistent device-resident scratch backing the capture-safe MoE LoRA // path. One LoraDevicePathBuffers per module (fc1, fc2, gated). All // at::Tensor members are allocated by ensureLoraDeviceScratch and reused @@ -1287,15 +1297,27 @@ class FusedMoeRunner : public torch::CustomClassHolder // ptr_data[req_id * 3 + 2] is the optional DoRA magnitude vector pointer; ignored here // (MoE+DoRA is rejected at load time, see tensorrt_llm/lora_manager.py). - auto const req_type = static_cast(req_types[req_id]); + // Validate the raw request type before trusting it. An unexpected + // value would otherwise fall into the CONTEXT branch and read an + // arbitrary context length, producing a negative/garbage repeat. + int32_t const req_type_raw = req_types[req_id]; + TORCH_CHECK(req_type_raw == static_cast(MoeLoraRequestType::kCONTEXT) + || req_type_raw == static_cast(MoeLoraRequestType::kGENERATION), + "MoE LoRA host_request_types[", req_id, "] must be 0 (context) or 1 (generation); got ", req_type_raw); + auto const req_type = static_cast(req_type_raw); + if (req_type == MoeLoraRequestType::kCONTEXT) + { + TORCH_CHECK(ctx_lens[req_id] >= 0, "MoE LoRA host_context_lengths[", req_id, + "] must be non-negative; got ", ctx_lens[req_id]); + } int64_t const repeat = (req_type == MoeLoraRequestType::kGENERATION) ? int64_t{1} : static_cast(ctx_lens[req_id]); // Guard the destination writes BEFORE producing them. expand_*_data // point at fixed-capacity pinned buffers sized for num_tokens, so a // malformed host_context_lengths (summing past num_tokens) must be a // clean error rather than an out-of-bounds write into pinned memory. - TORCH_CHECK(produced + repeat <= num_tokens, "MoE LoRA per-request expansion overran the ", num_tokens, - "-token buffer at request ", req_id, " (produced ", produced, " + ", repeat, + TORCH_CHECK(repeat >= 0 && produced + repeat <= num_tokens, "MoE LoRA per-request expansion overran the ", + num_tokens, "-token buffer at request ", req_id, " (produced ", produced, " + ", repeat, "). Check host_request_types / host_context_lengths against the op's token count."); for (int64_t i = 0; i < repeat; ++i) { @@ -1310,22 +1332,26 @@ class FusedMoeRunner : public torch::CustomClassHolder " tokens but op input has ", num_tokens, " tokens."); } - // Guard against reallocating MoE-LoRA scratch while a CUDA graph is being - // captured on `stream`. Reallocation would hand out fresh device/pinned - // addresses, silently invalidating the copies and kernels already recorded - // into earlier captured graphs (which keep replaying against the old - // addresses). Convert that silent corruption into a loud, actionable error. - // No-op when not capturing or when stream is null (warmup pre-sizing). + // Reallocating MoE-LoRA scratch hands out fresh addresses, which silently + // invalidates any CUDA graph that baked in the old ones. Reject reallocation + // both while capturing and after any capture has been observed, since an + // earlier graph keeps replaying. Callers invoke this only when reallocation + // is imminent. No-op before the first capture (e.g. warmup pre-sizing). void checkLoraReallocSafeDuringCapture(cudaStream_t stream, int64_t requested, int64_t current) const { - if (stream == nullptr || !tensorrt_llm::common::isCapturing(stream)) + bool const capturing = (stream != nullptr && tensorrt_llm::common::isCapturing(stream)); + if (capturing) + { + mLoraCaptureObserved = true; + } + if (!capturing && !mLoraCaptureObserved) { return; } TORCH_CHECK(false, "MoE LoRA scratch (current capacity ", current, ") is too small for ", requested, - " entries during CUDA graph capture. Growing it would invalidate addresses baked into " - "already-captured graphs. Run the device LoRA path eagerly through the worst-case shape before " - "capture so the scratch is pre-sized."); + capturing ? " entries during CUDA graph capture." : " entries after a CUDA graph capture was observed.", + " Growing it would invalidate addresses baked into already-captured graphs. Run the device LoRA path " + "eagerly through the worst-case shape before capture so the scratch is pre-sized."); } // Internal helper: (re)allocate the six pinned-host + six device tensor @@ -1537,6 +1563,27 @@ class FusedMoeRunner : public torch::CustomClassHolder int64_t const num_seqs = fc1_lora_ranks->size(0); bool const has_gated = is_gated_activation && gated_lora_ranks.has_value(); + // Every per-request rank must fit within lora_max_low_rank, which sizes + // both the lowrank workspace and the max-problem hints. A larger rank + // would make the device path build GEMM problems wider than the + // allocated scratch and write out of bounds, so reject it up front. + 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(); + 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); + } + // Ensure pinned/device buffers can hold num_tokens entries. // Idempotent at-or-below current capacity. if (num_tokens > mLoraHostBufCapacity) diff --git a/cpp/tests/unit_tests/kernels/moeLoraPointerExpandTest.cu b/cpp/tests/unit_tests/kernels/moeLoraPointerExpandTest.cu index 4ece53ffcb6a..78f61d424a77 100644 --- a/cpp/tests/unit_tests/kernels/moeLoraPointerExpandTest.cu +++ b/cpp/tests/unit_tests/kernels/moeLoraPointerExpandTest.cu @@ -103,6 +103,19 @@ T* deviceAllocZero(size_t count) return dev; } +// Like deviceAllocZero but pre-fills with a non-zero byte pattern. Simulates +// reused scratch holding stale values, so tests can verify the kernel actively +// zeroes ghost rows. +template +T* deviceAllocFilled(size_t count, int byte_pattern) +{ + T* dev = nullptr; + auto const bytes = count * sizeof(T); + TLLM_CUDA_CHECK(cudaMalloc(&dev, bytes)); + TLLM_CUDA_CHECK(cudaMemset(dev, byte_pattern, bytes)); + return dev; +} + template void deviceDownload(T* dev, std::vector& host) { @@ -149,10 +162,22 @@ protected: return p; } - // Run the kernel against ref and assert the device outputs match. + template + T* allocFilled(size_t count, int byte_pattern) + { + T* p = deviceAllocFilled(count, byte_pattern); + mAllocations.push_back(p); + return p; + } + + // Run the kernel against ref and assert the device outputs match. When + // prefill_garbage is true the output buffers start with a non-zero pattern, + // forcing the kernel to explicitly zero ghost rows for the comparison to + // pass. void runAndCompare(std::vector const& permuted_rows, std::vector 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, RefModule& fc1_ref, RefModule& fc2_ref, RefModule* gated_ref) + int64_t lora_dtype_bytes, RefModule& fc1_ref, RefModule& fc2_ref, RefModule* gated_ref, + bool prefill_garbage = false) { cpuExpand(permuted_rows, expert_first_token_offset, num_experts_per_node, start_expert, num_rows, expanded_num_rows, lora_dtype_bytes, fc1_ref, fc2_ref, gated_ref); @@ -167,8 +192,10 @@ protected: m.ptrs_src = upload(r.ptrs_src); m.dim_a = r.dim_a; m.dim_b = r.dim_b; - m.ranks_out = allocZero(expanded_num_rows); - m.ptrs_out = allocZero(expanded_num_rows * 2); + m.ranks_out = prefill_garbage ? allocFilled(expanded_num_rows, 0x7F) + : allocZero(expanded_num_rows); + m.ptrs_out = prefill_garbage ? allocFilled(expanded_num_rows * 2, 0x7F) + : allocZero(expanded_num_rows * 2); return m; }; @@ -322,4 +349,72 @@ TEST_F(MoeLoraPointerExpandTest, Fp32StrideBytes) expanded_num_rows, /*lora_dtype_bytes=*/4, fc1, fc2, /*gated=*/nullptr); } +// Ghost rows: expanded_num_rows exceeds the last valid expert offset, so the +// trailing rows have expert_idx == num_experts_per_node and must be zeroed by +// the kernel. Buffers are pre-filled with garbage to verify the kernel actively +// resets them. +TEST_F(MoeLoraPointerExpandTest, GhostRowsRemainZero) +{ + int32_t const num_experts_per_node = 3; + int32_t const start_expert = 0; + int64_t const num_rows = 4; + // expert_first_token_offset.back() == 6, but we run two extra ghost rows. + int64_t const expanded_num_rows = 8; + + std::vector permuted_rows = {0, 4, 1, 5, 2, 6, 0, 0}; + std::vector expert_first_token_offset = {0, 2, 4, 6}; + + RefModule fc1{}; + fc1.dim_a = 16; + fc1.dim_b = 32; + fc1.ranks_src = {2, 0, 4, 1}; + fc1.ptrs_src.resize(num_rows * 2); + for (int32_t s = 0; s < num_rows; ++s) + { + fc1.ptrs_src[2 * s + 0] = fakePtr(/*tag=*/1, s, /*side=*/0); + fc1.ptrs_src[2 * s + 1] = fakePtr(/*tag=*/1, s, /*side=*/1); + } + + RefModule fc2 = fc1; + fc2.dim_a = 32; + fc2.dim_b = 16; + + runAndCompare(permuted_rows, expert_first_token_offset, num_experts_per_node, start_expert, num_rows, + expanded_num_rows, /*lora_dtype_bytes=*/2, fc1, fc2, /*gated=*/nullptr, /*prefill_garbage=*/true); +} + +// num_experts_per_node above kMaxExpertsInSmem (1024) forces the kernel to take +// the global-memory expert-offset scan instead of the shared-memory path. +TEST_F(MoeLoraPointerExpandTest, ForceGlobalScanNumExperts1025) +{ + int32_t const num_experts_per_node = 1025; + int32_t const start_expert = 0; + int64_t const num_rows = 4; + int64_t const expanded_num_rows = 4; // top_k=1 + + std::vector permuted_rows = {0, 1, 2, 3}; + // All tokens land in expert 0; every other expert is empty. Offset array has + // num_experts_per_node + 1 == 1026 entries. + std::vector expert_first_token_offset(num_experts_per_node + 1, expanded_num_rows); + expert_first_token_offset[0] = 0; + + RefModule fc1{}; + fc1.dim_a = 8; + fc1.dim_b = 16; + fc1.ranks_src = {1, 2, 3, 4}; + fc1.ptrs_src.resize(num_rows * 2); + for (int32_t s = 0; s < num_rows; ++s) + { + fc1.ptrs_src[2 * s + 0] = fakePtr(/*tag=*/1, s, /*side=*/0); + fc1.ptrs_src[2 * s + 1] = fakePtr(/*tag=*/1, s, /*side=*/1); + } + + RefModule fc2 = fc1; + fc2.dim_a = 16; + fc2.dim_b = 8; + + runAndCompare(permuted_rows, expert_first_token_offset, num_experts_per_node, start_expert, num_rows, + expanded_num_rows, /*lora_dtype_bytes=*/2, fc1, fc2, /*gated=*/nullptr); +} + } // namespace diff --git a/cpp/tests/unit_tests/kernels/moeLoraProblemBuilderTest.cu b/cpp/tests/unit_tests/kernels/moeLoraProblemBuilderTest.cu index 9b47b4ee8328..e9a414e837ca 100644 --- a/cpp/tests/unit_tests/kernels/moeLoraProblemBuilderTest.cu +++ b/cpp/tests/unit_tests/kernels/moeLoraProblemBuilderTest.cu @@ -163,9 +163,12 @@ protected: return p; } + // When with_splitk is false, out.splitk_offsets is left null to exercise the + // kernel's null-offset branch (and the launch_count path that drops the +1 + // sentinel thread); the splitk_offsets comparison is then skipped. void runAndCompare(std::vector const& ranks, std::vector const& ptrs, int64_t input_base, int64_t lowrank_workspace, int64_t output_base, int64_t in_hidden_size, int64_t out_hidden_size, - int64_t max_lora_rank, int64_t dtype_bytes, int64_t splitk_slices) + int64_t max_lora_rank, int64_t dtype_bytes, int64_t splitk_slices, bool with_splitk = true) { auto const P = static_cast(ranks.size()); RefOutputs ref = cpuReference(ranks, ptrs, input_base, lowrank_workspace, output_base, in_hidden_size, @@ -189,7 +192,7 @@ protected: out.ldd_in = allocZero(P); out.ldb_out = allocZero(P); out.ldd_out = allocZero(P); - out.splitk_offsets = allocZero(P + 1); + out.splitk_offsets = with_splitk ? allocZero(P + 1) : nullptr; launchMoeLoraProblemBuilder(ranks_dev, ptrs_dev, reinterpret_cast(input_base), reinterpret_cast(lowrank_workspace), reinterpret_cast(output_base), P, in_hidden_size, @@ -234,7 +237,10 @@ protected: check_int64("ldd_in", out.ldd_in, ref.ldd_in); check_int64("ldb_out", out.ldb_out, ref.ldb_out); check_int64("ldd_out", out.ldd_out, ref.ldd_out); - check_int64("splitk_offsets", out.splitk_offsets, ref.splitk_offsets); + if (with_splitk) + { + check_int64("splitk_offsets", out.splitk_offsets, ref.splitk_offsets); + } } cudaStream_t mStream{}; @@ -328,4 +334,31 @@ TEST_F(MoeLoraProblemBuilderTest, BoundaryCases) } } +// splitk_offsets == nullptr: the caller does not need the split-K scratch +// offsets, so the kernel must skip the sentinel write and the per-row offset +// store while still producing all other arrays correctly. +TEST_F(MoeLoraProblemBuilderTest, NullSplitkOffsets) +{ + int64_t const in_hidden_size = 16; + int64_t const out_hidden_size = 32; + int64_t const max_lora_rank = 8; + int64_t const dtype_bytes = 2; + int64_t const splitk_slices = 4; + + int64_t const input_base = static_cast(0xA'0000'0000ull); + int64_t const lowrank_workspace = static_cast(0xB'0000'0000ull); + int64_t const output_base = static_cast(0xC'0000'0000ull); + + std::vector ranks = {2, 0, 4, 1, 8}; + std::vector ptrs; + for (int32_t i = 0; i < static_cast(ranks.size()); ++i) + { + ptrs.push_back(fakeAdapter(/*tag=*/4, i, /*side=*/0)); + ptrs.push_back(fakeAdapter(/*tag=*/4, i, /*side=*/1)); + } + + runAndCompare(ranks, ptrs, input_base, lowrank_workspace, output_base, in_hidden_size, out_hidden_size, + max_lora_rank, dtype_bytes, splitk_slices, /*with_splitk=*/false); +} + } // namespace