From 24a4c29344b4960300d4e06e691a7b42fffee6d1 Mon Sep 17 00:00:00 2001 From: RunguoLi Date: Mon, 21 Sep 2026 18:17:00 -0500 Subject: [PATCH] [#19527][fix] Account for beam width when sizing CUDA graph warmup KV requests With beam search, only KV blocks fully covered by the prompt are shared among beams; the partial last prompt block and every block allocated while appending tokens are allocated once per beam. The V1 KVCacheManager ignored this when sizing CUDA graph warmup dummy requests, so with a small KV cache pool the warmup could request more blocks than exist and abort startup with "No free block found. This shouldn't happen!". - get_num_available_tokens takes max_beam_width and returns a length such that every sequence up to it fits with per-beam allocation. - add_dummy_requests returns None instead of failing inside the block manager when beam-search dummy requests cannot fit. - The CUDA graph warmup passes max_beam_width; KVCacheManagerV2 accepts the argument (it only supports a beam width of 1). Signed-off-by: RunguoLi --- .../kv_cache/kv_cache_manager_v2.py | 10 +- .../_torch/pyexecutor/model_engine.py | 6 +- .../_torch/pyexecutor/resource_manager.py | 69 +++++++++++++- .../_torch/executor/test_resource_manager.py | 95 +++++++++++++++++++ 4 files changed, 174 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index cb8d99a9b8ea..66a9d8189dd1 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/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. diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 79f9fd7ff7b5..206d209dd286 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/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( diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index c755dda60358..4262a442907a 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/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]: diff --git a/tests/unittest/_torch/executor/test_resource_manager.py b/tests/unittest/_torch/executor/test_resource_manager.py index bded6c870e89..c3859a1c13b3 100644 --- a/tests/unittest/_torch/executor/test_resource_manager.py +++ b/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