[Perf][MRV2] Share token-to-request mappings across KV cache groups (#57102)

Signed-off-by: Juntian Liu <Juntianl777@gmail.com>
Co-authored-by: Woosuk Kwon <woosuk@inferact.ai>

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