[#19527][fix] Account for beam width when sizing CUDA graph warmup KV requests

tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py

@@ -3151,7 +3151,12 @@ def get_index_k_buffer(
         return full_view[:, 0]
 
     def get_num_available_tokens(
-        self, *, token_num_upper_bound: int, batch_size: int = 1, max_num_draft_tokens: int = 0
+        self,
+        *,
+        token_num_upper_bound: int,
+        batch_size: int = 1,
+        max_num_draft_tokens: int = 0,
+        max_beam_width: int = 1,
     ) -> int:
         """Clamp ``token_num_upper_bound`` to the allocatable token capacity.
 
@@ -3163,6 +3168,9 @@ def get_num_available_tokens(
         ``max_num_tokens``) stay consistent because a helix context forward
         replicates all tokens on every rank, so both bounds constrain the
         same request-length variable.
+
+        ``max_beam_width`` is accepted for interface parity with the V1
+        manager; V2 only supports a beam width of 1.
         """
         extra_tokens = self.num_extra_kv_tokens + max_num_draft_tokens
         # Token num upper bound is the maximum number of tokens that can be allocated in the kv cache manager.

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -3335,14 +3335,16 @@ def free_warmup_requests() -> None:
         available_tokens = kv_cache_manager.get_num_available_tokens(
             token_num_upper_bound=max_seq_len,
             batch_size=batch_size,
-            max_num_draft_tokens=_kv_draft)
+            max_num_draft_tokens=_kv_draft,
+            max_beam_width=self.max_beam_width)
 
         # Also consider draft KV cache capacity when it exists
         if draft_kv_cache_manager is not None:
             draft_available_tokens = draft_kv_cache_manager.get_num_available_tokens(
                 batch_size=batch_size,
                 token_num_upper_bound=max_seq_len,
-                max_num_draft_tokens=_kv_draft)
+                max_num_draft_tokens=_kv_draft,
+                max_beam_width=self.max_beam_width)
             available_tokens = min(available_tokens, draft_available_tokens)
 
         token_num = max(

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -1347,6 +1347,25 @@ def add_dummy_requests(
                 _populate_dummy_mrope_config(req, token_num, is_gen)
             requests.append(req)
 
+        # Beam search allocates most blocks once per beam, so the single-block
+        # check above does not guarantee the dummy requests fit. Skip padding
+        # instead of failing inside the block manager. VSWA pools are sized
+        # per window, which this full-attention count does not model.
+        if beam_width > 1 and batch_request_infos and not self.is_vswa:
+            num_appended_tokens = self.num_extra_kv_tokens + (_kv_draft
+                                                              if is_gen else 0)
+            num_required_blocks = sum(
+                self._get_num_blocks_for_dummy_request(
+                    token_num, num_appended_tokens, beam_width)
+                for _, token_num, _ in batch_request_infos)
+            if num_required_blocks > available_blocks:
+                logger.debug(
+                    f"[add_dummy_requests] {len(batch_request_infos)} dummy "
+                    f"requests with beam_width={beam_width} need "
+                    f"{num_required_blocks} blocks, only {available_blocks} "
+                    f"free; skipping.")
+                return None
+
         try:
             # Use add_sequence_batch for all dummy requests, then add extra tokens.
             # This must happen before is_gen state modifications below, which may
@@ -1839,18 +1858,62 @@ def get_num_kv_blocks(self, num_tokens: int) -> int:
     def get_num_available_tokens(self,
                                  token_num_upper_bound: int,
                                  max_num_draft_tokens: int = 0,
+                                 max_beam_width: int = 1,
                                  **kwargs) -> int:
+        """Return a token count such that one sequence of any length up to it
+        fits in the free blocks.
+
+        Args:
+            token_num_upper_bound: Upper bound on the returned token count.
+            max_num_draft_tokens: Draft tokens appended after the sequence.
+            max_beam_width: Beam width of the sequence. With beam search, only
+                blocks fully covered by the prompt are shared among beams; the
+                rest are allocated once per beam.
+        """
         free_blocks = self.get_num_free_blocks()
-        result = min(
-            token_num_upper_bound, free_blocks * self.tokens_per_block -
-            self.num_extra_kv_tokens - max_num_draft_tokens)
+        num_appended_tokens = self.num_extra_kv_tokens + max_num_draft_tokens
+        if max_beam_width > 1 and self.kv_cache_type != CacheTypeCpp.CROSS:
+            # Block usage is not monotonic in the sequence length (a
+            # block-aligned prompt shares all of its blocks), so bound it by
+            # the worst case: a partially filled last prompt block followed by
+            # the appended tokens, all allocated per beam.
+            max_blocks_per_beam = math.ceil(
+                (self.tokens_per_block - 1 + num_appended_tokens) /
+                self.tokens_per_block)
+            num_shared_blocks = free_blocks - max_beam_width * max_blocks_per_beam
+            capacity = (num_shared_blocks + 1) * self.tokens_per_block - 1
+        else:
+            capacity = free_blocks * self.tokens_per_block - num_appended_tokens
+        result = min(token_num_upper_bound, capacity)
         logger.debug(
             f"[get_num_available_tokens] free_blocks={free_blocks}, "
             f"tokens_per_block={self.tokens_per_block}, "
             f"num_extra_kv_tokens={self.num_extra_kv_tokens}, "
+            f"max_beam_width={max_beam_width}, "
             f"token_num_upper_bound={token_num_upper_bound}, result={result}")
         return result
 
+    def _get_num_blocks_for_dummy_request(self, token_num: int,
+                                          num_appended_tokens: int,
+                                          beam_width: int) -> int:
+        """Number of blocks ``add_dummy_requests`` allocates for one sequence
+        of ``token_num`` prompt tokens followed by ``num_appended_tokens``
+        tokens added one at a time.
+
+        Blocks fully covered by the prompt are shared among beams (for cross
+        KV, the partial last prompt block is shared too); every other block is
+        allocated once per beam.
+        """
+        num_blocks = math.ceil(
+            (token_num + num_appended_tokens) / self.tokens_per_block)
+        if beam_width == 1:
+            return num_blocks
+        if self.kv_cache_type == CacheTypeCpp.CROSS:
+            num_shared_blocks = math.ceil(token_num / self.tokens_per_block)
+        else:
+            num_shared_blocks = token_num // self.tokens_per_block
+        return num_shared_blocks + beam_width * (num_blocks - num_shared_blocks)
+
     def get_buffers(self,
                     layer_idx: int,
                     kv_layout: str = "NHD") -> Optional[torch.Tensor]:

tests/unittest/_torch/executor/test_resource_manager.py

@@ -1050,6 +1050,101 @@ def test_batch_cache_indices_honor_requested_blocks_for_beam0(self):
         finally:
             kv_cache_manager.shutdown()
 
+    @staticmethod
+    def _create_beam_search_kv_cache_manager(max_beam_width: int,
+                                             max_batch_size: int = 1):
+        # 32 blocks of 8 tokens.
+        return KVCacheManager(
+            kv_cache_config=KvCacheConfig(max_tokens=256,
+                                          enable_block_reuse=False),
+            kv_cache_type=tensorrt_llm.bindings.internal.batch_manager.
+            CacheType.SELF,
+            num_layers=2,
+            num_kv_heads=2,
+            head_dim=128,
+            tokens_per_block=8,
+            max_seq_len=256,
+            max_batch_size=max_batch_size,
+            max_beam_width=max_beam_width,
+            mapping=Mapping(),
+        )
+
+    def test_dummy_request_block_count_matches_beam_search_allocation(self):
+        """Prompt-covered blocks are shared among beams; the partial last
+        block is allocated once per beam."""
+        beam_width = 4
+        kv_cache_manager = self._create_beam_search_kv_cache_manager(beam_width)
+        try:
+            total_free = kv_cache_manager.get_num_free_blocks()
+            for request_id, token_num in enumerate([1, 8, 9, 100, 128]):
+                requests = kv_cache_manager.add_dummy_requests(
+                    [request_id], [token_num],
+                    is_gen=True,
+                    max_beam_width=beam_width)
+                self.assertIsNotNone(requests)
+                used_blocks = (total_free -
+                               kv_cache_manager.get_num_free_blocks())
+                self.assertEqual(
+                    used_blocks,
+                    kv_cache_manager._get_num_blocks_for_dummy_request(
+                        token_num, 0, beam_width), f"token_num={token_num}")
+                kv_cache_manager.free_resources(requests[0])
+        finally:
+            kv_cache_manager.shutdown()
+
+    def test_get_num_available_tokens_accounts_for_beam_width(self):
+        """Every length up to the reported capacity must fit with beam
+        search, including lengths that are not block aligned."""
+        beam_width = 4
+        kv_cache_manager = self._create_beam_search_kv_cache_manager(beam_width)
+        try:
+            self.assertEqual(kv_cache_manager.get_num_free_blocks(), 32)
+            self.assertEqual(
+                kv_cache_manager.get_num_available_tokens(
+                    token_num_upper_bound=1024), 256)
+            capacity = kv_cache_manager.get_num_available_tokens(
+                token_num_upper_bound=1024, max_beam_width=beam_width)
+            # 32 free blocks minus one per-beam tail block for each beam.
+            self.assertEqual(capacity, (32 - beam_width + 1) * 8 - 1)
+            for token_num in range(1, capacity + 1):
+                requests = kv_cache_manager.add_dummy_requests(
+                    [0], [token_num], is_gen=True, max_beam_width=beam_width)
+                self.assertIsNotNone(requests, f"token_num={token_num}")
+                kv_cache_manager.free_resources(requests[0])
+            self.assertEqual(kv_cache_manager.get_num_free_blocks(), 32)
+        finally:
+            kv_cache_manager.shutdown()
+
+    def test_add_dummy_requests_beam_search_returns_none_when_pool_too_small(
+            self):
+        """Dummy requests that cannot fit with beam search are skipped
+        instead of failing inside the block manager, and nothing leaks."""
+        beam_width = 4
+        kv_cache_manager = self._create_beam_search_kv_cache_manager(
+            beam_width, max_batch_size=16)
+        try:
+            total_free = kv_cache_manager.get_num_free_blocks()
+            # 31 shared blocks plus one tail block per beam: 35 > 32.
+            self.assertIsNone(
+                kv_cache_manager.add_dummy_requests([0], [255],
+                                                    is_gen=True,
+                                                    max_beam_width=beam_width))
+            # One per-beam block per request: 9 * 4 = 36 > 32.
+            self.assertIsNone(
+                kv_cache_manager.add_dummy_requests(list(range(9)),
+                                                    is_gen=True,
+                                                    max_beam_width=beam_width))
+            self.assertEqual(kv_cache_manager.get_num_free_blocks(), total_free)
+            # 8 * 4 = 32 blocks fits exactly.
+            requests = kv_cache_manager.add_dummy_requests(
+                list(range(8)), is_gen=True, max_beam_width=beam_width)
+            self.assertIsNotNone(requests)
+            self.assertEqual(kv_cache_manager.get_num_free_blocks(), 0)
+            for request in requests:
+                kv_cache_manager.free_resources(request)
+        finally:
+            kv_cache_manager.shutdown()
+
     def test_add_dummy_requests_failure_frees_partial_allocation(self):
         """A partial add_dummy_requests failure must free every block it
         allocated (TRTLLM-14903): leaked blocks on the minimal pool built for