@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0a2c2dc3f1

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, 15:15 UTC

fixed mypy warnings for files vllm/v1/attention with TEMPORARY workaround (#31465)

Signed-off-by: Zhuohao Yang <zy242@cornell.edu> Co-authored-by: Zhuohao Yang <zy242@cornell.edu> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>

Jack Yang committed Jan 6, 2026 at 23:08 UTC 0a2c2dc3f14620de291e67d6d26fca72b09f5617

18 files changed +140 -56

tools/pre_commit/mypy.py

+1 -1

@@ -41,6 +41,7 @@ FILES = [

41 "vllm/usage",

42 "vllm/utils",

43 "vllm/worker",

44 + "vllm/v1/attention",

45 "vllm/v1/core",

46 "vllm/v1/engine",

47 "vllm/v1/executor",

@@ -60,7 +61,6 @@ SEPARATE_GROUPS = [

61 "vllm/lora",

62 "vllm/model_executor",

63 # v1 related

63 - "vllm/v1/attention",

64 "vllm/v1/kv_offload",

65 "vllm/v1/spec_decode",

66 "vllm/v1/structured_output",

vllm/attention/backends/abstract.py

+9 -3

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 from abc import ABC, abstractmethod

5 +from enum import Enum

6 from typing import TYPE_CHECKING, ClassVar, Generic, Protocol, TypeVar, get_args

8 import torch

@@ -14,7 +15,7 @@ if TYPE_CHECKING:

15 from vllm.v1.attention.backends.utils import KVCacheLayoutType

17 -class AttentionType:

18 +class AttentionType(str, Enum):

19 """

20 Attention type.

21 Use string to be compatible with `torch.compile`.

@@ -193,7 +194,7 @@ class AttentionBackend(ABC):

194 head_size: int,

195 dtype: torch.dtype,

196 kv_cache_dtype: "CacheDType | None",

196 - block_size: int | None,

197 + block_size: int,

198 use_mla: bool,

199 has_sink: bool,

200 use_sparse: bool,

@@ -207,7 +208,7 @@ class AttentionBackend(ABC):

208 head_size: int,

209 dtype: torch.dtype,

210 kv_cache_dtype: "CacheDType | None",

210 - block_size: int | None,

211 + block_size: int,

212 use_mla: bool,

213 has_sink: bool,

214 use_sparse: bool,

@@ -290,6 +291,11 @@ class AttentionLayer(Protocol):

291

292

293 class AttentionImpl(ABC, Generic[T]):

294 + # Required attributes that all impls should have

295 + num_heads: int

296 + head_size: int

297 + scale: float

298 +

299 # Whether the attention impl can return the softmax lse for decode.

300 # Some features like decode context parallelism require the softmax lse.

301 can_return_lse_for_decode: bool = False

vllm/model_executor/layers/attention_layer_base.py

+3 -1

@@ -4,7 +4,7 @@

5 from abc import ABC, abstractmethod

7 -from vllm.attention.backends.abstract import AttentionBackend

7 +from vllm.attention.backends.abstract import AttentionBackend, AttentionImpl

8 from vllm.config import VllmConfig

9 from vllm.v1.kv_cache_interface import KVCacheSpec

@@ -18,6 +18,8 @@ class AttentionLayerBase(ABC):

18 from different layer types.

19 """

21 + impl: "AttentionImpl"

22 +

23 @abstractmethod

24 def get_attn_backend(self) -> type[AttentionBackend]:

25 """Get the attention backend class for this layer."""

vllm/v1/attention/backends/flash_attn.py

+34 -8

@@ -167,7 +167,7 @@ class FlashAttentionBackend(AttentionBackend):

167 head_size: int,

168 dtype: torch.dtype,

169 kv_cache_dtype: CacheDType | None,

170 - block_size: int,

170 + block_size: int | None,

171 use_mla: bool,

172 has_sink: bool,

173 use_sparse: bool,

@@ -354,7 +354,11 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

354 aot_schedule = False

355

356 max_num_splits = 0 # 0 means use FA3's heuristics, not CG compatible

357 - if self.use_full_cuda_graph and num_actual_tokens <= self.max_cudagraph_size:

357 + if (

358 + self.use_full_cuda_graph

359 + and self.max_cudagraph_size is not None

360 + and num_actual_tokens <= self.max_cudagraph_size

361 + ):

362 # NOTE(woosuk): Setting num_splits > 1 may increase the memory

363 # usage, because the intermediate buffers of size [num_splits,

364 # num_heads, num_tokens, head_size] are allocated. Therefore,

@@ -599,6 +603,9 @@ class FlashAttentionImpl(AttentionImpl):

603 We use torch's .expand() to avoid duplicating values

604 """

605 assert output is not None, "Output tensor must be provided."

606 + assert self.vllm_flash_attn_version is not None, (

607 + "FlashAttention version not detected."

608 + )

609

610 if output_scale is not None or output_block_scale is not None:

611 raise NotImplementedError(

@@ -697,6 +704,11 @@ class FlashAttentionImpl(AttentionImpl):

704 )

705 return output

706 else:

707 + sliding_window_size = (

708 + list(self.sliding_window)

709 + if self.sliding_window is not None

710 + else None

711 + )

712 flash_attn_varlen_func(

713 q=query[:num_actual_tokens],

714 k=key_cache,

@@ -709,7 +721,7 @@ class FlashAttentionImpl(AttentionImpl):

721 softmax_scale=self.scale,

722 causal=attn_metadata.causal,

723 alibi_slopes=self.alibi_slopes,

712 - window_size=self.sliding_window,

724 + window_size=sliding_window_size,

725 block_table=block_table,

726 softcap=self.logits_soft_cap,

727 scheduler_metadata=scheduler_metadata,

@@ -764,12 +776,19 @@ class FlashAttentionImpl(AttentionImpl):

776 k_descale: torch.Tensor | None = None,

777 v_descale: torch.Tensor | None = None,

778 ) -> torch.Tensor:

779 + assert self.vllm_flash_attn_version is not None, (

780 + "FlashAttention version not detected."

781 + )

782 +

783 cu_seqlens_q = attn_metadata.query_start_loc

784 max_seqlen_q = attn_metadata.max_query_len

785 block_table = attn_metadata.block_table

786

787 query = query.contiguous()

788 query_across_dcp = get_dcp_group().all_gather(query, dim=1)

789 + sliding_window_size = (

790 + list(self.sliding_window) if self.sliding_window is not None else None

791 + )

792 context_attn_out, context_lse = flash_attn_varlen_func(

793 q=query_across_dcp,

794 k=key_cache,

@@ -782,7 +801,7 @@ class FlashAttentionImpl(AttentionImpl):

801 softmax_scale=self.scale,

802 causal=False,

803 alibi_slopes=self.alibi_slopes,

785 - window_size=self.sliding_window,

804 + window_size=sliding_window_size,

805 block_table=block_table,

806 softcap=self.logits_soft_cap,

807 return_softmax_lse=True,

@@ -813,7 +832,7 @@ class FlashAttentionImpl(AttentionImpl):

832 softmax_scale=self.scale,

833 causal=attn_metadata.causal,

834 alibi_slopes=self.alibi_slopes,

816 - window_size=self.sliding_window,

835 + window_size=sliding_window_size,

836 softcap=self.logits_soft_cap,

837 return_softmax_lse=True,

838 fa_version=self.vllm_flash_attn_version,

@@ -850,6 +869,10 @@ class FlashAttentionImpl(AttentionImpl):

869 attn_metadata: Encoder attention metadata

870 layer: The attention layer

871 """

872 + assert self.vllm_flash_attn_version is not None, (

873 + "FlashAttention version not detected."

874 + )

875 +

876 # For encoder attention, process FP8 quantization if needed

877 if self.kv_cache_dtype.startswith("fp8"):

878 raise NotImplementedError(

@@ -868,6 +891,9 @@ class FlashAttentionImpl(AttentionImpl):

891 )

892

893 # Call flash attention directly on Q, K, V tensors

894 + sliding_window_size = (

895 + list(self.sliding_window) if self.sliding_window is not None else None

896 + )

897 flash_attn_varlen_func(

898 q=query,

899 k=key,

@@ -880,7 +906,7 @@ class FlashAttentionImpl(AttentionImpl):

906 softmax_scale=self.scale,

907 causal=False, # Encoder attention is bidirectional

908 alibi_slopes=self.alibi_slopes,

883 - window_size=self.sliding_window,

909 + window_size=sliding_window_size,

910 softcap=self.logits_soft_cap,

911 fa_version=self.vllm_flash_attn_version,

912 q_descale=layer._q_scale.expand(descale_shape),

@@ -1020,7 +1046,7 @@ def cascade_attention(

1046 max_seqlen_k=common_prefix_len,

1047 softmax_scale=softmax_scale,

1048 causal=False,

1023 - window_size=sliding_window,

1049 + window_size=list(sliding_window),

1050 block_table=block_table[:1],

1051 softcap=logits_soft_cap,

1052 return_softmax_lse=True,

@@ -1048,7 +1074,7 @@ def cascade_attention(

1074 max_seqlen_k=max_kv_len - common_prefix_len,

1075 softmax_scale=softmax_scale,

1076 causal=True,

1051 - window_size=sliding_window,

1077 + window_size=list(sliding_window),

1078 block_table=block_table[:, num_common_kv_blocks:],

1079 softcap=logits_soft_cap,

1080 return_softmax_lse=True,

vllm/v1/attention/backends/flash_attn_diffkv.py

+9 -1

@@ -113,6 +113,9 @@ class FlashAttentionDiffKVImpl(FlashAttentionImpl):

113 We use torch's .expand() to avoid duplicating values

114 """

115 assert output is not None, "Output tensor must be provided."

116 + assert self.vllm_flash_attn_version is not None, (

117 + "FlashAttention version not detected."

118 + )

119

120 if output_scale is not None or output_block_scale is not None:

121 raise NotImplementedError(

@@ -214,6 +217,11 @@ class FlashAttentionDiffKVImpl(FlashAttentionImpl):

217 )

218 return output

219 else:

220 + sliding_window_size = (

221 + list(self.sliding_window)

222 + if self.sliding_window is not None

223 + else None

224 + )

225 flash_attn_varlen_func(

226 q=query[:num_actual_tokens],

227 k=key_cache,

@@ -226,7 +234,7 @@ class FlashAttentionDiffKVImpl(FlashAttentionImpl):

234 softmax_scale=self.scale,

235 causal=attn_metadata.causal,

236 alibi_slopes=self.alibi_slopes,

229 - window_size=self.sliding_window,

237 + window_size=sliding_window_size,

238 block_table=block_table,

239 softcap=self.logits_soft_cap,

240 scheduler_metadata=scheduler_metadata,

vllm/v1/attention/backends/flashinfer.py

+6 -5

@@ -530,11 +530,12 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

530 self._decode_wrappers_cudagraph: dict[

531 int, BatchDecodeWithPagedKVCacheWrapper

532 ] = {}

533 - self._decode_cudagraph_max_bs = min(

534 - (1 + num_spec_tokens) * max_num_reqs,

535 - self.compilation_config.max_cudagraph_capture_size,

536 - )

537 -

533 + self._decode_cudagraph_max_bs = (1 + num_spec_tokens) * max_num_reqs

534 + if self.compilation_config.max_cudagraph_capture_size is not None:

535 + self._decode_cudagraph_max_bs = min(

536 + self._decode_cudagraph_max_bs,

537 + self.compilation_config.max_cudagraph_capture_size,

538 + )

539 try:

540 self.dcp_world_size = get_dcp_group().world_size

541 self.dcp_rank = get_dcp_group().rank_in_group

vllm/v1/attention/backends/flex_attention.py

+1 -1

@@ -215,7 +215,7 @@ def physical_to_logical_mapping(

215 )

216

217 # Only process valid blocks to avoid garbage values

218 - num_blocks_per_seq = cdiv(seq_lens, block_size)

218 + num_blocks_per_seq: torch.Tensor = cdiv(seq_lens, block_size)

219 mask = (

220 torch.arange(max_num_blocks, device=device)[None, :]

221 < num_blocks_per_seq[:, None]

vllm/v1/attention/backends/gdn_attn.py

+11 -4

@@ -75,8 +75,10 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]

75 self.compilation_config = vllm_config.compilation_config

76 self.speculative_config = vllm_config.speculative_config

77 self.kv_cache_spec = kv_cache_spec

78 +

79 if self.speculative_config:

79 - self.num_spec = self.speculative_config.num_speculative_tokens

80 + assert self.speculative_config.num_speculative_tokens is not None

81 + self.num_spec: int = self.speculative_config.num_speculative_tokens

82 else:

83 self.num_spec = 0

84 self.use_spec_decode = self.num_spec > 0

@@ -85,10 +87,15 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]

87 self.use_full_cuda_graph = (

88 self.compilation_config.cudagraph_mode.has_full_cudagraphs()

89 )

88 - self.decode_cudagraph_max_bs = min(

89 - self.vllm_config.scheduler_config.max_num_seqs * (self.num_spec + 1),

90 - self.compilation_config.max_cudagraph_capture_size,

90 +

91 + self.decode_cudagraph_max_bs = (

92 + self.vllm_config.scheduler_config.max_num_seqs * (self.num_spec + 1)

93 )

94 + if self.compilation_config.max_cudagraph_capture_size is not None:

95 + self.decode_cudagraph_max_bs = min(

96 + self.decode_cudagraph_max_bs,

97 + self.compilation_config.max_cudagraph_capture_size,

98 + )

100 self.spec_state_indices_tensor = torch.empty(

101 (self.decode_cudagraph_max_bs, self.num_spec + 1),

vllm/v1/attention/backends/mamba2_attn.py

+3 -2

@@ -123,10 +123,11 @@ class Mamba2AttentionMetadataBuilder(

123 device: torch.device,

124 ):

125 super().__init__(kv_cache_spec, layer_names, vllm_config, device)

126 - self.chunk_size = vllm_config.model_config.get_mamba_chunk_size()

127 - assert self.chunk_size is not None, (

126 + chunk_size = vllm_config.model_config.get_mamba_chunk_size()

127 + assert chunk_size is not None, (

128 "chunk_size needs to be set in the model config for Mamba2 models"

129 )

130 + self.chunk_size: int = chunk_size

131

132 def _compute_chunk_metadata(

133 self,

vllm/v1/attention/backends/mamba_attn.py

+12 -6

@@ -69,10 +69,12 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

70 assert isinstance(kv_cache_spec, MambaSpec)

71 self.compilation_config = vllm_config.compilation_config

72 - self.decode_cudagraph_max_bs = min(

73 - self.vllm_config.scheduler_config.max_num_seqs,

74 - self.compilation_config.max_cudagraph_capture_size,

75 - )

72 + self.decode_cudagraph_max_bs = self.vllm_config.scheduler_config.max_num_seqs

73 + if self.compilation_config.max_cudagraph_capture_size is not None:

74 + self.decode_cudagraph_max_bs = min(

75 + self.decode_cudagraph_max_bs,

76 + self.compilation_config.max_cudagraph_capture_size,

77 + )

79 if self.vllm_config.cache_config.enable_prefix_caching:

80 self.state_indices_tensor = torch.empty(

@@ -150,9 +152,13 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

152 cdiv(common_attn_metadata.seq_lens, mamba_block_size) - 1

153 )

154 # -1 in case it's non-computed and causes later issues with indexing

153 - block_idx_last_computed_token = block_idx_last_computed_token.clamp(min=0)

155 + block_idx_last_computed_token = torch.clamp(

156 + block_idx_last_computed_token, min=0

157 + )

158 # -1 in the case we have a padded request (0 seq-len)

155 - block_idx_last_scheduled_token = block_idx_last_scheduled_token.clamp(min=0)

159 + block_idx_last_scheduled_token = torch.clamp(

160 + block_idx_last_scheduled_token, min=0

161 + )

162

163 return (

164 block_idx_last_computed_token,

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

+1 -1

@@ -62,7 +62,7 @@ class AiterTritonMLAImpl(AiterMLAImpl):

62 k,

63 v,

64 softmax_scale=softmax_scale,

65 - return_lse=return_softmax_lse,

65 + return_softmax_lse=return_softmax_lse,

66 **kwargs,

67 )

68 # Transpose the LSE if Triton MHA is used:

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

+21 -12

@@ -202,6 +202,7 @@ from vllm._aiter_ops import rocm_aiter_ops

202 from vllm.attention.backends.abstract import (

203 AttentionBackend,

204 AttentionLayer,

205 + AttentionMetadata,

206 MLAAttentionImpl,

207 )

208 from vllm.attention.backends.utils import get_mla_dims

@@ -251,13 +252,15 @@ class QueryLenSupport(Enum):

252

253

254 try:

254 - from vllm.vllm_flash_attn import flash_attn_varlen_func

255 + from vllm.vllm_flash_attn import ( # type: ignore[attr-defined]

256 + flash_attn_varlen_func,

257 + )

258

259 is_vllm_fa = True

260 except ImportError:

261 # For rocm use upstream flash attention

262 if current_platform.is_rocm():

260 - from flash_attn import flash_attn_varlen_func

263 + from flash_attn import flash_attn_varlen_func # type: ignore[no-redef]

264 is_vllm_fa = False

265

266 try:

@@ -386,7 +389,7 @@ D = TypeVar("D", bound=MLACommonDecodeMetadata)

389

390

391 @dataclass

389 -class MLACommonMetadata(Generic[D]):

392 +class MLACommonMetadata(AttentionMetadata, Generic[D]):

393 """Metadata for MLACommon.

394

395 NOTE: Please read the comment at the top of the file before trying to

@@ -434,7 +437,7 @@ class MLACommonMetadata(Generic[D]):

437

438

439 M = TypeVar("M", bound=MLACommonMetadata)

437 -A = TypeVar("A")

440 +A = TypeVar("A", bound=AttentionMetadata)

441

442

443 def use_flashinfer_prefill() -> bool:

@@ -617,7 +620,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

620 self._fi_prefill_chunks: list[BatchPrefillWithRaggedKVCacheWrapper] = []

621

622 self._global_hyperparameters = infer_global_hyperparameters(

620 - get_per_layer_parameters(vllm_config, layer_names, MLACommonImpl)

623 + get_per_layer_parameters(vllm_config, layer_names, MLACommonImpl) # type: ignore[type-abstract]

624 )

625

626 if self._use_trtllm_ragged_prefill:

@@ -874,7 +877,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

877 )

878 # Note(qcs): The max local context lengths

879 # padded to `dcp_local_block_size`.

877 - padded_local_context_lens_cpu = (

880 + padded_local_context_lens_cpu: torch.Tensor = (

881 cdiv(

882 context_lens_cpu,

883 self.dcp_virtual_block_size,

@@ -1171,7 +1174,9 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):

1174 )

1175

1176 def get_and_maybe_dequant_weights(layer: LinearBase):

1174 - if not isinstance(layer.quant_method, UnquantizedLinearMethod):

1177 + if layer.quant_method is not None and not isinstance(

1178 + layer.quant_method, UnquantizedLinearMethod

1179 + ):

1180 # NOTE: This should only be used offline, since it's O(N^3)

1181 eye = torch.eye(

1182 layer.input_size_per_partition,

@@ -1327,12 +1332,14 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1332 # v with 0s to match the qk head dim for attention backends that do

1333 # not support different headdims

1334 # We don't need to pad V if we are on a hopper system with FA3

1335 + device_capability = current_platform.get_device_capability()

1336 self._pad_v = self.vllm_flash_attn_version is None or not (

1337 self.vllm_flash_attn_version == 3

1332 - and current_platform.get_device_capability()[0] == 9

1338 + and device_capability is not None

1339 + and device_capability[0] == 9

1340 )

1341

1335 - self.dcp_world_size: int | None = None

1342 + self.dcp_world_size: int = -1

1343

1344 self.chunked_prefill_workspace_size = (

1345 MLACommonMetadataBuilder.determine_chunked_prefill_workspace_size(

@@ -1583,7 +1590,9 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1590 )

1591

1592 def get_and_maybe_dequant_weights(layer: LinearBase):

1586 - if not isinstance(layer.quant_method, UnquantizedLinearMethod):

1593 + if layer.quant_method is not None and not isinstance(

1594 + layer.quant_method, UnquantizedLinearMethod

1595 + ):

1596 # NOTE: This should only be used offline, since it's O(N^3)

1597 eye = torch.eye(

1598 layer.input_size_per_partition,

@@ -1875,7 +1884,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1884 ) -> None:

1885 # TODO (zyongye): Prefill function here

1886 assert attn_metadata.prefill is not None

1878 - assert self.dcp_world_size is not None

1887 + assert self.dcp_world_size != -1

1888

1889 has_context = attn_metadata.prefill.chunked_context is not None

1890 kv_nope = self.kv_b_proj(kv_c_normed)[0].view(

@@ -1975,7 +1984,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1984 # same expert outputs.

1985 return output.fill_(0)

1986

1978 - if self.dcp_world_size is None:

1987 + if self.dcp_world_size == -1:

1988 self.dcp_world_size = get_dcp_group().world_size

1989

1990 fp8_attention = self.kv_cache_dtype.startswith("fp8")

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

+9 -2

@@ -33,7 +33,10 @@ from vllm.v1.attention.backends.mla.common import (

33 )

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 flash_attn_varlen_func, get_scheduler_metadata

36 +from vllm.vllm_flash_attn import ( # type: ignore[attr-defined]

37 + flash_attn_varlen_func,

38 + get_scheduler_metadata,

39 +)

41 logger = init_logger(__name__)

@@ -181,7 +184,11 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]

184

185 # For Flash Attention MLA + full cudagraph

186 max_num_splits = 0

184 - if self.use_full_cuda_graph and num_decode_tokens <= self.max_cudagraph_size:

187 + if (

188 + self.use_full_cuda_graph

189 + and self.max_cudagraph_size is not None

190 + and num_decode_tokens <= self.max_cudagraph_size

191 + ):

192 # NOTE(woosuk): Setting num_splits > 1 may increase the memory

193 # usage, because the intermediate buffers of size [num_splits,

194 # num_heads, num_tokens, head_size] are allocated. Therefore,

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

+4 -2

@@ -10,6 +10,7 @@ from vllm import _custom_ops as ops

10 from vllm.attention.backends.abstract import (

11 AttentionBackend,

12 AttentionLayer,

13 + AttentionMetadata,

14 MultipleOf,

15 )

16 from vllm.attention.backends.utils import get_mla_dims

@@ -124,7 +125,7 @@ class FlashMLASparseBackend(AttentionBackend):

125

126

127 @dataclass

127 -class FlashMLASparseMetadata:

128 +class FlashMLASparseMetadata(AttentionMetadata):

129 num_reqs: int

130 max_query_len: int

131 max_seq_len: int

@@ -718,7 +719,7 @@ class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):

719 )

720 self.softmax_scale = scale

721 assert indexer is not None

721 - self.topk_indices_buffer = indexer.topk_indices_buffer

722 + self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer

723 self.padding = 128 if current_platform.is_device_capability_family(100) else 64

724

725 if kv_cache_dtype == "fp8_ds_mla":

@@ -980,6 +981,7 @@ class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):

981 q = q[:num_actual_toks, ...]

982 k_c_normed = k_c_normed[:num_actual_toks, ...]

983 k_pe = k_pe[:num_actual_toks, ...]

984 + assert self.topk_indices_buffer is not None

985 topk_indices = self.topk_indices_buffer[:num_actual_toks]

986

987 q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

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

+2 -1

@@ -236,7 +236,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):

236 k=k,

237 v=v,

238 softmax_scale=softmax_scale,

239 - return_lse=return_softmax_lse,

239 + return_softmax_lse=return_softmax_lse,

240 **kwargs,

241 )

242

@@ -251,6 +251,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):

251 ) -> tuple[torch.Tensor, torch.Tensor | None]:

252 assert kv_c_and_k_pe_cache.numel() > 0

253 assert attn_metadata.decode is not None

254 + assert attn_metadata.decode.max_qo_len is not None

255

256 if type(q) is tuple:

257 q = torch.cat(q, dim=-1)

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

+4 -3

@@ -43,7 +43,7 @@ class ROCMAiterMLASparseBackend(AttentionBackend):

43 return "ROCM_AITER_MLA_SPARSE"

45 @staticmethod

46 - def get_metadata_cls() -> type[AttentionMetadata]:

46 + def get_metadata_cls() -> type["ROCMAiterMLASparseMetadata"]:

47 return ROCMAiterMLASparseMetadata

49 @staticmethod

@@ -74,7 +74,7 @@ class ROCMAiterMLASparseBackend(AttentionBackend):

76 @dataclass

77 -class ROCMAiterMLASparseMetadata:

77 +class ROCMAiterMLASparseMetadata(AttentionMetadata):

78 num_reqs: int

79 max_query_len: int

80 max_seq_len: int

@@ -223,7 +223,7 @@ class ROCMAiterMLASparseImpl(MLACommonBaseImpl[ROCMAiterMLASparseMetadata]):

223 )

