[3/N][Attention] Move AttentionMetadata-related code from utils.py to backend.py (vllm-project#32054)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

Author: MatthewBonanni

docs/design/cuda_graphs.md

@@ -149,7 +149,7 @@ The CUDA Graphs wrapper no longer manages the warm-up logic. The warm-up process
 
 ## CUDA Graphs Compatibility of Attention Backends
 
-To signal the CUDA Graphs compatibility of the attention backends, we introduce a new enum type [AttentionCGSupport][vllm.v1.attention.backends.utils.AttentionCGSupport], which is an enum type that tracks the capability of the attention backend to support CUDA Graphs. The value is sorted in the order of the capability, i.e., `ALWAYS`> `UNIFORM_BATCH`> `UNIFORM_SINGLE_TOKEN_DECODE`> `NEVER`.
+To signal the CUDA Graphs compatibility of the attention backends, we introduce a new enum type [AttentionCGSupport][vllm.v1.attention.backend.AttentionCGSupport], which is an enum type that tracks the capability of the attention backend to support CUDA Graphs. The value is sorted in the order of the capability, i.e., `ALWAYS`> `UNIFORM_BATCH`> `UNIFORM_SINGLE_TOKEN_DECODE`> `NEVER`.
 
 ```python
 class AttentionCGSupport(enum.Enum):

tests/v1/attention/test_attention_backends.py

@@ -23,10 +23,9 @@
     is_torch_equal_or_newer,
     set_random_seed,
 )
-from vllm.v1.attention.backend import AttentionType
+from vllm.v1.attention.backend import AttentionType, CommonAttentionMetadata
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.attention.backends.utils import (
-    CommonAttentionMetadata,
     set_kv_cache_layout,
 )
 from vllm.v1.kv_cache_interface import FullAttentionSpec

tests/v1/attention/test_mla_backends.py

@@ -22,10 +22,10 @@
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.utils.math_utils import cdiv
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
+from vllm.v1.attention.backend import CommonAttentionMetadata
 from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
 from vllm.v1.attention.backends.mla.common import QueryLenSupport
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
 from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported
 from vllm.v1.kv_cache_interface import FullAttentionSpec
 

tests/v1/attention/utils.py

@@ -18,12 +18,12 @@
     VllmConfig,
 )
 from vllm.config.model import ModelDType
-from vllm.v1.attention.backend import AttentionImpl
-from vllm.v1.attention.backends.registry import AttentionBackendEnum
-from vllm.v1.attention.backends.utils import (
+from vllm.v1.attention.backend import (
+    AttentionImpl,
     AttentionMetadataBuilder,
     CommonAttentionMetadata,
 )
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.kv_cache_interface import FullAttentionSpec
 
 

tests/v1/e2e/test_async_spec_decode.py

@@ -19,7 +19,7 @@ def sync_tracker():
     Fixture that patches CommonAttentionMetadata.seq_lens_cpu to detect
     lazy init syncs. Prints stack traces immediately when syncs occur.
     """
-    from vllm.v1.attention.backends.utils import CommonAttentionMetadata
+    from vllm.v1.attention.backend import CommonAttentionMetadata
 
     # Shared counter for cross-process communication (inherited by fork)
     sync_count = multiprocessing.Value("i", 0)

tests/v1/spec_decode/test_tree_attention.py

@@ -12,9 +12,9 @@
     try_get_attention_backend,
 )
 from vllm.config import ParallelConfig, SpeculativeConfig
+from vllm.v1.attention.backend import CommonAttentionMetadata
 from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
 
 if not is_flash_attn_varlen_func_available():
     pytest.skip(

vllm/model_executor/layers/attention/chunked_local_attention.py

@@ -8,11 +8,13 @@
 from vllm.config import CacheConfig
 from vllm.config.vllm import VllmConfig
 from vllm.model_executor.layers.quantization import QuantizationConfig
-from vllm.v1.attention.backend import AttentionBackend
-from vllm.v1.attention.backends.utils import (
+from vllm.v1.attention.backend import (
+    AttentionBackend,
     AttentionCGSupport,
     AttentionMetadataBuilder,
     CommonAttentionMetadata,
+)
+from vllm.v1.attention.backends.utils import (
     make_local_attention_virtual_batches,
     subclass_attention_backend,
 )

vllm/model_executor/layers/attention/cross_attention.py

@@ -14,9 +14,9 @@
     AttentionBackend,
     AttentionMetadata,
     AttentionType,
+    CommonAttentionMetadata,
 )
 from vllm.v1.attention.backends.utils import (
-    CommonAttentionMetadata,
     subclass_attention_backend,
 )
 from vllm.v1.attention.selector import get_attn_backend

vllm/model_executor/layers/attention/encoder_only_attention.py

@@ -12,9 +12,9 @@
     AttentionBackend,
     AttentionMetadata,
     AttentionType,
+    CommonAttentionMetadata,
 )
 from vllm.v1.attention.backends.utils import (
-    CommonAttentionMetadata,
     subclass_attention_backend,
 )
 from vllm.v1.attention.selector import get_attn_backend

vllm/model_executor/layers/attention/static_sink_attention.py

@@ -15,9 +15,9 @@
     AttentionBackend,
     AttentionMetadata,
     AttentionType,
+    CommonAttentionMetadata,
 )
 from vllm.v1.attention.backends.utils import (
-    CommonAttentionMetadata,
     subclass_attention_backend,
 )
 from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (