diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 4906ad52f5e5..2c70e6bc5e3e 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -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) @@ -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: + free_batch_resources() + except Exception as e: # noqa: BLE001 + logger.warning( + f"Failed to free warmup batch resources while unwinding: {e}" + ) + raise + 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): diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 50950d86bf15..bb437a1dcfcc 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -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) 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)