Merge branch 'main' into add_decoding_case

Author: LarryXFly

cpp/tensorrt_llm/nanobind/executor/executorConfig.cpp

@@ -477,7 +477,7 @@ void initConfigBindings(nb::module_& m)
             c.getExtendedRuntimePerfKnobConfig(), c.getDebugConfig(), c.getRecvPollPeriodMs(),
             c.getMaxSeqIdleMicroseconds(), c.getSpecDecConfig(), c.getGuidedDecodingConfig(),
             c.getAdditionalModelOutputs(), c.getCacheTransceiverConfig(), c.getGatherGenerationLogits(),
-            c.getPromptTableOffloading(), c.getEnableTrtOverlap());
+            c.getPromptTableOffloading(), c.getEnableTrtOverlap(), c.getFailFastOnAttentionWindowTooLarge());
         auto pickle_tuple = nb::make_tuple(cpp_states, nb::getattr(self, "__dict__"));
         return pickle_tuple;
     };
@@ -490,7 +490,7 @@ void initConfigBindings(nb::module_& m)
         }
 
         auto cpp_states = nb::cast<nb::tuple>(state[0]);
-        if (cpp_states.size() != 28)
+        if (cpp_states.size() != 29)
         {
             throw std::runtime_error("Invalid cpp_states!");
         }
@@ -525,7 +525,8 @@ void initConfigBindings(nb::module_& m)
             nb::cast<std::optional<tle::CacheTransceiverConfig>>(cpp_states[24]),             // CacheTransceiverConfig
             nb::cast<bool>(cpp_states[25]),                                                   // GatherGenerationLogits
             nb::cast<bool>(cpp_states[26]),                                                   // PromptTableOffloading
-            nb::cast<bool>(cpp_states[27])                                                    // EnableTrtOverlap
+            nb::cast<bool>(cpp_states[27]),                                                   // EnableTrtOverlap
+            nb::cast<bool>(cpp_states[28]) // FailFastOnAttentionWindowTooLarge
         );
 
         // Restore Python data
@@ -564,7 +565,8 @@ void initConfigBindings(nb::module_& m)
                  std::optional<tle::CacheTransceiverConfig>,             // CacheTransceiverConfig
                  bool,                                                   // GatherGenerationLogits
                  bool,                                                   // PromptTableOffloading
-                 bool                                                    // EnableTrtOverlap
+                 bool,                                                   // EnableTrtOverlap
+                 bool                                                    // FailFastOnAttentionWindowTooLarge
                  >(),
             nb::arg("max_beam_width") = 1, nb::arg("scheduler_config") = tle::SchedulerConfig(),
             nb::arg("kv_cache_config") = tle::KvCacheConfig(), nb::arg("enable_chunked_context") = false,
@@ -582,7 +584,7 @@ void initConfigBindings(nb::module_& m)
             nb::arg("spec_dec_config") = nb::none(), nb::arg("guided_decoding_config") = nb::none(),
             nb::arg("additional_model_outputs") = nb::none(), nb::arg("cache_transceiver_config") = nb::none(),
             nb::arg("gather_generation_logits") = false, nb::arg("mm_embedding_offloading") = false,
-            nb::arg("enable_trt_overlap") = false)
+            nb::arg("enable_trt_overlap") = false, nb::arg("fail_fast_on_attention_window_too_large") = false)
         .def_prop_rw("max_beam_width", &tle::ExecutorConfig::getMaxBeamWidth, &tle::ExecutorConfig::setMaxBeamWidth)
         .def_prop_rw("max_batch_size", &tle::ExecutorConfig::getMaxBatchSize, &tle::ExecutorConfig::setMaxBatchSize)
         .def_prop_rw("max_num_tokens", &tle::ExecutorConfig::getMaxNumTokens, &tle::ExecutorConfig::setMaxNumTokens)
@@ -632,6 +634,9 @@ void initConfigBindings(nb::module_& m)
             &tle::ExecutorConfig::setPromptTableOffloading)
         .def_prop_rw(
             "enable_trt_overlap", &tle::ExecutorConfig::getEnableTrtOverlap, &tle::ExecutorConfig::setEnableTrtOverlap)
+        .def_prop_rw("fail_fast_on_attention_window_too_large",
+            &tle::ExecutorConfig::getFailFastOnAttentionWindowTooLarge,
+            &tle::ExecutorConfig::setFailFastOnAttentionWindowTooLarge)
         .def("__getstate__", executorConfigGetState)
         .def("__setstate__", executorConfigSetState);
 }

examples/llm-api/quickstart_advanced.py

@@ -1,10 +1,10 @@
 import argparse
 
 from tensorrt_llm import LLM, SamplingParams
-from tensorrt_llm.llmapi import (CudaGraphConfig, DraftTargetDecodingConfig,
-                                 EagleDecodingConfig, KvCacheConfig, MoeConfig,
-                                 MTPDecodingConfig, NGramDecodingConfig,
-                                 TorchCompileConfig)
+from tensorrt_llm.llmapi import (AutoDecodingConfig, CudaGraphConfig,
+                                 DraftTargetDecodingConfig, EagleDecodingConfig,
+                                 KvCacheConfig, MoeConfig, MTPDecodingConfig,
+                                 NGramDecodingConfig, TorchCompileConfig)
 
 example_prompts = [
     "Hello, my name is",
@@ -181,6 +181,8 @@ def setup_llm(args, **kwargs):
             is_use_oldest=True,
             is_public_pool=True,
         )
+    elif spec_decode_algo == "AUTO":
+        spec_config = AutoDecodingConfig()
     else:
         spec_config = None
 