224 self.softmax_scale = scale

225 assert indexer is not None

226 - self.topk_indices_buffer = indexer.topk_indices_buffer

226 + self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer

227 self.is_fp8bmm_enabled = rocm_aiter_ops.is_fp8bmm_enabled()

228

229 def _forward_bf16_kv(

@@ -294,6 +294,7 @@ class ROCMAiterMLASparseImpl(MLACommonBaseImpl[ROCMAiterMLASparseMetadata]):

294 # Convert from (N, B, L) to (B, N, L)

295 ql_nope = ql_nope.transpose(0, 1)

296

297 + assert self.topk_indices_buffer is not None

298 topk_indices = self.topk_indices_buffer[:num_actual_toks]

299

300 topk_indices_global = triton_convert_req_index_to_global_index(

vllm/v1/attention/backends/tree_attn.py

+3 -1

@@ -155,7 +155,9 @@ class TreeAttentionMetadataBuilder(AttentionMetadataBuilder[TreeAttentionMetadat

155 self.block_size = kv_cache_spec.block_size

156

157 spec_config = vllm_config.speculative_config

158 - spec_token_tree = (spec := spec_config) and spec.speculative_token_tree

158 + spec_token_tree: str | None = None

159 + if spec := spec_config:

160 + spec_token_tree = spec.speculative_token_tree

161 tree_choices: list[tuple[int, ...]] = (

162 ast.literal_eval(spec_token_tree) if spec_token_tree is not None else [(0,)]

163 )

vllm/v1/attention/backends/utils.py

+7 -2

@@ -469,6 +469,7 @@ def get_kv_cache_layout():

469 # Format specified by the code.

470 global _KV_CACHE_LAYOUT_OVERRIDE

471

472 + cache_layout: Literal["NHD", "HND"] | None = None

473 if _KV_CACHE_LAYOUT_OVERRIDE is not None:

474 cache_layout = _KV_CACHE_LAYOUT_OVERRIDE

475 logger.info_once(

@@ -524,7 +525,11 @@ def get_per_layer_parameters(

525 to use during `plan`.

526 """

527

527 - layers = get_layers_from_vllm_config(vllm_config, AttentionLayerBase, layer_names)

528 + layers = get_layers_from_vllm_config(

529 + vllm_config,

530 + AttentionLayerBase, # type: ignore[type-abstract]

531 + layer_names,

532 + )

533 per_layer_params: dict[str, PerLayerParameters] = {}

534

535 for key, layer in layers.items():

@@ -1125,7 +1130,7 @@ class KVSharingFastPrefillMetadata(Protocol):

1130

1131 def create_fast_prefill_custom_backend(

1132 prefix: str,

1128 - underlying_attn_backend: AttentionBackend,

1133 + underlying_attn_backend: type[AttentionBackend],

1134 ) -> type[AttentionBackend]:

1135 underlying_builder = underlying_attn_backend.get_builder_cls()

1136