[TRTLLM-10232][feat] Support LoRA adapter for nemotron-h models

Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>

Author: Wanli-Jiang

cpp/include/tensorrt_llm/runtime/loraModule.h

@@ -53,6 +53,10 @@ class LoraModule
         kSHARED_EXPERT_H_TO_4H = 19,
         kSHARED_EXPERT_4H_TO_H = 20,
         kSHARED_EXPERT_GATE = 21,
+        kMAMBA_IN_PROJ = 22,
+        kMAMBA_OUT_PROJ = 23,
+        kMOE_LATENT_UP = 24,
+        kMOE_LATENT_DOWN = 25,
     };
 
     explicit constexpr LoraModule(ModuleType const& t, SizeType32 inDim, SizeType32 outDim, bool inDimFirst,
@@ -196,7 +200,7 @@ class LoraModule
     static std::vector<LoraModule> createLoraModules(std::vector<std::string> const& loraModuleNames,
         SizeType32 hiddenSize, SizeType32 mlpHiddenSize, SizeType32 numAttentionHeads, SizeType32 numKvAttentionHeads,
         SizeType32 attentionHeadSize, SizeType32 tpSize, SizeType32 numExperts, SizeType32 sharedExpertHiddenSize = 0,
-        SizeType32 moeHiddenSize = 0);
+        SizeType32 moeHiddenSize = 0, SizeType32 mambaInProjSize = 0, SizeType32 mambaInnerSize = 0);
 
     static ModuleType constexpr toModuleType(std::string_view const& name)
     {
@@ -244,6 +248,14 @@ class LoraModule
             return ModuleType::kSHARED_EXPERT_4H_TO_H;
         else if (name == "shared_expert_gate")
             return ModuleType::kSHARED_EXPERT_GATE;
+        else if (name == "mamba_in_proj")
+            return ModuleType::kMAMBA_IN_PROJ;
+        else if (name == "mamba_out_proj")
+            return ModuleType::kMAMBA_OUT_PROJ;
+        else if (name == "moe_latent_up")
+            return ModuleType::kMOE_LATENT_UP;
+        else if (name == "moe_latent_down")
+            return ModuleType::kMOE_LATENT_DOWN;
         else
             return ModuleType::kINVALID;
     }
@@ -274,6 +286,10 @@ class LoraModule
         case ModuleType::kSHARED_EXPERT_H_TO_4H: return "shared_expert_h_to_4h";
         case ModuleType::kSHARED_EXPERT_4H_TO_H: return "shared_expert_4h_to_h";
         case ModuleType::kSHARED_EXPERT_GATE: return "shared_expert_gate";
+        case ModuleType::kMAMBA_IN_PROJ: return "mamba_in_proj";
+        case ModuleType::kMAMBA_OUT_PROJ: return "mamba_out_proj";
+        case ModuleType::kMOE_LATENT_UP: return "moe_latent_up";
+        case ModuleType::kMOE_LATENT_DOWN: return "moe_latent_down";
         case ModuleType::kINVALID: return "INVALID";
         }
         return "INVALID";

cpp/include/tensorrt_llm/runtime/modelConfig.h

@@ -228,6 +228,47 @@ class ModelConfig
         return countLocalLayers(LayerType::kRECURRENT, pipelineParallelism, pipelineParallelismRank);
     }
 
