@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 20228cb851
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 10, 06:30 UTC

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

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

Matthew Bonanni committed Jan 12, 2026 at 12:13 UTC 20228cb8514ec2f8155b94fe0af204d839238280
37 files changed +374 -370
docs/design/cuda_graphs.md
+1 -1
@@ -149,7 +149,7 @@ The CUDA Graphs wrapper no longer manages the warm-up logic. The warm-up process
149
150 ## CUDA Graphs Compatibility of Attention Backends
151
152 -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`.
152 +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`.
153
154 ```python
155 class AttentionCGSupport(enum.Enum):
tests/v1/attention/test_attention_backends.py
+1 -2
@@ -23,10 +23,9 @@ from vllm.utils.torch_utils import (
23 is_torch_equal_or_newer,
24 set_random_seed,
25 )
26 -from vllm.v1.attention.backend import AttentionType
26 +from vllm.v1.attention.backend import AttentionType, CommonAttentionMetadata
27 from vllm.v1.attention.backends.registry import AttentionBackendEnum
28 from vllm.v1.attention.backends.utils import (
29 - CommonAttentionMetadata,
29 set_kv_cache_layout,
30 )
31 from vllm.v1.kv_cache_interface import FullAttentionSpec
tests/v1/attention/test_mla_backends.py
+1 -1
@@ -22,10 +22,10 @@ from vllm.config.vllm import set_current_vllm_config
22 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
23 from vllm.utils.math_utils import cdiv
24 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
25 +from vllm.v1.attention.backend import CommonAttentionMetadata
26 from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
27 from vllm.v1.attention.backends.mla.common import QueryLenSupport
28 from vllm.v1.attention.backends.registry import AttentionBackendEnum
28 -from vllm.v1.attention.backends.utils import CommonAttentionMetadata
29 from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported
30 from vllm.v1.kv_cache_interface import FullAttentionSpec
31
tests/v1/attention/utils.py
+3 -3
@@ -18,12 +18,12 @@ from vllm.config import (
18 VllmConfig,
19 )
20 from vllm.config.model import ModelDType
21 -from vllm.v1.attention.backend import AttentionImpl
22 -from vllm.v1.attention.backends.registry import AttentionBackendEnum
23 -from vllm.v1.attention.backends.utils import (
21 +from vllm.v1.attention.backend import (
22 + AttentionImpl,
23 AttentionMetadataBuilder,
24 CommonAttentionMetadata,
25 )
26 +from vllm.v1.attention.backends.registry import AttentionBackendEnum
27 from vllm.v1.kv_cache_interface import FullAttentionSpec
28
29
tests/v1/e2e/test_async_spec_decode.py
+1 -1
@@ -19,7 +19,7 @@ def sync_tracker():
19 Fixture that patches CommonAttentionMetadata.seq_lens_cpu to detect
20 lazy init syncs. Prints stack traces immediately when syncs occur.
21 """
22 - from vllm.v1.attention.backends.utils import CommonAttentionMetadata
22 + from vllm.v1.attention.backend import CommonAttentionMetadata
23
24 # Shared counter for cross-process communication (inherited by fork)
25 sync_count = multiprocessing.Value("i", 0)
tests/v1/spec_decode/test_tree_attention.py
+1 -1
@@ -12,9 +12,9 @@ from tests.v1.attention.utils import (
12 try_get_attention_backend,
13 )
14 from vllm.config import ParallelConfig, SpeculativeConfig
15 +from vllm.v1.attention.backend import CommonAttentionMetadata
16 from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available
17 from vllm.v1.attention.backends.registry import AttentionBackendEnum
17 -from vllm.v1.attention.backends.utils import CommonAttentionMetadata
18
19 if not is_flash_attn_varlen_func_available():
20 pytest.skip(
vllm/model_executor/layers/attention/chunked_local_attention.py
+4 -2
@@ -8,11 +8,13 @@ from vllm.attention.layer import Attention
8 from vllm.config import CacheConfig
9 from vllm.config.vllm import VllmConfig
10 from vllm.model_executor.layers.quantization import QuantizationConfig
11 -from vllm.v1.attention.backend import AttentionBackend
12 -from vllm.v1.attention.backends.utils import (
11 +from vllm.v1.attention.backend import (
12 + AttentionBackend,
13 AttentionCGSupport,
14 AttentionMetadataBuilder,
15 CommonAttentionMetadata,
16 +)
17 +from vllm.v1.attention.backends.utils import (
18 make_local_attention_virtual_batches,
19 subclass_attention_backend,
20 )
vllm/model_executor/layers/attention/cross_attention.py
+1 -1
@@ -14,9 +14,9 @@ from vllm.v1.attention.backend import (
14 AttentionBackend,
15 AttentionMetadata,
16 AttentionType,
17 + CommonAttentionMetadata,
18 )
19 from vllm.v1.attention.backends.utils import (
19 - CommonAttentionMetadata,
20 subclass_attention_backend,
21 )
22 from vllm.v1.attention.selector import get_attn_backend
vllm/model_executor/layers/attention/encoder_only_attention.py
+1 -1
@@ -12,9 +12,9 @@ from vllm.v1.attention.backend import (
12 AttentionBackend,
13 AttentionMetadata,
14 AttentionType,
15 + CommonAttentionMetadata,
16 )
17 from vllm.v1.attention.backends.utils import (
17 - CommonAttentionMetadata,
18 subclass_attention_backend,
19 )
20 from vllm.v1.attention.selector import get_attn_backend
vllm/model_executor/layers/attention/static_sink_attention.py
+1 -1
@@ -15,9 +15,9 @@ from vllm.v1.attention.backend import (
15 AttentionBackend,
16 AttentionMetadata,
17 AttentionType,
18 + CommonAttentionMetadata,
19 )
20 from vllm.v1.attention.backends.utils import (
20 - CommonAttentionMetadata,
21 subclass_attention_backend,
22 )
23 from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
vllm/model_executor/models/whisper_utils.py
+1 -1
@@ -16,10 +16,10 @@ from vllm.v1.attention.backend import (
16 AttentionBackend,
17 AttentionMetadata,
18 AttentionType,
19 + CommonAttentionMetadata,
20 )
21 from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
22 from vllm.v1.attention.backends.utils import (
22 - CommonAttentionMetadata,
23 subclass_attention_backend_with_overrides,
24 )
25 from vllm.v1.attention.selector import get_attn_backend
vllm/v1/attention/backend.py
+287
@@ -2,17 +2,22 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from abc import ABC, abstractmethod
5 +from dataclasses import dataclass
6 from enum import Enum
7 from typing import TYPE_CHECKING, ClassVar, Generic, Protocol, TypeVar, get_args
8
9 +import numpy as np
10 import torch
11 +from typing_extensions import deprecated
12
13 if TYPE_CHECKING:
14 + from vllm.config import VllmConfig
15 from vllm.config.cache import CacheDType
16 from vllm.model_executor.layers.linear import ColumnParallelLinear
17 from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
18 from vllm.platforms.interface import DeviceCapability
19 from vllm.v1.attention.backends.utils import KVCacheLayoutType
20 + from vllm.v1.kv_cache_interface import AttentionSpec
21
22
23 class AttentionType(str, Enum):
@@ -271,6 +276,288 @@ class AttentionMetadata:
276 T = TypeVar("T", bound=AttentionMetadata)
277
278
279 +@dataclass
280 +class CommonAttentionMetadata:
281 + """
282 + Per-batch attention metadata, shared across layers and backends.
283 + AttentionMetadataBuilder instances use it to construct per-layer metadata.
284 +
285 + For many of the tensors we keep both GPU and CPU versions.
286 + """
287 +
288 + query_start_loc: torch.Tensor
289 + query_start_loc_cpu: torch.Tensor
290 + """(batch_size + 1,), the start location of each request in query Tensor"""
291 +
292 + seq_lens: torch.Tensor
293 + """(batch_size,), the number of computed tokens for each request"""
294 +
295 + num_reqs: int
296 + """Number of requests"""
297 + # TODO(lucas): rename to num_tokens since it may be padded and this is misleading
298 + num_actual_tokens: int
299 + """Total number of tokens in batch"""
300 + max_query_len: int
301 + """Longest query in batch"""
302 + max_seq_len: int
303 + """Longest context length (may be an upper bound)"""
304 +
305 + block_table_tensor: torch.Tensor
306 + slot_mapping: torch.Tensor
307 +
308 + causal: bool = True
309 +
310 + # Needed by FastPrefillAttentionBuilder
311 + logits_indices_padded: torch.Tensor | None = None
312 + num_logits_indices: int | None = None
313 +
314 + # Needed by CrossAttentionBuilder
315 + encoder_seq_lens: torch.Tensor | None = None
316 + encoder_seq_lens_cpu: np.ndarray | None = None
317 +
318 + dcp_local_seq_lens: torch.Tensor | None = None
319 + dcp_local_seq_lens_cpu: torch.Tensor | None = None
320 + """Sequence lengths of the local rank in decode context parallelism world"""
321 +
322 + # WARNING: Deprecated fields. Will be removed in a future release (v0.15.0)
323 + _seq_lens_cpu: torch.Tensor | None = None
324 + _num_computed_tokens_cpu: torch.Tensor | None = None
325 +
326 + _num_computed_tokens_cache: torch.Tensor | None = None
327 +
328 + @property
329 + @deprecated(
330 + """
331 + Prefer using device seq_lens directly to avoid implicit H<>D sync.
332 + If a CPU copy is needed, use `seq_lens.cpu()` instead.
333 + Will be removed in a future release (v0.15.0)
334 + """
335 + )
336 + def seq_lens_cpu(self) -> torch.Tensor:
337 + if self._seq_lens_cpu is None:
338 + self._seq_lens_cpu = self.seq_lens.to("cpu")
339 + return self._seq_lens_cpu
340 +
341 + @property
342 + @deprecated(
343 + """
344 + Prefer using device seq_lens directly to avoid implicit H<>D sync which breaks full
345 + async scheduling. If a CPU copy is needed, it can be derived from
346 + query_start_loc_cpu and seq_lens.
347 + Will be removed in a future release (v0.15.0)
348 + """
349 + )
350 + def num_computed_tokens_cpu(self) -> torch.Tensor:
351 + if self._num_computed_tokens_cpu is None:
352 + query_seq_lens = (
353 + self.query_start_loc_cpu[1:] - self.query_start_loc_cpu[:-1]
354 + )
355 + self._num_computed_tokens_cpu = self.seq_lens_cpu - query_seq_lens
356 + return self._num_computed_tokens_cpu
357 +
358 + def compute_num_computed_tokens(self) -> torch.Tensor:
359 + """Compute num_computed_tokens on device (seq_lens - query_lens)."""
360 + if self._num_computed_tokens_cache is None:
361 + query_lens = self.query_start_loc[1:] - self.query_start_loc[:-1]
362 + self._num_computed_tokens_cache = self.seq_lens - query_lens
363 + return self._num_computed_tokens_cache
364 +
365 + # TODO(lucas): remove once we have FULL-CG spec-decode support
366 + def unpadded(
367 + self, num_actual_tokens: int, num_actual_reqs: int
368 + ) -> "CommonAttentionMetadata":
369 + maybe_slice_reqs = lambda x: x[:num_actual_reqs] if x is not None else None
370 + return CommonAttentionMetadata(
371 + query_start_loc=self.query_start_loc[: num_actual_reqs + 1],
372 + query_start_loc_cpu=self.query_start_loc_cpu[: num_actual_reqs + 1],
373 + seq_lens=self.seq_lens[:num_actual_reqs],
374 + _seq_lens_cpu=self._seq_lens_cpu[:num_actual_reqs]
375 + if self._seq_lens_cpu is not None
376 + else None,
377 + _num_computed_tokens_cpu=self._num_computed_tokens_cpu[:num_actual_reqs]
378 + if self._num_computed_tokens_cpu is not None
379 + else None,
380 + num_reqs=num_actual_reqs,
381 + num_actual_tokens=num_actual_tokens,
382 + max_query_len=self.max_query_len,
383 + max_seq_len=self.max_seq_len,
384 + block_table_tensor=self.block_table_tensor[:num_actual_reqs],
385 + slot_mapping=self.slot_mapping[:num_actual_tokens],
386 + causal=self.causal,
387 + logits_indices_padded=self.logits_indices_padded,
388 + num_logits_indices=self.num_logits_indices,
389 + encoder_seq_lens=maybe_slice_reqs(self.encoder_seq_lens),
390 + encoder_seq_lens_cpu=maybe_slice_reqs(self.encoder_seq_lens_cpu),
391 + dcp_local_seq_lens=maybe_slice_reqs(self.dcp_local_seq_lens),
392 + dcp_local_seq_lens_cpu=maybe_slice_reqs(self.dcp_local_seq_lens_cpu),
393 + )
394 +
395 +
396 +M = TypeVar("M")
397 +
398 +
399 +class AttentionCGSupport(Enum):
400 + """Constants for the cudagraph support of the attention backend
401 + Here we do not consider the cascade attention, as currently
402 + it is never cudagraph supported."""
403 +
404 + ALWAYS = 3
405 + """Cudagraph always supported; supports mixed-prefill-decode"""
406 + UNIFORM_BATCH = 2
407 + """Cudagraph supported for batches the only contain query lengths that are
408 + the same, this can be used for spec-decode
409 + i.e. "decodes" are 1 + num_speculative_tokens"""
410 + UNIFORM_SINGLE_TOKEN_DECODE = 1
411 + """Cudagraph supported for batches the only contain query_len==1 decodes"""
412 + NEVER = 0
413 + """NO cudagraph support"""
414 +
415 +
416 +class AttentionMetadataBuilder(ABC, Generic[M]):
417 + # Does this backend/builder support CUDA Graphs for attention (default: no).
418 + # Do not access directly. Call get_cudagraph_support() instead.
419 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.NEVER
420 + # Does this backend/builder reorder the batch?
421 + # If not, set this to None. Otherwise set it to the query
422 + # length that will be pulled into the front of the batch.
423 + reorder_batch_threshold: int | None = None
424 + # Does this backend/builder support updating the block table in existing
425 + # metadata
426 + supports_update_block_table: bool = False
427 +
428 + @abstractmethod
429 + def __init__(
430 + self,
431 + kv_cache_spec: "AttentionSpec",
432 + layer_names: list[str],
433 + vllm_config: "VllmConfig",
434 + device: torch.device,
435 + ):
436 + self.kv_cache_spec = kv_cache_spec
437 + self.layer_names = layer_names
438 + self.vllm_config = vllm_config
439 + self.device = device
440 +
441 + @classmethod
442 + def get_cudagraph_support(
443 + cls: type["AttentionMetadataBuilder"],
444 + vllm_config: "VllmConfig",
445 + kv_cache_spec: "AttentionSpec",
446 + ) -> AttentionCGSupport:
447 + """Get the cudagraph support level of this builder class."""
448 + return cls._cudagraph_support
449 +
450 + def _init_reorder_batch_threshold(
451 + self,
452 + reorder_batch_threshold: int | None = 1,
453 + supports_spec_as_decode: bool = False,
454 + supports_dcp_with_varlen: bool = False,
455 + ) -> None:
456 + self.reorder_batch_threshold = reorder_batch_threshold
457 + if self.reorder_batch_threshold is not None and supports_spec_as_decode:
458 + # If the backend supports spec-as-decode kernels, then we can set
459 + # the reorder_batch_threshold based on the number of speculative
460 + # tokens from the config.
461 + speculative_config = self.vllm_config.speculative_config
462 + if (
463 + speculative_config is not None
464 + and speculative_config.num_speculative_tokens is not None
465 + ):
466 + self.reorder_batch_threshold = max(
467 + self.reorder_batch_threshold,
468 + 1 + speculative_config.num_speculative_tokens,
469 + )
470 +
471 + if (
472 + self.vllm_config.parallel_config.decode_context_parallel_size > 1
473 + and not supports_dcp_with_varlen
474 + ):
475 + self.reorder_batch_threshold = 1
476 +
477 + @abstractmethod
478 + def build(
479 + self,
480 + common_prefix_len: int,
481 + common_attn_metadata: CommonAttentionMetadata,
482 + fast_build: bool = False,
483 + ) -> M:
484 + """
485 + Central method that builds attention metadata.
486 + Some builders (MLA) require reorder_batch to be called prior to build.
487 +
488 + Args:
489 + common_prefix_len: The length of the common prefix of the batch.
490 + common_attn_metadata: The common attention metadata.
491 + fast_build: The meta-data will prioritize speed of building over
492 + then speed at execution. Can be used for spec-decode where the
493 + result of a build call may only be used for few layers/iters.
494 + """
495 + raise NotImplementedError
496 +
497 + def update_block_table(
498 + self,
499 + metadata: M,
500 + blk_table: torch.Tensor,
501 + slot_mapping: torch.Tensor,
502 + ) -> M:
503 + """
504 + Update the block table for the attention metadata.
505 + Faster when theres multiple kv-cache groups that create virtually the
506 + same metadata but just with different block tables.
507 +
508 + Only needs to be implemented if supports_update_block_table is True.
509 + """
510 + raise NotImplementedError
511 +
512 + def build_for_cudagraph_capture(
513 + self, common_attn_metadata: CommonAttentionMetadata
514 + ) -> M:
515 + """
516 + Build attention metadata for CUDA graph capture. Uses build by default.
517 + Subclasses that override this method should call self.build or
518 + super().build_for_cudagraph_capture.
519 + """
520 + return self.build(
521 + common_prefix_len=0, common_attn_metadata=common_attn_metadata
522 + )
523 +
524 + def build_for_drafting(
525 + self,
526 + common_attn_metadata: CommonAttentionMetadata,
527 + draft_index: int,
528 + ) -> M:
529 + """
530 + Build attention metadata for draft model. Uses build by default.
531 +
532 + Args:
533 + common_attn_metadata: The common attention metadata.
534 + draft_index: The index of the current draft operation.
535 + When speculating a chain of tokens, this index refers to the
536 + draft attempt for the i-th token.
537 + For tree-based attention, this index instead refers to the
538 + draft attempt for the i-th level in the tree of tokens.
539 + """
540 + return self.build(
541 + common_prefix_len=0,
542 + common_attn_metadata=common_attn_metadata,
543 + fast_build=True,
544 + )
545 +
546 + def use_cascade_attention(
547 + self,
548 + common_prefix_len: int,
549 + query_lens: np.ndarray,
550 + num_query_heads: int,
551 + num_kv_heads: int,
552 + use_alibi: bool,
553 + use_sliding_window: bool,
554 + use_local_attention: bool,
555 + num_sms: int,
556 + dcp_world_size: int,
557 + ) -> bool:
558 + return False
559 +
560 +
561 class AttentionLayer(Protocol):
562 _q_scale: torch.Tensor
563 _k_scale: torch.Tensor
vllm/v1/attention/backends/cpu_attn.py
+2 -2
@@ -13,12 +13,12 @@ from vllm.v1.attention.backend import (
13 AttentionBackend,
14 AttentionImpl,
15 AttentionLayer,
16 + AttentionMetadataBuilder,
17 AttentionType,
18 + CommonAttentionMetadata,
19 is_quantized_kv_cache,
20 )
21 from vllm.v1.attention.backends.utils import (
20 - AttentionMetadataBuilder,
21 - CommonAttentionMetadata,
22 split_decodes_and_prefills,
23 )
24 from vllm.v1.kv_cache_interface import AttentionSpec, CrossAttentionSpec
vllm/v1/attention/backends/flash_attn.py
+3 -1
@@ -41,10 +41,12 @@ from vllm.model_executor.layers.batch_invariant import (
41 )
42 from vllm.platforms.interface import DeviceCapability
43 from vllm.utils.math_utils import cdiv
44 -from vllm.v1.attention.backends.utils import (
44 +from vllm.v1.attention.backend import (
45 AttentionCGSupport,
46 AttentionMetadataBuilder,
47 CommonAttentionMetadata,
48 +)
49 +from vllm.v1.attention.backends.utils import (
50 get_dcp_local_seq_lens,
51 get_kv_cache_layout,
52 )
vllm/v1/attention/backends/flashinfer.py
+3 -3
@@ -43,14 +43,14 @@ from vllm.utils.platform_utils import is_pin_memory_available
43 from vllm.utils.torch_utils import is_strictly_contiguous
44 from vllm.v1.attention.backend import (
45 AttentionBackend,
46 + AttentionCGSupport,
47 AttentionImpl,
48 + AttentionMetadataBuilder,
49 AttentionType,
50 + CommonAttentionMetadata,
51 MultipleOf,
52 )
53 from vllm.v1.attention.backends.utils import (
51 - AttentionCGSupport,
52 - AttentionMetadataBuilder,
53 - CommonAttentionMetadata,
54 KVCacheLayoutType,
55 get_dcp_local_seq_lens,
56 get_kv_cache_layout,
vllm/v1/attention/backends/flex_attention.py
+2 -4
@@ -32,12 +32,10 @@ from vllm.utils.torch_utils import is_torch_equal_or_newer
32 from vllm.v1.attention.backend import (
33 AttentionBackend,
34 AttentionImpl,
35 - AttentionType,
36 - is_quantized_kv_cache,
37 -)
38 -from vllm.v1.attention.backends.utils import (
35 AttentionMetadataBuilder,
36 + AttentionType,
37 CommonAttentionMetadata,
38 + is_quantized_kv_cache,
39 )
40 from vllm.v1.kv_cache_interface import AttentionSpec
41
vllm/v1/attention/backends/gdn_attn.py
+5 -3
@@ -7,12 +7,14 @@ from dataclasses import dataclass
7 import torch
8
9 from vllm.config import VllmConfig
10 -from vllm.v1.attention.backend import AttentionBackend
11 -from vllm.v1.attention.backends.utils import (
12 - PAD_SLOT_ID,
10 +from vllm.v1.attention.backend import (
11 + AttentionBackend,
12 AttentionCGSupport,
13 AttentionMetadataBuilder,
14 CommonAttentionMetadata,
15 +)
16 +from vllm.v1.attention.backends.utils import (
17 + PAD_SLOT_ID,
18 compute_causal_conv1d_metadata,
19 split_decodes_and_prefills,
20 )
vllm/v1/attention/backends/linear_attn.py
+3 -3
@@ -5,13 +5,13 @@ from dataclasses import dataclass
5 import torch
6
7 from vllm.config import VllmConfig
8 -from vllm.v1.attention.backend import AttentionBackend
9 -from vllm.v1.attention.backends.utils import (
8 +from vllm.v1.attention.backend import (
9 + AttentionBackend,
10 AttentionCGSupport,
11 AttentionMetadataBuilder,
12 CommonAttentionMetadata,
13 - split_decodes_and_prefills,
13 )
14 +from vllm.v1.attention.backends.utils import split_decodes_and_prefills
15 from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
16
17
vllm/v1/attention/backends/mamba2_attn.py
+1 -4
@@ -7,14 +7,11 @@ import torch
7
8 from vllm.config import VllmConfig
9 from vllm.utils.math_utils import cdiv
10 -from vllm.v1.attention.backend import AttentionBackend
10 +from vllm.v1.attention.backend import AttentionBackend, CommonAttentionMetadata
11 from vllm.v1.attention.backends.mamba_attn import (
12 BaseMambaAttentionMetadata,
13 BaseMambaAttentionMetadataBuilder,
14 )
15 -from vllm.v1.attention.backends.utils import (
16 - CommonAttentionMetadata,
17 -)
15 from vllm.v1.kv_cache_interface import AttentionSpec
16
17
vllm/v1/attention/backends/mamba_attn.py
+4 -2
@@ -10,11 +10,13 @@ import torch
10
11 from vllm.config import VllmConfig
12 from vllm.utils.math_utils import cdiv
13 -from vllm.v1.attention.backends.utils import (
14 - PAD_SLOT_ID,
13 +from vllm.v1.attention.backend import (
14 AttentionCGSupport,
15 AttentionMetadataBuilder,
16 CommonAttentionMetadata,
17 +)
18 +from vllm.v1.attention.backends.utils import (
19 + PAD_SLOT_ID,
20 compute_causal_conv1d_metadata,
21 split_decodes_and_prefills,
22 )
vllm/v1/attention/backends/mla/common.py
+2 -2
@@ -217,12 +217,12 @@ from vllm.v1.attention.backend import (
217 AttentionBackend,
218 AttentionLayer,
219 AttentionMetadata,
220 + AttentionMetadataBuilder,
221 + CommonAttentionMetadata,
222 MLAAttentionImpl,
223 )
224 from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
225 from vllm.v1.attention.backends.utils import (
224 - AttentionMetadataBuilder,
225 - CommonAttentionMetadata,
226 get_dcp_local_seq_lens,
227 get_per_layer_parameters,
228 infer_global_hyperparameters,
vllm/v1/attention/backends/mla/cutlass_mla.py
+1 -1
@@ -11,6 +11,7 @@ from vllm.config.cache import CacheDType
11 from vllm.logger import init_logger
12 from vllm.platforms.interface import DeviceCapability
13 from vllm.v1.attention.backend import (
14 + AttentionCGSupport,
15 AttentionLayer,
16 AttentionType,
17 MultipleOf,
@@ -22,7 +23,6 @@ from vllm.v1.attention.backends.mla.common import (
23 MLACommonMetadata,
24 MLACommonMetadataBuilder,
25 )
25 -from vllm.v1.attention.backends.utils import AttentionCGSupport
26
27 logger = init_logger(__name__)
28
vllm/v1/attention/backends/mla/flashattn_mla.py
+1 -1
@@ -14,6 +14,7 @@ from vllm.model_executor.layers.batch_invariant import (
14 )
15 from vllm.platforms.interface import DeviceCapability
16 from vllm.v1.attention.backend import (
17 + AttentionCGSupport,
18 AttentionLayer,
19 AttentionType,
20 MultipleOf,
@@ -31,7 +32,6 @@ from vllm.v1.attention.backends.mla.common import (
32 MLACommonMetadataBuilder,
33 QueryLenSupport,
34 )
34 -from vllm.v1.attention.backends.utils import AttentionCGSupport
35 from vllm.v1.kv_cache_interface import AttentionSpec
36 from vllm.vllm_flash_attn import ( # type: ignore[attr-defined]
37 flash_attn_varlen_func,
vllm/v1/attention/backends/mla/flashinfer_mla.py
+2 -1
@@ -10,6 +10,7 @@ from vllm.config.cache import CacheDType
10 from vllm.logger import init_logger
11 from vllm.platforms.interface import DeviceCapability
12 from vllm.v1.attention.backend import (
13 + AttentionCGSupport,
14 AttentionLayer,
15 AttentionType,
16 MultipleOf,
@@ -21,7 +22,7 @@ from vllm.v1.attention.backends.mla.common import (
22 MLACommonMetadataBuilder,
23 QueryLenSupport,
24 )
24 -from vllm.v1.attention.backends.utils import AttentionCGSupport, KVCacheLayoutType
25 +from vllm.v1.attention.backends.utils import KVCacheLayoutType
26
27 logger = init_logger(__name__)
28
vllm/v1/attention/backends/mla/flashmla.py
+6 -2
@@ -13,7 +13,12 @@ from vllm.model_executor.layers.batch_invariant import (
13 vllm_is_batch_invariant,
14 )
15 from vllm.platforms.interface import DeviceCapability
16 -from vllm.v1.attention.backend import AttentionLayer, AttentionType, MultipleOf
16 +from vllm.v1.attention.backend import (
17 + AttentionCGSupport,
18 + AttentionLayer,
19 + AttentionType,
20 + MultipleOf,
21 +)
22 from vllm.v1.attention.backends.mla.common import (
23 MLACommonBackend,
24 MLACommonDecodeMetadata,
@@ -23,7 +28,6 @@ from vllm.v1.attention.backends.mla.common import (
28 QueryLenSupport,
29 )
30 from vllm.v1.attention.backends.utils import (
26 - AttentionCGSupport,
31 reshape_attn_output_for_spec_decode,
32 reshape_query_for_spec_decode,
33 )
vllm/v1/attention/backends/mla/flashmla_sparse.py
+3 -3
@@ -16,15 +16,15 @@ from vllm.triton_utils import tl, triton
16 from vllm.utils.math_utils import cdiv
17 from vllm.v1.attention.backend import (
18 AttentionBackend,
19 + AttentionCGSupport,
20 AttentionLayer,
21 AttentionMetadata,
22 + AttentionMetadataBuilder,
23 + CommonAttentionMetadata,
24 MultipleOf,
25 )
26 from vllm.v1.attention.backends.mla.common import MLACommonBaseImpl, get_mla_dims
27 from vllm.v1.attention.backends.utils import (
25 - AttentionCGSupport,
26 - AttentionMetadataBuilder,
27 - CommonAttentionMetadata,
28 reshape_attn_output_for_spec_decode,
29 reshape_query_for_spec_decode,
30 split_decodes_and_prefills,
vllm/v1/attention/backends/mla/indexer.py
+3 -3
@@ -11,12 +11,12 @@ from vllm.platforms import current_platform
11 from vllm.utils.deep_gemm import get_paged_mqa_logits_metadata, is_deep_gemm_supported
12 from vllm.v1.attention.backend import (
13 AttentionBackend,
14 - MultipleOf,
15 -)
16 -from vllm.v1.attention.backends.utils import (
14 AttentionCGSupport,
15 AttentionMetadataBuilder,
16 CommonAttentionMetadata,
17 + MultipleOf,
18 +)
19 +from vllm.v1.attention.backends.utils import (
20 split_decodes_and_prefills,
21 split_prefill_chunks,
22 )
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+1 -2
@@ -8,7 +8,7 @@ import torch
8
9 from vllm._aiter_ops import rocm_aiter_ops
10 from vllm.config import VllmConfig
11 -from vllm.v1.attention.backend import AttentionLayer, MultipleOf
11 +from vllm.v1.attention.backend import AttentionCGSupport, AttentionLayer, MultipleOf
12 from vllm.v1.attention.backends.mla.common import (
13 MLACommonBackend,
14 MLACommonDecodeMetadata,
@@ -17,7 +17,6 @@ from vllm.v1.attention.backends.mla.common import (
17 MLACommonMetadataBuilder,
18 QueryLenSupport,
19 )
20 -from vllm.v1.attention.backends.utils import AttentionCGSupport
20 from vllm.v1.kv_cache_interface import AttentionSpec
21
22
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+3 -5
@@ -13,18 +13,16 @@ from vllm.config import VllmConfig
13 from vllm.logger import init_logger
14 from vllm.v1.attention.backend import (
15 AttentionBackend,
16 + AttentionCGSupport,
17 AttentionLayer,
18 AttentionMetadata,
19 + AttentionMetadataBuilder,
20 + CommonAttentionMetadata,
21 )
22 from vllm.v1.attention.backends.mla.common import MLACommonBaseImpl, get_mla_dims
23 from vllm.v1.attention.backends.mla.flashmla_sparse import (
24 triton_convert_req_index_to_global_index,
25 )
23 -from vllm.v1.attention.backends.utils import (
24 - AttentionCGSupport,
25 - AttentionMetadataBuilder,
26 - CommonAttentionMetadata,
27 -)
26 from vllm.v1.kv_cache_interface import AttentionSpec
27
28 if TYPE_CHECKING:
vllm/v1/attention/backends/rocm_aiter_fa.py
+3 -3
@@ -15,14 +15,14 @@ from vllm.utils.math_utils import cdiv
15 from vllm.utils.platform_utils import get_cu_count
16 from vllm.v1.attention.backend import (
17 AttentionBackend,
18 + AttentionCGSupport,
19 AttentionImpl,
20 + AttentionMetadataBuilder,
21 AttentionType,
22 + CommonAttentionMetadata,
23 MultipleOf,
24 )
25 from vllm.v1.attention.backends.utils import (
23 - AttentionCGSupport,
24 - AttentionMetadataBuilder,
25 - CommonAttentionMetadata,
26 split_decodes_prefills_and_extends,
27 )
28 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
vllm/v1/attention/backends/rocm_attn.py
+3 -5
@@ -16,16 +16,14 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
16 from vllm.platforms import current_platform
17 from vllm.v1.attention.backend import (
18 AttentionBackend,
19 + AttentionCGSupport,
20 AttentionImpl,
21 + AttentionMetadataBuilder,
22 AttentionType,
23 + CommonAttentionMetadata,
24 MultipleOf,
25 )
26 from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
24 -from vllm.v1.attention.backends.utils import (
25 - AttentionCGSupport,
26 - AttentionMetadataBuilder,
27 - CommonAttentionMetadata,
28 -)
27 from vllm.v1.attention.ops.chunked_prefill_paged_decode import (
28 chunked_prefill_paged_decode,
29 )
vllm/v1/attention/backends/tree_attn.py
+2 -2
@@ -14,12 +14,12 @@ from vllm.logger import init_logger
14 from vllm.v1.attention.backend import (
15 AttentionBackend,
16 AttentionImpl,
17 + AttentionMetadataBuilder,
18 AttentionType,
19 + CommonAttentionMetadata,
20 MultipleOf,
21 )
22 from vllm.v1.attention.backends.utils import (
21 - AttentionMetadataBuilder,
22 - CommonAttentionMetadata,
23 split_decodes_and_prefills,
24 )
25 from vllm.v1.attention.ops.triton_unified_attention import unified_attention
vllm/v1/attention/backends/triton_attn.py
+3 -5
@@ -19,14 +19,12 @@ from vllm.platforms.interface import DeviceCapability
19 from vllm.utils.math_utils import next_power_of_2
20 from vllm.v1.attention.backend import (
21 AttentionBackend,
22 - AttentionImpl,
23 - AttentionType,
24 - MultipleOf,
25 -)
26 -from vllm.v1.attention.backends.utils import (
22 AttentionCGSupport,
23 + AttentionImpl,
24 AttentionMetadataBuilder,
25 + AttentionType,
26 CommonAttentionMetadata,
27 + MultipleOf,
28 )
29 from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd
30 from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
vllm/v1/attention/backends/utils.py
+6 -289
@@ -1,16 +1,11 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import abc
4 -import enum
3 import functools
6 -from abc import abstractmethod
4 from collections.abc import Callable
5 from dataclasses import dataclass, field, fields, make_dataclass
6 from typing import (
7 TYPE_CHECKING,
8 Any,
12 - ClassVar,
13 - Generic,
9 Literal,
10 Protocol,
11 TypeVar,
@@ -19,7 +14,7 @@ from typing import (
14
15 import numpy as np
16 import torch
22 -from typing_extensions import deprecated, runtime_checkable
17 +from typing_extensions import runtime_checkable
18
19 from vllm.config import VllmConfig, get_layers_from_vllm_config
20 from vllm.utils.math_utils import cdiv
@@ -38,8 +33,9 @@ from vllm.v1.attention.backend import (
33 AttentionBackend,
34 AttentionImpl,
35 AttentionMetadata,
36 + AttentionMetadataBuilder,
37 + CommonAttentionMetadata,
38 )
42 -from vllm.v1.kv_cache_interface import AttentionSpec
39 from vllm.v1.worker.ubatch_utils import UBatchSlice
40
41 logger = init_logger(__name__)
@@ -53,123 +49,6 @@ def is_valid_kv_cache_layout(value: str) -> bool:
49 return value in get_args(KVCacheLayoutType)
50
51
56 -@dataclass
57 -class CommonAttentionMetadata:
58 - """
59 - Per-batch attention metadata, shared across layers and backends.
60 - AttentionMetadataBuilder instances use it to construct per-layer metadata.
61 -
62 - For many of the tensors we keep both GPU and CPU versions.
63 - """
64 -
65 - query_start_loc: torch.Tensor
66 - query_start_loc_cpu: torch.Tensor
67 - """(batch_size + 1,), the start location of each request in query Tensor"""
68 -
69 - seq_lens: torch.Tensor
70 - """(batch_size,), the number of computed tokens for each request"""
71 -
72 - num_reqs: int
73 - """Number of requests"""
74 - # TODO(lucas): rename to num_tokens since it may be padded and this is misleading
75 - num_actual_tokens: int
76 - """Total number of tokens in batch"""
77 - max_query_len: int
78 - """Longest query in batch"""
79 - max_seq_len: int
80 - """Longest context length (may be an upper bound)"""
81 -
82 - block_table_tensor: torch.Tensor
83 - slot_mapping: torch.Tensor
84 -
85 - causal: bool = True
86 -
87 - # Needed by FastPrefillAttentionBuilder
88 - logits_indices_padded: torch.Tensor | None = None
89 - num_logits_indices: int | None = None
90 -
91 - # Needed by CrossAttentionBuilder
92 - encoder_seq_lens: torch.Tensor | None = None
93 - encoder_seq_lens_cpu: np.ndarray | None = None
94 -
95 - dcp_local_seq_lens: torch.Tensor | None = None
96 - dcp_local_seq_lens_cpu: torch.Tensor | None = None
97 - """Sequence lengths of the local rank in decode context parallelism world"""
98 -
99 - # WARNING: Deprecated fields. Will be removed in a future release (v0.15.0)
100 - _seq_lens_cpu: torch.Tensor | None = None
101 - _num_computed_tokens_cpu: torch.Tensor | None = None
102 -
103 - _num_computed_tokens_cache: torch.Tensor | None = None
104 -
105 - @property
106 - @deprecated(
107 - """
108 - Prefer using device seq_lens directly to avoid implicit H<>D sync.
109 - If a CPU copy is needed, use `seq_lens.cpu()` instead.
110 - Will be removed in a future release (v0.15.0)
111 - """
112 - )
113 - def seq_lens_cpu(self) -> torch.Tensor:
114 - if self._seq_lens_cpu is None:
115 - self._seq_lens_cpu = self.seq_lens.to("cpu")
116 - return self._seq_lens_cpu
117 -
118 - @property
119 - @deprecated(
120 - """
121 - Prefer using device seq_lens directly to avoid implicit H<>D sync which breaks full
122 - async scheduling. If a CPU copy is needed, it can be derived from
123 - query_start_loc_cpu and seq_lens.
124 - Will be removed in a future release (v0.15.0)
125 - """
126 - )
127 - def num_computed_tokens_cpu(self) -> torch.Tensor:
128 - if self._num_computed_tokens_cpu is None:
129 - query_seq_lens = (
130 - self.query_start_loc_cpu[1:] - self.query_start_loc_cpu[:-1]
131 - )
132 - self._num_computed_tokens_cpu = self.seq_lens_cpu - query_seq_lens
133 - return self._num_computed_tokens_cpu
134 -
135 - def compute_num_computed_tokens(self) -> torch.Tensor:
136 - """Compute num_computed_tokens on device (seq_lens - query_lens)."""
137 - if self._num_computed_tokens_cache is None:
138 - query_lens = self.query_start_loc[1:] - self.query_start_loc[:-1]
139 - self._num_computed_tokens_cache = self.seq_lens - query_lens
140 - return self._num_computed_tokens_cache
141 -
142 - # TODO(lucas): remove once we have FULL-CG spec-decode support
143 - def unpadded(
144 - self, num_actual_tokens: int, num_actual_reqs: int
145 - ) -> "CommonAttentionMetadata":
146 - maybe_slice_reqs = lambda x: x[:num_actual_reqs] if x is not None else None
147 - return CommonAttentionMetadata(
148 - query_start_loc=self.query_start_loc[: num_actual_reqs + 1],
149 - query_start_loc_cpu=self.query_start_loc_cpu[: num_actual_reqs + 1],
150 - seq_lens=self.seq_lens[:num_actual_reqs],
151 - _seq_lens_cpu=self._seq_lens_cpu[:num_actual_reqs]
152 - if self._seq_lens_cpu is not None
153 - else None,
154 - _num_computed_tokens_cpu=self._num_computed_tokens_cpu[:num_actual_reqs]
155 - if self._num_computed_tokens_cpu is not None
156 - else None,
157 - num_reqs=num_actual_reqs,
158 - num_actual_tokens=num_actual_tokens,
159 - max_query_len=self.max_query_len,
160 - max_seq_len=self.max_seq_len,
161 - block_table_tensor=self.block_table_tensor[:num_actual_reqs],
162 - slot_mapping=self.slot_mapping[:num_actual_tokens],
163 - causal=self.causal,
164 - logits_indices_padded=self.logits_indices_padded,
165 - num_logits_indices=self.num_logits_indices,
166 - encoder_seq_lens=maybe_slice_reqs(self.encoder_seq_lens),
167 - encoder_seq_lens_cpu=maybe_slice_reqs(self.encoder_seq_lens_cpu),
168 - dcp_local_seq_lens=maybe_slice_reqs(self.dcp_local_seq_lens),
169 - dcp_local_seq_lens_cpu=maybe_slice_reqs(self.dcp_local_seq_lens_cpu),
170 - )
171 -
172 -
52 def slice_query_start_locs(
53 query_start_loc: torch.Tensor,
54 request_slice: slice,
@@ -299,171 +178,6 @@ def split_attn_metadata(
178 return results
179
180
302 -M = TypeVar("M")
303 -
304 -
305 -class AttentionCGSupport(enum.Enum):
306 - """Constants for the cudagraph support of the attention backend
307 - Here we do not consider the cascade attention, as currently
308 - it is never cudagraph supported."""
309 -
310 - ALWAYS = 3
311 - """Cudagraph always supported; supports mixed-prefill-decode"""
312 - UNIFORM_BATCH = 2
313 - """Cudagraph supported for batches the only contain query lengths that are
314 - the same, this can be used for spec-decode
315 - i.e. "decodes" are 1 + num_speculative_tokens"""
316 - UNIFORM_SINGLE_TOKEN_DECODE = 1
317 - """Cudagraph supported for batches the only contain query_len==1 decodes"""
318 - NEVER = 0
319 - """NO cudagraph support"""
320 -
321 -
322 -class AttentionMetadataBuilder(abc.ABC, Generic[M]):
323 - # Does this backend/builder support CUDA Graphs for attention (default: no).
324 - # Do not access directly. Call get_cudagraph_support() instead.
325 - _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.NEVER
326 - # Does this backend/builder reorder the batch?
327 - # If not, set this to None. Otherwise set it to the query
328 - # length that will be pulled into the front of the batch.
329 - reorder_batch_threshold: int | None = None
330 - # Does this backend/builder support updating the block table in existing
331 - # metadata
332 - supports_update_block_table: bool = False
333 -
334 - @abstractmethod
335 - def __init__(
336 - self,
337 - kv_cache_spec: AttentionSpec,
338 - layer_names: list[str],
339 - vllm_config: VllmConfig,
340 - device: torch.device,
341 - ):
342 - self.kv_cache_spec = kv_cache_spec
343 - self.layer_names = layer_names
344 - self.vllm_config = vllm_config
345 - self.device = device
346 -
347 - @classmethod
348 - def get_cudagraph_support(
349 - cls: type["AttentionMetadataBuilder"],
350 - vllm_config: VllmConfig,
351 - kv_cache_spec: AttentionSpec,
352 - ) -> AttentionCGSupport:
353 - """Get the cudagraph support level of this builder class."""
354 - return cls._cudagraph_support
355 -
356 - def _init_reorder_batch_threshold(
357 - self,
358 - reorder_batch_threshold: int | None = 1,
359 - supports_spec_as_decode: bool = False,
360 - supports_dcp_with_varlen: bool = False,
361 - ) -> None:
362 - self.reorder_batch_threshold = reorder_batch_threshold
363 - if self.reorder_batch_threshold is not None and supports_spec_as_decode:
364 - # If the backend supports spec-as-decode kernels, then we can set
365 - # the reorder_batch_threshold based on the number of speculative
366 - # tokens from the config.
367 - speculative_config = self.vllm_config.speculative_config
368 - if (
369 - speculative_config is not None
370 - and speculative_config.num_speculative_tokens is not None
371 - ):
372 - self.reorder_batch_threshold = max(
373 - self.reorder_batch_threshold,
374 - 1 + speculative_config.num_speculative_tokens,
375 - )
376 -
377 - if (
378 - self.vllm_config.parallel_config.decode_context_parallel_size > 1
379 - and not supports_dcp_with_varlen
380 - ):
381 - self.reorder_batch_threshold = 1
382 -
383 - @abstractmethod
384 - def build(
385 - self,
386 - common_prefix_len: int,
387 - common_attn_metadata: CommonAttentionMetadata,
388 - fast_build: bool = False,
389 - ) -> M:
390 - """
391 - Central method that builds attention metadata.
392 - Some builders (MLA) require reorder_batch to be called prior to build.
393 -
394 - Args:
395 - common_prefix_len: The length of the common prefix of the batch.
396 - common_attn_metadata: The common attention metadata.
397 - fast_build: The meta-data will prioritize speed of building over
398 - then speed at execution. Can be used for spec-decode where the
399 - result of a build call may only be used for few layers/iters.
400 - """
401 - raise NotImplementedError
402 -
403 - def update_block_table(
404 - self,
405 - metadata: M,
406 - blk_table: torch.Tensor,
407 - slot_mapping: torch.Tensor,
408 - ) -> M:
409 - """
410 - Update the block table for the attention metadata.
411 - Faster when theres multiple kv-cache groups that create virtually the
412 - same metadata but just with different block tables.
413 -
414 - Only needs to be implemented if supports_update_block_table is True.
415 - """
416 - raise NotImplementedError
417 -
418 - def build_for_cudagraph_capture(
419 - self, common_attn_metadata: CommonAttentionMetadata
420 - ) -> M:
421 - """
422 - Build attention metadata for CUDA graph capture. Uses build by default.
423 - Subclasses that override this method should call self.build or
424 - super().build_for_cudagraph_capture.
425 - """
426 - return self.build(
427 - common_prefix_len=0, common_attn_metadata=common_attn_metadata
428 - )
429 -
430 - def build_for_drafting(
431 - self,
432 - common_attn_metadata: CommonAttentionMetadata,
433 - draft_index: int,
434 - ) -> M:
435 - """
436 - Build attention metadata for draft model. Uses build by default.
437 -
438 - Args:
439 - common_attn_metadata: The common attention metadata.
440 - draft_index: The index of the current draft operation.
441 - When speculating a chain of tokens, this index refers to the
442 - draft attempt for the i-th token.
443 - For tree-based attention, this index instead refers to the
444 - draft attempt for the i-th level in the tree of tokens.
445 - """
446 - return self.build(
447 - common_prefix_len=0,
448 - common_attn_metadata=common_attn_metadata,
449 - fast_build=True,
450 - )
451 -
452 - def use_cascade_attention(
453 - self,
454 - common_prefix_len: int,
455 - query_lens: np.ndarray,
456 - num_query_heads: int,
457 - num_kv_heads: int,
458 - use_alibi: bool,
459 - use_sliding_window: bool,
460 - use_local_attention: bool,
461 - num_sms: int,
462 - dcp_world_size: int,
463 - ) -> bool:
464 - return False
465 -
466 -
181 @functools.lru_cache
182 def get_kv_cache_layout():
183 # Format specified by the code.
@@ -834,6 +548,9 @@ def make_kv_sharing_fast_prefill_common_attn_metadata(
548 return common_attn_metadata
549
550
551 +M = TypeVar("M")
552 +
553 +
554 def subclass_attention_backend(
555 name_prefix: str,
556 attention_backend_cls: type[AttentionBackend],
vllm/v1/spec_decode/eagle.py
+4 -4
@@ -26,16 +26,16 @@ from vllm.multimodal import MULTIMODAL_REGISTRY
26 from vllm.platforms import current_platform
27 from vllm.triton_utils import triton
28 from vllm.utils.platform_utils import is_pin_memory_available
29 +from vllm.v1.attention.backend import (
30 + AttentionMetadataBuilder,
31 + CommonAttentionMetadata,
32 +)
33 from vllm.v1.attention.backends.registry import AttentionBackendEnum
34 from vllm.v1.attention.backends.tree_attn import (
35 TreeAttentionMetadata,
36 TreeAttentionMetadataBuilder,
37 )
38 from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata
35 -from vllm.v1.attention.backends.utils import (
36 - AttentionMetadataBuilder,
37 - CommonAttentionMetadata,
38 -)
39 from vllm.v1.kv_cache_interface import KVCacheConfig
40 from vllm.v1.sample.metadata import SamplingMetadata
41 from vllm.v1.sample.sampler import _SAMPLING_EPS
vllm/v1/worker/gpu/attn_utils.py
+2 -2
@@ -7,8 +7,8 @@ import torch
7
8 from vllm.config import VllmConfig, get_layers_from_vllm_config
9 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
10 -from vllm.v1.attention.backend import AttentionBackend
11 -from vllm.v1.attention.backends.utils import (
10 +from vllm.v1.attention.backend import (
11 + AttentionBackend,
12 AttentionMetadataBuilder,
13 CommonAttentionMetadata,
14 )
vllm/v1/worker/gpu_model_runner.py
+3 -3
@@ -100,15 +100,15 @@ from vllm.utils.torch_utils import (
100 )
101 from vllm.v1.attention.backend import (
102 AttentionBackend,
103 + AttentionCGSupport,
104 AttentionMetadata,
105 + AttentionMetadataBuilder,
106 AttentionType,
107 + CommonAttentionMetadata,
108 MultipleOf,
109 )
110 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder
111 from vllm.v1.attention.backends.utils import (
109 - AttentionCGSupport,
110 - AttentionMetadataBuilder,
111 - CommonAttentionMetadata,
112 create_fast_prefill_custom_backend,
113 get_dcp_local_seq_lens,
114 reorder_batch_to_split_decodes_and_prefills,