Skip to content

[https://nvbugs/6539941][fix] Preserve primary warmup error when batch cleanup also fails #17152

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
trtllm-agent wants to merge 2 commits into
base: main
Choose a base branch
from
+20 −4
Open

[https://nvbugs/6539941][fix] Preserve primary warmup error when batch cleanup also fails #17152

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
82c4a6f
[nvbugs/6539941][fix] Stop warmup cleanup from replacing the primary …
trtllm-agent Aug 1, 2026
363d49c
[nvbugs/6539941][chore] Remove stale waiver after fix
trtllm-agent Aug 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 20 additions & 3 deletions tensorrt_llm/_torch/pyexecutor/model_engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2520,9 +2520,8 @@ def _release_batch_context(self, batch: Optional[ScheduledRequests],
ResourceManagerType.CROSS_KV_CACHE_MANAGER)
spec_resource_manager = resource_manager.get_resource_manager(
ResourceManagerType.SPEC_RESOURCE_MANAGER)
try:
yield batch
finally:

def free_batch_resources() -> None:
if batch is not None and kv_cache_manager is not None:
for req in batch.all_requests():
kv_cache_manager.free_resources(req)
Expand All @@ -2533,6 +2532,24 @@ def _release_batch_context(self, batch: Optional[ScheduledRequests],
if spec_resource_manager is not None:
spec_resource_manager.free_resources(req)

try:
yield batch
except BaseException:
# Freeing issues GPU work, so it raises again whenever the failure
# being unwound already left the CUDA context in a sticky error
# state. Letting that secondary error escape from a `finally` would
# *replace* the primary one, blaming the cache manager for a fault
# that actually happened in the model forward.
try:

@brnguyen2 brnguyen2 Aug 5, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: the outer handler is except BaseException but the inner swallow is except Exception. A BaseException-derived cleanup failure (a nested KeyboardInterrupt, or an assertion-free SystemExit from a C-extension abort path) would still escape and replace the primary error — exactly the case this is meant to prevent. Making the inner one except BaseException too closes the gap for free.

free_batch_resources()
except Exception as e: # noqa: BLE001

@brnguyen2 brnguyen2 Aug 5, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Worth being aware this isn't purely diagnostic on the recoverable path. _general_warmup_impl catches torch.OutOfMemoryError around this context manager (model_engine.py:1451) and continues to the next shape. Previously, a cleanup failure during an OOM unwind escaped as a non-OOM exception and hard-failed warmup; now it's swallowed at warning level and warmup continues with those KV blocks leaked, silently shrinking capacity for later shapes and for serving. That's arguably the better trade, but I'd make it logger.error and say so explicitly — e.g. f"Failed to free warmup batch resources while unwinding; {n} request(s) may leak KV blocks: {e}" — so a later "not enough KV cache space" is traceable back here.

logger.warning(
f"Failed to free warmup batch resources while unwinding: {e}"
)
raise
Comment thread
coderabbitai[bot] marked this conversation as resolved.
else:
free_batch_resources()

def _get_num_extra_decoding_steps(self) -> int:
"""Determines extra decoding steps needed for fused drafting loops."""
if isinstance(self.model, BaseDraftingLoopWrapper):
Expand Down
1 change: 0 additions & 1 deletion tests/integration/test_lists/waives.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -183,7 +183,6 @@ full:B300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=Fa
full:B300/accuracy/test_llm_api_pytorch.py::TestMistralLarge3_675B::test_nvfp4_4gpus[latency_moe_trtllm] SKIP (https://nvbugs/6529874)

@brnguyen2 brnguyen2 Aug 5, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This unwaives the B300 TestStep3_7::test_fp8_block_scales[tp_size=4-ep_size=4-mtp_nextn=3] case, but the PR states the kernel fault behind nvbugs/6539941 is untouched and wasn't reproducible here. Better error reporting doesn't make the test pass. Please keep the waiver (and update the bug with the now-correct first-order error), or cite the passing B300 runs of this exact node ID and close the NVBug — otherwise the entry comes back at the next triage sweep.

full:B300/accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_8gpus[attention_dp_off-trtllm] SKIP (https://nvbugs/6474894)
full:B300/accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_dummy_load_format SKIP (https://nvbugs/6525059)
full:B300/accuracy/test_llm_api_pytorch.py::TestStep3_7::test_fp8_block_scales[tp_size=4-ep_size=4-mtp_nextn=3] SKIP (https://nvbugs/6539941)
full:B300/llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks[cudagraph] SKIP (https://nvbugs/6475623)
full:DGX_B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Pro::test_gsm8k_full_accuracy SKIP (https://nvbugs/6571418)
full:DGX_B200/disaggregated/test_disaggregated.py::test_disaggregated_gpt_oss_120b_harmony[gpt_oss/gpt-oss-120b] SKIP (https://nvbugs/6594241)
Expand Down
Loading