From 24a132241124ffdc95ae84d7edb8185191494b7e Mon Sep 17 00:00:00 2001 From: William Zhang <133824995+2ez4bz@users.noreply.github.com> Date: Mon, 31 Aug 2026 23:01:36 -0700 Subject: [PATCH] [https://nvbugs/6689016][fix] Reuse FlashInfer plans across batch sizes * Why? Including the live generation batch size in FlashInfer plan keys created a separate eager wrapper and large block-table allocation for every observed batch size, exhausting device memory on long runs. * What? Key decode plans only by their required query width and static attention configuration. CUDA-graph metadata keeps private wrapper caches per captured batch size, while eager execution safely replans shared wrappers for each iteration Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com> --- .../_torch/attention/backends/flashinfer.py | 20 +++---- tests/integration/test_lists/waives.txt | 7 --- .../attention/test_flashinfer_attention.py | 52 +++++++++++++++---- 3 files changed, 48 insertions(+), 31 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/flashinfer.py b/tensorrt_llm/_torch/attention/backends/flashinfer.py index 3029c19226b6..199b92b35adf 100644 --- a/tensorrt_llm/_torch/attention/backends/flashinfer.py +++ b/tensorrt_llm/_torch/attention/backends/flashinfer.py @@ -205,11 +205,13 @@ class PlanParams: sm_scale: Optional[float] = None window_left: Optional[int] = None kv_pool_id: Optional[int] = None - # Decode wrappers own persistent graph-visible buffers and counters. Speculative query width - # and generation batch size must distinguish cache entries; ordinary single-token decode - # reuses max-sized graph buffers across generation batch sizes. + # Decode wrappers bake the speculative query width into their launch plan, so it must + # distinguish cache entries. The generation batch size deliberately does not: CUDA graph + # metadata is built per captured batch size and owns a private wrapper cache, while eager + # metadata replans its wrappers on every prepare(). Keying on it instead grew the eager cache + # -- and its per-wrapper block tables, which are always sized by max_num_requests -- once per + # live batch size observed. q_len_per_req: int = 1 - num_generations: int = 0 # NB: Some features (multi-item scoring) are only supported with the paged KV-cache wrapper. @@ -1909,15 +1911,6 @@ def plan(self, "FlashInfer decode requires a uniform query length per " f"request, but got {generation_seq_lens.tolist()}") - # Ordinary single-token decode reuses max-sized graph buffers as the - # generation batch changes. Speculative and draft paths retain the - # batch size because it affects graph-visible wrapper state. - plan_num_generations = (self.num_generations if - (q_len_per_req > 1 - or self._uses_full_generation_page_table - or self._is_shared_kv_draft_view - or self._is_separate_kv_draft_view) else 0) - plan_params = PlanParams( num_heads=num_heads, num_kv_heads=num_kv_heads, @@ -1932,7 +1925,6 @@ def plan(self, multi_item_params=self._multi_item_params, kv_pool_id=getattr(self, "_vswa_active_pool_id", None), q_len_per_req=q_len_per_req, - num_generations=plan_num_generations, ) return self._plan_with_params(plan_params, flashinfer_backend) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index e1b5e7de2f0d..e8e461a77db5 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -224,13 +224,6 @@ full:L40S/accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_ full:L40S/accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_multi_instance[MMLU] SKIP (https://nvbugs/6649818) full:L40S/accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_tp_pp_symmetric[GSM8K-tp1pp2] SKIP (https://nvbugs/6649818) full:L40S/accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_tp_pp_symmetric[MMLU-tp1pp2] SKIP (https://nvbugs/6649818) -full:L40S/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16_4gpus[tp2pp2-attn_backend=FLASHINFER-torch_compile=False] SKIP (https://nvbugs/6689016) -full:L40S/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16_4gpus[tp2pp2-attn_backend=FLASHINFER-torch_compile=True] SKIP (https://nvbugs/6689016) -full:L40S/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16_4gpus[tp4-attn_backend=FLASHINFER-torch_compile=False] SKIP (https://nvbugs/6689016) -full:L40S/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16_4gpus[tp4-attn_backend=FLASHINFER-torch_compile=True] SKIP (https://nvbugs/6689016) -full:L40S/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp2pp2-fp8kv=False-attn_backend=FLASHINFER-torch_compile=False] SKIP (https://nvbugs/6689016) -full:L40S/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=False-attn_backend=FLASHINFER-torch_compile=False] SKIP (https://nvbugs/6689016) -full:L40S/accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=False-attn_backend=FLASHINFER-torch_compile=True] SKIP (https://nvbugs/6689016) full:L40S/disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_genpp2[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6728119) full:L40S/disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_gentp2[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6728119) full:L40S/disaggregated/test_disaggregated.py::test_disaggregated_ctxtp2_genpp2[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6728119) diff --git a/tests/unittest/_torch/attention/test_flashinfer_attention.py b/tests/unittest/_torch/attention/test_flashinfer_attention.py index d518e02eaa63..d577f2dd6bf9 100644 --- a/tests/unittest/_torch/attention/test_flashinfer_attention.py +++ b/tests/unittest/_torch/attention/test_flashinfer_attention.py @@ -185,7 +185,7 @@ def test_generation_page_table_uses_reserved_block_count(self): ) manager.get_batch_cache_indices.assert_called_once_with([99]) - def test_decode_plan_cache_key_reuses_single_token_batches(self): + def test_decode_query_width_is_part_of_plan_params(self): if not torch.cuda.is_available(): self.skipTest("CUDA is required for FlashInfer metadata") @@ -270,19 +270,16 @@ def return_plan_params(plan_params, _flashinfer_backend): ) self.assertEqual(single_token_plan.q_len_per_req, 1) - self.assertEqual(single_token_plan.num_generations, 0) self.assertEqual(larger_single_token_plan.q_len_per_req, 1) - self.assertEqual(larger_single_token_plan.num_generations, 0) self.assertEqual(single_token_plan, larger_single_token_plan) - self.assertEqual(full_page_single_token_plan.num_generations, 2) - self.assertNotEqual(single_token_plan, full_page_single_token_plan) - self.assertEqual(shared_draft_single_token_plan.num_generations, 2) - self.assertNotEqual(single_token_plan, shared_draft_single_token_plan) + self.assertEqual(single_token_plan, full_page_single_token_plan) + self.assertEqual(single_token_plan, shared_draft_single_token_plan) self.assertEqual(multi_token_plan.q_len_per_req, 6) self.assertNotEqual(single_token_plan, multi_token_plan) - self.assertEqual(multi_token_plan.num_generations, 2) - self.assertEqual(larger_batch_plan.num_generations, 3) - self.assertNotEqual(multi_token_plan, larger_batch_plan) + # The live generation batch size deliberately does not key the cache. Eager metadata + # replans every wrapper on each prepare() and sizes their block tables by + # max_num_requests, so keying on it would allocate a fresh wrapper per batch size seen. + self.assertEqual(multi_token_plan, larger_batch_plan) metadata.seq_lens = torch.tensor([6, 5], dtype=torch.int32) with self.assertRaisesRegex(ValueError, "uniform query length"): @@ -296,6 +293,41 @@ def return_plan_params(plan_params, _flashinfer_backend): flashinfer_backend="trtllm-gen", ) + def test_cuda_graph_metadata_owns_a_private_plan_cache(self): + if not torch.cuda.is_available(): + self.skipTest("CUDA is required for FlashInfer metadata") + + metadata = FlashInferAttentionMetadata( + seq_lens=torch.tensor([1, 1], dtype=torch.int32), + num_contexts=0, + kv_cache_manager=None, + request_ids=[0, 1], + max_num_requests=4, + max_num_tokens=16, + ) + plan_params = PlanParams( + num_heads=32, + num_kv_heads=4, + head_dim=128, + q_dtype=torch.bfloat16, + kv_dtype=torch.bfloat16, + attention_mask_type=AttentionMaskType.causal, + ) + eager_wrappers = FlashInferWrappers(is_planned=True) + metadata._plan_params_to_wrappers[plan_params] = eager_wrappers + + graph_metadata = metadata.create_cuda_graph_metadata(2) + + # Wrappers belong to a metadata instance, and the runner builds one metadata per captured + # batch size. That partitioning -- not the cache key -- is what keeps a graph wrapper from + # being reused across batch sizes, so PlanParams need not carry the generation batch size. + self.assertTrue(graph_metadata.is_cuda_graph) + self.assertIsNot(graph_metadata._plan_params_to_wrappers, + metadata._plan_params_to_wrappers) + self.assertEqual(graph_metadata._plan_params_to_wrappers, {}) + self.assertIs(metadata._plan_params_to_wrappers[plan_params], + eager_wrappers) + def test_generation_page_table_keeps_logical_positions(self): if not torch.cuda.is_available(): self.skipTest("CUDA is required for FlashInfer metadata")