+    // Get the first LoRA layer index for a given PP rank.
+    // Distributes extra layers to lower ranks when num_lora_layers is not evenly divisible by PP size.
+    [[nodiscard]] SizeType32 getFirstLoraLayer(
+        SizeType32 pipelineParallelism = 1, SizeType32 pipelineParallelismRank = 0) const
+    {
+        TLLM_CHECK_WITH_INFO(pipelineParallelism > 0, "Invalid pipelineParallelism: %d", pipelineParallelism);
+        if (mNbLoraLayers > 0)
+        {
+            auto const numBaseLayers = mNbLoraLayers / pipelineParallelism;
+            auto const numExtraLayers = mNbLoraLayers % pipelineParallelism;
+            // If num_lora_layers % pp_size = n != 0, first n ranks get one extra layer
+            return pipelineParallelismRank * numBaseLayers + std::min(pipelineParallelismRank, numExtraLayers);
+        }
+        // Fall back to attention layer distribution
+        return countLowerRankLayers(LayerType::kATTENTION, pipelineParallelism, pipelineParallelismRank);
+    }
+
+    // Get number of layers that can have LoRA applied for the given PP rank.
+    // For hybrid models (e.g., Nemotron-H with Mamba + Attention), this may differ from num_attention_layers
+    // because LoRA can be applied to non-attention layers (e.g., Mamba in_proj/out_proj).
+    // Handles uneven PP splits by distributing extra layers to lower ranks.
+    [[nodiscard]] SizeType32 getNbLoraLayers(
+        SizeType32 pipelineParallelism = 1, SizeType32 pipelineParallelismRank = 0) const
+    {
+        TLLM_CHECK_WITH_INFO(pipelineParallelism > 0, "Invalid pipelineParallelism: %d", pipelineParallelism);
+        // If mNbLoraLayers is set (non-zero), use it with proper PP distribution
+        if (mNbLoraLayers > 0)
+        {
+            auto const numBaseLayers = mNbLoraLayers / pipelineParallelism;
+            auto const numExtraLayers = mNbLoraLayers % pipelineParallelism;
+            // If num_lora_layers % pp_size = n != 0, first n ranks get one extra layer
+            return numBaseLayers + (pipelineParallelismRank < numExtraLayers ? 1 : 0);
+        }
+        return getNbAttentionLayers(pipelineParallelism, pipelineParallelismRank);
+    }
+
+    void setNbLoraLayers(SizeType32 nbLoraLayers)
+    {
+        mNbLoraLayers = nbLoraLayers;
+    }
+
     [[nodiscard]] SizeType32 constexpr getNbHeads() const noexcept
     {
         return mNbHeads;
@@ -922,6 +963,8 @@ class ModelConfig
     std::vector<LoraModule> mLoraModules;
     SizeType32 mMlpHiddenSize;
     SizeType32 mMaxLoraRank;
+    // Number of layers that can have LoRA applied (for hybrid models this may be > num_attention_layers)
+    SizeType32 mNbLoraLayers{0};
 
     std::optional<RnnConfig> mRnnConfig;
 

cpp/tensorrt_llm/nanobind/bindings.cpp

@@ -220,7 +220,11 @@ NB_MODULE(TRTLLM_NB_MODULE, m)
         .value("MLP_GATE_UP", tr::LoraModule::ModuleType::kMLP_GATE_UP)
         .value("SHARED_EXPERT_H_TO_4H", tr::LoraModule::ModuleType::kSHARED_EXPERT_H_TO_4H)
         .value("SHARED_EXPERT_4H_TO_H", tr::LoraModule::ModuleType::kSHARED_EXPERT_4H_TO_H)
-        .value("SHARED_EXPERT_GATE", tr::LoraModule::ModuleType::kSHARED_EXPERT_GATE);
+        .value("SHARED_EXPERT_GATE", tr::LoraModule::ModuleType::kSHARED_EXPERT_GATE)
+        .value("MAMBA_IN_PROJ", tr::LoraModule::ModuleType::kMAMBA_IN_PROJ)
+        .value("MAMBA_OUT_PROJ", tr::LoraModule::ModuleType::kMAMBA_OUT_PROJ)
+        .value("MOE_LATENT_UP", tr::LoraModule::ModuleType::kMOE_LATENT_UP)
+        .value("MOE_LATENT_DOWN", tr::LoraModule::ModuleType::kMOE_LATENT_DOWN);
 
     nb::class_<tr::LoraModule>(m, "LoraModule")
         .def(nb::init<tr::LoraModule::ModuleType, SizeType32, SizeType32, bool, bool, SizeType32, SizeType32>(),
@@ -236,7 +240,8 @@ NB_MODULE(TRTLLM_NB_MODULE, m)
         .def_static("create_lora_modules", &tr::LoraModule::createLoraModules, nb::arg("lora_module_names"),
             nb::arg("hidden_size"), nb::arg("mlp_hidden_size"), nb::arg("num_attention_heads"),
             nb::arg("num_kv_attention_heads"), nb::arg("attention_head_size"), nb::arg("tp_size") = 1,
-            nb::arg("num_experts") = 0, nb::arg("shared_expert_hidden_size") = 0, nb::arg("moe_hidden_size") = 0);
+            nb::arg("num_experts") = 0, nb::arg("shared_expert_hidden_size") = 0, nb::arg("moe_hidden_size") = 0,
+            nb::arg("mamba_in_proj_size") = 0, nb::arg("mamba_inner_size") = 0);
 
     nb::class_<tc::QuantMode>(m, "QuantMode")
         .def_static("none", &tc::QuantMode::none)
@@ -342,6 +347,11 @@ NB_MODULE(TRTLLM_NB_MODULE, m)
         .def_prop_rw("lora_modules", &tr::ModelConfig::getLoraModules, &tr::ModelConfig::setLoraModules)
         .def_prop_rw("max_lora_rank", &tr::ModelConfig::getMaxLoraRank, &tr::ModelConfig::setMaxLoraRank)
         .def_prop_rw("mlp_hidden_size", &tr::ModelConfig::getMlpHiddenSize, &tr::ModelConfig::setMlpHiddenSize)
+        .def("num_lora_layers", &tr::ModelConfig::getNbLoraLayers, nb::arg("pipeline_parallelism") = 1,
+            nb::arg("pipeline_parallelism_rank") = 0)
+        .def("first_lora_layer", &tr::ModelConfig::getFirstLoraLayer, nb::arg("pipeline_parallelism") = 1,
+            nb::arg("pipeline_parallelism_rank") = 0)
+        .def("set_num_lora_layers", &tr::ModelConfig::setNbLoraLayers, nb::arg("num_lora_layers"))
         .def_prop_rw("size_per_head", &tr::ModelConfig::getSizePerHead, &tr::ModelConfig::setSizePerHead);
 
     nb::class_<tr::WorldConfig>(m, "WorldConfig")

cpp/tensorrt_llm/runtime/loraCache.cpp

@@ -454,8 +454,8 @@ SizeType32 LoraCache::determineNumPages(TaskIdType taskId) const
 SizeType32 LoraCache::determineNumPages(TensorPtr loraConfig) const
 {
     TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__);
-    auto const localNumLayers = mModelConfig.getNbAttentionLayers(
-        mWorldConfig.getPipelineParallelism(), mWorldConfig.getPipelineParallelRank());
+    auto const localNumLayers
+        = mModelConfig.getNbLoraLayers(mWorldConfig.getPipelineParallelism(), mWorldConfig.getPipelineParallelRank());
     auto const firstLayerId = mWorldConfig.getPipelineParallelRank() * localNumLayers;
     auto const lastLayerId = firstLayerId + localNumLayers;
 
@@ -579,8 +579,7 @@ std::vector<LoraCache::TaskLayerModuleConfig> LoraCache::copyToPages(TensorPtr s
     auto const tpRank = worldConfig.getTensorParallelRank();
     auto const ppSize = worldConfig.getPipelineParallelism();
     auto const ppRank = worldConfig.getPipelineParallelRank();
-    // TODO(oargov): why *attention* layers?
-    auto const localNumLayers = modelConfig.getNbAttentionLayers(ppSize, ppRank);
+    auto const localNumLayers = modelConfig.getNbLoraLayers(ppSize, ppRank);
     auto const firstLayerId = ppRank * localNumLayers;
     auto const lastLayerId = firstLayerId + localNumLayers;
 

cpp/tensorrt_llm/runtime/loraManager.cpp

@@ -72,7 +72,7 @@ void LoraManager::fillInputTensors(TensorPtr weightsPtrs, TensorPtr adapterSizes
 
     auto const ppSize = worldConfig.getPipelineParallelism();
     auto const ppRank = worldConfig.getPipelineParallelRank();
-    auto const localNumLayers = modelConfig.getNbAttentionLayers(ppSize, ppRank);
+    auto const localNumLayers = modelConfig.getNbLoraLayers(ppSize, ppRank);
     auto const firstLayerId = ppRank * localNumLayers;
 
     auto weightsPointersPtr = bufferCast<int64_t>(*weightsPtrs);
@@ -124,7 +124,7 @@ void LoraManager::insertInputTensors(TensorMap& inputTensors, TensorPtr weightsP
 {
     TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__);
     auto localNbLayers
-        = modelConfig.getNbAttentionLayers(worldConfig.getPipelineParallelism(), worldConfig.getPipelineParallelRank());
+        = modelConfig.getNbLoraLayers(worldConfig.getPipelineParallelism(), worldConfig.getPipelineParallelRank());
     auto firstLayerId = worldConfig.getPipelineParallelRank() * localNbLayers;
 
     for (auto const& [modId, mod] : mModuleIdToModule)

cpp/tensorrt_llm/runtime/loraModule.cpp

@@ -22,12 +22,16 @@ namespace tensorrt_llm::runtime
 std::vector<LoraModule> LoraModule::createLoraModules(std::vector<std::string> const& loraModuleNames,
     SizeType32 hiddenSize, SizeType32 mlpHiddenSize, SizeType32 numAttentionHeads, SizeType32 numKvAttentionHeads,
     SizeType32 attentionHeadSize, SizeType32 tpSize, SizeType32 numExperts, SizeType32 sharedExpertHiddenSize,
-    SizeType32 moeHiddenSize)
+    SizeType32 moeHiddenSize, SizeType32 mambaInProjSize, SizeType32 mambaInnerSize)
 {
     auto const hidden = hiddenSize * tpSize;
     auto const mlpHidden = mlpHiddenSize * tpSize;
     auto const sharedExpertHidden = sharedExpertHiddenSize > 0 ? sharedExpertHiddenSize * tpSize : mlpHidden;
     auto const moeHidden = moeHiddenSize > 0 ? moeHiddenSize * tpSize : mlpHidden;
+    // Mamba dimensions: in_proj outputs d_in_proj, out_proj inputs d_inner
+    // Fall back to mlpHidden if not specified (for backward compatibility)
+    auto const mambaInProj = mambaInProjSize > 0 ? mambaInProjSize * tpSize : mlpHidden;
+    auto const mambaInner = mambaInnerSize > 0 ? mambaInnerSize * tpSize : mlpHidden;
     auto const numHeads = numAttentionHeads * tpSize;
     auto const numKvHeads = numKvAttentionHeads * tpSize;
     auto const attnHeadSize = attentionHeadSize;
@@ -74,6 +78,12 @@ std::vector<LoraModule> LoraModule::createLoraModules(std::vector<std::string> c
         case ModuleType::kMOE_ROUTER: modules.emplace_back(t, hidden, numExperts, false, true, -1, -1); break;
         case ModuleType::kMLP_ROUTER: modules.emplace_back(t, hidden, 1, false, true, -1, -1); break;
         case ModuleType::kMLP_GATE_UP: modules.emplace_back(t, hidden, 2 * mlpHidden, false, true, -1, 0); break;
+        // Mamba modules: in_proj (hidden -> d_in_proj), out_proj (d_inner -> hidden)
+        case ModuleType::kMAMBA_IN_PROJ: modules.emplace_back(t, hidden, mambaInProj, false, true, -1, 0); break;
+        case ModuleType::kMAMBA_OUT_PROJ: modules.emplace_back(t, mambaInner, hidden, false, true, 1, -1); break;
+        // MoE latent projections: up expands to moe_hidden, down contracts back
+        case ModuleType::kMOE_LATENT_UP: modules.emplace_back(t, hidden, mlpHidden, false, true, -1, 0); break;
+        case ModuleType::kMOE_LATENT_DOWN: modules.emplace_back(t, mlpHidden, hidden, false, true, 1, -1); break;
         case ModuleType::kINVALID: throw std::runtime_error("Invalid LoRA module " + moduleName);
         }
     }

cpp/tensorrt_llm/runtime/loraUtils.cpp

@@ -84,7 +84,7 @@ void loraValidateRequestTensors(std::optional<std::uint64_t> const& optTaskId,
             ? config
             : ITensor::view(config, ITensor::makeShape({config->getShape().d[1], config->getShape().d[2]}));
 
-        SizeType32 nbModelLayers = modelConfig.getNbAttentionLayers();
+        SizeType32 nbModelLayers = modelConfig.getNbLoraLayers();
         TLLM_CHECK_WITH_INFO(weights->getDataType() == modelConfig.getDataType(),
             "Expected lora weights to be the same data type as base model");
 

examples/llm-api/quickstart_advanced.py

@@ -3,11 +3,13 @@
 import time
 
 from tensorrt_llm import LLM, SamplingParams
+from tensorrt_llm.executor.request import LoRARequest
 from tensorrt_llm.llmapi import (AttentionDpConfig, AutoDecodingConfig,
                                  CudaGraphConfig, DraftTargetDecodingConfig,
                                  Eagle3DecodingConfig, KvCacheConfig, MoeConfig,
                                  MTPDecodingConfig, NGramDecodingConfig,
                                  TorchCompileConfig)
+from tensorrt_llm.lora_helper import LoraConfig
 
 example_prompts = [
     "Hello, my name is",
@@ -198,6 +200,18 @@ def add_llm_args(parser):
     parser.add_argument('--relaxed_topk', type=int, default=1)
     parser.add_argument('--relaxed_delta', type=float, default=0.)
 
+    # LoRA
+    parser.add_argument('--lora_dir',
+                        type=str,
+                        default=None,
+                        help='Path to LoRA adapter directory.')
+    parser.add_argument(
+        '--max_lora_rank',
+        type=int,
+        default=None,
+        help='Maximum LoRA rank. If not specified, inferred from adapter config.'
+    )
+
     # HF
     parser.add_argument('--trust_remote_code',
                         default=False,
@@ -292,6 +306,18 @@ def setup_llm(args, **kwargs):
         batching_wait_iters=args.attention_dp_batching_wait_iters,
     )
 
+    lora_config = None
+    lora_request = None
+    if args.lora_dir:
+        max_lora_rank = args.max_lora_rank if args.max_lora_rank is not None else 64
+        lora_config = LoraConfig(lora_dir=[args.lora_dir],
+                                 max_lora_rank=max_lora_rank)
+        lora_request = LoRARequest(
+            lora_name="lora_adapter",
+            lora_int_id=0,  # First adapter ID
+            lora_path=args.lora_dir,
+        )
+
     llm = LLM(
         model=args.model_dir,
         backend='pytorch',
@@ -327,6 +353,7 @@ def setup_llm(args, **kwargs):
         gather_generation_logits=args.return_generation_logits,
         max_beam_width=args.max_beam_width,
         orchestrator_type=args.orchestrator_type,
+        lora_config=lora_config,
         **kwargs)
 
     use_beam_search = args.max_beam_width > 1
@@ -352,14 +379,14 @@ def setup_llm(args, **kwargs):
         use_beam_search=use_beam_search,
         additional_model_outputs=args.additional_model_outputs,
     )
-    return llm, sampling_params
+    return llm, sampling_params, lora_request
 
 
 def main():
     args = parse_arguments()
     prompts = args.prompt if args.prompt else example_prompts
 
-    llm, sampling_params = setup_llm(args)
+    llm, sampling_params, lora_request = setup_llm(args)
     new_prompts = []
     if args.apply_chat_template:
         for prompt in prompts:
@@ -369,7 +396,7 @@ def main():
                                                   tokenize=False,
                                                   add_generation_prompt=True))
         prompts = new_prompts
-    outputs = llm.generate(prompts, sampling_params)
+    outputs = llm.generate(prompts, sampling_params, lora_request=lora_request)
 
     for i, output in enumerate(outputs):
         prompt = output.prompt

tensorrt_llm/_torch/model_config.py

@@ -686,6 +686,12 @@ def get_bindings_model_config(self,
             num_kv_heads = num_key_value_heads // (attn_tp_size * attn_cp_size)
             model_config_cpp.set_num_kv_heads(num_kv_heads)
 
+        # For hybrid models (e.g., Nemotron-H with Mamba + Attention), LoRA can be applied
+        # to non-attention layers (e.g., Mamba in_proj/out_proj). Set num_lora_layers to
+        # total layers so the C++ LoRA validation accepts all layer indices.
+        if is_nemotron_hybrid(self.pretrained_config):
+            model_config_cpp.set_num_lora_layers(num_layers)
+
         mlp_hidden_size = None
         if self.pretrained_config.intermediate_size is not None:
             mlp_hidden_size = self.pretrained_config.intermediate_size // self.mapping.tp_size

tensorrt_llm/_torch/models/modeling_deepseekv3.py

@@ -43,6 +43,7 @@
     ConsumableWeightsDict
 from tensorrt_llm._utils import get_sm_version
 from tensorrt_llm.functional import PositionEmbeddingType
+from tensorrt_llm.lora_helper import LoraConfig
 from tensorrt_llm.mapping import Mapping
 from tensorrt_llm.models.modeling_utils import QuantConfig
 from tensorrt_llm.quantization.mode import QuantAlgo
@@ -852,8 +853,11 @@ def __init__(
         fuse_routing_kernel: bool = True,
         apply_routing: bool = False,
         moe_backend: str = 'CUTLASS',
+        lora_config: Optional[LoraConfig] = None,
     ):
         super().__init__()
+        self.hidden_size = hidden_size
+        self.num_experts = num_experts
         self.weight = nn.Parameter(torch.empty((num_experts, hidden_size),
                                                dtype=dtype),
                                    requires_grad=False)
@@ -877,11 +881,27 @@ def __init__(
             routed_scaling_factor=routed_scaling_factor,
             is_fused=fuse_routing_kernel)
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # LoRA for gate (router) - only create when LoRA is configured
+        from ..peft.lora.layer import LoraModuleType
+        self.gate_lora = (LoraLayer([LoraModuleType.MOE_ROUTER], [num_experts])
+                          if lora_config is not None else None)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        lora_params: Optional[dict] = None,
+        layer_idx: Optional[int] = None,
+    ) -> torch.Tensor:
         logits = torch.ops.trtllm.dsv3_router_gemm_op(hidden_states,
                                                       self.weight.t(),
                                                       bias=None,
                                                       out_dtype=torch.float32)
+        # Apply LoRA to gate (if LoRA is configured and weights are loaded)
+        if self.gate_lora is not None and bool(
+                lora_params) and layer_idx is not None:
+            lora_output = self.gate_lora(hidden_states, lora_params, layer_idx)
+            if lora_output is not None:
+                logits = logits + lora_output.to(logits.dtype)
         return logits
 
     def load_weights(self, weights: List[Dict]):