[https://nvbugs/6701493][fix] Latch warmup peak before KV cache size estimation resets it - #18669
[https://nvbugs/6701493][fix] Latch warmup peak before KV cache size estimation resets it#18669trtllm-agent wants to merge 1 commit into
Conversation
… 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>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. WalkthroughKV-cache capacity profiling now preserves CUDA memory usage from executor initialization warmup. It compares that value with the dummy-run peak before sizing the subsequent warmup replay. ChangesKV-cache profiling
Estimated code review effort: 2 (Simple) | ~10 minutes Merge Risk: ⚪ Minimal · up to KV-cache sizing now retains memory needed by executor warmup, reducing warmup-replay OOM risk without changing behavior when profiling already has the higher peak. No current merge-blocking risk remains. Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Full details: Description checkExplanation The description clearly explains the root cause, the fix, the affected failure, and the validation performed. It includes test planning and a bug link, although it does not reproduce the template's explicit PR Checklist section. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
1 participant
Summary
KvCacheCreator.estimate_max_kv_cache_size()callstorch.cuda.reset_peak_memory_stats()before its profiling pass, discarding the high-water mark left by warmup (torch.compile specialization, autotuning, CUDA-graph capture, pool pre-population) that already ran inPyExecutor.__init__. It then sizes the KV pool against only the text-only profiling dummy's peak, which is smaller than warmup's, so the pool absorbs memory that warmup will need again when it is replayed against the final cache. On Gemma3-27B-FP8 at the defaultfree_gpu_memory_fraction=0.9this handed out 42.38 GiB and left ~7.2 GiB of torch headroom — too little for MMLU's 128-row × 262144-vocab fp32 logits, so the third eval OOMed.torch.cuda.max_memory_allocated()intowarmup_peak_memoryimmediately before theempty_cache()/reset_peak_memory_stats()pair, then takemax(post-profiling peak, warmup_peak_memory)astorch_peak_memorywhen computing the available KV budget. This keeps the estimator honest about the true torch high-water mark instead of raising the model-side workaround (loweringfree_gpu_memory_fraction/max_batch_sizeper test, as the passing bf16 sibling does), so every model benefits without per-test tuning; the CnnDailymail rouge1 check confirms the change is accuracy-neutral (28.756 vs 28.778 pre-fix baseline, threshold 25.910).pytest tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestGemma3_27BInstruct::test_fp8_prequantized -vTest plan
Links
Reproduction comparison
Signature: tensorrt_llm.executor.utils.RequestError: CUDA out of memory. Tried to allocate 42.00 MiB. GPU 0 has a total capacity of 79.17 GiB of which 8.81 MiB is free. Process 362 has 518.00 MiB memory in use. Including non-PyTorch memory, this process has 78.56 GiB memory in use. 35.18 GiB allowed; Of the allocated memory 34.11 GiB is allocated by PyTorch, and 1.01 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segmen
Signature: tensorrt_llm.executor.utils.RequestError: CUDA out of memory. Tried to allocate 84.00 MiB. GPU 0 has a total capacity of 79.18 GiB of which 66.56 MiB is free. Process 433318 has 518.00 MiB memory in use. Including non-PyTorch memory, this process has 78.56 GiB memory in use. 35.21 GiB allowed; Of the allocated memory 34.07 GiB is allocated by PyTorch, and 1.03 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_se
Dev Engineer Review
QA Engineer Review
No test changes.