diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 50d562634250..b7ce063b83a3 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/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, diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py index 5aee38cabb35..990f0773e918 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py +++ b/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(