diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index eccfdbabcd..89f54d4d11 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -426,6 +426,7 @@ def build_attn_metadata( seq_lens_cpu_upper_bound = seq_lens_cpu_upper_bound[:num_reqs] attn_metadata: dict[str, Any] = {} + token_to_req_indices: torch.Tensor | None = None num_kv_cache_groups = len(kv_cache_config.kv_cache_groups) for i in range(num_kv_cache_groups): if not attn_groups[i]: @@ -472,6 +473,7 @@ def build_attn_metadata( mm_req_doc_ranges=mm_req_doc_ranges, rswa_prefix_lens=rswa_prefix_lens, req_idx=req_idx, + _token_to_req_indices_cache=token_to_req_indices, **common_attn_metadata_extra_kwargs, ) @@ -497,6 +499,7 @@ def build_attn_metadata( ) for layer_name in attn_group.layer_names: attn_metadata[layer_name] = metadata + token_to_req_indices = common_attn_metadata._token_to_req_indices_cache return attn_metadata