[https://nvbugs/6739116][fix] Reserve full-batch logits memory in KV cache estimation

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -1296,6 +1296,30 @@ def _get_multimodal_encoder_memory_reserve(self,
                 model.model_config.multimodal_config.encoder_cache_max_bytes)
         return unprofiled_output_bytes + cache_bytes
 
+    def _get_logits_memory_reserve(self) -> int:
+        """Return the eager-step logits working set absent from the measured peak.
+
+        The profiling forward is a context-only dummy that produces a single
+        logits row, while a serving step outside CUDA graphs produces one row
+        per scheduled request and draft token. ``LogitsProcessor`` keeps the
+        model-dtype logits alive while it materializes their float32 copy, so
+        both are resident at the peak of such a step.
+        """
+        if not self._mapping.is_last_pp_rank():
+            return 0
+        model = self._model_engine.model
+        vocab_size = getattr(model, "vocab_size_padded", None)
+        if not vocab_size:
+            return 0
+        tokens_per_step = (1 if self._speculative_config is None else
+                           self._speculative_config.tokens_per_gen_step)
+        num_rows = min(
+            self._max_num_tokens,
+            self._max_batch_size * self._max_beam_width * tokens_per_step)
+        bytes_per_row = vocab_size * (model.model_config.torch_dtype.itemsize +
+                                      torch.float32.itemsize)
+        return num_rows * bytes_per_row
+
     def _get_token_num_for_estimation(self) -> int:
         """Compute KV cache capacity required for estimate_max_kv_cache_tokens to succeed."""
         if 'cp_type' in self._mapping.cp_config:
@@ -1605,6 +1629,14 @@ def configure_kv_cache_capacity(self,
                 f"Reserving {mem_gb:.2f} GiB for multimodal encoder memory "
                 "not materialized by the profiling run.")
 
+        logits_memory_reserve = self._get_logits_memory_reserve()
+        peak_memory += logits_memory_reserve
+        if logits_memory_reserve > 0:
+            logger.info(
+                f"Reserving {logits_memory_reserve / GB:.2f} GiB for the "
+                "logits of a full-batch step, which the profiling run does "
+                "not materialize.")
+
         # calculate max memory from peak memory and free gpu memory fraction
         kv_cache_max_memory = self._cal_max_memory(peak_memory,
                                                    total_gpu_memory, fraction,

tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py

@@ -553,6 +553,60 @@ def test_reserve_adds_only_unprofiled_output_capacity():
     assert creator._get_multimodal_encoder_memory_reserve(profiled_output_bytes=400) == 112
 
 
+def _make_logits_reserve_creator(
+    *,
+    vocab_size_padded=1000,
+    max_batch_size=8,
+    max_num_tokens=64,
+    max_beam_width=1,
+    speculative_config=None,
+    mapping=None,
+) -> KvCacheCreator:
+    creator = object.__new__(KvCacheCreator)
+    creator._model_engine = SimpleNamespace(
+        model=SimpleNamespace(
+            vocab_size_padded=vocab_size_padded,
+            model_config=SimpleNamespace(torch_dtype=torch.bfloat16),
+        )
+    )
+    creator._mapping = mapping or Mapping()
+    creator._max_batch_size = max_batch_size
+    creator._max_num_tokens = max_num_tokens
+    creator._max_beam_width = max_beam_width
+    creator._speculative_config = speculative_config
+    return creator
+
+
+def test_logits_reserve_covers_full_batch_model_and_float_logits():
+    creator = _make_logits_reserve_creator()
+
+    # One row per request; bf16 logits plus their float32 copy.
+    assert creator._get_logits_memory_reserve() == 8 * 1000 * (2 + 4)
+
+
+def test_logits_reserve_is_bounded_by_max_num_tokens():
+    creator = _make_logits_reserve_creator(max_batch_size=512, max_num_tokens=64)
+
+    assert creator._get_logits_memory_reserve() == 64 * 1000 * (2 + 4)
+
+
+def test_logits_reserve_counts_beams_and_draft_tokens():
+    spec = SimpleNamespace(tokens_per_gen_step=4)
+    creator = _make_logits_reserve_creator(
+        max_batch_size=4, max_num_tokens=1024, max_beam_width=2, speculative_config=spec
+    )
+
+    assert creator._get_logits_memory_reserve() == 4 * 2 * 4 * 1000 * (2 + 4)
+
+
+def test_logits_reserve_skips_ranks_without_lm_head():
+    mapping = Mapping(world_size=2, rank=0, pp_size=2)
+    creator = _make_logits_reserve_creator(mapping=mapping)
+
+    assert not mapping.is_last_pp_rank()
+    assert creator._get_logits_memory_reserve() == 0
+
+
 # ---------------------------------------------------------------------------
 # VSWA hybrid attention pool-group scaling (Gemma4 hybrid MMMU Pro hang fix)
 # ---------------------------------------------------------------------------
@@ -1185,6 +1239,8 @@ def test_estimation_temporarily_uses_inferred_pool_sizing(
     # A bare Mock would auto-create the attribute; real engines set it to
     # None unless the model opted into MM item scheduling.
     model_engine.mm_encoder_output_budget_bytes = None
+    # Likewise for the vocabulary size that sizes the logits reserve.
+    model_engine.model.vocab_size_padded = None
     llm_args = Mock(cache_transceiver_config=None, enable_chunked_prefill=chunked_prefill)
 
     with patch.object(