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