From 8d068b7a726cb8fa6f2bb4ab17c18376fe09bde6 Mon Sep 17 00:00:00 2001
From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
Date: Thu, 3 Sep 2026 06:30:19 -0700
Subject: [PATCH] [https://nvbugs/6701493][fix] Respect the warmup peak when
 sizing the KV cache

configure_kv_cache_capacity() calls reset_peak_memory_stats() before
profiling, but PyExecutor.__init__ has already run the full warmup
(torch.compile specialization, autotuning, CUDA-graph capture and
memory-pool pre-population) by then. Resetting drops warmup's torch
high-water mark, so the KV pool is sized against the text-only
profiling dummy alone. Warmup is replayed against the final KV cache
right after the estimate is applied, so whichever of the two peaks is
larger has to fit -- and when warmup's is, the pool is oversized and
the replay OOMs.

For Gemma3-27B FP8 the profiling dummy reports 1.59 GiB of dynamic
activation while warmup actually peaks at 3.04 GiB, leaving the second
warmup pass ~3 GiB short and failing an 84 MiB allocation.

Latch the peak before the reset and take the max of the two, so the
estimate covers both passes. This is a no-op whenever the profiling
dummy already dominates.

Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/_util.py | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 6a8ee2021065..3852cd07eb8a 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -1156,6 +1156,14 @@ def configure_kv_cache_capacity(self,
 
         fraction = self._kv_cache_config.free_gpu_memory_fraction
 
+        # Warmup (torch.compile specialization, autotuning, CUDA-graph capture,
+        # memory-pool pre-population) already ran in PyExecutor.__init__ and is
+        # replayed against the final KV cache once this estimate is applied. Its
+        # high-water mark can exceed the text-only profiling dummy's, so latch it
+        # before reset_peak_memory_stats() drops it; otherwise the pool is sized
+        # against the smaller peak and warmup OOMs on the second pass.
+        warmup_peak_memory = torch.cuda.max_memory_allocated()
+
         torch.cuda.empty_cache()
         torch.cuda.reset_peak_memory_stats()
         end, total_gpu_memory = torch.cuda.mem_get_info()
@@ -1197,8 +1205,8 @@ def configure_kv_cache_capacity(self,
                         if response.has_error():
                             raise RuntimeError(response.error_msg)
 
-                torch_peak_memory = torch.cuda.memory_stats(
-                )["allocated_bytes.all.peak"]
+                torch_peak_memory = max(torch.cuda.max_memory_allocated(),
+                                        warmup_peak_memory)
 
                 # Release before measuring current usage so the retained
                 # embeddings count toward the peak but not the steady state.