[https://nvbugs/6689016][fix] Limit eager FlashInfer plan cache growth

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)
 

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)

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")