Refactor common attention metadata

Signed-off-by: luka <[email protected]>

Author: ProExpertProg

vllm/v1/attention/backends/cpu_attn.py

@@ -119,9 +119,12 @@ def reorder_batch(self, input_batch: InputBatch,
 
         return True
 
-    def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
-              common_prefix_len: int,
+    def build(self, common_prefix_len: int,
               common_attn_metadata: CommonAttentionMetadata):
+        num_reqs = common_attn_metadata.num_reqs
+        num_actual_tokens = common_attn_metadata.num_actual_tokens
+        max_query_len = common_attn_metadata.max_query_len
+
         runner = self.runner
         block_table = self.block_table
         seq_lens_np = runner.seq_lens_np[:num_reqs]

vllm/v1/attention/backends/flash_attn.py

@@ -341,9 +341,13 @@ def __init__(self, runner: "GPUModelRunner", kv_cache_spec: AttentionSpec,
         self.aot_sliding_window: Optional[tuple[int, int]] = None
 
     def build(
-        self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
-        common_prefix_len: int, common_attn_metadata: CommonAttentionMetadata
+        self, common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata
     ) -> FlashAttentionMetadata:
+        num_reqs = common_attn_metadata.num_reqs
+        num_actual_tokens = common_attn_metadata.num_actual_tokens
+        max_query_len = common_attn_metadata.max_query_len
+
         max_seq_len = int(self.runner.seq_lens_np[:num_reqs].max())
         query_start_loc = common_attn_metadata.query_start_loc
         seq_lens = common_attn_metadata.seq_lens

vllm/v1/attention/backends/flashinfer.py

@@ -400,9 +400,11 @@ def _plan(self, attn_metadata: FlashInferMetadata):
                     kv_data_type=attn_metadata.data_type,
                 )
 
-    def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
-              common_prefix_len: int,
+    def build(self, common_prefix_len: int,
               common_attn_metadata: CommonAttentionMetadata):
+        num_reqs = common_attn_metadata.num_reqs
+        num_actual_tokens = common_attn_metadata.num_actual_tokens
+
         assert self._num_decodes + self._num_prefills == num_reqs
         assert (self._num_decode_tokens +
                 self._num_prefill_tokens == num_actual_tokens)

vllm/v1/attention/backends/flex_attention.py

@@ -272,9 +272,12 @@ def __init__(self, runner: "GPUModelRunner", kv_cache_spec: AttentionSpec,
         self.kv_cache_spec = kv_cache_spec
         self.block_table = block_table
 
-    def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
-              common_prefix_len: int,
+    def build(self, common_prefix_len: int,
               common_attn_metadata: CommonAttentionMetadata):
+        num_reqs = common_attn_metadata.num_reqs
+        num_actual_tokens = common_attn_metadata.num_actual_tokens
+        max_query_len = common_attn_metadata.max_query_len
+
         max_seq_len = self.runner.seq_lens_np[:num_reqs].max()
         query_start_loc = common_attn_metadata.query_start_loc
         seq_lens = common_attn_metadata.seq_lens

vllm/v1/attention/backends/mla/common.py

@@ -459,26 +459,31 @@ def _build_decode(self, block_table_tensor: torch.Tensor,
         )
 
     def build_for_cudagraph_capture(
-            self, num_reqs: int, num_tokens: int,
-            common_attn_metadata: CommonAttentionMetadata) -> M:
+            self, common_attn_metadata: CommonAttentionMetadata) -> M:
         """
         This method builds the metadata for full cudagraph capture.
         Currently, only decode is supported for full cudagraphs with MLA.
         """
-        assert num_reqs == num_tokens, \
+        m = common_attn_metadata
+        assert m.num_reqs == m.num_actual_tokens, \
             "MLA only supports decode-only full CUDAGraph capture. " \
             "Make sure all cudagraph capture sizes <= max_num_seq."
 
+        m.max_query_len = 1  # decode-only
+
         # Update state usually set in reorder_batch.
-        self._num_decodes = num_tokens
-        self._num_decode_tokens = num_tokens
+        self._num_decodes = m.num_reqs
+        self._num_decode_tokens = m.num_actual_tokens
         self._num_prefills = 0
         self._num_prefill_tokens = 0
-        return self.build(num_tokens, num_tokens, 1, 0, common_attn_metadata)
+        return self.build(0, m)
 
-    def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
-              common_prefix_len: int,
+    def build(self, common_prefix_len: int,
               common_attn_metadata: CommonAttentionMetadata) -> M:
+        num_reqs = common_attn_metadata.num_reqs
+        num_actual_tokens = common_attn_metadata.num_actual_tokens
+        max_query_len = common_attn_metadata.max_query_len
+
         assert self._num_decodes + self._num_prefills == num_reqs
 
         # Note(simon): be careful about the CPU <> GPU memory movement in this

vllm/v1/attention/backends/utils.py

@@ -26,15 +26,21 @@ class CommonAttentionMetadata:
     """(batch_size,), the length of each request including both computed tokens
     and newly scheduled tokens"""
 
+    num_reqs: int
+    """Number of requests"""
+    num_actual_tokens: int
+    """Total number of tokens in batch"""
+    max_query_len: int
+    """Longest query in batch"""
+
 
 M = TypeVar("M")
 
 
 class AttentionMetadataBuilder(abc.ABC, Generic[M]):
 
     @abstractmethod
-    def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
-              common_prefix_len: int,
+    def build(self, common_prefix_len: int,
               common_attn_metadata: CommonAttentionMetadata) -> M:
         """
         Central method that builds attention metadata.
@@ -43,14 +49,14 @@ def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
         raise NotImplementedError
 
     def build_for_cudagraph_capture(
-            self, num_reqs: int, num_tokens: int,
-            common_attn_metadata: CommonAttentionMetadata) -> M:
+            self, common_attn_metadata: CommonAttentionMetadata) -> M:
         """
         Build attention metadata for CUDA graph capture. Uses build by default.
-        Subclasses that override this method should call self.build.
+        Subclasses that override this method should call self.build or
+        super().build_for_cudagraph_capture.
         """
-        return self.build(num_reqs, num_tokens, num_tokens, 0,
-                          common_attn_metadata)
+        return self.build(common_prefix_len=0,
+                          common_attn_metadata=common_attn_metadata)
 
     def use_cascade_attention(
         self,

vllm/v1/spec_decode/eagle.py

@@ -138,15 +138,17 @@ def propose(
             max_query_len = query_lens.max().item()
 
             common_attn_metadata = CommonAttentionMetadata(
-                query_start_loc=cu_num_tokens, seq_lens=seq_lens)
+                query_start_loc=cu_num_tokens,
+                seq_lens=seq_lens,
+                num_reqs=batch_size,
+                num_actual_tokens=num_tokens,
+                max_query_len=max_query_len,
+            )
 
             assert self.runner is not None
 
             # FIXME: need to consider multiple kv_cache_groups
             attn_metadata = self.runner.attn_metadata_builder.build(
-                num_reqs=batch_size,
-                num_actual_tokens=num_tokens,
-                max_query_len=max_query_len,
                 common_prefix_len=0,
                 common_attn_metadata=common_attn_metadata,
             )

vllm/v1/worker/gpu_model_runner.py

@@ -669,7 +669,12 @@ def _prepare_inputs(
         seq_lens = self.seq_lens[:num_reqs]
 
         common_attn_metadata = CommonAttentionMetadata(
-            query_start_loc=query_start_loc, seq_lens=seq_lens)
+            query_start_loc=query_start_loc,
+            seq_lens=seq_lens,
+            num_reqs=num_reqs,
+            num_actual_tokens=total_num_scheduled_tokens,
+            max_query_len=max_num_scheduled_tokens,
+        )
 
         attn_metadata: dict[str, Any] = {}
         # Prepare the attention metadata for each KV cache group and make layers
@@ -690,11 +695,9 @@ def _prepare_inputs(
 
             attn_metadata_i = (
                 self.attn_metadata_builders[kv_cache_group_id].build(
-                    num_reqs=num_reqs,
-                    num_actual_tokens=total_num_scheduled_tokens,
-                    max_query_len=max_num_scheduled_tokens,
                     common_prefix_len=common_prefix_len,
-                    common_attn_metadata=common_attn_metadata))
+                    common_attn_metadata=common_attn_metadata,
+                ))
             for layer_name in kv_cache_group_spec.layer_names:
                 attn_metadata[layer_name] = attn_metadata_i
 
@@ -1809,18 +1812,20 @@ def _dummy_run(
             seq_lens = self.seq_lens[:num_reqs]
 
             common_attn_metadata = CommonAttentionMetadata(
-                query_start_loc=query_start_loc, seq_lens=seq_lens)
+                query_start_loc=query_start_loc,
+                seq_lens=seq_lens,
+                num_reqs=num_reqs,
+                num_actual_tokens=num_tokens,
+                max_query_len=num_tokens,
+            )
 
             attn_metadata = {}
             for kv_cache_group_id, kv_cache_group_spec in enumerate(
                     self.kv_cache_config.kv_cache_groups):
 
                 attn_metadata_i = self.attn_metadata_builders[
                     kv_cache_group_id].build_for_cudagraph_capture(
-                        num_reqs=num_reqs,
-                        num_tokens=num_tokens,
-                        common_attn_metadata=common_attn_metadata,
-                    )
+                        common_attn_metadata)
                 for layer_name in kv_cache_group_spec.layer_names:
                     attn_metadata[layer_name] = attn_metadata_i