requirements.txt

@@ -61,3 +61,4 @@ etcd3
 blake3
 llguidance==0.7.29
 soundfile
+triton==3.3.1

tensorrt_llm/_torch/model_config.py

@@ -299,48 +299,6 @@ def get_bindings_model_config(self,
         num_heads = self.pretrained_config.num_attention_heads // (
             self.mapping.tp_size * self.mapping.cp_size)
 
-        # Handle both uniform and per-layer KV heads
-        num_kv_heads_per_layer = getattr(self.pretrained_config,
-                                         'num_kv_heads_per_layer', None)
-        if num_kv_heads_per_layer is not None:
-            # For models with per-layer KV heads, like nemotron-nas
-            kv_heads_per_layer_raw = num_kv_heads_per_layer
-            use_per_layer_kv_heads = True
-        else:
-            # Check if num_key_value_heads is a list (per-layer) or scalar (uniform)
-            num_kv_heads_raw = getattr(self.pretrained_config,
-                                       'num_key_value_heads', None)
-
-            if num_kv_heads_raw is not None and isinstance(
-                    num_kv_heads_raw, list):
-                # num_key_value_heads is a list - treat as per-layer KV heads
-                kv_heads_per_layer_raw = num_kv_heads_raw
-                use_per_layer_kv_heads = True
-            else:
-                # num_key_value_heads is scalar or None - treat as uniform KV heads
-                if num_kv_heads_raw is None:
-                    # For uniform models, check: num_key_value_heads (standard) -> num_query_groups (NeMo) -> num_attention_heads
-                    num_kv_heads_raw = getattr(
-                        self.pretrained_config, 'num_query_groups',
-                        self.pretrained_config.num_attention_heads)
-
-                num_kv_heads = num_kv_heads_raw // (self.mapping.tp_size *
-                                                    self.mapping.cp_size)
-                use_per_layer_kv_heads = False
-
-        if use_per_layer_kv_heads:
-            # TRT-LLM LoRA requires uniform KV heads across layers
-            if self.lora_config is not None and len(
-                    set(kv_heads_per_layer_raw)) > 1:
-                raise ValueError(
-                    f"TRT-LLM LoRA requires uniform KV heads across layers, "
-                    f"got: {kv_heads_per_layer_raw}")
-            # Apply TP/CP scaling to each layer
-            num_kv_heads_per_layer = [
-                kv_heads // (self.mapping.tp_size * self.mapping.cp_size)
-                for kv_heads in kv_heads_per_layer_raw
-            ]
-
         hidden_size = self.pretrained_config.hidden_size // self.mapping.tp_size
 
         model_config_cpp = ModelConfigCpp(
@@ -361,9 +319,18 @@ def get_bindings_model_config(self,
         else:
             model_config_cpp.tokens_per_block = tokens_per_block
 
-        if use_per_layer_kv_heads:
+        num_key_value_heads = getattr(self.pretrained_config,
+                                      "num_key_value_heads", num_heads)
+        if isinstance(num_key_value_heads, (list, tuple)):
+            # Per-layer KV heads (e.g., Nemotron-NAS, variable GQA models)
+            num_kv_heads_per_layer = [
+                kv_heads // (self.mapping.tp_size * self.mapping.cp_size)
+                for kv_heads in num_key_value_heads
+            ]
             model_config_cpp.num_kv_heads_per_layer = num_kv_heads_per_layer
         else:
+            num_kv_heads = num_key_value_heads // (self.mapping.tp_size *
+                                                   self.mapping.cp_size)
             model_config_cpp.set_num_kv_heads(num_kv_heads)
 
         mlp_hidden_size = None

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -451,18 +451,16 @@ def create_py_executor_instance(
 
         num_experts = _try_infer_num_experts(model_engine.model.model_config)
 
-        num_attn_layers = model_binding_config.num_attention_layers()
-        per_layer_kv_heads = [
-            model_binding_config.num_kv_heads(i) for i in range(num_attn_layers)
-        ]
-        num_kv_attention_heads = max(per_layer_kv_heads)
-        if len(set(per_layer_kv_heads)) > 1:
-            # NOTE: This code-path is currently untested and not validated. Can fail!
-            # This support is tracked in TRTLLM-6561
+        num_kv_attention_heads_per_layer = model_binding_config.num_kv_heads_per_layer
+        if max(num_kv_attention_heads_per_layer) != min(
+                num_kv_attention_heads_per_layer):
             logger.warning(
-                f"Non-uniform KV heads per layer detected, using max ({num_kv_attention_heads}) for LoRA. "
-                "This code-path is currently untested and not validated. May fail!"
+                "Defining LORA with per-layer KV heads is not supported for LORA, using the max number of KV heads per layer"
             )
+            num_kv_attention_heads = max(num_kv_attention_heads_per_layer)
+        else:
+            # all layers have the same number of KV heads
+            num_kv_attention_heads = num_kv_attention_heads_per_layer[0]
 
         lora_modules = LoraModule.create_lora_modules(
             lora_module_names=lora_config.lora_target_modules,

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -477,3 +477,17 @@ def executor_request_to_llm_request(
         py_multimodal_data=getattr(executor_request, "py_multimodal_data",
                                    None))
     return llm_request
+
+
+def get_draft_token_length(request: LlmRequest) -> int:
+    """Get the length of draft tokens for a given request.
+
+    Args:
+        request: The LlmRequest to get draft token length for
+
+    Returns:
+        The number of draft tokens, or 0 if no draft tokens exist
+    """
+    if request.py_draft_tokens is not None:
+        return len(request.py_draft_tokens)
+    return 0