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

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 14, 05:30 UTC

Remove V0 attention backends (#25351)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Woosuk Kwon committed Sep 21, 2025 at 16:03 UTC bc6e542d9fc400d7002853b3878c81b14fbdc998

28 files changed +142 -7375

examples/offline_inference/qwen_1m.py

-1

@@ -5,7 +5,6 @@ from urllib.request import urlopen

6 from vllm import LLM, SamplingParams

8 -os.environ["VLLM_ATTENTION_BACKEND"] = "DUAL_CHUNK_FLASH_ATTN"

8 os.environ["VLLM_ALLOW_LONG_MAX_MODEL_LEN"] = "1"

tests/compile/test_fusion_attn.py

+3 -2

@@ -334,8 +334,9 @@ else:

334 [7, 256, 533] if current_platform.is_cuda() else [8])

335 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])

336 @pytest.mark.parametrize("model_name, model_class", MODELS)

337 -@pytest.mark.parametrize("backend", [_Backend.FLASHINFER] if

338 - current_platform.is_cuda() else [_Backend.ROCM_FLASH])

337 +@pytest.mark.parametrize("backend",

338 + [_Backend.FLASHINFER] if current_platform.is_cuda()

339 + else [_Backend.TRITON_ATTN_VLLM_V1])

340 @pytest.mark.parametrize(

341 "split_attention",

342 [False, True] if current_platform.is_rocm() else [False])

tests/kernels/attention/test_attention.py

+3 -3

@@ -18,7 +18,7 @@ if not current_platform.is_rocm():

18 from xformers import ops as xops

19 from xformers.ops.fmha.attn_bias import BlockDiagonalCausalMask

21 - from vllm.attention.backends.xformers import _make_alibi_bias

21 + from tests.kernels.utils import make_alibi_bias

23 FLOAT32_BYTES = torch.finfo(torch.float).bits // 8

24 # This will change depending on the compute capability.

@@ -429,8 +429,8 @@ def test_multi_query_kv_attention(

429 alibi_bias = None

430 if use_alibi:

431 alibi_slopes = torch.randn(num_query_heads, dtype=torch.float)

432 - attn_bias = _make_alibi_bias(alibi_slopes, num_kv_heads, dtype,

433 - seq_lens)

432 + attn_bias = make_alibi_bias(alibi_slopes, num_kv_heads, dtype,

433 + seq_lens)

434 output = torch.empty_like(query)

435 start = 0

436 # Dynamic sequence length not supported with custom attn_bias.

tests/kernels/attention/test_attention_selector.py

@@ -67,6 +67,7 @@ def generate_params():

67 return params

70 +@pytest.mark.skip(reason="Skipped for now. Should be revisited.")

71 @pytest.mark.parametrize("device, name, use_mla, block_size",

72 generate_params())

73 def test_env(

tests/kernels/attention/test_prefix_prefill.py

+3 -3

@@ -11,7 +11,7 @@ import torch

11 from xformers import ops as xops

12 from xformers.ops.fmha.attn_bias import BlockDiagonalCausalFromBottomRightMask

14 -from vllm.attention.backends.xformers import _make_alibi_bias

14 +from tests.kernels.utils import make_alibi_bias

15 from vllm.attention.ops.chunked_prefill_paged_decode import (

16 chunked_prefill_paged_decode)

17 from vllm.attention.ops.prefix_prefill import context_attention_fwd

@@ -470,7 +470,7 @@ def test_contexted_kv_attention_alibi(

470 key = key.unsqueeze(0)

471 value = value.unsqueeze(0)

472

473 - attn_bias = _make_alibi_bias(alibi_slopes, num_kv_heads, dtype, seq_lens)

473 + attn_bias = make_alibi_bias(alibi_slopes, num_kv_heads, dtype, seq_lens)

474 output_ref = torch.empty_like(output)

475 seq_start = 0

476 query_start = 0

@@ -479,7 +479,7 @@ def test_contexted_kv_attention_alibi(

479 # FIXME(DefTruth): Because xformers does not support dynamic sequence

480 # lengths with custom attention bias, we process each prompt one by

481 # one. This is inefficient, especially when we have many short prompts.

482 - # modified from: vllm/attention/backends/xformers.py#L343

482 + # modified from: vllm/v1/attention/backends/xformers.py#L343

483 for i, (query_len, seq_len) in enumerate(zip(query_lens, seq_lens)):

484 seq_end = seq_start + seq_len

485 query_end = query_start + query_len

tests/kernels/attention/test_rocm_attention_selector.py

@@ -16,6 +16,7 @@ def clear_cache():

16 _cached_get_attn_backend.cache_clear()

19 +@pytest.mark.skip(reason="Skipped for now. Should be revisited.")

20 def test_selector(monkeypatch: pytest.MonkeyPatch):

21 with monkeypatch.context() as m:

22 m.setenv(STR_BACKEND_ENV_VAR, "ROCM_FLASH")

tests/kernels/utils.py

+56 -10

@@ -513,10 +513,6 @@ def make_backend(backend_name: str) -> AttentionBackend:

513 Construct the backend instance determined by the backend_name string

514 argument.

515

516 - "XFORMERS" -> construct xformers backend

517 -

518 - TODO: other backends

519 -

516 Note: at time of writing the Attention wrapper automatically selects

517 its own backend for Attention.forward(); so the backend instance which

518 you generate with this function is not meant to be used for *running*

@@ -528,18 +524,68 @@ def make_backend(backend_name: str) -> AttentionBackend:

524

525 * Backend instance

526 '''

531 - if backend_name == STR_XFORMERS_ATTN_VAL:

532 - # NOTE: xFormers backend cannot be imported for CPU and AMD GPUs.

533 - from vllm.attention.backends.xformers import XFormersBackend

534 - return XFormersBackend()

535 - elif backend_name == STR_FLASH_ATTN_VAL:

536 - from vllm.attention.backends.flash_attn import FlashAttentionBackend

527 + if backend_name in (STR_XFORMERS_ATTN_VAL, "XFORMERS_VLLM_V1"):

528 + from vllm.v1.attention.backends.xformers import (

529 + XFormersAttentionBackend)

530 + return XFormersAttentionBackend()

531 + if backend_name in (STR_FLASH_ATTN_VAL, "FLASH_ATTN_VLLM_V1"):

532 + from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend

533 return FlashAttentionBackend()

534 + if backend_name == "TRITON_ATTN_VLLM_V1":

535 + from vllm.v1.attention.backends.triton_attn import (

536 + TritonAttentionBackend)

537 + return TritonAttentionBackend()

538 + if backend_name == "FLEX_ATTENTION":

539 + from vllm.v1.attention.backends.flex_attention import (

540 + FlexAttentionBackend)

541 + return FlexAttentionBackend()

542 + if backend_name in ("TORCH_SDPA", "TORCH_SDPA_VLLM_V1"):

543 + from vllm.v1.attention.backends.cpu_attn import TorchSDPABackend

544 + return TorchSDPABackend()

545 + if backend_name == "FLASHINFER":

546 + from vllm.v1.attention.backends.flashinfer import FlashInferBackend

547 + return FlashInferBackend()

548

549 raise AssertionError(

550 f"Unrecognized backend_name {backend_name} for unit test")

551

552

553 +def make_alibi_bias(

554 + alibi_slopes: torch.Tensor,

555 + num_kv_heads: int,

556 + dtype: torch.dtype,

557 + seq_lens: list[int],

558 +) -> list[Any]:

559 + """Create ALiBi biases compatible with xFormers attention tests."""

560 + from xformers.ops.fmha.attn_bias import LowerTriangularMaskWithTensorBias

561 +

562 + if alibi_slopes is None:

563 + return [None for _ in seq_lens]

564 +

565 + attn_biases: list[Any] = []

566 + num_heads = alibi_slopes.shape[0]

567 + assert num_heads >= num_kv_heads, (

568 + "ALiBi slopes expect at least as many heads as KV heads")

569 +

570 + for seq_len in seq_lens:

571 + bias = torch.arange(seq_len, dtype=dtype, device=alibi_slopes.device)

572 + bias = bias[None, :] - bias[:, None]

573 +

574 + padded_len = (seq_len + 7) // 8 * 8

575 + bias_tensor = torch.empty(

576 + 1,

577 + num_heads,

578 + seq_len,

579 + padded_len,

580 + device=alibi_slopes.device,

581 + dtype=dtype,

582 + )[:, :, :, :seq_len].copy_(bias)

583 + bias_tensor.mul_(alibi_slopes[:, None, None])

584 + attn_biases.append(LowerTriangularMaskWithTensorBias(bias_tensor))

585 +

586 + return attn_biases

587 +

588 +

589 def _make_metadata_tensors(

590 seq_lens: Optional[list[int]],

591 context_lens: Optional[list[int]],

tests/models/test_initialization.py

+2 -3

@@ -78,9 +78,8 @@ def can_initialize(model_arch: str, monkeypatch: pytest.MonkeyPatch,

78 return

80 if model_arch in ("Phi4FlashForCausalLM", "MotifForCausalLM"):

81 - # Phi4FlashForCausalLM and MotifForCausalLM

82 - # only supports DIFFERENTIAL_FLASH_ATTN backend

83 - m.setenv("VLLM_ATTENTION_BACKEND", "DIFFERENTIAL_FLASH_ATTN")

81 + pytest.skip(

82 + "Differential Flash Attention backend has been removed.")

83 if model_arch == "GptOssForCausalLM":

84 # FIXME: A hack to bypass FA3 assertion because our CI's L4 GPU

85 # has cc==8.9 which hasn't supported FA3 yet. Remove this hack when

vllm/attention/backends/differential_flash_attn.py deleted

-931

@@ -1,931 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -"""" An implementation of https://arxiv.org/pdf/2410.05258 """

4 -from collections import defaultdict

5 -from dataclasses import dataclass

6 -from itertools import accumulate

7 -from typing import Any, Dict, List, Optional, Tuple, Type

8 -

9 -import torch

10 -from einops import rearrange

11 -

12 -from vllm import _custom_ops as ops

13 -# yapf conflicts with isort for this block

14 -# yapf: disable

15 -from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,

16 - AttentionLayer,

17 - AttentionMetadata,

18 - AttentionMetadataBuilder,

19 - AttentionType,

20 - is_quantized_kv_cache)

21 -from vllm.attention.backends.flash_attn import FlashAttentionBackend

22 -# yapf: enable

23 -from vllm.attention.backends.utils import (PAD_SLOT_ID, CommonAttentionState,

24 - compute_slot_mapping,

25 - compute_slot_mapping_start_idx,

26 - is_all_cross_attn_metadata_set,

27 - is_all_encoder_attn_metadata_set,

28 - is_block_tables_empty)

29 -from vllm.attention.utils.fa_utils import (flash_attn_supports_fp8,

30 - get_flash_attn_version)

31 -from vllm.logger import init_logger

32 -from vllm.multimodal import MultiModalPlaceholderMap

33 -from vllm.utils import async_tensor_h2d, make_tensor_with_pad

34 -from vllm.vllm_flash_attn import (flash_attn_varlen_func,

35 - flash_attn_with_kvcache)

36 -

37 -logger = init_logger(__name__)

38 -

39 -

40 -class DifferentialFlashAttentionBackend(AttentionBackend):

41 - accept_output_buffer = False

42 -

43 - @staticmethod

44 - def get_supported_head_sizes() -> List[int]:

45 - return [32, 64, 96, 128, 160, 192, 224, 256]

46 -

47 - @staticmethod

48 - def get_kv_cache_shape(

49 - num_blocks: int,

50 - block_size: int,

51 - num_kv_heads: int,

52 - head_size: int,

53 - ) -> Tuple[int, ...]:

54 - if block_size % 16 != 0:

55 - raise ValueError("Block size must be a multiple of 16.")

56 - assert num_kv_heads % 2 == 0, "num_kv_heads must be divisible by 2"

57 - return (2, 2, num_blocks, block_size, num_kv_heads // 2, head_size)

58 -

59 - @staticmethod

60 - def get_name() -> str:

61 - return "DIFFERENTIAL_FLASH_ATTN"

62 -

63 - @staticmethod

64 - def get_impl_cls() -> Type["DifferentialFlashAttentionImpl"]:

65 - return DifferentialFlashAttentionImpl

66 -

67 - @staticmethod

68 - def get_metadata_cls() -> Type["DifferentialFlashAttentionMetadata"]:

69 - return DifferentialFlashAttentionMetadata

70 -

71 - @staticmethod

72 - def get_builder_cls() -> Type["DifferentialFlashAttentionMetadataBuilder"]:

73 - return DifferentialFlashAttentionMetadataBuilder

74 -

75 - @staticmethod

76 - def get_state_cls() -> Type["CommonAttentionState"]:

77 - return CommonAttentionState

78 -

79 - @staticmethod

80 - def swap_blocks(

81 - src_kv_cache: torch.Tensor,

82 - dst_kv_cache: torch.Tensor,

83 - src_to_dst: torch.Tensor,

84 - ) -> None:

85 - src_key_cache = src_kv_cache[0]

86 - dst_key_cache = dst_kv_cache[0]

87 - ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dst)

88 - src_value_cache = src_kv_cache[1]

89 - dst_value_cache = dst_kv_cache[1]

90 - ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dst)

91 -

92 - @staticmethod

93 - def copy_blocks(

94 - kv_caches: List[torch.Tensor],

95 - src_to_dists: torch.Tensor,

96 - ) -> None:

97 - key_caches = [kv_cache[0] for kv_cache in kv_caches]

98 - value_caches = [kv_cache[1] for kv_cache in kv_caches]

99 -

100 - ops.copy_blocks(key_caches, value_caches, src_to_dists)

101 -

102 -

103 -@dataclass

104 -class DifferentialFlashAttentionMetadata(AttentionMetadata):

105 - """Metadata for FlashAttentionBackend.

106 -

107 - NOTE: Any python object stored here is not updated when it is

108 - cuda-graph replayed. If you have values that need to be changed

109 - dynamically, it should be stored in tensor. The tensor has to be

110 - updated from `CUDAGraphRunner.forward` API.

111 - """

112 - # (batch_size,). The sequence length per sequence. Sequence length means

113 - # the computed tokens + new tokens None if it is a decoding.

114 - seq_lens: Optional[List[int]]

115 - # seq_lens stored as a tensor.

116 - seq_lens_tensor: Optional[torch.Tensor]

117 -

118 - # NOTE(sang): Definition of context_len, query_len, and seq_len.

119 - # |---------- N-1 iteration --------|

120 - # |---------------- N iteration ---------------------|

121 - # |- tokenA -|......................|-- newTokens ---|

122 - # |---------- context_len ----------|

123 - # |-------------------- seq_len ---------------------|

124 - # |-- query_len ---|

125 -

126 - # Maximum sequence length among prefill batch. 0 if there are decoding

127 - # requests only.

128 - max_prefill_seq_len: int

129 - # Maximum sequence length among decode batch. 0 if there are prefill

130 - # requests only.

131 - max_decode_seq_len: int

132 - # (batch_size,) A tensor of context lengths (tokens that are computed

133 - # so far).

134 - context_lens_tensor: Optional[torch.Tensor]

135 -

136 - # (batch_size, max_blocks_per_seq).

137 - # Block addresses per sequence. (Seq id -> list of physical block)

138 - # E.g., [0, 1, 2] means tokens are stored in 0th, 1st, and 2nd blocks

139 - # in the kv cache. Each block can contain up to block_size tokens.

140 - # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph

141 - # captured.

142 - block_tables: Optional[torch.Tensor]

143 -

144 - # Whether or not if cuda graph is enabled.

145 - # Cuda-graph is currently enabled for decoding only.

146 - # TODO(woosuk): Move `use_cuda_graph` out since it's unrelated to attention.

147 -

148 - use_cuda_graph: bool

149 -

150 - # Maximum query length in the batch.

151 - max_query_len: Optional[int] = None

152 -

153 - # Max number of query tokens among request in the batch.

154 - max_decode_query_len: Optional[int] = None

155 -

156 - # (batch_size + 1,). The cumulative subquery lengths of the sequences in

157 - # the batch, used to index into subquery. E.g., if the subquery length

158 - # is [4, 6], it is [0, 4, 10].

159 - query_start_loc: Optional[torch.Tensor] = None

160 - # (batch_size + 1,). The cumulative sequence lengths of the sequences in

161 - # the batch, used to index into sequence. E.g., if the sequence length is

162 - # [4, 6], it is [0, 4, 10].

163 - seq_start_loc: Optional[torch.Tensor] = None

164 -

165 - _cached_prefill_metadata: Optional[

166 - "DifferentialFlashAttentionMetadata"] = None

167 - _cached_decode_metadata: Optional[

168 - "DifferentialFlashAttentionMetadata"] = None

169 -

170 - # Begin encoder attn & enc/dec cross-attn fields...

171 -

172 - # Encoder sequence lengths representation

173 - encoder_seq_lens: Optional[List[int]] = None

174 - encoder_seq_lens_tensor: Optional[torch.Tensor] = None

175 - # (batch_size + 1,). The cumulative sequence lengths of the sequences in

176 - # the batch, used to index into sequence. E.g., if the sequence length is

177 - # [4, 6], it is [0, 4, 10].

178 - encoder_seq_start_loc: Optional[torch.Tensor] = None

179 - # Maximum sequence length among encoder sequences

180 - max_encoder_seq_len: Optional[int] = None

181 - # Number of tokens input to encoder

182 - num_encoder_tokens: Optional[int] = None

183 -

184 - # Cross-attention memory-mapping data structures: slot mapping

185 - # and block tables

186 - cross_slot_mapping: Optional[torch.Tensor] = None

187 - cross_block_tables: Optional[torch.Tensor] = None

188 -

189 - # Cross-layer shared attention block tables

190 - cross_layer_shared_block_tables: Optional[torch.Tensor] = None

191 -

192 - @property

193 - def is_all_encoder_attn_metadata_set(self):

194 - '''

195 - All attention metadata required for encoder attention is set.

196 - '''

197 - return is_all_encoder_attn_metadata_set(self)

198 -

199 - @property

200 - def is_all_cross_attn_metadata_set(self):

201 - '''

202 - All attention metadata required for enc/dec cross-attention is set.

203 -

204 - Superset of encoder attention required metadata.

205 - '''

206 - return is_all_cross_attn_metadata_set(self)

207 -

208 - @property

209 - def prefill_metadata(

210 - self) -> Optional["DifferentialFlashAttentionMetadata"]:

211 - if self.num_prefills == 0:

212 - return None

213 -

214 - if self._cached_prefill_metadata is not None:

215 - return self._cached_prefill_metadata

216 -

217 - assert ((self.seq_lens is not None)

218 - or (self.encoder_seq_lens is not None))

219 - assert ((self.seq_lens_tensor is not None)

220 - or (self.encoder_seq_lens_tensor is not None))

221 -

222 - # Compute some attn_metadata fields which default to None

223 - query_start_loc = (None if self.query_start_loc is None else

224 - self.query_start_loc[:self.num_prefills + 1])

225 - slot_mapping = (None if self.slot_mapping is None else

226 - self.slot_mapping[:self.num_prefill_tokens])

227 - seq_lens = (None if self.seq_lens is None else

228 - self.seq_lens[:self.num_prefills])

229 - seq_lens_tensor = (None if self.seq_lens_tensor is None else

230 - self.seq_lens_tensor[:self.num_prefills])

231 - seq_start_loc = (None if self.seq_start_loc is None else

232 - self.seq_start_loc[:self.num_prefills + 1])

233 - context_lens_tensor = (None if self.context_lens_tensor is None else

234 - self.context_lens_tensor[:self.num_prefills])

235 - block_tables = (None if self.block_tables is None else

236 - self.block_tables[:self.num_prefills])

237 - cross_layer_shared_block_tables = (

238 - None if self.cross_layer_shared_block_tables is None else

239 - self.cross_layer_shared_block_tables[:self.num_prefills])

240 -

241 - self._cached_prefill_metadata = DifferentialFlashAttentionMetadata(

242 - num_prefills=self.num_prefills,

243 - num_prefill_tokens=self.num_prefill_tokens,

244 - num_decode_tokens=0,

245 - slot_mapping=slot_mapping,

246 - multi_modal_placeholder_index_maps=self.

247 - multi_modal_placeholder_index_maps,

248 - enable_kv_scales_calculation=self.enable_kv_scales_calculation,

249 - seq_lens=seq_lens,

250 - seq_lens_tensor=seq_lens_tensor,

251 - max_query_len=self.max_query_len,

252 - max_prefill_seq_len=self.max_prefill_seq_len,

253 - max_decode_query_len=0,

254 - max_decode_seq_len=0,

255 - query_start_loc=query_start_loc,

256 - seq_start_loc=seq_start_loc,

257 - context_lens_tensor=context_lens_tensor,

258 - block_tables=block_tables,

259 - cross_layer_shared_block_tables=cross_layer_shared_block_tables,

260 - use_cuda_graph=False,

261 - # Begin encoder & cross attn fields below...

262 - encoder_seq_lens=self.encoder_seq_lens,

263 - encoder_seq_lens_tensor=self.encoder_seq_lens_tensor,

264 - encoder_seq_start_loc=self.encoder_seq_start_loc,

265 - max_encoder_seq_len=self.max_encoder_seq_len,

266 - cross_slot_mapping=self.cross_slot_mapping,

267 - cross_block_tables=self.cross_block_tables)

268 - return self._cached_prefill_metadata

269 -

270 - @property

271 - def decode_metadata(

272 - self) -> Optional["DifferentialFlashAttentionMetadata"]:

273 - if self.num_decode_tokens == 0:

274 - return None

275 -

276 - if self._cached_decode_metadata is not None:

277 - return self._cached_decode_metadata

278 - assert ((self.seq_lens_tensor is not None)

279 - or (self.encoder_seq_lens_tensor is not None))

280 -

281 - # Compute some attn_metadata fields which default to None

282 - slot_mapping = (None if self.slot_mapping is None else

283 - self.slot_mapping[self.num_prefill_tokens:])

284 - seq_lens_tensor = (None if self.seq_lens_tensor is None else

285 - self.seq_lens_tensor[self.num_prefills:])

286 - block_tables = (None if self.block_tables is None else

287 - self.block_tables[self.num_prefills:])

288 - cross_layer_shared_block_tables = (

289 - None if self.cross_layer_shared_block_tables is None else

290 - self.cross_layer_shared_block_tables[self.num_prefills:])

291 - self._cached_decode_metadata = DifferentialFlashAttentionMetadata(

292 - num_prefills=0,

293 - num_prefill_tokens=0,

294 - num_decode_tokens=self.num_decode_tokens,

295 - slot_mapping=slot_mapping,

296 - multi_modal_placeholder_index_maps=None,

297 - enable_kv_scales_calculation=True,

298 - seq_lens=None,

299 - seq_lens_tensor=seq_lens_tensor,

300 - max_decode_query_len=self.max_decode_query_len,

301 - max_query_len=self.max_query_len,

302 - max_prefill_seq_len=0,

303 - max_decode_seq_len=self.max_decode_seq_len,

304 - # Batch may be composed of prefill|decodes, adjust query start

305 - # indices to refer to the start of decodes. E.g.

306 - # in tokens:[3 prefills|6 decodes], query_start_loc=[3,9] => [0,6].

307 - query_start_loc=(self.query_start_loc[self.num_prefills:] -

308 - self.query_start_loc[self.num_prefills])

309 - if self.query_start_loc is not None else None,

310 - seq_start_loc=self.seq_start_loc[self.num_prefills:]

311 - if self.seq_start_loc is not None else None,

312 - context_lens_tensor=None,

313 - block_tables=block_tables,

314 - cross_layer_shared_block_tables=cross_layer_shared_block_tables,

315 - use_cuda_graph=self.use_cuda_graph,

316 - # Begin encoder & cross attn fields below...

317 - encoder_seq_lens=self.encoder_seq_lens,

318 - encoder_seq_lens_tensor=self.encoder_seq_lens_tensor,

319 - encoder_seq_start_loc=self.encoder_seq_start_loc,

320 - max_encoder_seq_len=self.max_encoder_seq_len,

321 - cross_slot_mapping=self.cross_slot_mapping,

322 - cross_block_tables=self.cross_block_tables)

323 - return self._cached_decode_metadata

324 -

325 -

326 -class DifferentialFlashAttentionMetadataBuilder(

327 - AttentionMetadataBuilder[DifferentialFlashAttentionMetadata]):

328 -

329 - def __init__(self, input_builder):

330 - self.input_builder = input_builder

331 - self.runner = input_builder.runner

332 - self.sliding_window = input_builder.sliding_window

333 - self.block_size = input_builder.block_size

334 -

335 - def prepare(self):

336 - self.slot_mapping: List[int] = []

337 - self.prefill_seq_lens: List[int] = []

338 - self.context_lens: List[int] = []

339 - self.block_tables: List[List[int]] = []

340 - self.cross_layer_shared_block_tables: List[List[int]] = []

341 - self.curr_seq_lens: List[int] = []

342 - self.multimodal_placeholder_maps: Dict[

343 - str,

344 - MultiModalPlaceholderMap] = defaultdict(MultiModalPlaceholderMap)

345 - self.num_prefills = 0

346 - self.num_prefill_tokens = 0

347 - self.num_decode_tokens = 0

348 - self.has_prefix_cache_hit = False

349 -

350 - def _add_seq_group(self, inter_data, chunked_prefill_enabled: bool,

351 - prefix_cache_hit: bool):

352 - """Add a sequence group to the metadata. Specifically update/append

353 - 1. context length.

354 - 2. block table.

355 - 3. slot mapping.

356 - """

357 - # TODO: add support for chunked prefill and prefix caching.

358 - assert not chunked_prefill_enabled, \

359 - "chunked prefill is not supported for now"

360 - assert not prefix_cache_hit, "prefix caching is not supported for now"

361 -

362 - is_prompt = inter_data.is_prompt

363 - block_tables = inter_data.block_tables

364 -

365 - for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len,

366 - curr_sliding_window_block) in zip(

367 - inter_data.seq_ids, [len(t) for t in inter_data.input_tokens],

368 - inter_data.orig_seq_lens, inter_data.seq_lens,

369 - inter_data.query_lens, inter_data.context_lens,

370 - inter_data.curr_sliding_window_blocks):

371 - self.context_lens.append(context_len)

372 -

373 - if is_prompt:

374 - mm_maps = inter_data.multi_modal_placeholder_maps

375 - if mm_maps:

376 - for modality, placeholders in mm_maps.items():

377 - self.multimodal_placeholder_maps[modality].extend(

378 - placeholders)

379 -

380 - self.num_prefills += 1

381 - self.num_prefill_tokens += token_len

382 - self.prefill_seq_lens.append(seq_len)

383 - else:

384 - self.num_decode_tokens += query_len

385 - self.curr_seq_lens.append(curr_seq_len)

386 -

387 - # Compute block table.

388 - # TODO(sang): Combine chunked prefill and prefix caching by

389 - # only allowing multiple of block_size chunk size.

390 - # NOTE: This only works for oooooooxxx style attention.

391 - block_table = []

392 - if prefix_cache_hit:

393 - # NOTE(woosuk): For flash-attn, the block table should

394 - # include the entries for the incoming prefill tokens.

395 - block_table = block_tables[seq_id]

396 - elif ((chunked_prefill_enabled or not is_prompt)

397 - and block_tables is not None):

398 - if curr_sliding_window_block == 0:

399 - block_table = block_tables[seq_id]

400 - else:

401 - block_table = block_tables[seq_id][

402 - -curr_sliding_window_block:]

403 - self.block_tables.append(block_table)

404 -

405 - cross_layer_shared_block_table = []

406 - if prefix_cache_hit:

407 - cross_layer_shared_block_table = block_tables[seq_id]

408 - elif block_tables is not None:

409 - if curr_sliding_window_block == 0:

410 - cross_layer_shared_block_table = block_tables[seq_id]

411 - else:

412 - cross_layer_shared_block_table = block_tables[seq_id][

413 - -curr_sliding_window_block:]

414 - self.cross_layer_shared_block_tables.append(

415 - cross_layer_shared_block_table)

416 -

417 - # Compute slot mapping.

418 - is_profile_run = is_block_tables_empty(block_tables)

419 - start_idx = compute_slot_mapping_start_idx(is_prompt, query_len,

420 - context_len,

421 - self.sliding_window)

422 - compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id,

423 - seq_len, context_len, start_idx,

424 - self.block_size, inter_data.block_tables)

425 -

426 - def _get_graph_runner_block_tables(self, num_seqs: int,

427 - block_tables: List[List[int]],

428 - graph_block_tables) -> torch.Tensor:

429 - # The shape of graph_block_tables is

430 - # [max batch size, max context len // block size].

431 - # max_batch_size, max_blocks = self.runner.graph_block_tables.shape

432 - max_batch_size, max_blocks = graph_block_tables.shape

433 - assert max_batch_size >= num_seqs

434 -

435 - # graph_block_tables = self.runner.graph_block_tables[:num_seqs]

436 - graph_block_tables = graph_block_tables[:num_seqs]

437 - for i, block_table in enumerate(block_tables):

438 - if block_table:

439 - num_blocks = len(block_table)

440 - if num_blocks <= max_blocks:

441 - graph_block_tables[i, :num_blocks] = block_table

442 - else:

443 - # It may be possible to have more blocks allocated due

444 - # to lookahead slots of multi-step, however, they are

445 - # not used anyway, so can be safely ignored.

446 - graph_block_tables[

447 - i, :max_blocks] = block_table[:max_blocks]

448 -

449 - return torch.from_numpy(graph_block_tables).to(

450 - device=self.runner.device, non_blocking=True)

451 -

452 - def build(self, seq_lens: List[int], query_lens: List[int],

453 - cuda_graph_pad_size: int, batch_size: int):

454 - """Build attention metadata with on-device tensors.

455 -

456 - Args:

457 - seq_lens: The maybe padded sequence lengths of the input sequences.

458 - query_lens: The query lengths of the input sequences.

459 - cuda_graph_pad_size: The padding size for cuda graph.

460 - -1 if cuda graph is not used.

461 - batch_size: The maybe padded batch size.

462 - """

463 - prefix_cache_hit = any([

464 - inter_data.prefix_cache_hit

465 - for inter_data in self.input_builder.inter_data_list

466 - ])

467 - for inter_data in self.input_builder.inter_data_list:

468 - self._add_seq_group(inter_data,

469 - self.input_builder.chunked_prefill_enabled,

470 - prefix_cache_hit)

471 -

472 - device = self.runner.device

473 - use_captured_graph = cuda_graph_pad_size != -1

474 -

475 - max_query_len = max(query_lens)

476 - decode_query_lens = query_lens[self.num_prefills:]

477 - if len(decode_query_lens) > 0:

478 - max_decode_query_len = max(decode_query_lens)

479 - else:

480 - max_decode_query_len = 1

481 - max_prefill_seq_len = max(self.prefill_seq_lens, default=0)

482 - max_decode_seq_len = max(self.curr_seq_lens, default=0)

483 - num_decode_tokens = self.num_decode_tokens

484 - query_start_loc = list(accumulate(query_lens, initial=0))

485 - seq_start_loc = list(accumulate(seq_lens, initial=0))

486 -

487 - num_seqs = len(seq_lens)

488 - if use_captured_graph:

489 - self.slot_mapping.extend([PAD_SLOT_ID] * cuda_graph_pad_size)

490 - self.block_tables.extend([] * cuda_graph_pad_size)

491 -

492 - self.cross_layer_shared_block_tables.extend([] *

493 - cuda_graph_pad_size)

494 -

495 - num_decode_tokens = batch_size - self.num_prefill_tokens

496 - block_tables = self._get_graph_runner_block_tables(

497 - num_seqs, self.block_tables, self.runner.graph_block_tables)

498 - cross_layer_shared_block_tables = \

499 - self._get_graph_runner_block_tables(

500 - num_seqs, self.cross_layer_shared_block_tables,

501 - self.runner.cross_layer_shared_graph_block_tables)

502 - else:

503 - block_tables = make_tensor_with_pad(

504 - self.block_tables,

505 - pad=0,

506 - dtype=torch.int,

507 - device=device,

508 - )

509 - cross_layer_shared_block_tables = make_tensor_with_pad(

510 - self.cross_layer_shared_block_tables,

511 - pad=0,

512 - dtype=torch.int,

513 - device=device,

514 - )

515 - assert max_query_len > 0, ("query_lens: {}".format(query_lens))

516 -

517 - assert device is not None

518 - context_lens_tensor = async_tensor_h2d(self.context_lens, torch.int,

519 - device, self.runner.pin_memory)

520 - seq_lens_tensor = async_tensor_h2d(seq_lens, torch.int, device,

521 - self.runner.pin_memory)

522 - slot_mapping_tensor = async_tensor_h2d(self.slot_mapping, torch.long,

523 - device, self.runner.pin_memory)

524 - query_start_loc_tensor = async_tensor_h2d(query_start_loc, torch.int32,

525 - device,

526 - self.runner.pin_memory)

527 - seq_start_loc_tensor = async_tensor_h2d(seq_start_loc, torch.int32,

528 - device, self.runner.pin_memory)

529 - placeholder_index_maps = {

530 - modality: placeholder_map.index_map()

531 - for modality, placeholder_map in

532 - self.multimodal_placeholder_maps.items()

533 - }

534 -

535 - return DifferentialFlashAttentionMetadata(

536 - num_prefills=self.num_prefills,

537 - slot_mapping=slot_mapping_tensor,

538 - num_prefill_tokens=self.num_prefill_tokens,

539 - num_decode_tokens=num_decode_tokens,

540 - seq_lens=seq_lens,

541 - multi_modal_placeholder_index_maps=placeholder_index_maps,

542 - enable_kv_scales_calculation=True,

543 - seq_lens_tensor=seq_lens_tensor,

544 - max_query_len=max_query_len,

545 - max_decode_query_len=max_decode_query_len,

546 - max_prefill_seq_len=max_prefill_seq_len,

547 - max_decode_seq_len=max_decode_seq_len,

548 - query_start_loc=query_start_loc_tensor,

549 - seq_start_loc=seq_start_loc_tensor,

550 - context_lens_tensor=context_lens_tensor,

551 - block_tables=block_tables,

552 - cross_layer_shared_block_tables=cross_layer_shared_block_tables,

553 - use_cuda_graph=use_captured_graph,

554 - )

555 -

556 -

557 -class DifferentialFlashAttentionImpl(AttentionImpl):

558 - """

559 - If the input tensors contain prompt tokens, the layout is as follows:

560 - |<--------------- num_prefill_tokens ----------------->|

562 -

563 - Otherwise, the layout is as follows:

564 - |<----------------- num_decode_tokens ------------------>|

566 -

567 - Generation tokens can contain padding when cuda-graph is used.

568 - Currently, prompt tokens don't contain any padding.

569 -

570 - The prompts might have different lengths, while the generation tokens

571 - always have length 1.

572 -

573 - If chunked prefill is enabled, prefill tokens and decode tokens can be

574 - batched together in a flattened 1D query.

575 -

576 - |<----- num_prefill_tokens ---->|<------- num_decode_tokens --------->|

577 - |<-prefill_0->|...|<-prefill_N-1->|<--decode_0-->|...|<--decode_M-1-->|

578 -

579 - Currently, cuda graph is disabled for chunked prefill, meaning there's no

580 - padding between prefill and decode tokens.

581 - """

582 -

583 - def __init__(

584 - self,

585 - num_heads: int,

586 - head_size: int,

587 - scale: float,

588 - num_kv_heads: int,

589 - alibi_slopes: Optional[List[float]],

590 - sliding_window: Optional[int],

591 - kv_cache_dtype: str,

592 - logits_soft_cap: Optional[float] = None,

593 - attn_type: str = AttentionType.DECODER,

594 - kv_sharing_target_layer_name: Optional[str] = None,

595 - use_irope: bool = False,

596 - differential_flash_attention_config: Optional[Dict[str, Any]] = None,

597 - ) -> None:

598 - if differential_flash_attention_config is None:

599 - differential_flash_attention_config = {}

600 - self.differential_flash_attention_config = \

601 - differential_flash_attention_config

602 - self.used_shared_kv_cache = kv_sharing_target_layer_name is not None

603 - self.kv_sharing_target_layer_name = kv_sharing_target_layer_name

604 - if use_irope:

605 - logger.warning(

606 - "Using irope in V0 is not supported yet, it will fall back "

607 - "to global attention for long context.")

608 - self.num_heads = num_heads

609 - self.head_size = head_size

610 - self.scale = float(scale)

611 - self.num_kv_heads = num_kv_heads

612 - if alibi_slopes is not None:

613 - alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)

614 - self.alibi_slopes = alibi_slopes

615 - self.sliding_window = ((sliding_window - 1,

616 - 0) if sliding_window is not None else (-1, -1))

617 - self.kv_cache_dtype = kv_cache_dtype

618 - self.vllm_flash_attn_version = get_flash_attn_version(

619 - requires_alibi=self.alibi_slopes is not None)

620 - if is_quantized_kv_cache(self.kv_cache_dtype) and (

621 - not self.kv_cache_dtype.startswith("fp8")

622 - or not flash_attn_supports_fp8()):

623 - raise NotImplementedError(

624 - f"FlashAttention does not support {self.kv_cache_dtype} "

625 - "kv-cache on this device "

626 - f"(FA supports fp8 = {flash_attn_supports_fp8()}).")

627 - if logits_soft_cap is None:

628 - # In flash-attn, setting logits_soft_cap as 0 means no soft cap.

629 - logits_soft_cap = 0

630 - self.logits_soft_cap = logits_soft_cap

631 -

632 - assert self.num_heads % self.num_kv_heads == 0

633 - self.num_queries_per_kv = self.num_heads // self.num_kv_heads

634 -

635 - support_head_sizes = FlashAttentionBackend.get_supported_head_sizes()

636 - if head_size not in support_head_sizes:

637 - raise ValueError(

638 - f"Head size {head_size} is not supported by FlashAttention. "

639 - f"Supported head sizes are: {support_head_sizes}.")

640 - self.attn_type = attn_type

641 -

642 - self.lambda_full = None

643 - self.subln = self.differential_flash_attention_config["subln"]

644 -

645 - def split_heads(self, x):

646 - # split by num_heads, the stripe pattern is friendly to tensor parallel.

647 - x = rearrange(x, "... (H two) D -> ... H two D", two=2)

648 - x1 = x[..., 0, :]

649 - x2 = x[..., 1, :]

650 - return x1.contiguous(), x2.contiguous()

651 -

652 - def split_kv_cache(self, x):

653 - # split by num_heads, the stripe pattern is friendly to tensor parallel.

654 - if x.numel() == 0:

655 - return torch.empty(0), torch.empty(0)

656 -

657 - x1, x2 = x[0], x[1]

658 - return x1, x2

659 -

660 - def populate_kv_cache(self, layer: AttentionLayer, key: torch.Tensor,

661 - value: torch.Tensor, kv_cache: torch.Tensor,

662 - attn_metadata: DifferentialFlashAttentionMetadata):

663 - if kv_cache.numel() > 0 and key is not None and value is not None:

664 - updated_slot_mapping = attn_metadata.slot_mapping

665 - torch.ops._C_cache_ops.reshape_and_cache_flash(

666 - key,

667 - value,

668 - kv_cache[0],

669 - kv_cache[1],

670 - updated_slot_mapping.flatten(),

671 - self.kv_cache_dtype,

672 - layer._k_scale,

673 - layer._v_scale,

674 - )

675 -

676 - def forward_generate_kv_cache(

677 - self, query: torch.Tensor, key: Optional[torch.Tensor],

678 - value: Optional[torch.Tensor], k_cache: torch.Tensor,

679 - v_cache: torch.Tensor,

680 - attn_metadata: DifferentialFlashAttentionMetadata) -> torch.Tensor:

681 -

682 - head_size = self.head_size

683 - num_heads = self.num_heads // 2

684 - num_kv_heads = self.num_kv_heads // 2

685 -

686 - query = query.view(-1, num_heads, head_size)

687 - if key is not None:

688 - assert value is not None

689 - key = key.view(-1, num_kv_heads, head_size)

690 - value = value.view(-1, num_kv_heads, head_size)

691 - else:

692 - assert value is None

693 -

694 - num_prefill_tokens = attn_metadata.num_prefill_tokens

695 - num_decode_tokens = attn_metadata.num_decode_tokens

696 - assert key.shape[

697 - 0] == num_prefill_tokens + num_decode_tokens, "key shape mismatch"

698 - assert value.shape[

699 - 0] == num_prefill_tokens + num_decode_tokens, "value shape mismatch"

700 -

701 - output = torch.empty_like(query)

702 - # Query for decode. KV is not needed because it is already cached.

703 - decode_query = query[num_prefill_tokens:]

704 - # QKV for prefill.

705 - query = query[:num_prefill_tokens]

706 - if key is not None and value is not None:

707 - key = key[:num_prefill_tokens]

708 - value = value[:num_prefill_tokens]

709 -

710 - assert query.shape[0] == num_prefill_tokens, "query shape mismatch"

711 - assert decode_query.shape[

712 - 0] == num_decode_tokens, "decode query shape mismatch"

713 -

714 - if prefill_meta := attn_metadata.prefill_metadata:

715 - # Prompt run.

716 - if k_cache.numel() == 0 \

717 - or prefill_meta.block_tables is None \

718 - or prefill_meta.block_tables.numel() == 0:

719 - # normal attention

720 - prefill_output = flash_attn_varlen_func(

721 - q=query,

722 - k=key,

723 - v=value,

724 - cu_seqlens_q=prefill_meta.seq_start_loc,

725 - cu_seqlens_k=prefill_meta.seq_start_loc,

726 - max_seqlen_q=prefill_meta.max_prefill_seq_len,

727 - max_seqlen_k=prefill_meta.max_prefill_seq_len,

728 - softmax_scale=self.scale,

729 - causal=True,

730 - window_size=self.sliding_window,

731 - alibi_slopes=self.alibi_slopes,

732 - softcap=self.logits_soft_cap,

733 - fa_version=self.vllm_flash_attn_version,

734 - )

735 - assert prefill_output.shape == output[:

736 - num_prefill_tokens].shape

737 - output[:num_prefill_tokens] = prefill_output

738 - else:

739 - raise Exception("prefix caching not supported")

740 -

741 - if decode_meta := attn_metadata.decode_metadata:

742 - block_tables_arg = decode_meta.block_tables

743 - try:

744 - output[num_prefill_tokens:] = flash_attn_with_kvcache(

745 - q=decode_query.unsqueeze(1),

746 - k_cache=k_cache,

747 - v_cache=v_cache,

748 - block_table=block_tables_arg,

749 - cache_seqlens=decode_meta.seq_lens_tensor,

750 - softmax_scale=self.scale,

751 - causal=True,

752 - window_size=self.sliding_window,

753 - alibi_slopes=self.alibi_slopes,

754 - softcap=self.logits_soft_cap,

755 - fa_version=self.vllm_flash_attn_version,

756 - ).squeeze(1)

757 - except Exception as e:

758 - logger.error("Error in PagedAttention.forward_decode: %s",

759 - str(e))

760 - raise e

761 -

762 - # Reshape the output tensor.

763 - return output.view(-1, num_heads, head_size)

764 -

765 - def forward_with_kv_cache_only(

766 - self,

767 - query: torch.Tensor,

768 - k_cache: torch.Tensor,

769 - v_cache: torch.Tensor,

770 - attn_metadata: DifferentialFlashAttentionMetadata,

771 - ):

772 - if not attn_metadata.decode_metadata:

773 - block_tables_arg = attn_metadata.cross_layer_shared_block_tables

774 - else:

775 - block_tables_arg = attn_metadata.block_tables

776 -

777 - output = flash_attn_with_kvcache(

778 - q=query.unsqueeze(1),

779 - k_cache=k_cache,

780 - v_cache=v_cache,

781 - block_table=block_tables_arg,

782 - cache_seqlens=attn_metadata.seq_lens_tensor,

783 - softmax_scale=self.scale,

784 - causal=True,

785 - window_size=self.sliding_window,

786 - alibi_slopes=self.alibi_slopes,

787 - softcap=self.logits_soft_cap,

788 - fa_version=self.vllm_flash_attn_version,

789 - ).squeeze(1)

790 - return output

791 -

792 - def forward(

793 - self,

794 - layer: AttentionLayer,

795 - q: torch.Tensor,

796 - k: torch.Tensor,

797 - v: torch.Tensor,

798 - kv_cache: torch.Tensor,

799 - attn_metadata: DifferentialFlashAttentionMetadata,

800 - output: Optional[torch.Tensor] = None,

801 - output_scale: Optional[torch.Tensor] = None,

802 - output_block_scale: Optional[torch.Tensor] = None,

803 - ) -> torch.Tensor:

804 - """Forward pass with FlashAttention.

805 -

806 - Args:

807 - layer: Attention layer instance.

808 - q: Query tensor with shape = [num_tokens, num_heads, head_size]

809 - k: Key tensor with shape = [num_tokens, num_kv_heads, head_size]

810 - v: Value tensor with shape = [num_tokens, num_kv_heads, head_size]

811 - kv_cache: KV cache tensor with shape

812 - [2, num_blocks, block_size, num_kv_heads, head_size].

813 - NOTE: kv_cache will be an empty tensor with shape [0]

814 - for profiling run.

815 - attn_metadata: Metadata for attention.

816 - output: Output tensor with shape [num_tokens, num_heads, head_size]

817 - output_scale: Optional output scale tensor.

818 - output_block_scale: Optional output block scale tensor.

819 - NOTE: It in-place updates the output tensor.

820 - NOTE: FP8 quantization, flash-attn expect the size of

821 - {q,k,v}_descale to be (num_sequences, num_kv_heads).

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

823 - """

824 - if output_scale is not None or output_block_scale is not None:

825 - raise NotImplementedError(

826 - "fused output quantization is not yet supported"

827 - " for DifferentialFlashAttentionImpl")

828 -

829 - if self.lambda_full is None:

830 - self.lambda_init = self.differential_flash_attention_config[

831 - "lambda_init"]

832 - lambda_q1 = self.differential_flash_attention_config["lambda_q1"]

833 - lambda_k1 = self.differential_flash_attention_config["lambda_k1"]

834 - lambda_q2 = self.differential_flash_attention_config["lambda_q2"]

835 - lambda_k2 = self.differential_flash_attention_config["lambda_k2"]

836 - lambda_1 = torch.exp(

837 - torch.sum(lambda_q1 * lambda_k1, dim=-1).float()).type_as(q)

838 - lambda_2 = torch.exp(

839 - torch.sum(lambda_q2 * lambda_k2, dim=-1).float()).type_as(q)

840 - self.lambda_full = lambda_1 - lambda_2 + self.lambda_init

841 -

842 - if not self.used_shared_kv_cache: # need to generate kv-cache

843 - q = q.view(-1, self.num_heads, self.head_size)

844 - k = k.view(-1, self.num_kv_heads, self.head_size)

845 - v = v.view(-1, self.num_kv_heads, self.head_size)

846 -

847 - q1, q2 = self.split_heads(q)

848 - k1, k2 = self.split_heads(k)

849 - v1, v2 = self.split_heads(v)

850 -

851 - # kv_cache shape is (2, 2, num_blocks, block_size, num_kv_heads // 2, head_size) # noqa: E501

852 - # Split by half along the first dimension.

853 - kv_cache1, kv_cache2 = self.split_kv_cache(kv_cache)

854 - assert kv_cache1.is_contiguous(), "kv_cache1 is not contiguous"

855 - assert kv_cache2.is_contiguous(), "kv_cache2 is not contiguous"

856 -

857 - if kv_cache1.numel() != 0:

858 - self.populate_kv_cache(layer, k1, v1, kv_cache1, attn_metadata)

859 - self.populate_kv_cache(layer, k2, v2, kv_cache2, attn_metadata)

860 -

861 - key_cache1, value_cache1 = self.split_kv_cache(kv_cache1)

862 - key_cache2, value_cache2 = self.split_kv_cache(kv_cache2)

863 - else:

864 - key_cache1, value_cache1 = torch.empty(0), torch.empty(0)

865 - key_cache2, value_cache2 = torch.empty(0), torch.empty(0)

866 - attn11 = self.forward_generate_kv_cache(q1, k1, v1, key_cache1,

867 - value_cache1,

868 - attn_metadata)

869 - attn12 = self.forward_generate_kv_cache(q1, k1, v2, key_cache1,

870 - value_cache2,

871 - attn_metadata)

872 - attn11 = attn11.view(q1.shape)

873 - attn12 = attn12.view(q1.shape)

874 - attn1 = torch.cat([attn11, attn12], dim=-1)

875 -

876 - attn21 = self.forward_generate_kv_cache(q2, k2, v1, key_cache2,

877 - value_cache1,

878 - attn_metadata)

879 - attn22 = self.forward_generate_kv_cache(q2, k2, v2, key_cache2,

880 - value_cache2,

881 - attn_metadata)

882 - attn21 = attn21.view(q2.shape)

883 - attn22 = attn22.view(q2.shape)

884 - attn2 = torch.cat([attn21, attn22], dim=-1)

885 -

886 - attn = attn1 - self.lambda_full * attn2

887 - # attn shape (-1, self.num_heads // 2, 2 * self.head_dim)

888 - attn = self.subln(attn)

889 - attn = attn * (1 - self.lambda_init)

890 - # reshape back to 2 * num_head

891 - attn_output = rearrange(attn,

892 - "... H (two D) -> ... (H two) D",

893 - two=2)

894 -

895 - else: # reuse the kv cache, full attention

896 - q = q.view(-1, self.num_heads, self.head_size)

897 - q1, q2 = self.split_heads(q)

898 - # kv_cache shape is (2, num_blocks, block_size, num_kv_heads, head_size) # noqa: E501

899 - kv_cache1, kv_cache2 = self.split_kv_cache(kv_cache)

900 - key_cache1, value_cache1 = kv_cache1[0], kv_cache1[1]

901 - key_cache2, value_cache2 = kv_cache2[0], kv_cache2[1]

902 -

903 - attn11 = self.forward_with_kv_cache_only(q1, key_cache1,

904 - value_cache1,

905 - attn_metadata)

906 - attn12 = self.forward_with_kv_cache_only(q1, key_cache1,

907 - value_cache2,

908 - attn_metadata)

909 - attn11 = attn11.view(q1.shape)

910 - attn12 = attn12.view(q1.shape)

911 - attn1 = torch.cat([attn11, attn12], dim=-1)

912 -

913 - attn21 = self.forward_with_kv_cache_only(q2, key_cache2,

914 - value_cache1,

915 - attn_metadata)

916 - attn22 = self.forward_with_kv_cache_only(q2, key_cache2,

917 - value_cache2,

918 - attn_metadata)

919 - attn21 = attn21.view(q2.shape)

920 - attn22 = attn22.view(q2.shape)

921 - attn2 = torch.cat([attn21, attn22], dim=-1)

922 -

923 - attn = attn1 - self.lambda_full * attn2

924 - attn = self.subln(attn)

925 - attn = attn * (1 - self.lambda_init)

926 - # reshape back to 2 * num_head

927 - attn_output = rearrange(attn,

928 - "... H (two D) -> ... (H two) D",

929 - two=2)

930 - attn_output = attn_output.view(-1, self.num_heads * self.head_size)

931 - return attn_output

vllm/attention/backends/dual_chunk_flash_attn.py deleted

-1495

@@ -1,1495 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -"""Attention layer with Dual chunk flash attention and sparse attention.

4 -"""

5 -import math

6 -from dataclasses import dataclass

7 -from typing import Any, Dict, List, Optional, Tuple, Type

8 -

9 -import torch

10 -import torch.distributed

11 -import torch.nn.functional as F

12 -

13 -from vllm import _custom_ops as ops

14 -from vllm.attention.backends.abstract import AttentionLayer, AttentionType

15 -from vllm.attention.backends.flash_attn import (FlashAttentionBackend,

16 - FlashAttentionImpl,

17 - FlashAttentionMetadata,

18 - FlashAttentionMetadataBuilder)

19 -from vllm.distributed.parallel_state import get_tensor_model_parallel_rank

20 -from vllm.logger import init_logger

21 -from vllm.utils import async_tensor_h2d

22 -from vllm.vllm_flash_attn import (flash_attn_varlen_func,

23 - flash_attn_with_kvcache, sparse_attn_func)

24 -

25 -logger = init_logger(__name__)

26 -

27 -

28 -class DualChunkFlashAttentionBackend(FlashAttentionBackend):

29 -

30 - accept_output_buffer: bool = False

31 -

32 - @staticmethod

33 - def get_name() -> str:

34 - return "DUAL_CHUNK_FLASH_ATTN"

35 -

36 - @staticmethod

37 - def get_impl_cls() -> Type["DualChunkFlashAttentionImpl"]:

38 - return DualChunkFlashAttentionImpl

39 -

40 - @staticmethod

41 - def get_metadata_cls() -> Type["DualChunkFlashAttentionMetadata"]:

42 - return DualChunkFlashAttentionMetadata

43 -

44 - @staticmethod

45 - def get_builder_cls() -> Type["DualChunkFlashAttentionMetadataBuilder"]:

46 - return DualChunkFlashAttentionMetadataBuilder

47 -

48 -

49 -@dataclass

50 -class DualChunkFlashAttentionMetadata(FlashAttentionMetadata):

51 - # Block size of the paged kv cache.

52 - block_size: int = 16

53 -

54 - # Original max position embeddings.

55 - original_max_position_embeddings: int = 0

56 -

57 - # Chunk size

58 - chunk_size: int = 8192

59 -

60 - # Local size

61 - local_size: int = 1024

62 -

63 - # (batch_size,). The orig sequence length per sequence.

64 - orig_seq_lens: Optional[List[int]] = None

65 -

66 - # orig_seq_lens stored as a tensor.

67 - orig_seq_lens_tensor: Optional[torch.Tensor] = None

68 -

69 - # Length scaling factor

70 - scaling_factor: Optional[torch.Tensor] = None

71 -

72 - # (batch_size,). Sequence lengths for intra attention.

73 - seq_lens_intra: Optional[torch.Tensor] = None

74 -

75 - # Max sequence length for intra attention.

76 - max_seq_len_intra: Optional[int] = None

77 -

78 - # (batch_size, num_blocks). Block table for intra attention.

79 - block_tables_intra: Optional[torch.Tensor] = None

80 -

81 - # (batch_size,). Sequence lengths for succ attention.

82 - seq_lens_succ: Optional[torch.Tensor] = None

83 -

84 - # Max sequence length for succ attention.

85 - max_seq_len_succ: Optional[int] = None

86 -

87 - # (batch_size, num_blocks). Block table for succ attention.

88 - block_tables_succ: Optional[torch.Tensor] = None

89 -

90 - # (batch_size,). Sequence lengths for inter attention.

91 - seq_lens_inter: Optional[torch.Tensor] = None

92 -

93 - # Max sequence length for inter attention.

94 - max_seq_len_inter: Optional[int] = None

95 -

96 - _cached_prefill_metadata: Optional[

97 - "DualChunkFlashAttentionMetadata"] = None

98 - _cached_decode_metadata: Optional["DualChunkFlashAttentionMetadata"] = None

99 -

100 - @property

101 - def prefill_metadata(self) -> Optional["DualChunkFlashAttentionMetadata"]:

102 - if self.num_prefills == 0:

103 - return None

104 -

105 - if self._cached_prefill_metadata is not None:

106 - return self._cached_prefill_metadata

107 -

108 - prefill_metadata = super().prefill_metadata

109 - if prefill_metadata is None:

110 - return None

111 -

112 - prefill_metadata = DualChunkFlashAttentionMetadata(

113 - **prefill_metadata.asdict_zerocopy())

114 -

115 - prefill_metadata.orig_seq_lens = (

116 - None if self.orig_seq_lens is None else

117 - self.orig_seq_lens[:self.num_prefills])

118 - prefill_metadata.orig_seq_lens_tensor = (

119 - None if self.orig_seq_lens_tensor is None else

120 - self.orig_seq_lens_tensor[:self.num_prefills])

121 -

122 - if self.original_max_position_embeddings > 0:

123 - assert prefill_metadata.orig_seq_lens_tensor is not None

124 - prefill_metadata.scaling_factor = (

125 - 0.1 * torch.log(prefill_metadata.orig_seq_lens_tensor /

126 - self.original_max_position_embeddings) +

127 - 1.0).clip(min=1)

128 -

129 - self._cached_prefill_metadata = prefill_metadata

130 - return prefill_metadata

131 -

132 - @property

133 - def decode_metadata(self) -> Optional["DualChunkFlashAttentionMetadata"]:

134 - if self.num_decode_tokens == 0:

135 - return None

136 -

137 - if self._cached_decode_metadata is not None:

138 - return self._cached_decode_metadata

139 -

140 - decode_metadata = super().decode_metadata

141 - if decode_metadata is None:

142 - return None

143 -

144 - decode_metadata = DualChunkFlashAttentionMetadata(

145 - **decode_metadata.asdict_zerocopy())

146 -

147 - decode_metadata.orig_seq_lens_tensor = (

148 - None if self.orig_seq_lens_tensor is None else

149 - self.orig_seq_lens_tensor[self.num_prefills:])

150 -

151 - assert decode_metadata.orig_seq_lens_tensor is not None

152 - assert decode_metadata.block_tables is not None

153 -

154 - cache_seq_lens = decode_metadata.orig_seq_lens_tensor

155 - chunk_len = self.chunk_size - self.local_size

156 - chunk_num_curr = (cache_seq_lens - 1) // chunk_len

157 - batch_size = decode_metadata.num_decode_tokens

158 -

159 - if self.original_max_position_embeddings > 0:

160 - decode_metadata.scaling_factor = (0.1 * torch.log(

161 - cache_seq_lens / self.original_max_position_embeddings) +

162 - 1.0).clip(min=1)

163 -

164 - seq_lens_intra = cache_seq_lens - chunk_num_curr * chunk_len

165 - max_seq_len_intra = seq_lens_intra.max().item()

166 - decode_metadata.seq_lens_intra = seq_lens_intra

167 - decode_metadata.max_seq_len_intra = max_seq_len_intra

168 -

169 - block_tables_intra = torch.zeros(

170 - batch_size,

171 - (max_seq_len_intra - 1) // self.block_size + 1,

172 - dtype=decode_metadata.block_tables.dtype,

173 - device=decode_metadata.block_tables.device,

174 - )

175 - for i in range(batch_size):

176 - st = chunk_num_curr[i] * chunk_len // self.block_size

177 - ed = min(

178 - st + (max_seq_len_intra - 1) // self.block_size + 1,

179 - (cache_seq_lens[i] - 1) // self.block_size + 1,

180 - )

181 - block_tables_intra[i, :ed -

182 - st] = decode_metadata.block_tables[i, st:ed]

183 - decode_metadata.block_tables_intra = block_tables_intra

184 -

185 - seq_lens_succ = (chunk_num_curr -

186 - (chunk_num_curr - 1).clip(min=0)) * chunk_len

187 - max_seq_len_succ = seq_lens_succ.max().item()

188 - decode_metadata.seq_lens_succ = seq_lens_succ

189 - decode_metadata.max_seq_len_succ = max_seq_len_succ

190 - if max_seq_len_succ:

191 - block_tables_succ = torch.zeros(

192 - batch_size,

193 - (max_seq_len_succ - 1) // self.block_size + 1,

194 - dtype=decode_metadata.block_tables.dtype,

195 - device=decode_metadata.block_tables.device,

196 - )

197 - for i in range(batch_size):

198 - start = ((chunk_num_curr[i] - 1).clip(min=0) * chunk_len //

199 - self.block_size)

200 - end = min(

201 - start + (max_seq_len_succ - 1) // self.block_size + 1,

202 - (cache_seq_lens[i] - 1) // self.block_size + 1,

203 - )

204 - block_tables_succ[

205 - i, :end - start] = decode_metadata.block_tables[i,

206 - start:end]

207 - decode_metadata.block_tables_succ = block_tables_succ

208 -

209 - seq_lens_inter = (chunk_num_curr - 1).clip(min=0) * chunk_len

210 - max_seq_len_inter = seq_lens_inter.max().item()

211 - decode_metadata.seq_lens_inter = seq_lens_inter

212 - decode_metadata.max_seq_len_inter = max_seq_len_inter

213 -

214 - self._cached_decode_metadata = decode_metadata

215 - return decode_metadata

216 -

217 -

218 -class DualChunkFlashAttentionMetadataBuilder(FlashAttentionMetadataBuilder):

219 -

220 - def prepare(self):

221 - super().prepare()

222 - self.orig_seq_lens: List[int] = []

223 -

224 - def _add_seq_group(self, inter_data, chunked_prefill_enabled: bool,

225 - prefix_cache_hit: bool):

226 - super()._add_seq_group(inter_data, chunked_prefill_enabled,

227 - prefix_cache_hit)

228 - for prompt_len, seq_len in zip(inter_data.prompt_lens,

229 - inter_data.seq_lens):

230 - self.orig_seq_lens.append(max(prompt_len, seq_len))

231 -

232 - def build(self, seq_lens: List[int], query_lens: List[int],

233 - cuda_graph_pad_size: int, batch_size: int):

234 - attn_metadata = super().build(seq_lens, query_lens,

235 - cuda_graph_pad_size, batch_size)

236 - attn_metadata = DualChunkFlashAttentionMetadata(

237 - **attn_metadata.asdict_zerocopy())

238 -

239 - device = self.runner.device

240 - attn_metadata.orig_seq_lens = self.orig_seq_lens

241 - attn_metadata.orig_seq_lens_tensor = async_tensor_h2d(

242 - self.orig_seq_lens, torch.int, device, self.runner.pin_memory)

243 -

244 - attn_metadata.block_size = self.runner.block_size

245 - dual_chunk_attn_config = getattr(self.runner.model_config.hf_config,

246 - "dual_chunk_attention_config", {})

247 - attn_metadata.original_max_position_embeddings = \

248 - dual_chunk_attn_config.get("original_max_position_embeddings", 0)

249 - attn_metadata.chunk_size = dual_chunk_attn_config.get(

250 - "chunk_size", 8192)

251 - attn_metadata.local_size = dual_chunk_attn_config.get(

252 - "local_size", 1024)

253 -

254 - return attn_metadata

255 -

256 -

257 -class DualChunkFlashAttentionImpl(FlashAttentionImpl):

258 - """

259 - If the input tensors contain prompt tokens, the layout is as follows:

260 - |<--------------- num_prefill_tokens ----------------->|

262 - Otherwise, the layout is as follows:

263 - |<----------------- num_decode_tokens ------------------>|

265 - Generation tokens can contain padding when cuda-graph is used.

266 - Currently, prompt tokens don't contain any padding.

267 - The prompts might have different lengths, while the generation tokens

268 - always have length 1.

269 - If chunked prefill is enabled, prefill tokens and decode tokens can be

270 - batched together in a flattened 1D query.

271 - |<----- num_prefill_tokens ---->|<------- num_decode_tokens --------->|

272 - |<-prefill_0->|...|<-prefill_N-1->|<--decode_0-->|...|<--decode_M-1-->|

273 - Currently, cuda graph is disabled for chunked prefill, meaning there's no

274 - padding between prefill and decode tokens.

275 - """

276 -

277 - def __init__(

278 - self,

279 - num_heads: int,

280 - head_size: int,

281 - scale: float,

282 - num_kv_heads: int,

283 - alibi_slopes: Optional[List[float]],

284 - sliding_window: Optional[int],

285 - kv_cache_dtype: str,

286 - logits_soft_cap: Optional[float] = None,

287 - attn_type: str = AttentionType.DECODER,

288 - kv_sharing_target_layer_name: Optional[str] = None,

289 - layer_idx: int = -1,

290 - dual_chunk_attention_config: Optional[Dict[str, Any]] = None,

291 - ) -> None:

292 - if kv_sharing_target_layer_name is not None:

293 - raise NotImplementedError("KV sharing is not supported in V0 "

294 - "DUAL_CHUNK_FLASH_ATTN backend.")

295 - self.num_heads = num_heads

296 - self.head_size = head_size

297 - self.scale = float(scale)

298 - self.num_kv_heads = num_kv_heads

299 - if alibi_slopes is not None:

300 - alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)

301 - self.alibi_slopes = alibi_slopes

302 - self.sliding_window = ((sliding_window, sliding_window)

303 - if sliding_window is not None else (-1, -1))

304 - self.kv_cache_dtype = kv_cache_dtype

305 -

306 - self.num_queries_per_kv = self.num_heads // self.num_kv_heads

307 - if sliding_window is not None:

308 - # NOTE(woosuk): flash-attn's sliding window does not work with

309 - # paged KV cache.

310 - raise ValueError(

311 - "Sliding window is not supported in FlashAttention.")

312 -

313 - support_head_sizes = (

314 - DualChunkFlashAttentionBackend.get_supported_head_sizes())

315 -

316 - if head_size not in support_head_sizes:

317 - raise ValueError(

318 - f"Head size {head_size} is not supported by FlashAttention. "

319 - f"Supported head sizes are: {support_head_sizes}.")

320 -

321 - assert dual_chunk_attention_config is not None

322 - self.chunk_size = dual_chunk_attention_config.get("chunk_size", 8192)

323 - self.local_size = dual_chunk_attention_config.get("local_size", 1024)

324 - self.original_max_position_embeddings = dual_chunk_attention_config.get(

325 - "original_max_position_embeddings", 0)

326 - self.sparse_attention_config = dual_chunk_attention_config.get(

327 - "sparse_attention_config", None)

328 - if not self.sparse_attention_config:

329 - logger.warning_once("Sparse attention will not be enabled as "

330 - "sparse attention config is not provided.")

331 - self.sparse_attention_enabled = dual_chunk_attention_config.get(

332 - "sparse_attention_enabled", self.sparse_attention_config

333 - is not None)

334 - self.sparse_attention_threshold = dual_chunk_attention_config.get(

335 - "sparse_attention_threshold", 32768)

336 - self.sparse_attention_last_q = dual_chunk_attention_config.get(

337 - "sparse_attention_last_q", 64)

338 - self.layer_idx = layer_idx

339 - self.dual_chunk_attention_config = dual_chunk_attention_config

340 -

341 - if self.sparse_attention_config:

342 - self.sparse_attention_config = {

343 - int(i): j

344 - for i, j in self.sparse_attention_config[

345 - self.layer_idx].items()

346 - }

347 - start_head = self.num_heads * get_tensor_model_parallel_rank()

348 - end_head = start_head + self.num_heads

349 - self.sparse_attention_config = [

350 - self.sparse_attention_config[i]

351 - for i in range(start_head, end_head)

352 - ]

353 -

354 - if self.sparse_attention_enabled:

355 - self.arange = torch.arange(self.sparse_attention_last_q,

356 - device="cuda")

357 - self.last_q_mask = (self.arange[None, None, :, None]

358 - >= self.arange[None, None, None, :])

359 -

360 - def forward( # type: ignore

361 - self,

362 - layer: AttentionLayer,

363 - query: torch.Tensor,

364 - key: torch.Tensor,

365 - value: torch.Tensor,

366 - kv_cache: torch.Tensor,

367 - attn_metadata: DualChunkFlashAttentionMetadata,

368 - output: Optional[torch.Tensor] = None,

369 - output_scale: Optional[torch.Tensor] = None,

370 - output_block_scale: Optional[torch.Tensor] = None,

371 - ) -> torch.Tensor:

372 - """Forward pass with DualChunkFlashAttention.

373 - Args:

374 - query: shape = [num_tokens, num_heads * head_size]

375 - query_succ: shape = [num_tokens, num_heads * head_size]

376 - query_inter: shape = [num_tokens, num_heads * head_size]

377 - key: shape = [num_tokens, num_kv_heads * head_size]

378 - value: shape = [num_tokens, num_kv_heads * head_size]

379 - kv_cache = [2, num_blocks, block_size, num_kv_heads * head_size]

380 - attn_metadata: Metadata for attention.

381 - Returns:

382 - shape = [num_tokens, num_heads * head_size]

383 - """

384 - assert output is None, "Output tensor not supported for DualChunk"

385 -

386 - if output_scale is not None or output_block_scale is not None:

387 - raise NotImplementedError(

388 - "fused output quantization is not yet supported"

389 - " for FlashAttentionImpl")

390 -

391 - (

392 - query,

393 - query_succ,

394 - query_inter,

395 - query_succ_critical,

396 - query_inter_critical,

397 - ) = torch.split(query, query.shape[-1] // 5, dim=-1)

398 -

399 - assert (

400 - query_succ is not None and query_inter is not None

401 - ), "query_succ and query_inter are required in Dual Chunk Attention."

402 -

403 - num_tokens, hidden_size = query.shape

404 -

405 - # Reshape the query, key, and value tensors.

406 - query = query.view(-1, self.num_heads, self.head_size)

407 - query_succ = query_succ.view(-1, self.num_heads, self.head_size)

408 - query_inter = query_inter.view(-1, self.num_heads, self.head_size)

409 - query_succ_critical = query_succ_critical.view(-1, self.num_heads,

410 - self.head_size)

411 - query_inter_critical = query_inter_critical.view(

412 - -1, self.num_heads, self.head_size)

413 - key = key.view(-1, self.num_kv_heads, self.head_size)

414 - value = value.view(-1, self.num_kv_heads, self.head_size)

415 -

416 - if self.original_max_position_embeddings > 0:

417 - if prefill_meta := attn_metadata.prefill_metadata:

418 - assert prefill_meta.scaling_factor is not None

419 - assert prefill_meta.query_start_loc is not None

420 - assert prefill_meta.orig_seq_lens is not None

421 - current_start = 0

422 - query_start_loc_cpu = prefill_meta.query_start_loc.cpu()

423 - for i in range(len(prefill_meta.orig_seq_lens)):

424 - current_end = (current_start +

425 - (query_start_loc_cpu[i + 1] -

426 - query_start_loc_cpu[i]).item())

427 - key[current_start:current_end].mul_(

428 - prefill_meta.scaling_factor[i])

429 - current_start = current_end

430 - assert current_end <= attn_metadata.num_prefill_tokens

431 - if decode_meta := attn_metadata.decode_metadata:

432 - assert decode_meta.scaling_factor is not None

433 - scaling_factor = decode_meta.scaling_factor

434 - key[attn_metadata.num_prefill_tokens:].mul_(

435 - scaling_factor.unsqueeze(-1).unsqueeze(-1))

436 -

437 - if kv_cache is not None and kv_cache.numel() > 0:

438 - key_cache = kv_cache[0]

439 - value_cache = kv_cache[1]

440 -

441 - # Reshape the input keys and values and store them in the cache.

442 - # If kv_cache is not provided, the new key and value tensors are

443 - # not cached. This happens during the initial memory profiling run.

444 - ops.reshape_and_cache_flash(

445 - key,

446 - value,

447 - key_cache,

448 - value_cache,

449 - attn_metadata.slot_mapping.flatten(),

450 - self.kv_cache_dtype,

451 - layer._k_scale,

452 - layer._v_scale,

453 - )

454 -

455 - num_prefill_tokens = attn_metadata.num_prefill_tokens

456 - num_decode_tokens = attn_metadata.num_decode_tokens

457 - assert key.shape[0] == num_prefill_tokens + num_decode_tokens

458 - assert value.shape[0] == num_prefill_tokens + num_decode_tokens

459 - output = torch.empty_like(query)

460 -

461 - # Query for decode. KV is not needed because it is already cached.

462 - decode_query = query[num_prefill_tokens:]

463 - decode_query_succ = query_succ[num_prefill_tokens:]

464 - decode_query_inter = query_inter[num_prefill_tokens:]

465 -

466 - # QKV for prefill.

467 - query = query[:num_prefill_tokens]

468 - query_succ = query_succ[:num_prefill_tokens]

469 - query_inter = query_inter[:num_prefill_tokens]

470 - query_succ_critical = query_succ_critical[:num_prefill_tokens]

471 - query_inter_critical = query_inter_critical[:num_prefill_tokens]

472 - key = key[:num_prefill_tokens]

473 - value = value[:num_prefill_tokens]

474 - assert query.shape[0] == num_prefill_tokens

475 - assert decode_query.shape[0] == num_decode_tokens

476 -

477 - if prefill_meta := attn_metadata.prefill_metadata:

478 - # Prompt run.

479 - if (kv_cache is None or prefill_meta.block_tables is None

480 - or prefill_meta.block_tables.numel() == 0):

481 - # normal attention, called during the profiling run.

482 - out = flash_attn_varlen_func(

483 - q=query,

484 - k=key,

485 - v=value,

486 - cu_seqlens_q=prefill_meta.seq_start_loc,

487 - cu_seqlens_k=prefill_meta.seq_start_loc,

488 - max_seqlen_q=prefill_meta.max_prefill_seq_len,

489 - max_seqlen_k=prefill_meta.max_prefill_seq_len,

490 - softmax_scale=self.scale,

491 - causal=True,

492 - window_size=self.sliding_window,

493 - alibi_slopes=self.alibi_slopes,

494 - )

495 - assert output[:num_prefill_tokens].shape == out.shape

496 - output[:num_prefill_tokens] = out

497 - else:

498 - # prefix-enabled attention

499 - assert prefill_meta.seq_lens is not None

500 - assert prefill_meta.orig_seq_lens is not None

501 - output[:num_prefill_tokens] = (

502 - self._dual_chunk_flash_attn_prefill(

503 - q=query,

504 - q_succ=query_succ,

505 - q_inter=query_inter,

506 - q_succ_critical=query_succ_critical,

507 - q_inter_critical=query_inter_critical,

508 - k=key_cache,

509 - v=value_cache,

510 - cu_seqlens_q=prefill_meta.query_start_loc,

511 - cu_seqlens_k=prefill_meta.seq_start_loc,

512 - orig_seq_lens=prefill_meta.orig_seq_lens,

513 - scaling_factor=prefill_meta.scaling_factor,

514 - softmax_scale=self.scale,

515 - causal=True,

516 - window_size=(-1, -1),

517 - alibi_slopes=self.alibi_slopes,

518 - block_table=prefill_meta.block_tables,

519 - chunk_size=self.chunk_size,

520 - local_size=self.local_size,

521 - ))

522 -

523 - if decode_meta := attn_metadata.decode_metadata:

524 - # Decoding run.

525 - output[num_prefill_tokens:] = (

526 - self._dual_chunk_flash_attn_decoding(

527 - decode_query.unsqueeze(1),

528 - decode_query_succ.unsqueeze(1),

529 - decode_query_inter.unsqueeze(1),

530 - key_cache,

531 - value_cache,

532 - block_table=decode_meta.block_tables,

533 - cache_seqlens=decode_meta.seq_lens_tensor,

534 - softmax_scale=self.scale,

535 - causal=True,

536 - alibi_slopes=self.alibi_slopes,

537 - chunk_size=self.chunk_size,

538 - local_size=self.local_size,

539 - original_max_position_embeddings=self.

540 - original_max_position_embeddings,

541 - decode_meta=decode_meta,

542 - ).squeeze(1))

543 - # Reshape the output tensor.

544 - return output.view(num_tokens, hidden_size)

545 -

546 - def _dual_chunk_flash_attn_prefill(

547 - self,

548 - q,

549 - q_succ,

550 - q_inter,

551 - q_succ_critical,

552 - q_inter_critical,

553 - k,

554 - v,

555 - cu_seqlens_q,

556 - cu_seqlens_k,

557 - orig_seq_lens: List[int],

558 - scaling_factor: torch.Tensor,

559 - softmax_scale: float,

560 - causal: Optional[bool] = True,

561 - window_size: Tuple[int, int] = (-1, -1),

562 - alibi_slopes: Optional[torch.Tensor] = None,

563 - block_table: Optional[torch.Tensor] = None,

564 - chunk_size: int = 8192,

565 - local_size: int = 1024,

566 - ):

567 - if alibi_slopes is not None:

568 - raise ValueError(

569 - "Dual Chunk Attention does not support alibi_slopes")

570 - if not causal:

571 - raise ValueError(

572 - "Dual Chunk Attention does not support causal=False")

573 - if window_size != (-1, -1):

574 - raise ValueError(

575 - "Dual Chunk Attention does not support window_size")

576 -

577 - cu_seqlens_q_cpu = cu_seqlens_q.cpu().tolist()

578 - cu_seqlens_k_cpu = cu_seqlens_k.cpu().tolist()

579 - all_outputs = []

580 -

581 - for i in range(0, len(cu_seqlens_q_cpu) - 1):

582 - qs = cu_seqlens_q_cpu[i]

583 - qe = cu_seqlens_q_cpu[i:i + 2][-1]

584 - ks = cu_seqlens_k_cpu[i]

585 - ke = cu_seqlens_k_cpu[i:i + 2][-1]

586 -

587 - current_q = q[qs:qe]

588 - current_q_succ = q_succ[qs:qe]

589 - current_q_inter = q_inter[qs:qe]

590 - current_q_succ_critical = q_succ_critical[qs:qe]

591 - current_q_inter_critical = q_inter_critical[qs:qe]

592 -

593 - if block_table is None:

594 - current_k = k[ks:ke]

595 - current_v = v[ks:ke]

596 - current_block_table = None

597 - current_orig_seq_len = orig_seq_lens[i]

598 - else:

599 - current_block_table = block_table[i]

600 - current_orig_seq_len = orig_seq_lens[i]

601 - current_k = k

602 - current_v = v

603 - sparse_attn_enabled = (self.sparse_attention_enabled

604 - and current_orig_seq_len

605 - > self.sparse_attention_threshold)

606 -

607 - if current_q.shape[0] == 0:

608 - continue

609 -

610 - if current_k.shape[0] == 0:

611 - all_outputs.append(

612 - torch.zeros(

613 - (current_q.shape[0], current_q.shape[1], v.shape[2]),

614 - device=q.device,

615 - dtype=q.dtype,

616 - ))

617 - continue

618 -

619 - current_output = torch.empty_like(current_q)

620 - group_size = int(current_q.size(-2) / current_k.size(-2))

621 -

622 - if sparse_attn_enabled:

623 - num_device_q_heads = current_q.size(-2)

624 - heads_vertical_size = torch.empty(size=(num_device_q_heads, ),

625 - dtype=torch.int32)

626 - heads_slash_size = torch.empty(size=(num_device_q_heads, ),

627 - dtype=torch.int32)

628 - for head_id in range(current_q.size(-2)):

629 - (

630 - ty,

631 - vertical_size,

632 - slash_size,

633 - _,

634 - ) = self.sparse_attention_config[head_id]

635 - assert ty == "vertical_and_slash", "only support slash mode"

636 -

637 - if vertical_size == 30:

638 - vertical_size += 100

639 - heads_vertical_size[head_id] = vertical_size

640 - heads_slash_size[head_id] = slash_size

641 -

642 - current_output = self._dual_chunk_flash_attn_prefill_func(

643 - current_q, # allheads

644 - current_q_succ,

645 - current_q_inter,

646 - current_q_succ_critical,

647 - current_q_inter_critical,

648 - current_k,

649 - current_v,

650 - current_block_table,

651 - softmax_scale,

652 - chunk_size,

653 - local_size,

654 - scaling_factor[i].item(),

655 - ke - ks,

656 - sparse_attn_enabled=sparse_attn_enabled,

657 - heads_vertical_size=heads_vertical_size,

658 - heads_slash_size=heads_slash_size,

659 - group_size=group_size)

660 - else:

661 - for head_id in range(current_q.size(-2)):

662 - # (seq_len, num_heads, head_size)

663 - current_q_head = current_q[:, head_id, :].unsqueeze(1)

664 - current_q_succ_head = \

665 - current_q_succ[:, head_id, :].unsqueeze(1)

666 - current_q_inter_head = \

667 - current_q_inter[:, head_id, :].unsqueeze(1)

668 - current_q_succ_head_critical = \

669 - current_q_succ_critical[:, head_id, :].unsqueeze(1)

670 - current_q_inter_head_critical = \

671 - current_q_inter_critical[:, head_id, :].unsqueeze(1)

672 - if block_table is not None:

673 - current_k_head = current_k[..., head_id //

674 - group_size, :].unsqueeze(2)

675 - current_v_head = current_v[..., head_id //

676 - group_size, :].unsqueeze(2)

677 -

678 - else:

679 - current_k_head = current_k[:, head_id, :].unsqueeze(1)

680 - current_v_head = current_v[:, head_id, :].unsqueeze(1)

681 -

682 - current_out = self._dual_chunk_flash_attn_prefill_func(

683 - current_q_head,

684 - current_q_succ_head,

685 - current_q_inter_head,

686 - current_q_succ_head_critical,

687 - current_q_inter_head_critical,

688 - current_k_head,

689 - current_v_head,

690 - current_block_table,

691 - softmax_scale,

692 - chunk_size,

693 - local_size,

694 - scaling_factor[i].item(),

695 - ke - ks,

696 - sparse_attn_enabled=sparse_attn_enabled,

697 - )

698 - current_output[:, head_id:head_id + 1, :] = current_out

699 - all_outputs.append(current_output)

700 - return torch.cat(all_outputs, dim=0)

701 -

702 - def _dual_chunk_flash_attn_prefill_func(

703 - self,

704 - q,

705 - q_succ,

706 - q_inter,

707 - q_succ_critical,

708 - q_inter_critical,

709 - k,

710 - v,

711 - block_table,

712 - softmax_scale: float,

713 - chunk_size: int,

714 - local_size: int,

715 - scaling_factor: float,

716 - k_length: int,

717 - sparse_attn_enabled: Optional[bool] = True,

718 - heads_vertical_size=None,

719 - heads_slash_size=None,

720 - group_size=None,

721 - ):

722 - flash_results = []

723 - chunk_len = chunk_size - local_size

724 -

725 - if block_table is not None:

726 - block_size = v.shape[1]

727 - if chunk_len % block_size != 0:

728 - raise ValueError("chunk_len must be divisible by block_size.")

729 - else:

730 - block_size = 1

731 -

732 - if self.original_max_position_embeddings > 0:

733 - softmax_scale = softmax_scale * scaling_factor

734 -

735 - begin = k_length - q.shape[0]

736 - while begin < k_length:

737 - flash_per_chunk = []

738 -

739 - prev_chunk_end_pos = (begin // chunk_len) * chunk_len

740 - next_chunk_end_pos = prev_chunk_end_pos + chunk_len

741 - end = min(next_chunk_end_pos, k_length)

742 - qbegin = begin - (k_length - q.shape[0])

743 - qend = end - (k_length - q.shape[0])

744 -

745 - qk_chunks = []

746 - q_states_intra = q[qbegin:qend]

747 - # choose critical token

748 - if block_table is not None:

749 - block_tables_intra = _get_block(block_table, block_size,

750 - prev_chunk_end_pos, end)

751 - k_states_intra = k[block_tables_intra].view(

752 - -1, *k.shape[-2:])[:(end - prev_chunk_end_pos)]

753 - v_states_intra = v[block_tables_intra].view(

754 - -1, *v.shape[-2:])[:(end - prev_chunk_end_pos)]

755 - else:

756 - block_tables_intra = None

757 - k_states_intra = k[prev_chunk_end_pos:end]

758 - v_states_intra = v[prev_chunk_end_pos:end]

759 -

760 - if sparse_attn_enabled:

761 - last_q_size = min(qend - qbegin, self.sparse_attention_last_q)

762 - _, num_device_k_heads, head_dim = k_states_intra.shape

763 - k_states_intra = (k_states_intra.unsqueeze(2).repeat(

764 - 1, 1, group_size,

765 - 1).reshape(-1, num_device_k_heads * group_size, head_dim))

766 - v_states_intra = (v_states_intra.unsqueeze(2).repeat(

767 - 1, 1, group_size,

768 - 1).reshape(-1, num_device_k_heads * group_size, head_dim))

769 - qk_chunks.append(

770 - (q_states_intra.transpose(0, 1)[:, -last_q_size:] *

771 - softmax_scale) @ k_states_intra.permute(1, 2, 0))

772 -

773 - if prev_chunk_end_pos - chunk_len >= 0:

774 - q_states_succ = q_succ[qbegin:qend]

775 - q_states_succ_critical = q_succ_critical[qbegin:qend]

776 - if block_table is not None:

777 - block_tables_succ = _get_block(

778 - block_table, block_size,

779 - prev_chunk_end_pos - chunk_len, prev_chunk_end_pos)

780 - k_states_succ = k[block_tables_succ].view(

781 - -1, *k.shape[-2:])[:chunk_len]

782 - v_states_succ = v[block_tables_succ].view(

783 - -1, *v.shape[-2:])[:chunk_len]

784 - else:

785 - k_states_succ = k[prev_chunk_end_pos -

786 - chunk_len:prev_chunk_end_pos]

787 - v_states_succ = v[prev_chunk_end_pos -

788 - chunk_len:prev_chunk_end_pos]

789 -

790 - if sparse_attn_enabled:

791 - k_states_succ = (k_states_succ.unsqueeze(2).repeat(

792 - 1, 1, group_size,

793 - 1).reshape(-1, num_device_k_heads * group_size,

794 - head_dim))

795 - v_states_succ = (v_states_succ.unsqueeze(2).repeat(

796 - 1, 1, group_size,

797 - 1).reshape(-1, num_device_k_heads * group_size,

798 - head_dim))

799 - qk_chunks.append((q_states_succ_critical.transpose(

800 - 0, 1)[:, -last_q_size:] * softmax_scale)

801 - @ k_states_succ.permute(1, 2, 0))

802 -

803 - if prev_chunk_end_pos - chunk_len * 2 >= 0:

804 - q_states_inter = q_inter[qbegin:qend]

805 - q_states_inter_critical = q_inter_critical[qbegin:qend]

806 - if block_table is not None:

807 - block_tables_inter = _get_block(

808 - block_table, block_size, 0,

809 - prev_chunk_end_pos - chunk_len)

810 - k_states_inter = k[block_tables_inter].view(

811 - -1, *k.shape[-2:])[:(prev_chunk_end_pos - chunk_len)]

812 - v_states_inter = v[block_tables_inter].view(

813 - -1, *v.shape[-2:])[:(prev_chunk_end_pos - chunk_len)]

814 - else:

815 - k_states_inter = k[:prev_chunk_end_pos - chunk_len]

816 - v_states_inter = v[:prev_chunk_end_pos - chunk_len]

817 -

818 - if sparse_attn_enabled:

819 - k_states_inter = (k_states_inter.unsqueeze(2).repeat(

820 - 1, 1, group_size,

821 - 1).reshape(-1, num_device_k_heads * group_size,

822 - head_dim))

823 - v_states_inter = (v_states_inter.unsqueeze(2).repeat(

824 - 1, 1, group_size,

825 - 1).reshape(-1, num_device_k_heads * group_size,

826 - head_dim))

827 - qk_chunks.append((q_states_inter_critical.transpose(

828 - 0, 1)[:, -last_q_size:] * softmax_scale)

829 - @ k_states_inter.permute(1, 2, 0))

830 -

831 - if sparse_attn_enabled:

832 - reversed_qk = qk_chunks[::-1]

833 - qk = torch.cat(reversed_qk, dim=-1)

834 -

835 - qk[:, :, -last_q_size:] = torch.where(

836 - self.last_q_mask[..., -last_q_size:,

837 - -last_q_size:].to(qk.device),

838 - qk[:, :, -last_q_size:], -torch.inf)

839 - qk = F.softmax(qk, dim=-1, dtype=torch.float32)

840 -

841 - vertical = qk.sum(-2, keepdim=True)

842 - vertical[..., :30] = torch.inf

843 -

844 - # Avoid sorting by using the min/max ints to fill the indexer

845 - # buffers.

846 - int32_max = torch.iinfo(torch.int32).max

847 - int32_min = torch.iinfo(torch.int32).min

848 - n_heads = qk.size()[0]

849 - max_slash_topk = torch.max(heads_slash_size).item()

850 - max_vertical_topk = torch.max(heads_vertical_size).item()

851 - # store each head's slash topk, vertical topk

852 - vertical = vertical.reshape((n_heads, -1))

853 - # prevent out of range when prompt size < max_vertical_topk

854 - max_vertical_topk = min(vertical.shape[-1], max_vertical_topk)

855 - vertical_topk_buffer = torch.topk(vertical, max_vertical_topk,

856 - -1).indices

857 - slash_topk_buffer = torch.empty(size=(n_heads, max_slash_topk),

858 - dtype=torch.int64,

859 - device=qk.device)

860 - for head_i in range(n_heads):

861 - # (nqheads=1, lastq, k_len)

862 - head_score = qk[head_i:head_i + 1, :, :]

863 - slash_scores = _sum_all_diagonal_matrix(head_score)

864 - if head_score.size(1) != 1:

865 - # drop right up corner

866 - slash_scores = slash_scores[..., :-last_q_size + 1]

867 - slash_scores[..., -100:] = torch.inf

868 -

869 - head_slash_size = heads_slash_size[head_i]

870 - head_slash_size = min(head_slash_size, vertical.size(-1))

871 - slash_topk = torch.topk(slash_scores, head_slash_size,

872 - -1).indices

873 - #（nheads, max_topk）

874 - slash_topk_buffer[head_i, :head_slash_size] = slash_topk

875 -

876 - # reset heads topk

877 - heads_slash_size[head_i] = head_slash_size

878 - heads_vertical_size[head_i] = min(

879 - heads_vertical_size[head_i], max_vertical_topk)

880 -

881 - # store

882 - vertical_buffer = torch.full((n_heads, max_vertical_topk),

883 - int32_max,

884 - dtype=torch.int64,

885 - device=q.device)

886 - slash_buffer = torch.full((n_heads, max_slash_topk),

887 - int32_min,

888 - dtype=torch.int64,

889 - device=q.device)

890 - succ_vertical_buffer = torch.full((n_heads, max_vertical_topk),

891 - int32_max,

892 - dtype=torch.int64,

893 - device=q.device)

894 - succ_slash_buffer = torch.full((n_heads, max_slash_topk),

895 - int32_min,

896 - dtype=torch.int64,

897 - device=q.device)

898 - inter_vertical_buffer = torch.full(

899 - (n_heads, max_vertical_topk),

900 - int32_max,

901 - dtype=torch.int64,

902 - device=q.device)

903 - inter_slash_buffer = torch.full((n_heads, max_slash_topk),

904 - int32_min,

905 - dtype=torch.int64,

906 - device=q.device)

907 -

908 - vertical_size_buffer = torch.empty(size=(n_heads, ),

909 - dtype=torch.int32,

910 - device=q.device)

911 - slash_sizes_buffer = torch.empty(size=(n_heads, ),

912 - dtype=torch.int32,

913 - device=q.device)

914 - succ_vertical_size_buffer = torch.empty(size=(n_heads, ),

915 - dtype=torch.int32,

916 - device=q.device)

917 - succ_slash_sizes_buffer = torch.empty(size=(n_heads, ),

918 - dtype=torch.int32,

919 - device=q.device)

920 - inter_vertical_size_buffer = torch.empty(size=(n_heads, ),

921 - dtype=torch.int32,

922 - device=q.device)

923 - inter_slash_sizes_buffer = torch.empty(size=(n_heads, ),

924 - dtype=torch.int32,

925 - device=q.device)

926 -

927 - for head_i in range(n_heads):

928 - vertical_topk = vertical_topk_buffer[

929 - head_i, :heads_vertical_size[head_i]]

930 - # intra

931 - intra_vertical_indices = vertical_topk[

932 - vertical_topk >=

933 - prev_chunk_end_pos] - prev_chunk_end_pos

934 - if intra_vertical_indices.nelement() == 0:

935 - intra_vertical_indices = torch.cat([

936 - intra_vertical_indices,

937 - torch.arange(0,

938 - k_states_intra.size(0),

939 - max(1,

940 - k_states_intra.size(0) / 5),

941 - dtype=torch.int32,

942 - device=intra_vertical_indices.device)

943 - ])

944 - slash_topk = slash_topk_buffer[

945 - head_i, :heads_slash_size[head_i]]

946 - intra_slash_indices = (

947 - (qk.size(-1) - 1) -

948 - slash_topk[slash_topk >= prev_chunk_end_pos])

949 - # fill buffer

950 - v_count = intra_vertical_indices.nelement()

951 - s_count = intra_slash_indices.nelement()

952 - vertical_size_buffer[head_i] = v_count

953 - slash_sizes_buffer[head_i] = s_count

954 - vertical_buffer[head_i, :v_count].copy_(

955 - intra_vertical_indices)

956 - slash_buffer[head_i, :s_count].copy_(intra_slash_indices)

957 - # succ

958 - if prev_chunk_end_pos - chunk_len >= 0:

959 - succ_vertical_indices = vertical_topk[

960 - (vertical_topk < prev_chunk_end_pos)

961 - & (vertical_topk >= prev_chunk_end_pos -

962 - chunk_len)] - (prev_chunk_end_pos - chunk_len)

963 - # TODO: support no vertical

964 - if succ_vertical_indices.nelement() == 0:

965 - succ_vertical_indices = torch.cat([

966 - succ_vertical_indices,

967 - torch.arange(

968 - 0,

969 - k_states_succ.size(0),

970 - max(1,

971 - k_states_succ.size(0) / 5),

972 - dtype=torch.int32,

973 - device=intra_vertical_indices.device)

974 - ])

975 - succ_slash_indices = (

976 - (prev_chunk_end_pos + (qend - qbegin) - 1) -

977 - slash_topk[((slash_topk >=

978 - (prev_chunk_end_pos - chunk_len)) &

979 - (slash_topk < (prev_chunk_end_pos +

980 - (qend - qbegin))))])

981 - if succ_slash_indices.nelement() == 0:

982 - succ_slash_indices = torch.cat([

983 - succ_slash_indices,

984 - torch.arange(

985 - 0,

986 - k_states_succ.size(0),

987 - max(1,

988 - k_states_succ.size(0) / 5),

989 - dtype=torch.int32,

990 - device=intra_vertical_indices.device)

991 - ])

992 - # fill buffer

993 - v_count = succ_vertical_indices.nelement()

994 - s_count = succ_slash_indices.nelement()

995 - succ_vertical_size_buffer[head_i] = v_count

996 - succ_slash_sizes_buffer[head_i] = s_count

997 - succ_vertical_buffer[head_i, :v_count].copy_(

998 - succ_vertical_indices)

999 - succ_slash_buffer[head_i, :s_count].copy_(

1000 - succ_slash_indices)

1001 -

1002 - if prev_chunk_end_pos - 2 * chunk_len >= 0:

1003 - inter_vertical_indices = vertical_topk[

1004 - vertical_topk < prev_chunk_end_pos - chunk_len]

1005 -

1006 - if inter_vertical_indices.nelement() == 0:

1007 - inter_vertical_indices = torch.cat([

1008 - inter_vertical_indices,

1009 - torch.arange(

1010 - 0,

1011 - k_states_inter.size(0),

1012 - max(1,

1013 - k_states_inter.size(0) / 5),

1014 - dtype=torch.int32,

1015 - device=intra_vertical_indices.device)

1016 - ])

1017 - inter_slash_indices = (

1018 - (prev_chunk_end_pos - chunk_len +

1019 - (qend - qbegin) - 1) -

1020 - slash_topk[slash_topk < (prev_chunk_end_pos -

1021 - chunk_len +

1022 - (qend - qbegin))])

1023 - if inter_slash_indices.nelement() == 0:

1024 - inter_slash_indices = torch.cat([

1025 - inter_slash_indices,

1026 - torch.arange(

1027 - 0,

1028 - k_states_inter.size(0),

1029 - max(1,

1030 - k_states_inter.size(0) / 5),

1031 - dtype=torch.int32,

1032 - device=intra_vertical_indices.device)

1033 - ])

1034 - # fill buffer

1035 - v_count = inter_vertical_indices.nelement()

1036 - s_count = inter_slash_indices.nelement()

1037 - inter_vertical_size_buffer[head_i] = v_count

1038 - inter_slash_sizes_buffer[head_i] = s_count

1039 - inter_vertical_buffer[head_i, :v_count].copy_(

1040 - inter_vertical_indices)

1041 - inter_slash_buffer[head_i, :s_count].copy_(

1042 - inter_slash_indices)

1043 - else:

1044 - intra_vertical_indices, intra_slash_indices = None, None

1045 - succ_vertical_indices, succ_slash_indices = None, None

1046 - inter_vertical_indices, inter_slash_indices = None, None

1047 -

1048 - if sparse_attn_enabled:

1049 - flash_result = self._do_flash_attn(

1050 - q_states_intra,

1051 - k_states_intra,

1052 - v_states_intra,

1053 - softmax_scale=softmax_scale,

1054 - causal=True,

1055 - stage="intra",

1056 - vertical_indices=vertical_buffer,

1057 - slash_indices=slash_buffer,

1058 - vertical_indices_count=vertical_size_buffer,

1059 - slash_indices_count=slash_sizes_buffer,

1060 - mergehead_softmax_scale=softmax_scale,

1061 - sparse_attn_enabled=sparse_attn_enabled)

1062 - else:

1063 - flash_result = self._do_flash_attn(

1064 - q_states_intra,

1065 - k_states_intra,

1066 - v_states_intra,

1067 - softmax_scale=softmax_scale,

1068 - causal=True,

1069 - stage="intra",

1070 - vertical_indices=intra_vertical_indices,

1071 - slash_indices=intra_slash_indices,

1072 - sparse_attn_enabled=sparse_attn_enabled)

1073 - flash_per_chunk.append(flash_result)

1074 -

1075 - if prev_chunk_end_pos - chunk_len >= 0:

1076 - if sparse_attn_enabled:

1077 - flash_result = self._do_flash_attn(

1078 - q_states_succ,

1079 - k_states_succ,

1080 - v_states_succ,

1081 - softmax_scale=softmax_scale,

1082 - causal=False,

1083 - stage="succ",

1084 - vertical_indices=succ_vertical_buffer,

1085 - slash_indices=succ_slash_buffer,

1086 - vertical_indices_count=succ_vertical_size_buffer,

1087 - slash_indices_count=succ_slash_sizes_buffer,

1088 - mergehead_softmax_scale=softmax_scale,

1089 - sparse_attn_enabled=sparse_attn_enabled)

1090 - else:

1091 - flash_result = self._do_flash_attn(

1092 - q_states_succ,

1093 - k_states_succ,

1094 - v_states_succ,

1095 - softmax_scale=softmax_scale,

1096 - causal=False,

1097 - stage="succ",

1098 - vertical_indices=succ_vertical_indices,

1099 - slash_indices=succ_slash_indices,

1100 - sparse_attn_enabled=sparse_attn_enabled)

1101 - flash_per_chunk.append(flash_result)

1102 -

1103 - if prev_chunk_end_pos - chunk_len * 2 >= 0:

1104 - if sparse_attn_enabled:

1105 - flash_result = self._do_flash_attn(

1106 - q_states_inter,

1107 - k_states_inter,

1108 - v_states_inter,

1109 - softmax_scale=softmax_scale,

1110 - causal=False,

1111 - stage="inter",

1112 - vertical_indices=inter_vertical_buffer,

1113 - slash_indices=inter_slash_buffer,

1114 - vertical_indices_count=inter_vertical_size_buffer,

1115 - slash_indices_count=inter_slash_sizes_buffer,

1116 - mergehead_softmax_scale=softmax_scale,

1117 - sparse_attn_enabled=sparse_attn_enabled)

1118 - else:

1119 - flash_result = self._do_flash_attn(

1120 - q_states_inter,

1121 - k_states_inter,

1122 - v_states_inter,

1123 - softmax_scale=softmax_scale,

1124 - causal=False,

1125 - stage="inter",

1126 - vertical_indices=inter_vertical_indices,

1127 - slash_indices=inter_slash_indices,

1128 - sparse_attn_enabled=sparse_attn_enabled)

1129 - flash_per_chunk.append(flash_result)

1130 -

1131 - flash_results.append(flash_per_chunk)

1132 - begin = end

1133 -

1134 - attn_output = self._merge_attn_outputs(flash_results)

1135 - del flash_results

1136 - return attn_output

1137 -

1138 - def _do_flash_attn(

1139 - self,

1140 - query_states: torch.Tensor,

1141 - key_states: torch.Tensor,

1142 - value_states: torch.Tensor,

1143 - softmax_scale: float,

1144 - causal: bool = True,

1145 - max_seqlen_k: Optional[int] = None,

1146 - stage: str = "intra",

1147 - vertical_indices: Optional[torch.Tensor] = None,

1148 - slash_indices: Optional[torch.Tensor] = None,

1149 - vertical_indices_count: Optional[torch.Tensor] = None,

1150 - slash_indices_count: Optional[torch.Tensor] = None,

1151 - mergehead_softmax_scale: Optional[float] = None,

1152 - sparse_attn_enabled: Optional[bool] = False,

1153 - ):

1154 - if max_seqlen_k is None:

1155 - max_seqlen_k = key_states.shape[0]

1156 -

1157 - q_len = query_states.shape[0]

1158 - q_heads = query_states.shape[1]

1159 - h_dim = query_states.shape[-1]

1160 -

1161 - if sparse_attn_enabled:

1162 - assert slash_indices is not None

1163 - if stage == "intra":

1164 - assert causal

1165 - else:

1166 - assert not causal

1167 -

1168 - query_states = query_states.unsqueeze(0).transpose(1, 2)

1169 - key_states = key_states.unsqueeze(0).transpose(1, 2)

1170 - value_states = value_states.unsqueeze(0).transpose(1, 2)

1171 -

1172 - q = query_states

1173 - k = key_states

1174 - v = value_states

1175 -

1176 - if (vertical_indices_count is not None and \

1177 - slash_indices_count is not None):

1178 - assert mergehead_softmax_scale is not None

1179 -

1180 - res, s_lse = _vertical_slash_sparse_attention(

1181 - q,

1182 - k,

1183 - v,

1184 - vertical_indices,

1185 - slash_indices,

1186 - mergehead_softmax_scale,

1187 - causal=causal,

1188 - stage=stage,

1189 - vertical_indices_count=vertical_indices_count,

1190 - slash_indices_count=slash_indices_count)

1191 - res = res.view(q_heads, q_len,

1192 - h_dim).transpose(0, 1) # (qlen,nhead,h_dim)

1193 - s_lse = s_lse.view(

1194 - q_heads, q_len,

1195 - 1).squeeze(-1).unsqueeze(0).float() # (1, nhead,qlen)

1196 - else:

1197 - res, s_lse = _vertical_slash_sparse_attention(q,

1198 - k,

1199 - v,

1200 - vertical_indices,

1201 - slash_indices,

1202 - softmax_scale,

1203 - causal=causal,

1204 - stage=stage)

1205 - res = res.view(q_len, q_heads, h_dim)

1206 - s_lse = s_lse.view(q_len, q_heads, 1).transpose(0, 2).float()

1207 - return res, s_lse

1208 -

1209 - output, softmax_lse = flash_attn_varlen_func(

1210 - q=query_states,

1211 - k=key_states,

1212 - v=value_states,

1213 - softmax_scale=softmax_scale,

1214 - cu_seqlens_q=torch.tensor([0, query_states.shape[0]],

1215 - dtype=torch.int32,

1216 - device=query_states.device),

1217 - max_seqlen_q=query_states.shape[0],

1218 - cu_seqlens_k=torch.tensor([0, max_seqlen_k],

1219 - dtype=torch.int32,

1220 - device=query_states.device),

1221 - max_seqlen_k=max_seqlen_k,

1222 - causal=causal,

1223 - return_softmax_lse=True,

1224 - )

1225 - softmax_lse = softmax_lse.view(q_len, q_heads, 1).transpose(0,

1226 - 2).float()

1227 - return output, softmax_lse

1228 -

1229 - def _merge_attn_outputs(

1230 - self,

1231 - flash_results: List[List[Tuple[torch.Tensor, torch.Tensor]]],

1232 - return_lse: Optional[bool] = False,

1233 - ) -> torch.Tensor:

1234 - attn_outputs_all = []

1235 - logits_all = []

1236 -

1237 - for flash_per_chunk in flash_results:

1238 - if len(flash_per_chunk) == 1:

1239 - attn_outputs_all.append(flash_per_chunk[0][0])

1240 - if return_lse:

1241 - logits_all.append(flash_per_chunk[0][1])

1242 - continue

1243 -

1244 - attn_outputs = torch.stack([

1245 - flash_attn_output[0] for flash_attn_output in flash_per_chunk

1246 - ])

1247 - logits = torch.stack([

1248 - flash_attn_output[1] for flash_attn_output in flash_per_chunk

1249 - ])

1250 - logits = logits.to(torch.float32)

1251 -

1252 - if return_lse:

1253 - max_val = torch.max(logits, dim=0).values

1254 - diff = torch.abs(logits[0] - logits[1])

1255 - log_sum_exp = max_val + torch.log1p(torch.exp(-diff))

1256 - logits_all.append(log_sum_exp)

1257 -

1258 - max_logits = torch.max(logits, dim=0).values

1259 - stable_logits = logits - max_logits.unsqueeze(0)

1260 - lse_s = torch.exp(stable_logits).detach()

1261 - lse_sum = torch.sum(lse_s, dim=0)

1262 - lse_s /= lse_sum

1263 - attn_outputs *= lse_s.unsqueeze(-1).transpose(2, 3).squeeze(1)

1264 - attn_outputs_all.append(attn_outputs.sum(dim=0))

1265 -

1266 - if return_lse:

1267 - return (torch.cat(attn_outputs_all,

1268 - dim=0), torch.cat(logits_all, dim=-1))

1269 - else:

1270 - return torch.cat(attn_outputs_all, dim=0)

1271 -

1272 - def _dual_chunk_flash_attn_decoding(

1273 - self,

1274 - query: torch.Tensor,

1275 - query_succ: torch.Tensor,

1276 - query_inter: torch.Tensor,

1277 - key_cache: torch.Tensor,

1278 - value_cache: torch.Tensor,

1279 - block_table: torch.Tensor,

1280 - cache_seqlens: torch.Tensor,

1281 - softmax_scale: float,

1282 - causal: bool,

1283 - alibi_slopes: Optional[torch.Tensor],

1284 - chunk_size: int,

1285 - local_size: int,

1286 - original_max_position_embeddings: int,

1287 - decode_meta: DualChunkFlashAttentionMetadata,

1288 - ):

1289 - if not causal:

1290 - raise ValueError(

1291 - "Dual Chunk Attention does not support causal=False")

1292 -

1293 - block_size = value_cache.shape[1]

1294 - chunk_len = chunk_size - local_size

1295 - if chunk_len % block_size != 0:

1296 - raise ValueError("chunk_len must be divisible by block_size.")

1297 - if original_max_position_embeddings > 0:

1298 - assert decode_meta.scaling_factor is not None

1299 - scaling_factor = decode_meta.scaling_factor

1300 - query = (query * scaling_factor.view(-1, 1, 1, 1)).to(

1301 - query.dtype

1302 - ) # possible for numerical issue, need to fused in the kernel

1303 - query_succ = (query_succ * scaling_factor.view(-1, 1, 1, 1)).to(

1304 - query.dtype)

1305 - query_inter = (query_inter * scaling_factor.view(-1, 1, 1, 1)).to(

1306 - query.dtype)

1307 - outputs_list = []

1308 - softmax_lses_list = []

1309 -

1310 - # intra-attention

1311 - intra_output, intra_softmax_lse = (

1312 - self._dual_chunk_flash_attn_decoding_with_exp_sums(

1313 - query,

1314 - key_cache,

1315 - value_cache,

1316 - decode_meta.block_tables_intra,

1317 - decode_meta.seq_lens_intra,

1318 - softmax_scale,

1319 - alibi_slopes,

1320 - causal=False,

1321 - ))

1322 - outputs_list.append(intra_output)

1323 - softmax_lses_list.append(intra_softmax_lse)

1324 -

1325 - # succ-attention

1326 - if decode_meta.max_seq_len_succ:

1327 - succ_output, succ_softmax_lse = (

1328 - self._dual_chunk_flash_attn_decoding_with_exp_sums(

1329 - query_succ,

1330 - key_cache,

1331 - value_cache,

1332 - decode_meta.block_tables_succ,

1333 - decode_meta.seq_lens_succ,

1334 - softmax_scale,

1335 - alibi_slopes,

1336 - causal=False,

1337 - ))

1338 - outputs_list.append(succ_output)

1339 - softmax_lses_list.append(succ_softmax_lse)

1340 -

1341 - # inter-attention

1342 - if decode_meta.max_seq_len_inter:

1343 - inter_output, inter_softmax_lse = (

1344 - self._dual_chunk_flash_attn_decoding_with_exp_sums(

1345 - query_inter,

1346 - key_cache,

1347 - value_cache,

1348 - block_table[:, :decode_meta.max_seq_len_inter],

1349 - decode_meta.seq_lens_inter,

1350 - softmax_scale,

1351 - alibi_slopes,

1352 - causal=False,

1353 - ))

1354 - outputs_list.append(inter_output)

1355 - softmax_lses_list.append(inter_softmax_lse)

1356 - outputs = torch.stack(outputs_list, dim=0)

1357 - del outputs_list

1358 - softmax_lses = torch.stack(softmax_lses_list, dim=0).to(torch.float32)

1359 - del softmax_lses_list

1360 - max_logits = torch.max(softmax_lses, dim=0).values

1361 - stable_logits = softmax_lses - max_logits.unsqueeze(0)

1362 - lse_s = torch.exp(stable_logits).detach()

1363 - lse_sum = torch.sum(lse_s, dim=0)

1364 - lse_s /= lse_sum

1365 - outputs *= lse_s.unsqueeze(-1).transpose(2, 3)

1366 - return outputs.sum(0)

1367 -

1368 - def _dual_chunk_flash_attn_decoding_with_exp_sums(

1369 - self,

1370 - query: torch.Tensor,

1371 - key_cache: torch.Tensor,

1372 - value_cache: torch.Tensor,

1373 - block_table: torch.Tensor,

1374 - cache_seqlens: torch.Tensor,

1375 - softmax_scale: float,

1376 - alibi_slopes: Optional[torch.Tensor],

1377 - causal: bool,

1378 - ):

1379 - out, softmax_lse = flash_attn_with_kvcache(

1380 - q=query,

1381 - k_cache=key_cache,

1382 - v_cache=value_cache,

1383 - block_table=block_table,

1384 - cache_seqlens=cache_seqlens,

1385 - softmax_scale=softmax_scale,

1386 - alibi_slopes=alibi_slopes,

1387 - causal=causal,

1388 - return_softmax_lse=True,

1389 - )

1390 - mask = (cache_seqlens == 0)

1391 - out[mask] = 0

1392 - softmax_lse[mask] = -float("inf")

1393 - return out, softmax_lse

1394 -

1395 -

1396 -def _vertical_slash_sparse_attention(

1397 - query: torch.Tensor, # [BATCH, N_HEADS, N_CTX, D_HEAD]

1398 - key: torch.Tensor, # [BATCH, N_HEADS, N_KV_CTX, D_HEAD]

1399 - value: torch.Tensor, # [BATCH, N_HEADS, N_KV_CTX, D_HEAD]

1400 - v_idx: torch.Tensor, # [BATCH, N_HEADS, NNZ_V]

1401 - s_idx: torch.Tensor, # [BATCH, N_HEADS, NNZ_S]

1402 - softmax_scale: float,

1403 - causal: bool = True,

1404 - stage: str = "intra",

1405 - block_size_M: int = 64,

1406 - block_size_N: int = 64,

1407 - vertical_indices_count: torch.Tensor = None, # [N_HEADS,]

1408 - slash_indices_count: torch.Tensor = None,

1409 -):

1410 - if stage == "intra":

1411 - assert causal

1412 - else:

1413 - assert not causal

1414 -

1415 - batch_size, num_heads, context_size, head_dim = query.shape

1416 - _, _, kv_seq_len, _ = key.shape

1417 -

1418 - if head_dim not in [16, 32, 64, 128, 256, 512]:

1419 - target_dim = 2**math.ceil(math.log2(head_dim)) - head_dim

1420 - query = F.pad(query, [0, target_dim, 0, 0, 0, 0, 0, 0])

1421 - key = F.pad(key, [0, target_dim, 0, 0, 0, 0, 0, 0])

1422 - value = F.pad(value, [0, target_dim, 0, 0, 0, 0, 0, 0])

1423 -

1424 - v_idx = v_idx.to(torch.int32).reshape(

1425 - (batch_size, num_heads, -1)).sort(dim=-1, descending=False)[0]

1426 - s_idx = s_idx.to(torch.int32).reshape(

1427 - (batch_size, num_heads, -1)).sort(dim=-1, descending=True)[0]

1428 - q_seqlens = torch.tensor([context_size],

1429 - dtype=torch.int32,

1430 - device=query.device)

1431 - kv_seqlens = torch.tensor([kv_seq_len],

1432 - dtype=torch.int32,

1433 - device=query.device)

1434 -

1435 - if vertical_indices_count is not None and slash_indices_count is not None:

1436 - (

1437 - block_count,

1438 - block_offset,

1439 - column_count,

1440 - column_index,

1441 - ) = ops.convert_vertical_slash_indexes_mergehead(

1442 - q_seqlens, kv_seqlens, v_idx, s_idx, vertical_indices_count,

1443 - slash_indices_count, context_size, block_size_M, block_size_N,

1444 - causal)

1445 - else:

1446 - (

1447 - block_count,

1448 - block_offset,

1449 - column_count,

1450 - column_index,

1451 - ) = ops.convert_vertical_slash_indexes(q_seqlens, kv_seqlens, v_idx,

1452 - s_idx, context_size,

1453 - block_size_M, block_size_N,

1454 - causal)

1455 -

1456 - q = query.transpose(1, 2).contiguous()

1457 - k = key.transpose(1, 2).contiguous()

1458 - v = value.transpose(1, 2).contiguous()

1459 - out, lse = sparse_attn_func(

1460 - q,

1461 - k,

1462 - v,

1463 - block_count,

1464 - block_offset,

1465 - column_count,

1466 - column_index,

1467 - causal=causal,

1468 - softmax_scale=softmax_scale,

1469 - return_softmax_lse=True,

1470 - )

1471 - out = out.transpose(1, 2).contiguous()

1472 - softmax_lse = lse.reshape(*lse.shape, 1)

1473 - return (out[..., :context_size, :head_dim],

1474 - softmax_lse[..., :context_size, :])

1475 -

1476 -

1477 -def _sum_all_diagonal_matrix(mat: torch.tensor):

1478 - h, n, m = mat.shape

1479 - # Zero matrix used for padding

1480 - zero_mat = torch.zeros((h, n, n), device=mat.device)

1481 - # pads the matrix on left and right

1482 - mat_padded = torch.cat((zero_mat, mat, zero_mat), -1)

1483 - # Change the strides

1484 - mat_strided = mat_padded.as_strided((1, n, n + m),

1485 - (n * (2 * n + m), 2 * n + m + 1, 1))

1486 - # Sums the resulting matrix's columns

1487 - sum_diags = torch.sum(mat_strided, 1)

1488 - return sum_diags[:, 1:] # drop left bottom corner

1489 -

1490 -

1491 -def _get_block(block_table: torch.Tensor, block_size: int, begin: int,

1492 - end: int):

1493 - begin_block = begin // block_size

1494 - end_block = (end - 1) // block_size + 1

1495 - return block_table[begin_block:end_block]

vllm/attention/backends/flash_attn.py deleted

-929

@@ -1,929 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -"""Attention layer with FlashAttention."""

4 -from collections import defaultdict

5 -from dataclasses import dataclass

6 -from itertools import accumulate

7 -from typing import Dict, List, Optional, Tuple, Type

8 -

9 -import torch

10 -

11 -from vllm import _custom_ops as ops

12 -# yapf conflicts with isort for this block

13 -# yapf: disable

14 -from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,

15 - AttentionLayer,

16 - AttentionMetadata,

17 - AttentionMetadataBuilder,

18 - AttentionType,

19 - is_quantized_kv_cache)

20 -# yapf: enable

21 -from vllm.attention.backends.utils import (

22 - PAD_SLOT_ID, CommonAttentionState, compute_slot_mapping,

23 - compute_slot_mapping_start_idx, get_num_prefill_decode_query_kv_tokens,

24 - get_seq_len_block_table_args, is_all_cross_attn_metadata_set,

25 - is_all_encoder_attn_metadata_set, is_block_tables_empty)

26 -from vllm.attention.utils.fa_utils import (flash_attn_supports_fp8,

27 - get_flash_attn_version)

28 -from vllm.logger import init_logger

29 -from vllm.multimodal import MultiModalPlaceholderMap

30 -from vllm.utils import async_tensor_h2d, make_tensor_with_pad

31 -from vllm.vllm_flash_attn import (flash_attn_varlen_func,

32 - flash_attn_with_kvcache)

33 -

34 -logger = init_logger(__name__)

35 -

36 -

37 -class FlashAttentionBackend(AttentionBackend):

38 -

39 - accept_output_buffer: bool = True

40 -

41 - @staticmethod

42 - def get_supported_head_sizes() -> List[int]:

43 - return [32, 64, 96, 128, 160, 192, 224, 256]

44 -

45 - @staticmethod

46 - def get_name() -> str:

47 - return "FLASH_ATTN"

48 -

49 - @staticmethod

50 - def get_impl_cls() -> Type["FlashAttentionImpl"]:

51 - return FlashAttentionImpl

52 -

53 - @staticmethod

54 - def get_metadata_cls() -> Type["AttentionMetadata"]:

55 - return FlashAttentionMetadata

56 -

57 - @staticmethod

58 - def get_builder_cls() -> Type["FlashAttentionMetadataBuilder"]:

59 - return FlashAttentionMetadataBuilder

60 -

61 - @staticmethod

62 - def get_state_cls() -> Type["CommonAttentionState"]:

63 - return CommonAttentionState

64 -

65 - @staticmethod

66 - def get_kv_cache_shape(

67 - num_blocks: int,

68 - block_size: int,

69 - num_kv_heads: int,

70 - head_size: int,

71 - ) -> Tuple[int, ...]:

72 - if block_size % 16 != 0:

73 - raise ValueError("Block size must be a multiple of 16.")

74 - return (2, num_blocks, block_size, num_kv_heads, head_size)

75 -

76 - @staticmethod

77 - def swap_blocks(

78 - src_kv_cache: torch.Tensor,

79 - dst_kv_cache: torch.Tensor,

80 - src_to_dst: torch.Tensor,

81 - ) -> None:

82 - src_key_cache = src_kv_cache[0]

83 - dst_key_cache = dst_kv_cache[0]

84 - ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dst)

85 - src_value_cache = src_kv_cache[1]

86 - dst_value_cache = dst_kv_cache[1]

87 - ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dst)

88 -

89 - @staticmethod

90 - def copy_blocks(

91 - kv_caches: List[torch.Tensor],

92 - src_to_dists: torch.Tensor,

93 - ) -> None:

94 - key_caches = [kv_cache[0] for kv_cache in kv_caches]

95 - value_caches = [kv_cache[1] for kv_cache in kv_caches]

96 -

97 - ops.copy_blocks(key_caches, value_caches, src_to_dists)

98 -

99 -

100 -@dataclass

101 -class FlashAttentionMetadata(AttentionMetadata):

102 - """Metadata for FlashAttentionBackend.

103 -

104 - NOTE: Any python object stored here is not updated when it is

105 - cuda-graph replayed. If you have values that need to be changed

106 - dynamically, it should be stored in tensor. The tensor has to be

107 - updated from `CUDAGraphRunner.forward` API.

108 - """

109 - # (batch_size,). The sequence length per sequence. Sequence length means

110 - # the computed tokens + new tokens None if it is a decoding.

111 - seq_lens: Optional[List[int]]

112 - # seq_lens stored as a tensor.

113 - seq_lens_tensor: Optional[torch.Tensor]

114 -

115 - # NOTE(sang): Definition of context_len, query_len, and seq_len.

116 - # |---------- N-1 iteration --------|

117 - # |---------------- N iteration ---------------------|

118 - # |- tokenA -|......................|-- newTokens ---|

119 - # |---------- context_len ----------|

120 - # |-------------------- seq_len ---------------------|

121 - # |-- query_len ---|

122 -

123 - # Maximum sequence length among prefill batch. 0 if there are decoding

124 - # requests only.

125 - max_prefill_seq_len: int

126 - # Maximum sequence length among decode batch. 0 if there are prefill

127 - # requests only.

128 - max_decode_seq_len: int

129 - # (batch_size,) A tensor of context lengths (tokens that are computed

130 - # so far).

131 - context_lens_tensor: Optional[torch.Tensor]

132 -

133 - # (batch_size, max_blocks_per_seq).

134 - # Block addresses per sequence. (Seq id -> list of physical block)

135 - # E.g., [0, 1, 2] means tokens are stored in 0th, 1st, and 2nd blocks

136 - # in the kv cache. Each block can contain up to block_size tokens.

137 - # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph

138 - # captured.

139 - block_tables: Optional[torch.Tensor]

140 -

141 - # Whether or not if cuda graph is enabled.

142 - # Cuda-graph is currently enabled for decoding only.

143 - # TODO(woosuk): Move `use_cuda_graph` out since it's unrelated to attention.

144 -

145 - use_cuda_graph: bool

146 -

147 - # Maximum query length in the batch.

148 - max_query_len: Optional[int] = None

149 -

150 - # Max number of query tokens among request in the batch.

151 - max_decode_query_len: Optional[int] = None

152 -

153 - # (batch_size + 1,). The cumulative subquery lengths of the sequences in

154 - # the batch, used to index into subquery. E.g., if the subquery length

155 - # is [4, 6], it is [0, 4, 10].

156 - query_start_loc: Optional[torch.Tensor] = None

157 - # (batch_size + 1,). The cumulative sequence lengths of the sequences in

158 - # the batch, used to index into sequence. E.g., if the sequence length is

159 - # [4, 6], it is [0, 4, 10].

160 - seq_start_loc: Optional[torch.Tensor] = None

161 -

162 - _cached_prefill_metadata: Optional["FlashAttentionMetadata"] = None

163 - _cached_decode_metadata: Optional["FlashAttentionMetadata"] = None

164 -

165 - # Begin encoder attn & enc/dec cross-attn fields...

166 -

167 - # Encoder sequence lengths representation

168 - encoder_seq_lens: Optional[List[int]] = None

169 - encoder_seq_lens_tensor: Optional[torch.Tensor] = None

170 - # (batch_size + 1,). The cumulative sequence lengths of the sequences in

171 - # the batch, used to index into sequence. E.g., if the sequence length is

172 - # [4, 6], it is [0, 4, 10].

173 - encoder_seq_start_loc: Optional[torch.Tensor] = None

174 - # Maximum sequence length among encoder sequences

175 - max_encoder_seq_len: Optional[int] = None

176 - # Number of tokens input to encoder

177 - num_encoder_tokens: Optional[int] = None

178 -

179 - # Cross-attention memory-mapping data structures: slot mapping

180 - # and block tables

181 - cross_slot_mapping: Optional[torch.Tensor] = None

182 - cross_block_tables: Optional[torch.Tensor] = None

183 -

184 - @property

185 - def is_all_encoder_attn_metadata_set(self):

186 - '''

187 - All attention metadata required for encoder attention is set.

188 - '''

189 - return is_all_encoder_attn_metadata_set(self)

190 -

191 - @property

192 - def is_all_cross_attn_metadata_set(self):

193 - '''

194 - All attention metadata required for enc/dec cross-attention is set.

195 -

196 - Superset of encoder attention required metadata.

197 - '''

198 - return is_all_cross_attn_metadata_set(self)

199 -

200 - @property

201 - def prefill_metadata(self) -> Optional["FlashAttentionMetadata"]:

202 - if self.num_prefills == 0:

203 - return None

204 -

205 - if self._cached_prefill_metadata is not None:

206 - return self._cached_prefill_metadata

207 -

208 - assert ((self.seq_lens is not None)

209 - or (self.encoder_seq_lens is not None))

210 - assert ((self.seq_lens_tensor is not None)

211 - or (self.encoder_seq_lens_tensor is not None))

212 -

213 - # Compute some attn_metadata fields which default to None

214 - query_start_loc = (None if self.query_start_loc is None else

215 - self.query_start_loc[:self.num_prefills + 1])

216 - slot_mapping = (None if self.slot_mapping is None else

217 - self.slot_mapping[:self.num_prefill_tokens])

218 - seq_lens = (None if self.seq_lens is None else

219 - self.seq_lens[:self.num_prefills])

220 - seq_lens_tensor = (None if self.seq_lens_tensor is None else

221 - self.seq_lens_tensor[:self.num_prefills])

222 - seq_start_loc = (None if self.seq_start_loc is None else

223 - self.seq_start_loc[:self.num_prefills + 1])

224 - context_lens_tensor = (None if self.context_lens_tensor is None else

225 - self.context_lens_tensor[:self.num_prefills])

226 - block_tables = (None if self.block_tables is None else

227 - self.block_tables[:self.num_prefills])

228 -

229 - self._cached_prefill_metadata = FlashAttentionMetadata(

230 - num_prefills=self.num_prefills,

231 - num_prefill_tokens=self.num_prefill_tokens,

232 - num_decode_tokens=0,

233 - slot_mapping=slot_mapping,

234 - multi_modal_placeholder_index_maps=self.

235 - multi_modal_placeholder_index_maps,

236 - enable_kv_scales_calculation=self.enable_kv_scales_calculation,

237 - seq_lens=seq_lens,

238 - seq_lens_tensor=seq_lens_tensor,

239 - max_query_len=self.max_query_len,

240 - max_prefill_seq_len=self.max_prefill_seq_len,

241 - max_decode_query_len=0,

242 - max_decode_seq_len=0,

243 - query_start_loc=query_start_loc,

244 - seq_start_loc=seq_start_loc,

245 - context_lens_tensor=context_lens_tensor,

246 - block_tables=block_tables,

247 - use_cuda_graph=False,

248 - # Begin encoder & cross attn fields below...

249 - encoder_seq_lens=self.encoder_seq_lens,

250 - encoder_seq_lens_tensor=self.encoder_seq_lens_tensor,

251 - encoder_seq_start_loc=self.encoder_seq_start_loc,

252 - max_encoder_seq_len=self.max_encoder_seq_len,

253 - cross_slot_mapping=self.cross_slot_mapping,

254 - cross_block_tables=self.cross_block_tables)

255 - return self._cached_prefill_metadata

256 -

257 - @property

258 - def decode_metadata(self) -> Optional["FlashAttentionMetadata"]:

259 - if self.num_decode_tokens == 0:

260 - return None

261 -

262 - if self._cached_decode_metadata is not None:

263 - return self._cached_decode_metadata

264 - assert ((self.seq_lens_tensor is not None)

265 - or (self.encoder_seq_lens_tensor is not None))

266 -

267 - # Compute some attn_metadata fields which default to None

268 - slot_mapping = (None if self.slot_mapping is None else

269 - self.slot_mapping[self.num_prefill_tokens:])

270 - seq_lens_tensor = (None if self.seq_lens_tensor is None else

271 - self.seq_lens_tensor[self.num_prefills:])

272 - block_tables = (None if self.block_tables is None else

273 - self.block_tables[self.num_prefills:])

274 -

275 - self._cached_decode_metadata = FlashAttentionMetadata(

276 - num_prefills=0,

277 - num_prefill_tokens=0,

278 - num_decode_tokens=self.num_decode_tokens,

279 - slot_mapping=slot_mapping,

280 - multi_modal_placeholder_index_maps=None,

281 - enable_kv_scales_calculation=True,

282 - seq_lens=None,

283 - seq_lens_tensor=seq_lens_tensor,

284 - max_decode_query_len=self.max_decode_query_len,

285 - max_query_len=self.max_query_len,

286 - max_prefill_seq_len=0,

287 - max_decode_seq_len=self.max_decode_seq_len,

288 - # Batch may be composed of prefill|decodes, adjust query start

289 - # indices to refer to the start of decodes. E.g.

290 - # in tokens:[3 prefills|6 decodes], query_start_loc=[3,9] => [0,6].

291 - query_start_loc=(self.query_start_loc[self.num_prefills:] -

292 - self.query_start_loc[self.num_prefills])

293 - if self.query_start_loc is not None else None,

294 - seq_start_loc=self.seq_start_loc[self.num_prefills:]

295 - if self.seq_start_loc is not None else None,

296 - context_lens_tensor=None,

297 - block_tables=block_tables,

298 - use_cuda_graph=self.use_cuda_graph,

299 - # Begin encoder & cross attn fields below...

300 - encoder_seq_lens=self.encoder_seq_lens,

301 - encoder_seq_lens_tensor=self.encoder_seq_lens_tensor,

302 - encoder_seq_start_loc=self.encoder_seq_start_loc,

303 - max_encoder_seq_len=self.max_encoder_seq_len,

304 - cross_slot_mapping=self.cross_slot_mapping,

305 - cross_block_tables=self.cross_block_tables)

306 - return self._cached_decode_metadata

307 -

308 -

309 -class FlashAttentionMetadataBuilder(

310 - AttentionMetadataBuilder[FlashAttentionMetadata]):

311 -

312 - def __init__(self, input_builder):

313 - self.input_builder = input_builder

314 - self.runner = input_builder.runner

315 - self.sliding_window = input_builder.sliding_window

316 - self.block_size = input_builder.block_size

317 -

318 - def prepare(self):

319 - self.slot_mapping: List[int] = []

320 - self.prefill_seq_lens: List[int] = []

321 - self.context_lens: List[int] = []

322 - self.block_tables: List[List[int]] = []

323 - self.curr_seq_lens: List[int] = []

324 - self.multimodal_placeholder_maps: Dict[

325 - str,

326 - MultiModalPlaceholderMap] = defaultdict(MultiModalPlaceholderMap)

327 - self.num_prefills = 0

328 - self.num_prefill_tokens = 0

329 - self.num_decode_tokens = 0

330 - self.has_prefix_cache_hit = False

331 -

332 - def _add_seq_group(self, inter_data, chunked_prefill_enabled: bool,

333 - prefix_cache_hit: bool):

334 - """Add a sequence group to the metadata. Specifically update/append

335 - 1. context length.

336 - 2. block table.

337 - 3. slot mapping.

338 - """

339 - is_prompt = inter_data.is_prompt

340 - block_tables = inter_data.block_tables

341 -

342 - for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len,

343 - curr_sliding_window_block) in zip(

344 - inter_data.seq_ids, [len(t) for t in inter_data.input_tokens],

345 - inter_data.orig_seq_lens, inter_data.seq_lens,

346 - inter_data.query_lens, inter_data.context_lens,

347 - inter_data.curr_sliding_window_blocks):

348 - self.context_lens.append(context_len)

349 -

350 - if is_prompt:

351 - mm_maps = inter_data.multi_modal_placeholder_maps

352 - if mm_maps:

353 - for modality, placeholders in mm_maps.items():

354 - self.multimodal_placeholder_maps[modality].extend(

355 - placeholders)

356 -

357 - self.num_prefills += 1

358 - self.num_prefill_tokens += token_len

359 - self.prefill_seq_lens.append(seq_len)

360 - else:

361 - self.num_decode_tokens += query_len

362 - self.curr_seq_lens.append(curr_seq_len)

363 -

364 - # Compute block table.

365 - # TODO(sang): Combine chunked prefill and prefix caching by

366 - # only allowing multiple of block_size chunk size.

367 - # NOTE: This only works for oooooooxxx style attention.

368 - block_table = []

369 - if prefix_cache_hit:

370 - # NOTE(woosuk): For flash-attn, the block table should

371 - # include the entries for the incoming prefill tokens.

372 - block_table = block_tables[seq_id]

373 - elif ((chunked_prefill_enabled or not is_prompt)

374 - and block_tables is not None):

375 - if curr_sliding_window_block == 0:

376 - block_table = block_tables[seq_id]

377 - else:

378 - block_table = block_tables[seq_id][

379 - -curr_sliding_window_block:]

380 - self.block_tables.append(block_table)

381 -

382 - # Compute slot mapping.

383 - is_profile_run = is_block_tables_empty(block_tables)

384 - start_idx = compute_slot_mapping_start_idx(is_prompt, query_len,

385 - context_len,

386 - self.sliding_window)

387 - compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id,

388 - seq_len, context_len, start_idx,

389 - self.block_size, inter_data.block_tables)

390 -

391 - def _get_graph_runner_block_tables(

392 - self, num_seqs: int,

393 - block_tables: List[List[int]]) -> torch.Tensor:

394 - # The shape of graph_block_tables is

395 - # [max batch size, max context len // block size].

396 - max_batch_size, max_blocks = self.runner.graph_block_tables.shape

397 - assert max_batch_size >= num_seqs

398 -

399 - graph_block_tables = self.runner.graph_block_tables[:num_seqs]

400 - for i, block_table in enumerate(block_tables):

401 - if block_table:

402 - num_blocks = len(block_table)

403 - if num_blocks <= max_blocks:

404 - graph_block_tables[i, :num_blocks] = block_table

405 - else:

406 - # It may be possible to have more blocks allocated due

407 - # to lookahead slots of multi-step, however, they are

408 - # not used anyway, so can be safely ignored.

409 - graph_block_tables[

410 - i, :max_blocks] = block_table[:max_blocks]

411 -

412 - return torch.from_numpy(graph_block_tables).to(

413 - device=self.runner.device, non_blocking=True)

414 -

415 - def build(self, seq_lens: List[int], query_lens: List[int],

416 - cuda_graph_pad_size: int, batch_size: int):

417 - """Build attention metadata with on-device tensors.

418 -

419 - Args:

420 - seq_lens: The maybe padded sequence lengths of the input sequences.

421 - query_lens: The query lengths of the input sequences.

422 - cuda_graph_pad_size: The padding size for cuda graph.

423 - -1 if cuda graph is not used.

424 - batch_size: The maybe padded batch size.

425 - """

426 - prefix_cache_hit = any([

427 - inter_data.prefix_cache_hit

428 - for inter_data in self.input_builder.inter_data_list

429 - ])

430 - for inter_data in self.input_builder.inter_data_list:

431 - self._add_seq_group(inter_data,

432 - self.input_builder.chunked_prefill_enabled,

433 - prefix_cache_hit)

434 -

435 - device = self.runner.device

436 - use_captured_graph = cuda_graph_pad_size != -1

437 -

438 - max_query_len = max(query_lens)

439 - decode_query_lens = query_lens[self.num_prefills:]

440 - if len(decode_query_lens) > 0:

441 - max_decode_query_len = max(decode_query_lens)

442 - else:

443 - max_decode_query_len = 1

444 - max_prefill_seq_len = max(self.prefill_seq_lens, default=0)

445 - max_decode_seq_len = max(self.curr_seq_lens, default=0)

446 - num_decode_tokens = self.num_decode_tokens

447 - query_start_loc = list(accumulate(query_lens, initial=0))

448 - seq_start_loc = list(accumulate(seq_lens, initial=0))

449 -

450 - num_seqs = len(seq_lens)

451 - if use_captured_graph:

452 - self.slot_mapping.extend([PAD_SLOT_ID] * cuda_graph_pad_size)

453 - self.block_tables.extend([] * cuda_graph_pad_size)

454 - num_decode_tokens = batch_size - self.num_prefill_tokens

455 - block_tables = self._get_graph_runner_block_tables(

456 - num_seqs, self.block_tables)

457 - else:

458 - block_tables = make_tensor_with_pad(

459 - self.block_tables,

460 - pad=0,

461 - dtype=torch.int,

462 - device=device,

463 - )

464 - assert max_query_len > 0, ("query_lens: {}".format(query_lens))

465 -

466 - assert device is not None

467 - context_lens_tensor = async_tensor_h2d(self.context_lens, torch.int,

468 - device, self.runner.pin_memory)

469 - seq_lens_tensor = async_tensor_h2d(seq_lens, torch.int, device,

470 - self.runner.pin_memory)

471 - slot_mapping_tensor = async_tensor_h2d(self.slot_mapping, torch.long,

472 - device, self.runner.pin_memory)

473 - query_start_loc_tensor = async_tensor_h2d(query_start_loc, torch.int32,

474 - device,

475 - self.runner.pin_memory)

476 - seq_start_loc_tensor = async_tensor_h2d(seq_start_loc, torch.int32,

477 - device, self.runner.pin_memory)

478 - placeholder_index_maps = {

479 - modality: placeholder_map.index_map()

480 - for modality, placeholder_map in

481 - self.multimodal_placeholder_maps.items()

482 - }

483 -

484 - return FlashAttentionMetadata(

485 - num_prefills=self.num_prefills,

486 - slot_mapping=slot_mapping_tensor,

487 - num_prefill_tokens=self.num_prefill_tokens,

488 - num_decode_tokens=num_decode_tokens,

489 - seq_lens=seq_lens,

490 - multi_modal_placeholder_index_maps=placeholder_index_maps,

491 - enable_kv_scales_calculation=True,

492 - seq_lens_tensor=seq_lens_tensor,

493 - max_query_len=max_query_len,

494 - max_decode_query_len=max_decode_query_len,

495 - max_prefill_seq_len=max_prefill_seq_len,

496 - max_decode_seq_len=max_decode_seq_len,

497 - query_start_loc=query_start_loc_tensor,

498 - seq_start_loc=seq_start_loc_tensor,

499 - context_lens_tensor=context_lens_tensor,

500 - block_tables=block_tables,

501 - use_cuda_graph=use_captured_graph,

502 - )

503 -

504 -

505 -class FlashAttentionImpl(AttentionImpl):

506 - """

507 - If the input tensors contain prompt tokens, the layout is as follows:

508 - |<--------------- num_prefill_tokens ----------------->|

510 -

511 - Otherwise, the layout is as follows:

512 - |<----------------- num_decode_tokens ------------------>|

514 -

515 - Generation tokens can contain padding when cuda-graph is used.

516 - Currently, prompt tokens don't contain any padding.

517 -

518 - The prompts might have different lengths, while the generation tokens

519 - always have length 1.

520 -

521 - If chunked prefill is enabled, prefill tokens and decode tokens can be

522 - batched together in a flattened 1D query.

523 -

524 - |<----- num_prefill_tokens ---->|<------- num_decode_tokens --------->|

525 - |<-prefill_0->|...|<-prefill_N-1->|<--decode_0-->|...|<--decode_M-1-->|

526 -

527 - Currently, cuda graph is disabled for chunked prefill, meaning there's no

528 - padding between prefill and decode tokens.

529 - """

530 -

531 - def __init__(

532 - self,

533 - num_heads: int,

534 - head_size: int,

535 - scale: float,

536 - num_kv_heads: int,

537 - alibi_slopes: Optional[List[float]],

538 - sliding_window: Optional[int],

539 - kv_cache_dtype: str,

540 - logits_soft_cap: Optional[float] = None,

541 - attn_type: str = AttentionType.DECODER,

542 - kv_sharing_target_layer_name: Optional[str] = None,

543 - use_irope: bool = False,

544 - ) -> None:

545 - if kv_sharing_target_layer_name is not None:

546 - raise NotImplementedError("KV sharing is not supported in V0 "

547 - "FLASH_ATTN backend.")

548 - if use_irope:

549 - logger.warning(

550 - "Using irope in V0 is not supported yet, it will fall back "

551 - "to global attention for long context.")

552 - self.num_heads = num_heads

553 - self.head_size = head_size

554 - self.scale = float(scale)

555 - self.num_kv_heads = num_kv_heads

556 - if alibi_slopes is not None:

557 - alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)

558 - self.alibi_slopes = alibi_slopes

559 - self.sliding_window = ((sliding_window - 1,

560 - 0) if sliding_window is not None else (-1, -1))

561 - self.kv_cache_dtype = kv_cache_dtype

562 - self.vllm_flash_attn_version = get_flash_attn_version(

563 - requires_alibi=self.alibi_slopes is not None)

564 - if is_quantized_kv_cache(self.kv_cache_dtype) and (

565 - not self.kv_cache_dtype.startswith("fp8")

566 - or not flash_attn_supports_fp8()):

567 - raise NotImplementedError(

568 - f"FlashAttention does not support {self.kv_cache_dtype} "

569 - "kv-cache on this device "

570 - f"(FA supports fp8 = {flash_attn_supports_fp8()}).")

571 - if logits_soft_cap is None:

572 - # In flash-attn, setting logits_soft_cap as 0 means no soft cap.

573 - logits_soft_cap = 0

574 - self.logits_soft_cap = logits_soft_cap

575 -

576 - self.num_queries_per_kv = self.num_heads // self.num_kv_heads

577 -

578 - support_head_sizes = FlashAttentionBackend.get_supported_head_sizes()

579 - if head_size not in support_head_sizes:

580 - raise ValueError(

581 - f"Head size {head_size} is not supported by FlashAttention. "

582 - f"Supported head sizes are: {support_head_sizes}.")

583 - self.attn_type = attn_type

584 -

585 - def forward(

586 - self,

587 - layer: AttentionLayer,

588 - query: torch.Tensor,

589 - key: torch.Tensor,

590 - value: torch.Tensor,

591 - kv_cache: torch.Tensor,

592 - attn_metadata: FlashAttentionMetadata,

593 - output: Optional[torch.Tensor] = None,

594 - output_scale: Optional[torch.Tensor] = None,

595 - output_block_scale: Optional[torch.Tensor] = None,

596 - ) -> torch.Tensor:

597 - """Forward pass with FlashAttention.

598 -

599 - Args:

600 - query: shape = [num_tokens, num_heads, head_size]

601 - key: shape = [num_tokens, num_kv_heads, head_size]

602 - value: shape = [num_tokens, num_kv_heads, head_size]

603 - output: shape = [num_tokens, num_heads, head_size]

604 - kv_cache: KV cache tensor with shape

605 - [2, num_blocks, block_size, num_kv_heads, head_size].

606 - NOTE: kv_cache will be an empty tensor with shape [0]

607 - for profiling run.

608 - attn_metadata: Metadata for attention.

609 - NOTE: It in-place updates the output tensor.

610 - NOTE: FP8 quantization, flash-attn expect the size of

611 - {q,k,v}_descale to be (num_sequences, num_kv_heads).

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

613 - """

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

615 -

616 - if output_scale is not None or output_block_scale is not None:

617 - raise NotImplementedError(

618 - "fused output quantization is not yet supported"

619 - " for FlashAttentionImpl")

620 -

621 - # NOTE(woosuk): FlashAttention2 does not support FP8 KV cache.

622 - if not flash_attn_supports_fp8() or output.dtype != torch.bfloat16:

623 - assert (

624 - layer._k_scale_float == 1.0 and layer._v_scale_float == 1.0), (

625 - "key/v_scale is only supported in FlashAttention 3 with "

626 - "base dtype bfloat16")

627 -

628 - attn_type = self.attn_type

629 - if (attn_type == AttentionType.ENCODER

630 - and (not attn_metadata.is_all_encoder_attn_metadata_set)):

631 - raise AttributeError("Encoder attention requires setting "

632 - "encoder metadata attributes.")

633 - elif (attn_type == AttentionType.ENCODER_DECODER

634 - and (not attn_metadata.is_all_cross_attn_metadata_set)):

635 - raise AttributeError("Encoder/decoder cross-attention "

636 - "requires setting cross-attention "

637 - "metadata attributes.")

638 -

639 - kv_cache_dtype: str = self.kv_cache_dtype

640 - softmax_scale: float = self.scale

641 - window_size = self.sliding_window

642 - alibi_slopes: Optional[torch.Tensor] = self.alibi_slopes

643 - logits_soft_cap: Optional[float] = self.logits_soft_cap

644 - fp8_attention = kv_cache_dtype.startswith("fp8")

645 -

646 - if fp8_attention and not flash_attn_supports_fp8():

647 - raise NotImplementedError(

648 - "FlashAttention does not support FP8 kv-cache on this device.")

649 -

650 - if kv_cache.numel() > 0:

651 - key_cache = kv_cache[0]

652 - value_cache = kv_cache[1]

653 - # We skip updating the KV cache under two conditions:

654 - # a. When the Attention Type is ENCODER. In this phase, we compute

655 - # only the encoder attention without updating the cache.

656 - # b. When both Key and Value are None. This occurs during

657 - # cross-attention computation in the decoding phase, where the

658 - # KV cache is already populated with the cross-attention

659 - # tensor. Thus, we skip cache updates during this time.

660 - if (attn_type != AttentionType.ENCODER) and (key is not None) and (

661 - value is not None):

662 - if attn_type == AttentionType.ENCODER_DECODER:

663 - # Update cross-attention KV cache (prefill-only)

664 - updated_slot_mapping = attn_metadata.cross_slot_mapping

665 - else:

666 - # Update self-attention KV cache (prefill/decode)

667 - updated_slot_mapping = attn_metadata.slot_mapping

668 -

669 - # Reshape the input keys and values and store them in the cache.

670 - # If kv_cache is not provided, the new key and value tensors are

671 - # not cached. This happens during the initial memory

672 - # profiling run.

673 - torch.ops._C_cache_ops.reshape_and_cache_flash(

674 - key,

675 - value,

676 - kv_cache[0],

677 - kv_cache[1],

678 - updated_slot_mapping.flatten(), # type: ignore[union-attr]

679 - kv_cache_dtype,

680 - layer._k_scale,

681 - layer._v_scale,

682 - )

683 -

684 - if fp8_attention:

685 - kv_cache = kv_cache.view(torch.float8_e4m3fn)

686 - key_cache = key_cache.view(torch.float8_e4m3fn)

687 - value_cache = value_cache.view(torch.float8_e4m3fn)

688 -

689 - if fp8_attention:

690 - num_tokens, num_heads, head_size = query.shape

691 - query, _ = ops.scaled_fp8_quant(

692 - query.reshape(

693 - (num_tokens, num_heads * head_size)).contiguous(),

694 - layer._q_scale)

695 - query = query.reshape((num_tokens, num_heads, head_size))

696 -

697 - (num_prefill_query_tokens, num_prefill_kv_tokens,

698 - num_decode_query_tokens) = \

699 - get_num_prefill_decode_query_kv_tokens(attn_metadata, attn_type)

700 - decode_query = query[num_prefill_query_tokens:]

701 - decode_output = output[num_prefill_query_tokens:]

702 - # QKV for prefill.

703 - query = query[:num_prefill_query_tokens]

704 - prefill_output = output[:num_prefill_query_tokens]

705 - assert query.shape[0] == num_prefill_query_tokens

706 - assert decode_query.shape[0] == num_decode_query_tokens

707 -

708 - if prefill_meta := attn_metadata.prefill_metadata:

709 - # Prompt run.

710 - if (kv_cache.numel() == 0 or prefill_meta.block_tables is None

711 - or prefill_meta.block_tables.numel() == 0):

712 - # normal attention

713 - # When block_tables are not filled, it means q and k are the

714 - # prompt, and they have the same length.

715 - q_seq_start_loc, q_seq_len, k_seq_start_loc, k_seq_len = \

716 - _get_query_key_seq_metadata(prefill_meta, True, attn_type)

717 -

718 - key = key[:num_prefill_kv_tokens]

719 - value = value[:num_prefill_kv_tokens]

720 -

721 - if fp8_attention:

722 - num_kv_tokens, num_kv_heads, head_size = key.shape

723 -

724 - key, _ = ops.scaled_fp8_quant(

725 - key.reshape((num_kv_tokens,

726 - num_kv_heads * head_size)).contiguous(),

727 - layer._k_scale)

728 - key = key.reshape((num_kv_tokens, num_kv_heads, head_size))

729 -

730 - value, _ = ops.scaled_fp8_quant(

731 - value.reshape((num_kv_tokens,

732 - num_kv_heads * head_size)).contiguous(),

733 - layer._v_scale)

734 - value = value.reshape(

735 - (num_kv_tokens, num_kv_heads, head_size))

736 -

737 - descale_shape = (q_seq_start_loc.shape[0] - 1, key.shape[1])

738 - flash_attn_varlen_func(

739 - q=query,

740 - k=key,

741 - v=value,

742 - cu_seqlens_q=q_seq_start_loc,

743 - cu_seqlens_k=k_seq_start_loc,

744 - max_seqlen_q=q_seq_len,

745 - max_seqlen_k=k_seq_len,

746 - softmax_scale=softmax_scale,

747 - causal=_get_causal_option(attn_type),

748 - window_size=window_size,

749 - alibi_slopes=alibi_slopes,

750 - softcap=logits_soft_cap,

751 - out=prefill_output,

752 - fa_version=self.vllm_flash_attn_version,

753 - q_descale=layer._q_scale.expand(descale_shape),

754 - k_descale=layer._k_scale.expand(descale_shape),

755 - v_descale=layer._v_scale.expand(descale_shape),

756 - )

757 - else:

758 - # prefix-enabled attention

759 - assert attn_type == AttentionType.DECODER, (

760 - "Only decoder-only models support prefix caching")

761 - assert prefill_meta.seq_lens is not None

762 - assert prefill_meta.query_start_loc is not None

763 - max_seq_len = max(prefill_meta.seq_lens)

764 - descale_shape = (prefill_meta.query_start_loc.shape[0] - 1,

765 - key.shape[1])

766 - flash_attn_varlen_func( # noqa

767 - q=query,

768 - k=key_cache,

769 - v=value_cache,

770 - cu_seqlens_q=prefill_meta.query_start_loc,

771 - max_seqlen_q=prefill_meta.max_query_len,

772 - seqused_k=prefill_meta.seq_lens_tensor,

773 - max_seqlen_k=max_seq_len,

774 - softmax_scale=softmax_scale,

775 - causal=True,

776 - window_size=window_size,

777 - alibi_slopes=alibi_slopes,

778 - block_table=prefill_meta.block_tables,

779 - softcap=logits_soft_cap,

780 - out=prefill_output,

781 - fa_version=self.vllm_flash_attn_version,

782 - q_descale=layer._q_scale.expand(descale_shape),

783 - k_descale=layer._k_scale.expand(descale_shape),

784 - v_descale=layer._v_scale.expand(descale_shape),

785 - )

786 -

787 - if decode_meta := attn_metadata.decode_metadata:

788 - # Decoding run.

789 - # Use flash_attn_varlen_func kernel for speculative decoding

790 - # because different queries might have different lengths.

791 -

792 - assert decode_meta.max_decode_query_len is not None

793 - # use only for actual varlen decoding

794 - if decode_meta.max_decode_query_len > 1:

795 - assert attn_type == AttentionType.DECODER, (

796 - "Only decoder-only models support max_decode_query_len > 1"

797 - )

798 - assert decode_meta.query_start_loc is not None

799 - descale_shape = (decode_meta.query_start_loc.shape[0] - 1,

800 - key.shape[1])

801 - flash_attn_varlen_func(

802 - q=decode_query,

803 - k=key_cache,

804 - v=value_cache,

805 - cu_seqlens_q=decode_meta.query_start_loc,

806 - max_seqlen_q=decode_meta.max_decode_query_len,

807 - seqused_k=decode_meta.seq_lens_tensor,

808 - max_seqlen_k=decode_meta.max_decode_seq_len,

809 - softmax_scale=softmax_scale,

810 - causal=True,

811 - window_size=window_size,

812 - alibi_slopes=alibi_slopes,

813 - softcap=logits_soft_cap,

814 - block_table=decode_meta.block_tables,

815 - out=decode_output,

816 - fa_version=self.vllm_flash_attn_version,

817 - q_descale=layer._q_scale.expand(descale_shape),

818 - k_descale=layer._k_scale.expand(descale_shape),

819 - v_descale=layer._v_scale.expand(descale_shape),

820 - )

821 - else:

822 - # Use flash_attn_with_kvcache for normal decoding.

823 - (

824 - seq_lens_arg,

825 - _,

826 - block_tables_arg,

827 - ) = get_seq_len_block_table_args(decode_meta, False, attn_type)

828 - descale_shape = (seq_lens_arg.shape[0], key_cache.shape[-2])

829 - flash_attn_with_kvcache(

830 - q=decode_query.unsqueeze(1),

831 - k_cache=key_cache,

832 - v_cache=value_cache,

833 - block_table=block_tables_arg,

834 - cache_seqlens=seq_lens_arg,

835 - softmax_scale=softmax_scale,

836 - causal=True,

837 - window_size=window_size,

838 - alibi_slopes=alibi_slopes,

839 - softcap=logits_soft_cap,

840 - out=decode_output.unsqueeze(1),

841 - fa_version=self.vllm_flash_attn_version,

842 - q_descale=layer._q_scale.expand(descale_shape),

843 - k_descale=layer._k_scale.expand(descale_shape),

844 - v_descale=layer._v_scale.expand(descale_shape),

845 - )

846 - return output

847 -

848 -

849 -def _get_query_key_seq_metadata(

850 - attn_metadata: FlashAttentionMetadata,

851 - is_prompt: bool,

852 - attn_type: str,

853 -) -> tuple:

854 - """

855 - Returns sequence metadata for key and query based on the specified

856 - attention type and whether input is a prompt.

857 -

858 - This function computes the starting locations and maximum sequence lengths

859 - for key and query sequences for different attention types.

860 -

861 - Args:

862 - attn_metadata: The attention metadata object

863 - is_prompt (bool): A flag indicating if the input is a prompt

864 - attn_type (AttentionType): The type of attention being used.

865 -

866 - Returns:

867 - tuple: A tuple containing four integers:

868 - - Starting location for the query sequence.

869 - - Maximum sequence length for the query sequence.

870 - - Starting location for the key sequence.

871 - - Maximum sequence length for the key sequence.

872 -

873 - Raises:

874 - AttributeError: If an invalid attention type is provided.

875 - """

876 - if attn_type == AttentionType.DECODER:

877 - # Decoder self-attention

878 - # Choose max_seq_len based on whether we are in prompt_run

879 - if is_prompt:

880 - max_seq_len = attn_metadata.max_prefill_seq_len

881 - else:

882 - max_seq_len = attn_metadata.max_decode_seq_len

883 - return (attn_metadata.seq_start_loc, max_seq_len,

884 - attn_metadata.seq_start_loc, max_seq_len)

885 -

886 - elif attn_type == AttentionType.ENCODER_DECODER:

887 - # This is cross attention between the where the key

888 - # is the precomputed encoder attention and query

889 - # is the input sequence.

890 - # Choose query max length based on whether it is prompt

891 - # or not.

892 - if is_prompt:

893 - max_seq_len = attn_metadata.max_prefill_seq_len

894 - else:

895 - max_seq_len = attn_metadata.max_decode_seq_len

896 - return (attn_metadata.seq_start_loc, max_seq_len,

897 - attn_metadata.encoder_seq_start_loc,

898 - attn_metadata.max_encoder_seq_len)

899 - elif attn_type == AttentionType.ENCODER:

900 - # For encoder attention both the query and the key are same i.e. the

901 - # encoder sequence.

902 - return (attn_metadata.encoder_seq_start_loc,

903 - attn_metadata.max_encoder_seq_len,

904 - attn_metadata.encoder_seq_start_loc,

905 - attn_metadata.max_encoder_seq_len)

906 - elif attn_type == AttentionType.ENCODER_ONLY:

907 - assert is_prompt, "Should not have decode for encoder only model."

908 - return (attn_metadata.seq_start_loc, attn_metadata.max_prefill_seq_len,

909 - attn_metadata.seq_start_loc, attn_metadata.max_prefill_seq_len)

910 - else:

911 - raise AttributeError(f"Invalid attention type {str(attn_type)}")

912 -

913 -

914 -def _get_causal_option(attn_type: str) -> bool:

915 - """

916 - Determine whether the given attention type is suitable for causal

917 - attention mechanisms.

918 -

919 - Args:

920 - attn_type (AttentionType): The type of attention being evaluated

921 -

922 - Returns:

923 - bool: Returns `True` if the attention type is suitable for causal

924 - attention (i.e., not encoder, encoder-only, or encoder-decoder),

925 - otherwise returns `False`.

926 - """

927 - return not (attn_type == AttentionType.ENCODER

928 - or attn_type == AttentionType.ENCODER_ONLY

929 - or attn_type == AttentionType.ENCODER_DECODER)

vllm/attention/backends/flashmla.py deleted

-227

@@ -1,227 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -from contextlib import contextmanager

5 -from dataclasses import dataclass

6 -from typing import List, Optional, Tuple, Type

7 -

8 -import torch

9 -

10 -from vllm.attention.backends.abstract import (AttentionType,

11 - is_quantized_kv_cache)

12 -from vllm.attention.backends.mla.common import (MLACommonBackend,

13 - MLACommonImpl,

14 - MLACommonMetadata,

15 - MLACommonMetadataBuilder,

16 - MLACommonState)

17 -from vllm.attention.ops.flashmla import (flash_mla_with_kvcache,

18 - get_mla_metadata,

19 - is_flashmla_supported)

20 -

21 -

22 -class FlashMLABackend(MLACommonBackend):

23 -

24 - @staticmethod

25 - def get_name() -> str:

26 - return "FLASHMLA"

27 -

28 - @staticmethod

29 - def get_impl_cls() -> Type["FlashMLAImpl"]:

30 - return FlashMLAImpl

31 -

32 - @staticmethod

33 - def get_metadata_cls() -> Type["FlashMLAMetadata"]:

34 - return FlashMLAMetadata

35 -

36 - @staticmethod

37 - def get_builder_cls() -> Type["FlashMLAMetadataBuilder"]:

38 - return FlashMLAMetadataBuilder

39 -

40 - @staticmethod

41 - def get_state_cls() -> Type["FlashMLAState"]:

42 - return FlashMLAState

43 -

44 -

45 -@dataclass

46 -class FlashMLAMetadata(MLACommonMetadata):

47 - decode_tile_scheduler_metadata: Optional[Tuple[torch.Tensor,

48 - torch.Tensor]] = None

49 - decode_num_splits: Optional[torch.Tensor] = None

50 -

51 - @property

52 - def decode_metadata(self):

53 - decode_metadata = super().decode_metadata

54 - # TODO: cache assignment?

55 - if decode_metadata is not None:

56 - decode_metadata.decode_tile_scheduler_metadata=\

57 - self.decode_tile_scheduler_metadata

58 - decode_metadata.decode_num_splits=\

59 - self.decode_num_splits

60 - return decode_metadata

61 -

62 -

63 -class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):

64 -

65 - def __init__(self, *args, **kwargs):

66 - super().__init__(*args, **kwargs)

67 -

68 - self.num_q_heads = self.runner.model_config.get_num_attention_heads(

69 - self.runner.parallel_config)

70 -

71 - def build(self, seq_lens: List[int], query_lens: List[int],

72 - cuda_graph_pad_size: int, batch_size: int):

73 - m = super().build(seq_lens, query_lens, cuda_graph_pad_size,

74 - batch_size)

75 -

76 - if m.num_decode_tokens > 0:

77 - m.decode_tile_scheduler_metadata, m.decode_num_splits = \

78 - get_mla_metadata(

79 - m.seq_lens_tensor[m.num_prefills:],

80 - self.num_q_heads,

81 - 1, # MQA for the decode path

82 - )

83 -

84 - return m

85 -

86 -

87 -class FlashMLAState(MLACommonState[FlashMLAMetadata]):

88 -

89 - def __init__(self, *args, **kwds):

90 - super().__init__(*args, **kwds)

91 -

92 - self.num_q_heads = self.runner.model_config.get_num_attention_heads(

93 - self.runner.parallel_config)

94 -

95 - @contextmanager

96 - def graph_capture(self, max_batch_size: int):

97 - # Run a dummy `get_mla_metadata` so we can get the right shapes

98 - self._graph_decoder_tile_scheduler_metadata, \

99 - self._graph_decode_num_splits = get_mla_metadata(

100 - torch.ones(

101 - max_batch_size, dtype=torch.int32, device=self.runner.device),

102 - self.num_q_heads,

103 - 1, # MQA for the decode path

104 - )

105 -

106 - with super().graph_capture(max_batch_size):

107 - yield

108 -

109 - del self._graph_decoder_tile_scheduler_metadata

110 - del self._graph_decode_num_splits

111 -

112 - def graph_capture_get_metadata_for_batch(

113 - self, batch_size: int, is_encoder_decoder_model: bool = False):

114 - metadata = super().graph_capture_get_metadata_for_batch(

115 - batch_size, is_encoder_decoder_model)

116 - assert metadata.num_decode_tokens > 0

117 -

118 - decoder_tile_scheduler_metadata, decode_num_splits = get_mla_metadata(

119 - self._graph_seq_lens[:batch_size],

120 - self.num_q_heads,

121 - 1, # MQA for the decode path

122 - )

123 -

124 - self._graph_decoder_tile_scheduler_metadata.copy_(

125 - decoder_tile_scheduler_metadata)

126 - self._graph_decode_num_splits[:batch_size + 1].copy_(decode_num_splits)

127 -

128 - metadata.decode_tile_scheduler_metadata=\

129 - self._graph_decoder_tile_scheduler_metadata

130 - metadata.decode_num_splits=\

131 - self._graph_decode_num_splits[:batch_size + 1]

132 -

133 - return metadata

134 -

135 - def get_graph_input_buffers(self,

136 - attn_metadata,

137 - is_encoder_decoder_model: bool = False):

138 - input_buffers = super().get_graph_input_buffers(

139 - attn_metadata, is_encoder_decoder_model)

140 - input_buffers["decode_tile_scheduler_metadata"] = \

141 - attn_metadata.decode_metadata.decode_tile_scheduler_metadata

142 - input_buffers["decode_num_splits"] = \

143 - attn_metadata.decode_metadata.decode_num_splits

144 -

145 - return input_buffers

146 -

147 - def prepare_graph_input_buffers(self,

148 - input_buffers,

149 - attn_metadata,

150 - is_encoder_decoder_model: bool = False):

151 - super().prepare_graph_input_buffers(input_buffers, attn_metadata,

152 - is_encoder_decoder_model)

153 -

154 - input_buffers["decode_tile_scheduler_metadata"].copy_(

155 - attn_metadata.decode_metadata.decode_tile_scheduler_metadata)

156 - input_buffers["decode_num_splits"].copy_(

157 - attn_metadata.decode_metadata.decode_num_splits)

158 -

159 -

160 -class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):

161 -

162 - def __init__(

163 - self,

164 - num_heads: int,

165 - head_size: int,

166 - scale: float,

167 - num_kv_heads: int,

168 - alibi_slopes: Optional[List[float]],

169 - sliding_window: Optional[int],

170 - kv_cache_dtype: str,

171 - logits_soft_cap: Optional[float],

172 - attn_type: str,

173 - kv_sharing_target_layer_name: Optional[str] = None,

174 - # MLA Specific Arguments

175 - **mla_args) -> None:

176 - super().__init__(num_heads, head_size, scale, num_kv_heads,

177 - alibi_slopes, sliding_window, kv_cache_dtype,

178 - logits_soft_cap, attn_type,

179 - kv_sharing_target_layer_name, **mla_args)

180 -

181 - is_supported, reason = is_flashmla_supported()

182 - assert is_supported, reason

183 -

184 - unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]

185 - if any(unsupported_features):

186 - raise NotImplementedError(

187 - "FlashMLAImpl does not support one of the following: "

188 - "alibi_slopes, sliding_window, logits_soft_cap")

189 -

190 - if attn_type != AttentionType.DECODER:

191 - raise NotImplementedError("Encoder self-attention and "

192 - "encoder/decoder cross-attention "

193 - "are not implemented for "

194 - "FlashMLAImpl")

195 -

196 - if is_quantized_kv_cache(self.kv_cache_dtype):

197 - raise NotImplementedError(

198 - "FlashMLA with FP8 KV cache not yet supported")

199 -

200 - def _forward_decode(

201 - self,

202 - q_nope: torch.Tensor,

203 - q_pe: torch.Tensor,

204 - kv_c_and_k_pe_cache: torch.Tensor,

205 - attn_metadata: FlashMLAMetadata,

206 - ) -> torch.Tensor:

207 - assert kv_c_and_k_pe_cache.numel() > 0

208 -

209 - decode_meta = attn_metadata.decode_metadata

210 - assert decode_meta is not None

211 -

212 - q = torch.cat([q_nope, q_pe], dim=-1)\

213 - .unsqueeze(1) # Add seqlen dim of 1 (decode)

214 -

215 - o, _ = flash_mla_with_kvcache(

216 - q=q,

217 - k_cache=kv_c_and_k_pe_cache.unsqueeze(-2), # Add head dim of 1

218 - block_table=decode_meta.block_tables,

219 - cache_seqlens=decode_meta.seq_lens_tensor,

220 - head_dim_v=self.kv_lora_rank,

221 - tile_scheduler_metadata=decode_meta.decode_tile_scheduler_metadata,

222 - num_splits=decode_meta.decode_num_splits,

223 - softmax_scale=self.scale,

224 - causal=True,

225 - )

226 -

227 - return self._v_up_proj(o)

vllm/attention/backends/mla/__init__.py

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

-1305

@@ -1,1305 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -"""

4 -# MLA Common Components

5 -

6 -This file implements common components for MLA implementations.

7 -

8 -First we define:

9 -

10 -Sq as Q sequence length

11 -Skv as KV sequence length

12 -

13 -MLA has two possible ways of computing, a data-movement friendly approach and a

14 -compute friendly approach, we generally want to use the compute friendly

15 -approach for "prefill" (i.e. the ratio Sq / Skv is "small", is near 1)

16 -and the data-movement friendly approach for "decode" (i.e. the ratio

17 -Sq / Skv is "large").

18 -

19 -NOTE what we deem small and large is currently determined by if its labelled

20 -prefill or decode by the scheduler, but this is something we should probably

21 -tune.

22 -

23 -Main reference: DeepseekV2 paper, and FlashInfer Implementation

24 -(https://arxiv.org/abs/2405.04434 and https://github.com/flashinfer-ai/flashinfer/pull/551).

25 -

26 -Deepseek's MLA attention works the following way:

27 -* Use a single latent vector to represent the per-token entry of the KV cache.

28 -* For decode (i.e. the memory friendly approach) the attention "simulates" a

29 -multi-head attention, while the compute is similar to multi-query attention.

30 -

31 -Below is example of both paths assuming batchsize = 1

32 -

33 -## More Extent Definitions:

34 -

35 -C Context length, `Skv - Sq`

36 -H hidden size

37 -N number of attention heads

38 -Lq latent dimension for Q 1536 in DSV3

39 -Lkv latent dimension for K/V 512 in DSV3

40 -P nope dimension, no rope. 128 in DSV3

41 -R rope dimension, goes through rope. 64 in DSV3

42 -V V head dim. 128 in DSV3

43 -

44 -## Vector/Matrix Definitions

45 -

46 -h_t hidden states (input to attention) shape [Sq, H]

47 -q_c latent/compressed Q shape [Sq, Lq]

48 -q_nope uncompressed Q (no-rope) shape [Sq, N, P]

49 -q_pe uncompressed Q (rope) shape [Sq, N, R]

50 -kv_c latent/compressed KV shape [Skv, Lkv]

51 -k_pe decoupled k position embeddings shape [Skv, R]

52 -new_kv_c new kv_c from current iter shape [Sq, Lkv]

53 -new_k_pe new k_pe from current iter shape [Sq, R]

54 -cache_kv_c cached k_c from previous iters shape [C, Lkv]

55 -cache_k_pe cached k_pe from previous iters shape [C, R]

56 -W_DQ project h_t to q_c shape [H, Lq]

57 -W_UQ project q_c to q_nope shape [Lq, N * P]

58 -W_QR project q_c to q_pe shape [Lq, N * R]

59 -W_DKV project h_t to kv_c shape [H, Lkv]

60 -W_UK project kv_c to k_nope shape [Lkv, N, P]

61 -W_KR project h_t to k_pe shape [H, R]

62 -W_UV project kv_c to v shape [Lkv, N, V]

63 -W_O project v to h_t shape [N * V, H]

64 -

65 -

66 -## Compute Friendly Approach (i.e. "_forward_prefill"):

67 -

68 -q_c = h_t @ W_DQ

69 -q_nope = (q_c @ W_UQ).view(Sq, N, P)

70 -q_pe = RoPE(q_c @ W_QR).view(Sq, N, R)

71 -new_kv_c = h_t @ W_DKV

72 -new_k_pe = RoPE(h_t @ W_KR)

73 -kv_c = torch.cat([new_kv_c, cache_kv_c], dim=0)

74 -k_pe = torch.cat([new_k_pe, cache_k_pe], dim=0)

75 -k_nope = (kv_c @ W_UK.view(Lkv, N * P)).view(Skv, N, P)

76 -v = (kv_c @ W_UV.view(Lkv, N * V)).view(Skv, N, V)

77 -

78 -// MHA with QK headdim = P + R

79 -// V headdim = V

80 -// spda_o shape [Sq, N, V]

81 -spda_o = scaled_dot_product_attention(

82 - torch.cat([q_nope, q_pe], dim=-1),

83 - torch.cat([k_nope, k_pe.unsqueeze(1).expand(-1, N, -1)], dim=-1),

84 - v

85 -)

86 -return spda_o @ W_O

87 -

88 -NOTE: in the actual code,

89 - `kv_b_proj` is [W_UK; W_UV] concatenated per head

90 - `q_b_proj` is [W_UQ; W_QR] concatenated per head

91 - `out_proj` is W_O

92 -

93 -

94 -## Data-Movement Friendly Approach (i.e. "_forward_decode"):

95 -

96 -Runtime

97 -q_c = h_t @ W_DQ

98 -q_nope = (q_c @ W_UQ).view(-1, N, P)

99 -ql_nope = einsum("snh,lnh->snl", q, W_UK)

100 -q_pe = RoPE(q_c @ W_QR).view(Sq, N, R)

101 -new_kv_c = h_t @ W_DKV

102 -new_k_pe = RoPE(h_t @ W_KR)

103 -kv_c = torch.cat([new_kv_c, cache_kv_c], dim=0)

104 -k_pe = torch.cat([new_k_pe, cache_k_pe], dim=0)

105 -

106 -// MQA with QK headdim = Lkv + R

107 -// V headdim = Lkv

108 -// spda_o shape [Sq, N, Lkv]

109 -// NOTE: this is less compute-friendly since Lkv > P

110 -// but is more data-movement friendly since its MQA vs MHA

111 -spda_o = scaled_dot_product_attention(

112 - torch.cat([ql_nope, q_pe], dim=-1),

113 - torch.cat([kv_c, k_pe], dim=-1),

114 - kv_c

115 -)

116 -

117 -o = einsum("snl,lnv->snv", spda_o.reshape(-1, N, Lkv), W_UV)

118 -return o.view(-1, N * V) @ self.num_heads @ W_O

119 -

120 -

121 -## Chunked Prefill

122 -

123 -For chunked prefill we want to use the compute friendly algorithm. We are

124 -assuming sufficiently large Sq / Skv ratio, in the future may want to switch to

125 -the data-movement friendly approach if the chunk (i.e. `Sq`) is small.

126 -

127 -However, the compute-friendly approach can potentially run out of memory if Skv

128 -is large due to: `k_nope = (kv_c @ W_UK).view(Skv, N, P)`

129 -

130 -To mitigate this, we chunk the computation of attention with respect to the

131 -current context (i.e. `cache_kv_c` and `cache_k_pe`) so that we can used a

132 -fixed workspace size.

133 -

134 -The chunked prefill approach is as follows:

135 -

136 -MCC Max chunk of context to process per iter, computed dynamically,

137 - used to bound the memory usage

138 -

139 -q_c = h_t @ W_DQ

140 -q_nope = (q_c @ W_UQ).view(Sq, N, P)

141 -q_pe = RoPE(q_c @ W_QR).view(Sq, N, R)

142 -new_kv_c = h_t @ W_DKV

143 -new_k_pe = RoPE(h_t @ W_KR)

144 -new_k_nope = (new_kv_c @ W_UK.view(Lkv, N * P)).view(Sq, N, P)

145 -new_v = (new_kv_c @ W_UV.view(Lkv, N * V)).view(Sq, N, V)

146 -

147 -// MHA between queries and new KV

148 -// with QK headdim = P + R

149 -// V headdim = V

150 -// curr_o shape [Sq, N, V]

151 -// curr_lse shape [N, Sq], this is just order FA returns

152 -curr_o, curr_lse = scaled_dot_product_attention(

153 - torch.cat([q_nope, q_pe], dim=-1),

154 - torch.cat([new_k_nope, new_k_pe.unsqueeze(1).expand(-1, N, -1)], dim=-1),

155 - new_v,

156 - casual=True,

157 - return_softmax_lse=True

158 -)

159 -

160 -// Compute attention with the already existing context

161 -for chunk_idx in range(cdiv(C, MCC)):

162 - chunk_start = chunk_idx * MCC

163 - chunk_end = min(chunk_start + MCC, C)

164 - Sc = chunk_end - chunk_start

165 - cache_kv_c_chunk = cache_kv_c[chunk_start:chunk_end]

166 - cache_k_pe_chunk = cache_k_pe[chunk_start:chunk_end]

167 - cache_k_nope_chunk = (cache_kv_c_chunk @ W_UK).view(-1, N, P)

168 - cache_v_chunk = (cache_kv_c_chunk @ W_UV).view(-1, N, V)

169 -

170 - chunk_o, chunk_lse = scaled_dot_product_attention(

171 - torch.cat([q_nope, q_pe], dim=-1),

172 - torch.cat([cache_k_nope_chunk,

173 - cache_k_pe_chunk.unsqueeze(1).expand(-1, N, -1)],

174 - dim=-1),

175 - cache_v_chunk,

176 - casual=False,

177 - return_softmax_lse=True

178 - )

179 -

180 - curr_o, curr_lse = merge_attn_states(

181 - suffix_output=curr_o,

182 - suffix_lse=curr_lse,

183 - prefix_output=chunk_o,

184 - prefix_lse=chunk_lse,

185 - )

186 -

187 -return curr_o @ W_O

188 -"""

189 -

190 -import functools

191 -from abc import abstractmethod

192 -from collections import defaultdict

193 -from contextlib import contextmanager

194 -from dataclasses import dataclass

195 -from itertools import accumulate

196 -from typing import Any, Dict, Generic, List, Optional, Tuple, Type, TypeVar

197 -

198 -import torch

199 -

200 -from vllm import _custom_ops as ops

201 -from vllm import envs

202 -from vllm.attention.backends.abstract import (AttentionBackend, AttentionLayer,

203 - AttentionMetadata,

204 - AttentionMetadataBuilder,

205 - AttentionState, MLAAttentionImpl)

206 -from vllm.attention.backends.utils import (PAD_SLOT_ID, compute_slot_mapping,

207 - compute_slot_mapping_start_idx,

208 - is_block_tables_empty)

209 -from vllm.attention.ops.merge_attn_states import merge_attn_states

210 -from vllm.attention.utils.fa_utils import get_flash_attn_version

211 -from vllm.model_executor.layers.linear import (ColumnParallelLinear,

212 - LinearBase,

213 - UnquantizedLinearMethod)

214 -from vllm.multimodal import MultiModalPlaceholderMap

215 -from vllm.platforms import current_platform

216 -from vllm.triton_utils import HAS_TRITON

217 -from vllm.utils import async_tensor_h2d, cdiv, make_tensor_with_pad, round_down

218 -

219 -if HAS_TRITON:

220 - from vllm.attention.ops.triton_flash_attention import triton_attention

221 -else:

222 - triton_attention = None

223 -

224 -try:

225 - from vllm.vllm_flash_attn import flash_attn_varlen_func

226 - is_vllm_fa = True

227 -except ImportError:

228 - is_vllm_fa = False

229 - try:

230 - # For rocm use upstream flash attention

231 - from flash_attn import flash_attn_varlen_func

232 - except ImportError:

233 - flash_attn_varlen_func = None

234 -

235 -is_hip = current_platform.is_rocm()

236 -

237 -

238 -class MLACommonBackend(AttentionBackend):

239 -

240 - @staticmethod

241 - def get_name() -> str:

242 - return "TRITON_MLA"

243 -

244 - @staticmethod

245 - def get_metadata_cls() -> Type["AttentionMetadata"]:

246 - return MLACommonMetadata

247 -

248 - @staticmethod

249 - def get_builder_cls() -> Type["MLACommonMetadataBuilder"]:

250 - return MLACommonMetadataBuilder

251 -

252 - @staticmethod

253 - def get_state_cls() -> Type["MLACommonState"]:

254 - return MLACommonState

255 -

256 - @staticmethod

257 - def get_kv_cache_shape(

258 - num_blocks: int,

259 - block_size: int,

260 - num_kv_heads: int, # assumed to be 1 for MLA

261 - head_size: int,

262 - ) -> Tuple[int, ...]:

263 - return (num_blocks, block_size, head_size)

264 -

265 - @staticmethod

266 - def swap_blocks(

267 - src_kv_cache: torch.Tensor,

268 - dst_kv_cache: torch.Tensor,

269 - src_to_dst: torch.Tensor,

270 - ) -> None:

271 - ops.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dst)

272 -

273 - @staticmethod

274 - def copy_blocks(

275 - kv_caches: List[torch.Tensor],

276 - src_to_dists: torch.Tensor,

277 - ) -> None:

278 - ops.copy_blocks_mla(kv_caches, src_to_dists)

279 -

280 - @staticmethod

281 - def get_supported_head_sizes() -> List[int]:

282 - return [576]

283 -

284 -

285 -T = TypeVar("T", bound="MLACommonMetadata")

286 -

287 -

288 -class MLACommonState(AttentionState, Generic[T]):

289 -

290 - def __init__(self, runner):

291 - self.runner = runner

292 - self._is_graph_capturing = False

293 -

294 - scheduler_config = runner.scheduler_config

295 - self.model_config = runner.model_config

296 - cache_config = runner.cache_config

297 -

298 - self.chunked_prefill_enabled = scheduler_config.chunked_prefill_enabled

299 - self.enable_prefix_caching = cache_config.enable_prefix_caching

300 -

301 - if self.chunked_prefill_enabled or self.enable_prefix_caching:

302 - self.context_chunk_workspace_size = min(

303 - # Max sure there is enough for 8 full length request or at least

304 - # 4 pages of cache per request

305 - max(

306 - 8 * self.model_config.max_model_len, 4 *

307 - scheduler_config.max_num_seqs * cache_config.block_size),

308 - # For long-context models try not to over-allocate limiting

309 - # kv-cache space, limiting it to 64k tokens,

310 - # which would result in the workspace being:

311 - # 2*(576)*(64*1024) = 144mb

312 - # (assuming 576 MLA head dim, and fp16)

313 - # which would result in up-projected context being

314 - # 2*(192*128)*(64*1024) = 3gb

315 - # (assuming 192 QK head dim, 128 heads, and fp16)

316 - 128 * 1024)

317 - assert self.context_chunk_workspace_size >= \

318 - scheduler_config.max_num_seqs * cache_config.block_size

319 -

320 - @contextmanager

321 - def graph_capture(self, max_batch_size: int):

322 - self._is_graph_capturing = True

323 -

324 - self._graph_slot_mapping = torch.full((max_batch_size, ),

325 - PAD_SLOT_ID,

326 - dtype=torch.long,

327 - device=self.runner.device)

328 - self._graph_seq_lens = torch.ones(max_batch_size,

329 - dtype=torch.int32,

330 - device=self.runner.device)

331 - self._graph_block_tables = torch.from_numpy(

332 - self.runner.graph_block_tables).to(device=self.runner.device)

333 -

334 - self._positions = torch.zeros((max_batch_size, ),

335 - dtype=torch.long,

336 - device=self.runner.device)

337 -

338 - yield

339 -

340 - self._is_graph_capturing = False

341 - del self._graph_slot_mapping

342 - del self._graph_seq_lens

343 - del self._graph_block_tables

344 - del self._positions

345 -

346 - def graph_clone(self, batch_size: int):

347 - assert self._is_graph_capturing

348 - return self.__class__(self.runner)

349 -

350 - def graph_capture_get_metadata_for_batch(

351 - self,

352 - batch_size: int,

353 - is_encoder_decoder_model: bool = False) -> T:

354 - assert self._is_graph_capturing

355 -

356 - attn_metadata = self.runner.attn_backend.make_metadata(

357 - multi_modal_placeholder_index_maps=None,

358 - enable_kv_scales_calculation=False,

359 - use_cuda_graph=True,

360 - num_prefills=0,

361 - num_prefill_tokens=0,

362 - num_decode_tokens=batch_size,

363 - slot_mapping=self._graph_slot_mapping[:batch_size],

364 - seq_lens=None,

365 - seq_lens_tensor=self._graph_seq_lens[:batch_size],

366 - max_query_len=1,

367 - max_decode_query_len=1,

368 - max_prefill_seq_len=0,

369 - max_decode_seq_len=self.runner.max_seq_len_to_capture,

370 - query_start_loc=None,

371 - seq_start_loc=None,

372 - context_lens_tensor=None,

373 - block_tables=self._graph_block_tables[:batch_size],

374 - head_dim=self.runner.model_config.get_head_size())

375 -

376 - if is_encoder_decoder_model:

377 - raise NotImplementedError(

378 - "MLACommonState does not support encoder/decoder yet")

379 -

380 - return attn_metadata

381 -

382 - def get_graph_input_buffers(self,

383 - attn_metadata,

384 - is_encoder_decoder_model: bool = False):

385 - input_buffers = {

386 - "slot_mapping": attn_metadata.slot_mapping,

387 - "seq_lens_tensor": attn_metadata.decode_metadata.seq_lens_tensor,

388 - "block_tables": attn_metadata.decode_metadata.block_tables,

389 - }

390 - if is_encoder_decoder_model:

391 - raise NotImplementedError(

392 - "MLACommonState does not support encoder/decoder yet")

393 -

394 - return input_buffers

395 -

396 - def prepare_graph_input_buffers(self,

397 - input_buffers,

398 - attn_metadata,

399 - is_encoder_decoder_model: bool = False):

400 - input_buffers["seq_lens_tensor"].copy_(

401 - attn_metadata.decode_metadata.seq_lens_tensor, non_blocking=True)

402 - input_buffers["block_tables"].copy_(

403 - attn_metadata.decode_metadata.block_tables, non_blocking=True)

404 - if is_encoder_decoder_model:

405 - raise NotImplementedError(

406 - "TritonMLAState does not support encoder/decoder yet")

407 -

408 - def begin_forward(self, model_input):

409 - if self.chunked_prefill_enabled or self.enable_prefix_caching:

410 - if not hasattr(self, "context_chunk_workspace"):

411 - # not self.runner.device does not return the correct device

412 - # for this process, (init_device sets the correct device but

413 - # only on the Worker). The only way Ive figured out to get the

414 - # correct device is to allocate the workspace on the first call

415 - # to begin_forward and use the device of the input tokens

416 - assert model_input.input_tokens is not None

417 - self.context_chunk_workspace = torch.empty(

418 - (self.context_chunk_workspace_size,

419 - self.model_config.get_head_size()),

420 - dtype=self.model_config.dtype,

421 - device=model_input.input_tokens.device,

422 - )

423 -

424 - model_input.attn_metadata.context_chunk_workspace = \

425 - self.context_chunk_workspace

426 -

427 -

428 -@dataclass

429 -class MLACommonMetadata(AttentionMetadata):

430 - """Metadata for MLACommon.

431 -

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

433 - understand this class

434 -

435 - NOTE: Any python object stored here is not updated when it is

436 - cuda-graph replayed. If you have values that need to be changed

437 - dynamically, it should be stored in tensor. The tensor has to be

438 - updated from `CUDAGraphRunner.forward` API.

439 - """

440 - # Whether or not if cuda graph is enabled.

441 - # Cuda-graph is currently enabled for decoding only.

442 - # TODO(woosuk): Move `use_cuda_graph` out since it's unrelated to attention.

443 - use_cuda_graph: bool

444 -

445 - # NOTE(sang): Definition of context_len, query_len, and seq_len.

446 - # |---------- N-1 iteration --------|

447 - # |---------------- N iteration ---------------------|

448 - # |- tokenA -|......................|-- newTokens ---|

449 - # |---------- context_len ----------|

450 - # |-------------------- seq_len ---------------------|

451 - # |-- query_len ---|

452 -

453 - # (batch_size,). The sequence length per sequence. Sequence length means

454 - # the computed tokens + new tokens None if it is a decoding.

455 - seq_lens: Optional[List[int]]

456 - # seq_lens stored as a tensor.

457 - seq_lens_tensor: Optional[torch.Tensor]

458 -

459 - # Maximum sequence length among prefill batch. 0 if there are decoding

460 - # requests only.

461 - max_prefill_seq_len: int

462 - # Maximum sequence length among decode batch. 0 if there are prefill

463 - # requests only.

464 - max_decode_seq_len: int

465 - # (batch_size,) A tensor of context lengths (tokens that are computed

466 - # so far).

467 - context_lens_tensor: Optional[torch.Tensor]

468 -

469 - # (batch_size, max_blocks_per_seq).

470 - # Block addresses per sequence. (Seq id -> list of physical block)

471 - # E.g., [0, 1, 2] means tokens are stored in 0th, 1st, and 2nd blocks

472 - # in the kv cache. Each block can contain up to block_size tokens.

473 - # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph

474 - # captured.

475 - block_tables: Optional[torch.Tensor]

476 -

477 - # Maximum query length in the batch.

478 - max_query_len: Optional[int] = None

479 -

480 - # Max number of query tokens among request in the batch.

481 - max_decode_query_len: Optional[int] = None

482 -

483 - # (batch_size + 1,). The cumulative subquery lengths of the sequences in

484 - # the batch, used to index into subquery. E.g., if the subquery length

485 - # is [4, 6], it is [0, 4, 10].

486 - query_start_loc: Optional[torch.Tensor] = None

487 - # (batch_size + 1,). The cumulative sequence lengths of the sequences in

488 - # the batch, used to index into sequence. E.g., if the sequence length is

489 - # [4, 6], it is [0, 4, 10].

490 - seq_start_loc: Optional[torch.Tensor] = None

491 -

492 - _cached_prefill_metadata: Optional[Any] = None

493 - _cached_decode_metadata: Optional[Any] = None

494 -

495 - num_prefill_tokens: int

496 -

497 - # The dimension of the attention heads

498 - head_dim: Optional[int] = None

499 -

500 - # Used when chunked prefill is enabled to simulate worst case workspace

501 - # allocations, hopefully to avoid going OOM

502 - is_profile_run: bool = False

503 -

504 - # New for MLA (compared to FlashAttention)

505 - # For chunked prefill

506 - context_chunk_cu_seq_lens: Optional[torch.Tensor] = None

507 - context_chunk_starts: Optional[torch.Tensor] = None

508 - context_chunk_seq_tot: Optional[List[int]] = None

509 - context_chunk_max_seq_lens: Optional[List[int]] = None

510 - # Set by MLAAttentionState in `begin_forward` so it doesn't get broadcasted

511 - context_chunk_workspace: Optional[torch.Tensor] = None

512 -

513 - def __post_init__(self):

514 - supported_head_sizes = MLACommonBackend.get_supported_head_sizes()

515 - if self.head_dim is not None and self.head_dim \

516 - not in supported_head_sizes:

517 - raise ValueError(

518 - f"Only {supported_head_sizes} are supported for head_dim,",

519 - f" received {self.head_dim}.")

520 -

521 - @property

522 - def prefill_metadata(self):

523 - if self.num_prefills == 0:

524 - return None

525 -

526 - if self._cached_prefill_metadata is not None:

527 - return self._cached_prefill_metadata

528 -

529 - assert self.seq_lens is not None

530 - assert self.seq_lens_tensor is not None

531 -

532 - # Compute some attn_metadata fields which default to None

533 - query_start_loc = (None if self.query_start_loc is None else

534 - self.query_start_loc[:self.num_prefills + 1])

535 - slot_mapping = (None if self.slot_mapping is None else

536 - self.slot_mapping[:self.num_prefill_tokens])

537 - seq_lens = (None if self.seq_lens is None else

538 - self.seq_lens[:self.num_prefills])

539 - seq_lens_tensor = (None if self.seq_lens_tensor is None else

540 - self.seq_lens_tensor[:self.num_prefills])

541 - seq_start_loc = (None if self.seq_start_loc is None else

542 - self.seq_start_loc[:self.num_prefills + 1])

543 - context_lens_tensor = (None if self.context_lens_tensor is None else

544 - self.context_lens_tensor[:self.num_prefills])

545 - block_tables = (None if self.block_tables is None else

546 - self.block_tables[:self.num_prefills])

547 -

548 - self._cached_prefill_metadata = self.__class__(

549 - # Required by ModelRunner

550 - use_cuda_graph=False, # Not Attention Related

551 - # Required by Attention Metadata

552 - num_prefills=self.num_prefills,

553 - num_prefill_tokens=self.num_prefill_tokens,

554 - num_decode_tokens=0,

555 - slot_mapping=slot_mapping,

556 - # Required by Attention Metadata (not used)

557 - multi_modal_placeholder_index_maps=None,

558 - enable_kv_scales_calculation=False,

559 - # MLACommonMetadata

560 - seq_lens=seq_lens,

561 - seq_lens_tensor=seq_lens_tensor,

562 - max_query_len=self.max_query_len,

563 - max_prefill_seq_len=self.max_prefill_seq_len,

564 - max_decode_query_len=0,

565 - max_decode_seq_len=0,

566 - query_start_loc=query_start_loc,

567 - seq_start_loc=seq_start_loc,

568 - context_lens_tensor=context_lens_tensor,

569 - block_tables=block_tables,

570 - head_dim=self.head_dim,

571 - is_profile_run=self.is_profile_run,

572 - # MLACommonMetadata Chunk prefill specific

573 - context_chunk_cu_seq_lens=self.context_chunk_cu_seq_lens,

574 - context_chunk_starts=self.context_chunk_starts,

575 - context_chunk_seq_tot=self.context_chunk_seq_tot,

576 - context_chunk_max_seq_lens=self.context_chunk_max_seq_lens,

577 - )

578 - return self._cached_prefill_metadata

579 -

580 - @property

581 - def decode_metadata(self):

582 - if self.num_decode_tokens == 0:

583 - return None

584 -

585 - if self._cached_decode_metadata is not None:

586 - return self._cached_decode_metadata

587 - assert self.seq_lens_tensor is not None

588 -

589 - # Compute some attn_metadata fields which default to None

590 - slot_mapping = (None if self.slot_mapping is None else

591 - self.slot_mapping[self.num_prefill_tokens:])

592 - seq_lens_tensor = (None if self.seq_lens_tensor is None else

593 - self.seq_lens_tensor[self.num_prefills:])

594 - block_tables = (None if self.block_tables is None else

595 - self.block_tables[self.num_prefills:])

596 -

597 - self._cached_decode_metadata = self.__class__(

598 - # Required by ModelRunner

599 - use_cuda_graph=self.use_cuda_graph, # Not Attention Related

600 - # Required by Attention Metadata

601 - num_prefills=0,

602 - num_prefill_tokens=0,

603 - num_decode_tokens=self.num_decode_tokens,

604 - slot_mapping=slot_mapping,

605 - # Required by Attention Metadata (not used)

606 - multi_modal_placeholder_index_maps=None,

607 - enable_kv_scales_calculation=False,

608 - # MLACommonMetadata

609 - seq_lens=None,

610 - seq_lens_tensor=seq_lens_tensor,

611 - max_decode_query_len=self.max_decode_query_len,

612 - max_query_len=self.max_query_len,

613 - max_prefill_seq_len=0,

614 - max_decode_seq_len=self.max_decode_seq_len,

615 - # Batch may be composed of prefill|decodes, adjust query start

616 - # indices to refer to the start of decodes. E.g.

617 - # in tokens:[3 prefills|6 decodes], query_start_loc=[3,9] => [0,6].

618 - query_start_loc=(self.query_start_loc[self.num_prefills:] -

619 - self.query_start_loc[self.num_prefills])

620 - if self.query_start_loc is not None else None,

621 - seq_start_loc=self.seq_start_loc[self.num_prefills:]

622 - if self.seq_start_loc is not None else None,

623 - context_lens_tensor=None,

624 - block_tables=block_tables,

625 - head_dim=self.head_dim,

626 - is_profile_run=self.is_profile_run)

627 - return self._cached_decode_metadata

628 -

629 -

630 -class MLACommonMetadataBuilder(AttentionMetadataBuilder[T], Generic[T]):

631 - """

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

633 - understand this class

634 - """

635 - BLOCK_TABLE_EXTENDER: list[list[int]] = []

636 -

637 - def __init__(self, input_builder):

638 - self.input_builder = input_builder

639 - self.runner = input_builder.runner

640 - self.sliding_window = input_builder.sliding_window

641 - self.block_size = input_builder.block_size

642 - self.chunked_prefill_enabled = \

643 - self.runner.scheduler_config.chunked_prefill_enabled

644 - self.enable_prefix_caching = \

645 - self.runner.cache_config.enable_prefix_caching

646 -

647 - if self.chunked_prefill_enabled or self.enable_prefix_caching:

648 - attn_state = self.input_builder.runner.attn_state

649 - self.context_chunk_workspace_size = \

650 - attn_state.context_chunk_workspace_size

651 - self.page_size = self.runner.block_size

652 -

653 - def prepare(self):

654 - self.slot_mapping: List[int] = []

655 - self.prefill_seq_lens: List[int] = []

656 - self.context_lens: List[int] = []

657 - self.block_tables: List[List[int]] = []

658 - self.curr_seq_lens: List[int] = []

659 - self.multimodal_placeholder_maps: Dict[

660 - str,

661 - MultiModalPlaceholderMap] = defaultdict(MultiModalPlaceholderMap)

662 - self.num_prefills = 0

663 - self.num_prefill_tokens = 0

664 - self.num_decode_tokens = 0

665 - self.has_prefix_cache_hit = False

666 -

667 - def _add_seq_group(self, inter_data, chunked_prefill_enabled: bool,

668 - prefix_cache_hit: bool):

669 - """Add a sequence group to the metadata. Specifically update/append

670 - 1. context length.

671 - 2. block table.

672 - 3. slot mapping.

673 - """

674 - is_prompt = inter_data.is_prompt

675 - block_tables = inter_data.block_tables

676 -

677 - for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len,

678 - curr_sliding_window_block) in zip(

679 - inter_data.seq_ids, [len(t) for t in inter_data.input_tokens],

680 - inter_data.orig_seq_lens, inter_data.seq_lens,

681 - inter_data.query_lens, inter_data.context_lens,

682 - inter_data.curr_sliding_window_blocks):

683 - self.context_lens.append(context_len)

684 - if is_prompt:

685 - self.num_prefills += 1

686 - self.num_prefill_tokens += token_len

687 - self.prefill_seq_lens.append(seq_len)

688 - else:

689 - self.num_decode_tokens += query_len

690 - self.curr_seq_lens.append(curr_seq_len)

691 -

692 - # Compute block table.

693 - # TODO(sang): Combine chunked prefill and prefix caching by

694 - # only allowing multiple of block_size chunk size.

695 - # NOTE: This only works for oooooooxxx style attention.

696 - block_table = []

697 - if prefix_cache_hit:

698 - # NOTE(woosuk): For flash-attn, the block table should

699 - # include the entries for the incoming prefill tokens.

700 - block_table = block_tables[seq_id]

701 - elif ((chunked_prefill_enabled or not is_prompt)

702 - and block_tables is not None):

703 - if curr_sliding_window_block == 0:

704 - block_table = block_tables[seq_id]

705 - else:

706 - block_table = block_tables[seq_id][

707 - -curr_sliding_window_block:]

708 - self.block_tables.append(block_table)

709 -

710 - # Compute slot mapping.

711 - is_profile_run = is_block_tables_empty(block_tables)

712 - start_idx = compute_slot_mapping_start_idx(is_prompt, query_len,

713 - context_len,

714 - self.sliding_window)

715 - compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id,

716 - seq_len, context_len, start_idx,

717 - self.block_size, inter_data.block_tables)

718 -

719 - def _get_graph_runner_block_tables(

720 - self, num_seqs: int,

721 - block_tables: List[List[int]]) -> torch.Tensor:

722 - # The shape of graph_block_tables is

723 - # [max batch size, max context len // block size].

724 - max_batch_size, max_blocks = self.runner.graph_block_tables.shape

725 - assert max_batch_size >= num_seqs

726 -

727 - graph_block_tables = self.runner.graph_block_tables[:num_seqs]

728 - for i, block_table in enumerate(block_tables):

729 - if block_table:

730 - num_blocks = len(block_table)

731 - if num_blocks <= max_blocks:

732 - graph_block_tables[i, :num_blocks] = block_table

733 - else:

734 - # It may be possible to have more blocks allocated due

735 - # to lookahead slots of multi-step, however, they are

736 - # not used anyway, so can be safely ignored.

737 - graph_block_tables[

738 - i, :max_blocks] = block_table[:max_blocks]

739 -

740 - return torch.from_numpy(graph_block_tables).to(

741 - device=self.runner.device, non_blocking=True)

742 -

743 - def build(self, seq_lens: List[int], query_lens: List[int],

744 - cuda_graph_pad_size: int, batch_size: int):

745 - """Build attention metadata with on-device tensors.

746 -

747 - Args:

748 - seq_lens: The maybe padded sequence lengths of the input sequences.

749 - query_lens: The query lengths of the input sequences.

750 - cuda_graph_pad_size: The padding size for cuda graph.

751 - -1 if cuda graph is not used.

752 - batch_size: The maybe padded batch size.

753 - """

754 - prefix_cache_hit = any([

755 - inter_data.prefix_cache_hit

756 - for inter_data in self.input_builder.inter_data_list

757 - ])

758 -

759 - for inter_data in self.input_builder.inter_data_list:

760 - self._add_seq_group(inter_data,

761 - self.input_builder.chunked_prefill_enabled,

762 - prefix_cache_hit)

763 -

764 - device = self.runner.device

765 - use_captured_graph = cuda_graph_pad_size != -1

766 -

767 - max_query_len = max(query_lens)

768 - decode_query_lens = query_lens[self.num_prefills:]

769 - if len(decode_query_lens) > 0:

770 - max_decode_query_len = max(decode_query_lens)

771 - else:

772 - max_decode_query_len = 1

773 - max_prefill_seq_len = max(self.prefill_seq_lens, default=0)

774 - max_decode_seq_len = max(self.curr_seq_lens, default=0)

775 - num_decode_tokens = self.num_decode_tokens

776 - query_start_loc = list(accumulate(query_lens, initial=0))

777 - seq_start_loc = list(accumulate(seq_lens, initial=0))

778 -

779 - num_seqs = len(seq_lens)

780 - if use_captured_graph:

781 - self.slot_mapping.extend([PAD_SLOT_ID] * cuda_graph_pad_size)

782 - self.block_tables.extend(self.__class__.BLOCK_TABLE_EXTENDER *

783 - cuda_graph_pad_size)

784 - num_decode_tokens = batch_size - self.num_prefill_tokens

785 -

786 - block_tables = self._get_graph_runner_block_tables(

787 - num_seqs, self.block_tables)

788 - else:

789 - block_tables = make_tensor_with_pad(

790 - self.block_tables,

791 - pad=0,

792 - dtype=torch.int,

793 - device=device,

794 - )

795 - assert max_query_len > 0, ("query_lens: {}".format(query_lens))

796 -

797 - assert device is not None

798 - context_lens_tensor = async_tensor_h2d(self.context_lens, torch.int,

799 - device, self.runner.pin_memory)

800 - seq_lens_tensor = async_tensor_h2d(seq_lens, torch.int, device,

801 - self.runner.pin_memory)

802 - slot_mapping_tensor = async_tensor_h2d(self.slot_mapping, torch.long,

803 - device, self.runner.pin_memory)

804 - query_start_loc_tensor = async_tensor_h2d(query_start_loc, torch.int32,

805 - device,

806 - self.runner.pin_memory)

807 - seq_start_loc_tensor = async_tensor_h2d(seq_start_loc, torch.int32,

808 - device, self.runner.pin_memory)

809 -

810 - context_chunk_cu_seq_lens = None

811 - context_chunk_starts = None

812 - context_chunk_seq_tot = None

813 - context_chunk_max_seq_lens = None

814 -

815 - if (self.chunked_prefill_enabled or self.enable_prefix_caching) \

816 - and self.num_prefills > 0 \

817 - and context_lens_tensor is not None \

818 - and context_lens_tensor[:self.num_prefills].max() > 0:

819 -

820 - # NOTE: it is recommended you read the `Chunked Prefill` section in

821 - # the comment at the top of the file before trying to understand

822 - # the following code

823 -

824 - num_prefills_with_context = \

825 - (context_lens_tensor[:self.num_prefills] > 0).sum().item()

826 -

827 - # currently we allocate an equal amount of workspace for each

828 - # prefill in the batch, we could probably use a more advanced

829 - # algorithm here and allocate more workspace to prefills with

830 - # longer context lengths

831 - max_context_chunk = \

832 - self.context_chunk_workspace_size // num_prefills_with_context

833 -

834 - # align max_context_chunk to page_size by rounding down,

835 - # currently the `gather_and_maybe_dequant_cache` kernel cannot

836 - # handle `context_chunk_starts` that are not aligned to page_size

837 - max_context_chunk = round_down(max_context_chunk, self.page_size)

838 - assert max_context_chunk > 0

839 - num_chunks = cdiv(context_lens_tensor.max(), max_context_chunk)

840 -

841 - # if `max_context_chunk = 256`, `num_chunks = 3`, and

842 - # `num_prefills_with_context = 4`, create a tensor that looks like

843 - # [[0, 0, 0, 0], [256, 256, 256, 256], [512, 512, 512, 512]]

844 - context_chunk_starts = \

845 - torch.arange(num_chunks, device=device, dtype=torch.int32)\

846 - .unsqueeze(1).expand(-1, self.num_prefills)\

847 - * max_context_chunk

848 - chunk_ends = torch.min(context_lens_tensor[:self.num_prefills]\

849 - .unsqueeze(0), context_chunk_starts + max_context_chunk)

850 - chunk_seq_lens = (chunk_ends - context_chunk_starts).clamp(min=0)

851 - _context_chunk_cu_seq_lens = chunk_seq_lens.cumsum(dim=1).to(

852 - torch.int32)

853 - zero = torch.zeros(num_chunks, dtype=torch.int32, device=device)\

854 - .unsqueeze(-1)

855 - context_chunk_cu_seq_lens = \

856 - torch.cat([zero, _context_chunk_cu_seq_lens], dim=1)

857 - context_chunk_max_seq_lens = \

858 - chunk_seq_lens.max(dim=1).values.tolist()

859 - context_chunk_seq_tot = chunk_seq_lens.sum(dim=1).tolist()

860 - assert max(context_chunk_seq_tot) <= \

861 - self.context_chunk_workspace_size

862 -

863 - return self.runner.attn_backend.make_metadata(

864 - # Required by ModelRunner

865 - use_cuda_graph=use_captured_graph, # Not Attention Related

866 - # Required by Attention Metadata

867 - num_prefills=self.num_prefills,

868 - slot_mapping=slot_mapping_tensor,

869 - num_prefill_tokens=self.num_prefill_tokens,

870 - num_decode_tokens=num_decode_tokens,

871 - # Required by Attention Metadata (not used)

872 - multi_modal_placeholder_index_maps=None, # Not Attention Related

873 - enable_kv_scales_calculation=False,

874 - # MLACommonMetadata

875 - seq_lens=seq_lens,

876 - seq_lens_tensor=seq_lens_tensor,

877 - max_query_len=max_query_len,

878 - max_decode_query_len=max_decode_query_len,

879 - max_prefill_seq_len=max_prefill_seq_len,

880 - max_decode_seq_len=max_decode_seq_len,

881 - query_start_loc=query_start_loc_tensor,

882 - seq_start_loc=seq_start_loc_tensor,

883 - context_lens_tensor=context_lens_tensor,

884 - block_tables=block_tables,

885 - head_dim=self.runner.model_config.get_head_size(),

886 - is_profile_run=self.runner.in_profile_run,

887 - # MLACommonMetadata Chunk prefill specific

888 - context_chunk_cu_seq_lens=context_chunk_cu_seq_lens,

889 - context_chunk_starts=context_chunk_starts,

890 - context_chunk_seq_tot=context_chunk_seq_tot,

891 - context_chunk_max_seq_lens=context_chunk_max_seq_lens,

892 - )

893 -

894 -

895 -class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):

896 - """

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

898 - understand this class

899 - """

900 -

901 - def __init__(

902 - self,

903 - num_heads: int,

904 - head_size: int,

905 - scale: float,

906 - num_kv_heads: int,

907 - alibi_slopes: Optional[List[float]],

908 - sliding_window: Optional[int],

909 - kv_cache_dtype: str,

910 - logits_soft_cap: Optional[float],

911 - attn_type: str,

912 - kv_sharing_target_layer_name: Optional[str],

913 - # MLA Specific Arguments

914 - q_lora_rank: Optional[int],

915 - kv_lora_rank: int,

916 - qk_nope_head_dim: int,

917 - qk_rope_head_dim: int,

918 - qk_head_dim: int,

919 - v_head_dim: int,

920 - kv_b_proj: ColumnParallelLinear,

921 - ) -> None:

922 - if kv_sharing_target_layer_name is not None:

923 - raise NotImplementedError("KV sharing not supported in V0.")

924 - self.num_heads = num_heads

925 - self.head_size = head_size

926 - self.scale = float(scale)

927 - self.num_kv_heads = num_kv_heads

928 - self.kv_cache_dtype = kv_cache_dtype

929 -

930 - self.q_lora_rank = q_lora_rank

931 - self.kv_lora_rank = kv_lora_rank

932 - self.qk_nope_head_dim = qk_nope_head_dim

933 - self.qk_rope_head_dim = qk_rope_head_dim

934 - self.qk_head_dim = qk_head_dim

935 - self.v_head_dim = v_head_dim

936 - self.kv_b_proj = kv_b_proj

937 -

938 - self.triton_fa_func = triton_attention

939 - # Handle the differences between the flash_attn_varlen from flash_attn

940 - # and the one from vllm_flash_attn. The former is used on RoCM and the

941 - # latter has an additional parameter to control FA2 vs FA3

942 - self.flash_attn_varlen_func = flash_attn_varlen_func

943 - self.vllm_flash_attn_version = get_flash_attn_version()

944 - if self.vllm_flash_attn_version is not None:

945 - self.flash_attn_varlen_func = \

946 - functools.partial(flash_attn_varlen_func,

947 - fa_version=self.vllm_flash_attn_version)

948 -

949 - # For MLA the v head dim is smaller than qk head dim so we pad out

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

951 - # not support different headdims

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

953 - self._pad_v = self.vllm_flash_attn_version is None or not (

954 - self.vllm_flash_attn_version == 3

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

956 -

957 - def _flash_attn_varlen_diff_headdims(self, q, k, v, softmax_scale,

958 - return_softmax_lse, **kwargs):

959 - maybe_padded_v = v

960 - if self._pad_v:

961 - maybe_padded_v = torch.nn.functional.pad(

962 - v, [0, q.shape[-1] - v.shape[-1]], value=0)

963 -

964 - if is_hip and envs.VLLM_USE_TRITON_FLASH_ATTN \

965 - and not return_softmax_lse:

966 - attn_out = self.triton_fa_func(

967 - q,

968 - k,

969 - maybe_padded_v,

970 - None, # output

971 - kwargs["cu_seqlens_q"],

972 - kwargs["cu_seqlens_k"],

973 - kwargs["max_seqlen_q"],

974 - kwargs["max_seqlen_k"],

975 - kwargs["causal"],

976 - softmax_scale,

977 - None, # bias

978 - )

979 - elif is_vllm_fa:

980 - attn_out = self.flash_attn_varlen_func(

981 - q=q,

982 - k=k,

983 - v=maybe_padded_v,

984 - return_softmax_lse=return_softmax_lse,

985 - softmax_scale=softmax_scale,

986 - **kwargs,

987 - )

988 - else:

989 - # Use return_attn_probs instead of return_softmax_lse for RoCM

990 - attn_out = self.flash_attn_varlen_func(

991 - q=q,

992 - k=k,

993 - v=maybe_padded_v,

994 - return_attn_probs=return_softmax_lse,

995 - softmax_scale=softmax_scale,

996 - **kwargs,

997 - )

998 -

999 - # Unpack the output if there is multiple results,

1000 - # triton always returns (output, softmax_lse),

1001 - # vllm_flash_attn returns (output, softmax_lse) when

1002 - # `return_softmax_lse = True`

1003 - # flash_attn (RoCM) returns (output, softmax_lse, ...) when

1004 - # `return_attn_probs = True`

1005 - rest = None

1006 - if isinstance(attn_out, tuple):

1007 - attn_out, *rest = attn_out

1008 -

1009 - # Remain consistent with old `flash_attn_varlen_func` where there

1010 - # is only one output tensor if `return_softmax_lse` is False.

1011 - if return_softmax_lse:

1012 - assert rest is not None

1013 - return attn_out, rest[0]

1014 - return attn_out

1015 -

1016 - def _v_up_proj(self, x):

1017 - # Convert from (B, N, L) to (N, B, L)

1018 - x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)

1019 - # Multiply (N, B, L) x (N, L, V) -> (N, B, V)

1020 - x = torch.bmm(x, self.W_UV)

1021 - # Convert from (N, B, V) to (B, N * V)

1022 - return x.transpose(0, 1).reshape(-1, self.num_heads * self.v_head_dim)

1023 -

1024 - def process_weights_after_loading(self, act_dtype: torch.dtype):

1025 -

1026 - def get_layer_weight(layer):

1027 - WEIGHT_NAMES = ("weight", "qweight", "weight_packed")

1028 - for attr in WEIGHT_NAMES:

1029 - if hasattr(layer, attr):

1030 - return getattr(layer, attr)

1031 - raise AttributeError(

1032 - f"Layer '{layer}' has no recognized weight attribute:"

1033 - f" {WEIGHT_NAMES}.")

1034 -

1035 - def get_and_maybe_dequant_weights(layer: LinearBase):

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

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

1038 - eye = torch.eye(layer.input_size_per_partition,

1039 - dtype=act_dtype,

1040 - device=get_layer_weight(layer).device)

1041 - dequant_weights = layer.quant_method.apply(layer,

1042 - eye,

1043 - bias=None)

1044 - del eye

1045 - # standardize to (output, input)

1046 - return dequant_weights.T

1047 - return layer.weight

1048 -

1049 - # we currently do not have quantized bmm's which are needed for

1050 - # `W_UV` and `W_UK_T`, we just store fp16/bf16 copies and perform

1051 - # the bmm's in 16-bit, the extra memory overhead of this is fairly low

1052 - kv_b_proj_weight = get_and_maybe_dequant_weights(self.kv_b_proj).T

1053 - assert kv_b_proj_weight.shape == (

1054 - self.kv_lora_rank,

1055 - self.num_heads * (self.qk_nope_head_dim + self.v_head_dim)), (

1056 - f"{kv_b_proj_weight.shape=}, "

1057 - f"{self.kv_lora_rank=}, "

1058 - f"{self.num_heads=}, "

1059 - f"{self.qk_nope_head_dim=}, "

1060 - f"{self.v_head_dim=}")

1061 - kv_b_proj_weight = kv_b_proj_weight.view(

1062 - self.kv_lora_rank,

1063 - self.num_heads,

1064 - self.qk_nope_head_dim + self.v_head_dim,

1065 - )

1066 -

1067 - W_UK, W_UV = kv_b_proj_weight.split(

1068 - [self.qk_nope_head_dim, self.v_head_dim], dim=-1)

1069 -

1070 - # Convert from (L, N, V) to (N, L, V)

1071 - self.W_UV = W_UV.transpose(0, 1)

1072 - # Convert from (L, N, P) to (N, P, L)

1073 - self.W_UK_T = W_UK.permute(1, 2, 0)

1074 -

1075 - def _compute_prefill_context(

1076 - self,

1077 - q: torch.Tensor,

1078 - kv_c_and_k_pe_cache: torch.Tensor,

1079 - attn_metadata: MLACommonMetadata,

1080 - k_scale: torch.Tensor,

1081 - ):

1082 - prefill_metadata = attn_metadata.prefill_metadata

1083 - assert prefill_metadata is not None

1084 - assert prefill_metadata.context_chunk_seq_tot is not None

1085 - assert prefill_metadata.context_chunk_cu_seq_lens is not None

1086 - assert prefill_metadata.context_chunk_starts is not None

1087 - assert prefill_metadata.context_chunk_max_seq_lens is not None

1088 - assert prefill_metadata.context_lens_tensor is not None

1089 -

1090 - output = None

1091 - iters = len(prefill_metadata.context_chunk_seq_tot)

1092 -

1093 - # Fetch from attn_metadata directly, since it late bound by

1094 - # MLAAttentionState, grabbing it directly `attn_metadata` can avoid

1095 - # any weirdness around prefill_metadata caching

1096 - assert attn_metadata.context_chunk_workspace is not None

1097 - workspace = attn_metadata.context_chunk_workspace

1098 -

1099 - for i in range(iters):

1100 - toks = prefill_metadata.context_chunk_seq_tot[i]

1101 -

1102 - ops.gather_and_maybe_dequant_cache(

1103 - src_cache=kv_c_and_k_pe_cache,

1104 - dst=workspace,

1105 - block_table=prefill_metadata.block_tables,

1106 - cu_seq_lens=prefill_metadata.context_chunk_cu_seq_lens[i],

1107 - batch_size=prefill_metadata.num_prefills,

1108 - kv_cache_dtype=self.kv_cache_dtype,

1109 - scale=k_scale,

1110 - seq_starts=prefill_metadata.context_chunk_starts[i],

1111 - )

1112 -

1113 - kv_c_normed = workspace[:toks]\

1114 - [..., :self.kv_lora_rank]

1115 - k_pe = workspace[:toks]\

1116 - [..., self.kv_lora_rank:].unsqueeze(1)

1117 -

1118 - kv_nope = self.kv_b_proj(kv_c_normed)[0].view( \

1119 - -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)

1120 - k_nope, v = kv_nope\

1121 - .split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)

1122 -

1123 - k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))),

1124 - dim=-1)

1125 -

1126 - attn_output, attn_softmax_lse = \

1127 - self._flash_attn_varlen_diff_headdims(

1128 - q=q,

1129 - k=k,

1130 - v=v,

1131 - cu_seqlens_q=prefill_metadata.query_start_loc,

1132 - cu_seqlens_k=prefill_metadata.context_chunk_cu_seq_lens[i],

1133 - max_seqlen_q=prefill_metadata.max_query_len,

1134 - max_seqlen_k=prefill_metadata.context_chunk_max_seq_lens[i],

1135 - softmax_scale=self.scale,

1136 - causal=False, # Context is unmasked

1137 - return_softmax_lse=True,

1138 - )

1139 -

1140 - if output is None:

1141 - output = attn_output

1142 - output_lse = attn_softmax_lse

1143 - else:

1144 - output_tmp = torch.empty_like(output)

1145 - output_lse_tmp = torch.empty_like(output_lse)

1146 - merge_attn_states(

1147 - output=output_tmp,

1148 - output_lse=output_lse_tmp,

1149 - prefix_output=output,

1150 - prefix_lse=output_lse,

1151 - suffix_output=attn_output,

1152 - suffix_lse=attn_softmax_lse,

1153 - )

1154 - output = output_tmp

1155 - output_lse = output_lse_tmp

1156 -

1157 - return output, output_lse

1158 -

1159 - def _forward_prefill(

1160 - self,

1161 - q: torch.Tensor,

1162 - kv_c_normed: torch.Tensor,

1163 - k_pe: torch.Tensor,

1164 - kv_c_and_k_pe_cache: torch.Tensor,

1165 - attn_metadata: MLACommonMetadata,

1166 - k_scale: torch.Tensor,

1167 - ) -> torch.Tensor:

1168 -

1169 - prefill_metadata = attn_metadata.prefill_metadata

1170 - assert prefill_metadata is not None

1171 -

1172 - has_context = prefill_metadata.context_lens_tensor is not None \

1173 - and prefill_metadata.context_lens_tensor.max() > 0

1174 -

1175 - kv_nope = self.kv_b_proj(kv_c_normed)[0].view(\

1176 - -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)

1177 - k_nope, v = kv_nope\

1178 - .split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)

1179 -

1180 - k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)

1181 -

1182 - output = self._flash_attn_varlen_diff_headdims(

1183 - q=q,

1184 - k=k,

1185 - v=v,

1186 - cu_seqlens_q=prefill_metadata.query_start_loc,

1187 - cu_seqlens_k=prefill_metadata.query_start_loc,

1188 - max_seqlen_q=prefill_metadata.max_prefill_seq_len,

1189 - max_seqlen_k=prefill_metadata.max_prefill_seq_len,

1190 - softmax_scale=self.scale,

1191 - causal=True,

1192 - return_softmax_lse=has_context,

1193 - )

1194 -

1195 - if has_context:

1196 - # ROCm flash_attn_varlen_func will return 3 objects instead of 2

1197 - suffix_output, suffix_lse = output

1198 - context_output, context_lse = self._compute_prefill_context( \

1199 - q, kv_c_and_k_pe_cache, attn_metadata, k_scale)

1200 -

1201 - output = torch.empty_like(suffix_output)

1202 - merge_attn_states(

1203 - output=output,

1204 - prefix_output=context_output,

1205 - prefix_lse=context_lse,

1206 - suffix_output=suffix_output,

1207 - suffix_lse=suffix_lse,

1208 - )

1209 -

1210 - # unpad if necessary

1211 - if self._pad_v:

1212 - output = output[..., :v.shape[-1]]

1213 -

1214 - return output.flatten(start_dim=-2)

1215 -

1216 - @abstractmethod

1217 - def _forward_decode(

1218 - self,

1219 - ql_nope: torch.Tensor,

1220 - q_pe: torch.Tensor,

1221 - kv_c_and_k_pe_cache: torch.Tensor,

1222 - attn_metadata: T,

1223 - ) -> torch.Tensor:

1224 - raise NotImplementedError

1225 -

1226 - def forward(

1227 - self,

1228 - layer: AttentionLayer,

1229 - q: torch.Tensor, # query in unified attn

1230 - k_c_normed: torch.Tensor, # key in unified attn

1231 - k_pe: torch.Tensor, # value in unified attn

1232 - kv_cache: torch.Tensor,

1233 - attn_metadata: T,

1234 - output: Optional[torch.Tensor] = None,

1235 - output_scale: Optional[torch.Tensor] = None,

1236 - output_block_scale: Optional[torch.Tensor] = None,

1237 - ) -> torch.Tensor:

1238 - if output is not None:

1239 - raise NotImplementedError(

1240 - "output is not yet supported for MLAImplBase")

1241 -

1242 - if output_scale is not None or output_block_scale is not None:

1243 - raise NotImplementedError(

1244 - "fused output quantization is not yet supported"

1245 - " for MLAImplBase")

1246 -

1247 - if attn_metadata.is_profile_run and \

1248 - attn_metadata.context_chunk_workspace is not None:

1249 - # During the profile run try to simulate to worse case output size

1250 - # for `self.kv_b_proj(kv_c_normed)` in `_compute_prefill_context`

1251 - # since this can be large

1252 - _ = torch.empty(

1253 - (attn_metadata.context_chunk_workspace.shape[0],

1254 - self.num_heads, self.qk_nope_head_dim + self.v_head_dim),

1255 - device=k_c_normed.device,

1256 - dtype=k_c_normed.dtype,

1257 - )

1258 -

1259 - has_decode = attn_metadata.decode_metadata is not None

1260 - has_prefill = attn_metadata.prefill_metadata is not None

1261 -

1262 - num_prefill_tokens: int = attn_metadata.num_prefill_tokens

1263 - q = q.view(-1, self.num_heads, self.qk_head_dim)

1264 -

1265 - decode_q = q[num_prefill_tokens:]

1266 -

1267 - prefill_q = q[:num_prefill_tokens]

1268 - prefill_k_pe = k_pe[:num_prefill_tokens]

1269 - prefill_k_c_normed = k_c_normed[:num_prefill_tokens]

1270 -

1271 - # write the latent and rope to kv cache

1272 - if kv_cache.numel() > 0:

1273 - ops.concat_and_cache_mla(

1274 - k_c_normed,

1275 - k_pe.squeeze(1),

1276 - kv_cache,

1277 - attn_metadata.slot_mapping.flatten(),

1278 - kv_cache_dtype=self.kv_cache_dtype,

1279 - scale=layer._k_scale,

1280 - )

1281 -

1282 - output = torch.empty(attn_metadata.num_prefill_tokens +

1283 - attn_metadata.num_decode_tokens,

1284 - self.v_head_dim * self.num_heads,

1285 - device=q.device,

1286 - dtype=q.dtype)

1287 - if has_prefill:

1288 - output[:num_prefill_tokens] = self._forward_prefill(

1289 - prefill_q, prefill_k_c_normed, prefill_k_pe, kv_cache,

1290 - attn_metadata, layer._k_scale)

1291 -

1292 - if has_decode:

1293 - decode_q_nope, decode_q_pe = decode_q.split(

1294 - [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

1295 - # Convert from (B, N, P) to (N, B, P)

1296 - decode_q_nope = decode_q_nope.transpose(0, 1)

1297 - # Multiply (N, B, P) x (N, P, L) -> (N, B, L)

1298 - decode_ql_nope = torch.bmm(decode_q_nope, self.W_UK_T)

1299 - # Convert from (N, B, L) to (B, N, L)

1300 - decode_ql_nope = decode_ql_nope.transpose(0, 1)

1301 -

1302 - output[num_prefill_tokens:] = self._forward_decode(

1303 - decode_ql_nope, decode_q_pe, kv_cache, attn_metadata)

1304 -

1305 - return output

vllm/attention/backends/rocm_aiter_mla.py deleted

-407

@@ -1,407 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -from contextlib import contextmanager

5 -from dataclasses import dataclass

6 -from typing import Optional, Type, Union

7 -

8 -import torch

9 -

10 -import vllm.envs as envs

11 -from vllm.attention.backends.mla.common import (MLACommonBackend,

12 - MLACommonImpl,

13 - MLACommonMetadata,

14 - MLACommonMetadataBuilder,

15 - MLACommonState)

16 -from vllm.attention.backends.utils import (compute_slot_mapping,

17 - compute_slot_mapping_start_idx,

18 - is_block_tables_empty)

19 -from vllm.attention.ops.rocm_aiter_mla import (aiter_mla_decode_fwd,

20 - get_aiter_mla_metadata)

21 -

22 -

23 -def is_aiter_mla_enabled() -> bool:

24 - return envs.VLLM_ROCM_USE_AITER \

25 - and envs.VLLM_ROCM_USE_AITER_MLA

26 -

27 -

28 -class AiterMLABackend(MLACommonBackend):

29 -

30 - @staticmethod

31 - def get_name() -> str:

32 - return "ROCM_AITER_MLA"

33 -

34 - @staticmethod

35 - def get_impl_cls() -> Type["AiterMLAImpl"]:

36 - return AiterMLAImpl

37 -

38 - @staticmethod

39 - def get_metadata_cls() -> Type["AiterMLAMetadata"]:

40 - return AiterMLAMetadata

41 -

42 - @staticmethod

43 - def get_builder_cls() -> Type["AiterMLAMetadataBuilder"]:

44 - return AiterMLAMetadataBuilder

45 -

46 - @staticmethod

47 - def get_state_cls() -> Type["AiterMLAState"]:

48 - return AiterMLAState

49 -

50 -

51 -@dataclass

52 -class AiterMLAMetadata(MLACommonMetadata):

53 - # The following 5 tensors are for current version of AITER MLA

54 - block_table_bound: Optional[torch.Tensor] = None

55 - # The indptr of the paged kv cache, shape: [batch_size + 1]

56 - paged_kv_indptr: Optional[torch.Tensor] = None

57 - # The page indices of the paged kv cache

58 - paged_kv_indices: Optional[torch.Tensor] = None

59 - # The number of entries in the last page of each request in

60 - # the paged kv cache, shape: [batch_size]

61 - paged_kv_last_page_lens: Optional[torch.Tensor] = None

62 -

63 - # This is just to make new AITER MLA API work

64 - # -- MTP support is not added yet.

65 - qo_indptr: Optional[torch.Tensor] = None

66 -

67 - @property

68 - def prefill_metadata(self):

69 - prefill_metadata = super().prefill_metadata

70 - self._cached_prefill_metadata = prefill_metadata

71 -

72 - if prefill_metadata is not None:

73 - prefill_metadata.paged_kv_indptr = self.paged_kv_indptr

74 - prefill_metadata.paged_kv_indices = self.paged_kv_indices

75 - prefill_metadata\

76 - .paged_kv_last_page_lens = self.paged_kv_last_page_lens

77 - prefill_metadata.block_table_bound = self.block_table_bound

78 - prefill_metadata.qo_indptr = self.qo_indptr

79 -

80 - # update the cache

81 - self._cached_prefill_metadata = self.__class__(

82 - **prefill_metadata.__dict__)

83 -

84 - return self._cached_prefill_metadata

85 -

86 - @property

87 - def decode_metadata(self):

88 - decode_metadata = super().decode_metadata

89 -

90 - self._cached_decode_metadata = decode_metadata

91 -

92 - if decode_metadata is not None:

93 - decode_metadata.paged_kv_indptr = self.paged_kv_indptr

94 - decode_metadata.paged_kv_indices = self.paged_kv_indices

95 - decode_metadata\

96 - .paged_kv_last_page_lens = self.paged_kv_last_page_lens

97 - decode_metadata.block_table_bound = self.block_table_bound

98 - decode_metadata.qo_indptr = self.qo_indptr

99 -

100 - # update the cache

101 - self._cached_decode_metadata = self.__class__(

102 - **decode_metadata.__dict__)

103 -

104 - return self._cached_decode_metadata

105 -

106 -

107 -class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

108 - BLOCK_TABLE_EXTENDER: list[list[int]] = [[]]

109 -

110 - def __init__(self, input_builder):

111 - super().__init__(input_builder)

112 - assert self.block_size == 1, "AITER MLA requires only block size 1."

113 -

114 - def prepare(self):

115 - super().prepare()

116 - self.paged_kv_indices: list[int] = []

117 - self.paged_kv_indptr: list[int] = [0]

118 - self.paged_kv_last_page_lens: list[int] = []

119 - self.total_blocks = 0

120 - self.qo_indptr: list[int] = [0]

121 -

122 - def _add_seq_group(self, inter_data, chunked_prefill_enabled: bool,

123 - prefix_cache_hit: bool):

124 - """Add a sequence group to the metadata. Specifically update/append

125 - 1. context length.

126 - 2. block table.

127 - 3. slot mapping.

128 - """

129 - is_prompt = inter_data.is_prompt

130 - block_tables = inter_data.block_tables

131 -

132 - for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len,

133 - curr_sliding_window_block) in zip(

134 - inter_data.seq_ids, [len(t) for t in inter_data.input_tokens],

135 - inter_data.orig_seq_lens, inter_data.seq_lens,

136 - inter_data.query_lens, inter_data.context_lens,

137 - inter_data.curr_sliding_window_blocks):

138 - self.context_lens.append(context_len)

139 - if is_prompt:

140 - self.num_prefills += 1

141 - self.num_prefill_tokens += token_len

142 - self.prefill_seq_lens.append(seq_len)

143 - else:

144 - self.num_decode_tokens += query_len

145 - self.curr_seq_lens.append(curr_seq_len)

146 -

147 - # Compute block table.

148 - # TODO(sang): Combine chunked prefill and prefix caching by

149 - # only allowing multiple of block_size chunk size.

150 - # NOTE: This only works for oooooooxxx style attention.

151 - block_table = []

152 - if prefix_cache_hit:

153 - # NOTE(woosuk): For flash-attn, the block table should

154 - # include the entries for the incoming prefill tokens.

155 - block_table = block_tables[seq_id]

156 - elif ((chunked_prefill_enabled or not is_prompt)

157 - and block_tables is not None):

158 - if curr_sliding_window_block == 0:

159 - block_table = block_tables[seq_id]

160 - else:

161 - block_table = block_tables[seq_id][

162 - -curr_sliding_window_block:]

163 - self.block_tables.append(block_table)

164 -

165 - # Compute slot mapping.

166 - is_profile_run = is_block_tables_empty(block_tables)

167 - start_idx = compute_slot_mapping_start_idx(is_prompt, query_len,

168 - context_len,

169 - self.sliding_window)

170 - compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id,

171 - seq_len, context_len, start_idx,

172 - self.block_size, inter_data.block_tables)

173 - if is_profile_run:

174 - return

175 -

176 - # Update paged_kv_* tensors only for non-profile run

177 - block_table = block_tables[seq_id]

178 - self._update_paged_kv_tensors(block_table, seq_len)

179 -

180 - def _update_paged_kv_tensors(self, block_table: list[int], seq_len: int):

181 - # Get the number of valid blocks based on sequence length.

182 - # If seq_len = 16, block_size = 16,

183 - # block_table_bound is 1 with 1 valid block.

184 - # If seq_len = 15, block_size = 16,

185 - # block_table_bound is 0 + 1 with 1 valid block.

186 - self.total_blocks += len(block_table)

187 - block_table_bound = seq_len // self.block_size + 1 \

188 - if seq_len % self.block_size != 0 \

189 - else seq_len // self.block_size

190 - self.paged_kv_indices.extend(block_table[:block_table_bound])

191 - self.paged_kv_indptr.append(self.paged_kv_indptr[-1] +

192 - block_table_bound)

193 - self.qo_indptr.append(self.qo_indptr[-1] + 1)

194 -

195 - last_page_len = seq_len % self.block_size

196 - if last_page_len == 0:

197 - last_page_len = self.block_size

198 - self.paged_kv_last_page_lens.append(last_page_len)

199 -

200 - def build(self, seq_lens: list[int], query_lens: list[int],

201 - cuda_graph_pad_size: int, batch_size: int) -> AiterMLAMetadata:

202 - metadata = super().build(seq_lens, query_lens, cuda_graph_pad_size,

203 - batch_size)

204 - device = self.runner.device

205 - use_captured_graph = cuda_graph_pad_size != -1

206 -

207 - if use_captured_graph:

208 - last_paged_kv_indptr = self.paged_kv_indptr[-1]

209 - self.paged_kv_indptr.extend([last_paged_kv_indptr] *

210 - cuda_graph_pad_size)

211 - self.paged_kv_last_page_lens.extend([0] * cuda_graph_pad_size)

212 - last_qo_indptr = self.qo_indptr[-1]

213 - self.qo_indptr.extend([last_qo_indptr] * cuda_graph_pad_size)

214 -

215 - # For current version of AITER MLA

216 - if len(self.paged_kv_indptr) > 0:

217 - # extend to the maximum number of blocks as returned by the

218 - # scheduler

219 - self.paged_kv_indices.extend(

220 - [0] * (self.total_blocks - len(self.paged_kv_indices)))

221 - paged_kv_indices_tensor = torch.tensor(self.paged_kv_indices,

222 - device=device,

223 - dtype=torch.int)

224 - paged_kv_indptr_tensor = torch.tensor(self.paged_kv_indptr,

225 - device=device,

226 - dtype=torch.int)

227 - paged_kv_last_page_lens_tensor = torch.tensor(

228 - self.paged_kv_last_page_lens, device=device, dtype=torch.int)

229 - block_table_bound_tensor = torch.zeros(len(self.paged_kv_indptr) -

230 - 1,

231 - device=device,

232 - dtype=torch.int)

233 -

234 - qo_indptr = torch.tensor(self.qo_indptr,

235 - device=device,

236 - dtype=torch.int)

237 - else:

238 - paged_kv_indices_tensor = None

239 - paged_kv_indptr_tensor = None

240 - paged_kv_last_page_lens_tensor = None

241 - block_table_bound_tensor = None

242 - qo_indptr = None

243 -

244 - metadata.paged_kv_indptr = paged_kv_indptr_tensor

245 - metadata.paged_kv_indices = paged_kv_indices_tensor

246 - metadata.paged_kv_last_page_lens = paged_kv_last_page_lens_tensor

247 - metadata.block_table_bound = block_table_bound_tensor

248 - metadata.qo_indptr = qo_indptr

249 -

250 - return metadata

251 -

252 -

253 -class AiterMLAState(MLACommonState[AiterMLAMetadata]):

254 -

255 - @contextmanager

256 - def graph_capture(self, max_batch_size: int):

257 - kv_indices, kv_indptr, last_page_lens, qo_indptr = \

258 - get_aiter_mla_metadata(

259 - max_batch_size=max_batch_size,

260 - block_size=self.runner.block_size,

261 - max_block_per_batch=\

262 - self.runner.get_max_block_per_batch(),

263 - device=self.runner.device)

264 - self._paged_kv_indices_tensor = kv_indices

265 - self._paged_kv_indptr_tensor = kv_indptr

266 - self._paged_kv_last_page_lens_tensor = last_page_lens

267 - self._qo_indptr_tensor = qo_indptr

268 -

269 - with super().graph_capture(max_batch_size):

270 - yield

271 -

272 - del self._paged_kv_indices_tensor

273 - del self._paged_kv_indptr_tensor

274 - del self._paged_kv_last_page_lens_tensor

275 - del self._qo_indptr_tensor

276 -

277 - def graph_capture_get_metadata_for_batch(

278 - self,

279 - batch_size: int,

280 - is_encoder_decoder_model: bool = False) -> AiterMLAMetadata:

281 -

282 - metadata = super().graph_capture_get_metadata_for_batch(

283 - batch_size, is_encoder_decoder_model)

284 -

285 - paged_kv_indptr = self._paged_kv_indptr_tensor[:batch_size + 1]

286 - paged_kv_indices = self._paged_kv_indices_tensor

287 - paged_kv_last_page_lens = self._paged_kv_last_page_lens_tensor[:

288 - batch_size]

289 - qo_indptr = self._qo_indptr_tensor[:batch_size + 1]

290 -

291 - metadata.paged_kv_indptr = paged_kv_indptr

292 - metadata.paged_kv_indices = paged_kv_indices

293 - metadata.paged_kv_last_page_lens = paged_kv_last_page_lens

294 - metadata.qo_indptr = qo_indptr

295 -

296 - return metadata

297 -

298 - def get_graph_input_buffers(self,

299 - attn_metadata: AiterMLAMetadata,

300 - is_encoder_decoder_model: bool = False):

301 - input_buffers = super().get_graph_input_buffers(

302 - attn_metadata, is_encoder_decoder_model)

303 - input_buffers[

304 - 'paged_kv_indptr'] = attn_metadata.decode_metadata.paged_kv_indptr

305 - input_buffers[

306 - "paged_kv_indices"] = attn_metadata.\

307 - decode_metadata.paged_kv_indices

308 - input_buffers[

309 - "paged_kv_last_page_lens"] = attn_metadata.\

310 - decode_metadata.paged_kv_last_page_lens

311 - input_buffers['qo_indptr'] = attn_metadata.qo_indptr

312 -

313 - return input_buffers

314 -

315 - def prepare_graph_input_buffers(self,

316 - input_buffers,

317 - attn_metadata: AiterMLAMetadata,

318 - is_encoder_decoder_model: bool = False):

319 - super().prepare_graph_input_buffers(input_buffers, attn_metadata,

320 - is_encoder_decoder_model)

321 -

322 - num_total_blocks = attn_metadata.decode_metadata.paged_kv_indices.shape[

323 - 0]

324 - input_buffers["paged_kv_indptr"].copy_(

325 - attn_metadata.decode_metadata.paged_kv_indptr, non_blocking=True)

326 - input_buffers["paged_kv_indices"][:num_total_blocks].copy_(

327 - attn_metadata.decode_metadata.paged_kv_indices, non_blocking=True)

328 - input_buffers["paged_kv_last_page_lens"].copy_(

329 - attn_metadata.decode_metadata.paged_kv_last_page_lens,

330 - non_blocking=True)

331 - input_buffers["qo_indptr"].copy_(

332 - attn_metadata.decode_metadata.qo_indptr, non_blocking=True)

333 -

334 -

335 -class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):

336 -

337 - def __init__(

338 - self,

339 - num_heads: int,

340 - head_size: int,

341 - scale: float,

342 - num_kv_heads: int,

343 - alibi_slopes: Optional[list[float]],

344 - sliding_window: Optional[int],

345 - kv_cache_dtype: str,

346 - logits_soft_cap: Optional[float],

347 - attn_type: str,

348 - kv_sharing_target_layer_name: Optional[str],

349 - # MLA Specific Arguments

350 - **mla_args) -> None:

351 - super().__init__(num_heads, head_size, scale, num_kv_heads,

352 - alibi_slopes, sliding_window, kv_cache_dtype,

353 - logits_soft_cap, attn_type,

354 - kv_sharing_target_layer_name, **mla_args)

355 -

356 - unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]

357 - if any(unsupported_features):

358 - raise NotImplementedError(

359 - "Aiter MLA does not support one of the following: "

360 - "alibi_slopes, sliding_window, logits_soft_cap")

361 -

362 - from aiter import flash_attn_varlen_func

363 - self.flash_attn_varlen_func = flash_attn_varlen_func

364 -

365 - def _flash_attn_varlen_diff_headdims(

366 - self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,

367 - softmax_scale: float, return_softmax_lse: bool,

368 - **kwargs) -> Union[tuple[torch.Tensor, ...], torch.Tensor]:

369 - output = self.flash_attn_varlen_func(

370 - q,

371 - k,

372 - v,

373 - **kwargs,

374 - )

375 -

376 - return output

377 -

378 - def _forward_decode(

379 - self,

380 - q_nope: torch.Tensor,

381 - q_pe: torch.Tensor,

382 - kv_c_and_k_pe_cache: torch.Tensor,

383 - attn_metadata: AiterMLAMetadata,

384 - ) -> torch.Tensor:

385 - assert kv_c_and_k_pe_cache.numel() > 0

386 -

387 - decode_meta = attn_metadata.decode_metadata

388 - assert decode_meta is not None

389 - B = q_nope.shape[0]

390 -

391 - q = torch.cat([q_nope, q_pe], dim=-1)

392 - o = torch.empty(B,

393 - self.num_heads,

394 - self.kv_lora_rank,

395 - dtype=q.dtype,

396 - device=q.device)

397 -

398 - kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2)

399 -

400 - aiter_mla_decode_fwd(q, kv_buffer, o, self.scale,

401 - attn_metadata.qo_indptr,

402 - attn_metadata.max_query_len,

403 - attn_metadata.paged_kv_indptr,

404 - attn_metadata.paged_kv_indices,

405 - attn_metadata.paged_kv_last_page_lens)

406 -

407 - return self._v_up_proj(o)

vllm/attention/backends/rocm_flash_attn.py deleted

-953

@@ -1,953 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -"""Attention layer ROCm GPUs."""

4 -import itertools

5 -from dataclasses import dataclass

6 -from functools import cache

7 -from typing import List, Optional, Tuple, Type

8 -

9 -import torch

10 -

11 -import vllm.envs as envs

12 -from vllm import _custom_ops as ops

13 -from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,

14 - AttentionLayer,

15 - AttentionMetadata, AttentionType)

16 -from vllm.attention.backends.utils import (CommonAttentionState,

17 - CommonMetadataBuilder)

18 -from vllm.attention.ops.paged_attn import (PagedAttention,

19 - PagedAttentionMetadata)

20 -from vllm.config import get_current_vllm_config

21 -from vllm.logger import init_logger

22 -from vllm.model_executor.layers.quantization.utils.quant_utils import (

23 - QuantKey, kFp8StaticTensorSym)

24 -from vllm.platforms import current_platform

25 -

26 -logger = init_logger(__name__)

27 -_PARTITION_SIZE_ROCM = 256

28 -

29 -

30 -@cache

31 -def is_rocm_aiter_paged_attn_enabled() -> bool:

32 - return envs.VLLM_ROCM_USE_AITER_PAGED_ATTN \

33 - and envs.VLLM_ROCM_USE_AITER \

34 -

35 -

36 -@cache

37 -def _get_paged_attn_module() -> PagedAttention:

38 - """

39 - Initializes the appropriate PagedAttention module from `attention/ops`,

40 - which is used as helper function

41 - by `ROCmFlashAttentionImpl` and `ROCmFlashAttentionBackend`.

42 -

43 - The choice of attention module depends on whether

44 - AITER paged attention is enabled:

45 - - If enabled, `ROCmFlashAttentionImpl` uses `AITERPagedAttention`.

46 - - Otherwise, it defaults to using the original `PagedAttention`.

47 - """

48 - if is_rocm_aiter_paged_attn_enabled():

49 - # Import AITERPagedAttention only when the flag is enabled

50 - from vllm.attention.ops.rocm_aiter_paged_attn import (

51 - AITERPagedAttention)

52 - return AITERPagedAttention()

53 - return PagedAttention()

54 -

55 -

56 -class ROCmFlashAttentionBackend(AttentionBackend):

57 - accept_output_buffer: bool = True

58 -

59 - @staticmethod

60 - def get_name() -> str:

61 - return "ROCM_FLASH"

62 -

63 - @staticmethod

64 - def get_impl_cls() -> Type["ROCmFlashAttentionImpl"]:

65 - return ROCmFlashAttentionImpl

66 -

67 - @staticmethod

68 - def get_metadata_cls() -> Type["AttentionMetadata"]:

69 - return ROCmFlashAttentionMetadata

70 -

71 - @staticmethod

72 - def get_builder_cls() -> Type["ROCmFlashAttentionMetadataBuilder"]:

73 - return ROCmFlashAttentionMetadataBuilder

74 -

75 - @staticmethod

76 - def get_state_cls() -> Type["CommonAttentionState"]:

77 - return CommonAttentionState

78 -

79 - @staticmethod

80 - def get_kv_cache_shape(

81 - num_blocks: int,

82 - block_size: int,

83 - num_kv_heads: int,

84 - head_size: int,

85 - ) -> Tuple[int, ...]:

86 - paged_attn = _get_paged_attn_module()

87 - return paged_attn.get_kv_cache_shape(num_blocks, block_size,

88 - num_kv_heads, head_size)

89 -

90 - @staticmethod

91 - def swap_blocks(

92 - src_kv_cache: torch.Tensor,

93 - dst_kv_cache: torch.Tensor,

94 - src_to_dst: torch.Tensor,

95 - ) -> None:

96 - paged_attn = _get_paged_attn_module()

97 - paged_attn.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dst)

98 -

99 - @staticmethod

100 - def copy_blocks(

101 - kv_caches: List[torch.Tensor],

102 - src_to_dists: torch.Tensor,

103 - ) -> None:

104 - paged_attn = _get_paged_attn_module()

105 - paged_attn.copy_blocks(kv_caches, src_to_dists)

106 -

107 -

108 -@dataclass

109 -class ROCmFlashAttentionMetadata(AttentionMetadata, PagedAttentionMetadata):

110 - """Metadata for FlashAttentionBackend.

111 -

112 - NOTE: Any python object stored here is not updated when it is

113 - cuda-graph replayed. If you have values that need to be changed

114 - dynamically, it should be stored in tensor. The tensor has to be

115 - updated from `CUDAGraphRunner.forward` API.

116 - """

117 - # (batch_size,). The sequence length per sequence. Sequence length means

118 - # the computed tokens + new tokens None if it is a decoding.

119 - seq_lens: Optional[List[int]]

120 - # seq_lens stored as a tensor.

121 - seq_lens_tensor: Optional[torch.Tensor]

122 - # Maximum sequence length among prefill batch. 0 if there are decoding

123 - # requests only.

124 - max_prefill_seq_len: int

125 - # Maximum sequence length among decode batch. 0 if there are prefill

126 - # requests only.

127 - max_decode_seq_len: int

128 -

129 - # Whether or not if cuda graph is enabled.

130 - # Cuda-graph is currently enabled for decoding only.

131 - # TODO(woosuk): Move `use_cuda_graph` out since it's unrelated to attention.

132 - use_cuda_graph: bool

133 -

134 - # NOTE(sang): Definition of context_len, query_len, and seq_len.

135 - # |---------- N-1 iteration --------|

136 - # |---------------- N iteration ---------------------|

137 - # |- tokenA -|......................|-- newTokens ---|

138 - # |---------- context_len ----------|

139 - # |-------------------- seq_len ----------------------|

140 - # |-- query_len ---|

141 -

142 - # Maximum query length in the batch. None for decoding.

143 - max_query_len: Optional[int] = None

144 - # (batch_size + 1,). The cumulative subquery lengths of the sequences in

145 - # the batch, used to index into subquery. E.g., if the subquery length

146 - # is [4, 6], it is [0, 4, 10].

147 - query_start_loc: Optional[torch.Tensor] = None

148 - # (batch_size + 1,). The cumulative sequence lengths of the sequences in

149 - # the batch, used to index into sequence. E.g., if the sequence length is

150 - # [4, 6], it is [0, 4, 10].

151 - seq_start_loc: Optional[torch.Tensor] = None

152 - # (batch_size,) A tensor of context lengths (tokens that are computed

153 - # so far).

154 - context_lens_tensor: Optional[torch.Tensor] = None

155 -

156 - # Max number of query tokens among request in the batch.

157 - max_decode_query_len: Optional[int] = None

158 -

159 - _cached_prefill_metadata: Optional["ROCmFlashAttentionMetadata"] = None

160 - _cached_decode_metadata: Optional["ROCmFlashAttentionMetadata"] = None

161 -

162 - # Begin encoder attn & enc/dec cross-attn fields...

163 -

164 - # Encoder sequence lengths representation

165 - encoder_seq_lens: Optional[List[int]] = None

166 - encoder_seq_lens_tensor: Optional[torch.Tensor] = None

167 -

168 - # Maximum sequence length among encoder sequences

169 - max_encoder_seq_len: Optional[int] = None

170 -

171 - # Number of tokens input to encoder

172 - num_encoder_tokens: Optional[int] = None

173 -

174 - # Cross-attention memory-mapping data structures: slot mapping

175 - # and block tables

176 - cross_slot_mapping: Optional[torch.Tensor] = None

177 - cross_block_tables: Optional[torch.Tensor] = None

178 -

179 - @property

180 - def prefill_metadata(self) -> Optional["ROCmFlashAttentionMetadata"]:

181 - if self.num_prefills == 0:

182 - return None

183 -

184 - if self._cached_prefill_metadata is not None:

185 - return self._cached_prefill_metadata

186 -

187 - assert self.seq_lens is not None

188 - assert self.seq_lens_tensor is not None

189 - assert self.block_tables is not None

190 -

191 - self._cached_prefill_metadata = ROCmFlashAttentionMetadata(

192 - num_prefills=self.num_prefills,

193 - num_prefill_tokens=self.num_prefill_tokens,

194 - num_decode_tokens=0,

195 - slot_mapping=self.slot_mapping[:self.num_prefill_tokens],

196 - multi_modal_placeholder_index_maps=self.

197 - multi_modal_placeholder_index_maps,

198 - enable_kv_scales_calculation=self.enable_kv_scales_calculation,

199 - seq_lens=self.seq_lens[:self.num_prefills],

200 - seq_lens_tensor=self.seq_lens_tensor[:self.num_prefills],

201 - max_query_len=self.max_query_len,

202 - max_prefill_seq_len=self.max_prefill_seq_len,

203 - max_decode_seq_len=0,

204 - query_start_loc=None if self.query_start_loc is None else

205 - self.query_start_loc[:self.num_prefills + 1],

206 - seq_start_loc=None if self.seq_start_loc is None else

207 - self.seq_start_loc[:self.num_prefills + 1],

208 - context_lens_tensor=None if self.context_lens_tensor is None else

209 - self.context_lens_tensor[:self.num_prefills],

210 - block_tables=self.block_tables[:self.num_prefills],

211 - use_cuda_graph=False,

212 - # Begin encoder & cross attn fields below...

213 - encoder_seq_lens=self.encoder_seq_lens,

214 - encoder_seq_lens_tensor=self.encoder_seq_lens_tensor,

215 - max_encoder_seq_len=self.max_encoder_seq_len,

216 - cross_slot_mapping=self.cross_slot_mapping,

217 - cross_block_tables=self.cross_block_tables)

218 - return self._cached_prefill_metadata

219 -

220 - @property

221 - def decode_metadata(self) -> Optional["ROCmFlashAttentionMetadata"]:

222 - if self.num_decode_tokens == 0:

223 - return None

224 -

225 - if self._cached_decode_metadata is not None:

226 - return self._cached_decode_metadata

227 - assert self.block_tables is not None

228 - assert self.seq_lens_tensor is not None

229 -

230 - self._cached_decode_metadata = ROCmFlashAttentionMetadata(

231 - num_prefills=0,

232 - num_prefill_tokens=0,

233 - num_decode_tokens=self.num_decode_tokens,

234 - slot_mapping=self.slot_mapping[self.num_prefill_tokens:],

235 - multi_modal_placeholder_index_maps=None,

236 - enable_kv_scales_calculation=True,

237 - seq_lens=None,

238 - seq_lens_tensor=self.seq_lens_tensor[self.num_prefills:],

239 - max_query_len=None,

240 - max_prefill_seq_len=0,

241 - max_decode_seq_len=self.max_decode_seq_len,

242 - query_start_loc=None,

243 - seq_start_loc=None,

244 - context_lens_tensor=None,

245 - block_tables=self.block_tables[self.num_prefills:],

246 - use_cuda_graph=self.use_cuda_graph,

247 - # Begin encoder & cross attn fields below...

248 - encoder_seq_lens=self.encoder_seq_lens,

249 - encoder_seq_lens_tensor=self.encoder_seq_lens_tensor,

250 - max_encoder_seq_len=self.max_encoder_seq_len,

251 - cross_slot_mapping=self.cross_slot_mapping,

252 - cross_block_tables=self.cross_block_tables)

253 - # Batch may be composed of prefill|decodes, adjust query start indices

254 - # to refer to the start of decodes when the two are split apart.

255 - # E.g. in tokens:[3 prefills|6 decodes], query_start_loc=[3,9] => [0,6].

256 - if self._cached_decode_metadata.query_start_loc is not None:

257 - qs = self._cached_decode_metadata.query_start_loc

258 - self._cached_decode_metadata.query_start_loc = qs - qs[0]

259 - return self._cached_decode_metadata

260 -

261 -

262 -class ROCmFlashAttentionMetadataBuilder(

263 - CommonMetadataBuilder[ROCmFlashAttentionMetadata]):

264 -

265 - _metadata_cls = ROCmFlashAttentionMetadata

266 -

267 -

268 -def _make_alibi_bias(alibi_slopes: torch.Tensor,

269 - dtype: torch.dtype,

270 - seq_lens: Optional[List[int]],

271 - make_attn_mask: bool = True) -> List[torch.Tensor]:

272 - attn_biases = []

273 - if seq_lens:

274 - for seq_len in seq_lens:

275 - bias = torch.arange(seq_len, dtype=dtype)

276 - # NOTE(zhuohan): HF uses

277 - # `bias = bias[None, :].repeat(seq_len, 1)`

278 - # here. We find that both biases give the same results, but

279 - # the bias below more accurately follows the original ALiBi

280 - # paper.

281 - bias = bias[None, :] - bias[:, None]

282 -

283 - num_heads = alibi_slopes.shape[0]

284 - bias = bias[None, :].repeat(

285 - (num_heads, 1, 1)).to(alibi_slopes.device)

286 - bias.mul_(alibi_slopes[:, None, None])

287 - if make_attn_mask:

288 - inf_mask = torch.empty(

289 - (1, seq_len, seq_len),

290 - dtype=bias.dtype).fill_(-torch.inf).triu_(diagonal=1).to(

291 - alibi_slopes.device)

292 - attn_biases.append((bias + inf_mask).to(dtype))

293 - else:

294 - attn_biases.append(bias.to(dtype))

295 -

296 - return attn_biases

297 -

298 -

299 -def _get_seq_len_block_table_args(

300 - attn_metadata: ROCmFlashAttentionMetadata,

301 - attn_type: str,

302 -) -> tuple:

303 - '''

304 - The particular choice of sequence-length

305 - attributes which should be extracted from attn_metadata is dependent

306 - on the type of attention operation.

307 -

308 - Decoder attn -> select entirely decoder self-attention-related fields

309 - Encoder/decoder cross-attn -> select encoder sequence lengths

310 - Encoder attn -> select encoder sequence lengths fields

311 - Encoder-only attn -> select prefill sequence lengths with

312 - bidirectional attention

313 -

314 - Arguments:

315 -

316 - * attn_metadata: Attention metadata structure associated with attention op

317 - * attn_type: encoder attention, decoder self-attention,

318 - encoder/decoder cross-attention, encoder-only

319 -

320 - Returns:

321 -

322 - * Appropriate sequence-lengths tensors for query and key

323 - * Appropriate max sequence-length scalar

324 - * Causal masking flag

325 - '''

326 -

327 - if attn_type == AttentionType.ENCODER:

328 - assert attn_metadata.encoder_seq_lens is not None

329 - assert attn_metadata.encoder_seq_lens_tensor is not None

330 - query_seq_start_loc = torch.tensor(

331 - list(itertools.accumulate([0] + attn_metadata.encoder_seq_lens)),

332 - device=attn_metadata.encoder_seq_lens_tensor.device,

333 - dtype=attn_metadata.encoder_seq_lens_tensor.dtype)

334 - causal_mask = False

335 -

336 - # No block tables associated with encoder attention

337 - return (query_seq_start_loc, attn_metadata.max_encoder_seq_len,

338 - query_seq_start_loc, attn_metadata.max_encoder_seq_len,

339 - attn_metadata.encoder_seq_lens, causal_mask)

340 -

341 - elif attn_type == AttentionType.ENCODER_ONLY:

342 - # For encoder-only models, we use the prefill sequence lengths

343 - assert attn_metadata.seq_lens is not None

344 - assert attn_metadata.seq_lens_tensor is not None

345 - query_seq_start_loc = torch.tensor(

346 - list(itertools.accumulate([0] + attn_metadata.seq_lens)),

347 - device=attn_metadata.seq_lens_tensor.device,

348 - dtype=attn_metadata.seq_lens_tensor.dtype)

349 - max_seq_len = attn_metadata.max_prefill_seq_len

350 - # Encoder-only models typically use bidirectional attention

351 - causal_mask = False

352 -

353 - return (query_seq_start_loc, max_seq_len, query_seq_start_loc,

354 - max_seq_len, attn_metadata.seq_lens, causal_mask)

355 -

356 - elif attn_type == AttentionType.DECODER:

357 - # Decoder self-attention

358 - # Choose max_seq_len based on whether we are in prompt_run

359 - assert attn_metadata.seq_lens is not None

360 - assert attn_metadata.seq_lens_tensor is not None

361 - query_seq_start_loc = torch.tensor(

362 - list(itertools.accumulate([0] + attn_metadata.seq_lens)),

363 - device=attn_metadata.seq_lens_tensor.device,

364 - dtype=attn_metadata.seq_lens_tensor.dtype)

365 - max_seq_len = attn_metadata.max_prefill_seq_len

366 - causal_mask = True

367 -

368 - return (query_seq_start_loc, max_seq_len, query_seq_start_loc,

369 - max_seq_len, attn_metadata.seq_lens, causal_mask)

370 - elif attn_type == AttentionType.ENCODER_DECODER:

371 - assert attn_metadata.seq_lens is not None

372 - assert attn_metadata.encoder_seq_lens_tensor is not None

373 - query_start_loc = torch.tensor(

374 - list(itertools.accumulate([0] + attn_metadata.seq_lens)),

375 - device=attn_metadata.encoder_seq_lens_tensor.device,

376 - dtype=attn_metadata.encoder_seq_lens_tensor.dtype)

377 -

378 - assert attn_metadata.encoder_seq_lens is not None

379 - assert attn_metadata.seq_lens_tensor is not None

380 - key_seq_start_loc = torch.tensor(

381 - list(itertools.accumulate([0] + attn_metadata.encoder_seq_lens)),

382 - device=attn_metadata.seq_lens_tensor.device,

383 - dtype=attn_metadata.seq_lens_tensor.dtype)

384 - causal_mask = False

385 -

386 - # Enc/dec cross-attention KVs match encoder sequence length;

387 - # cross-attention utilizes special "cross" block tables

388 - return (query_start_loc, attn_metadata.max_prefill_seq_len,

389 - key_seq_start_loc, attn_metadata.max_encoder_seq_len,

390 - attn_metadata.seq_lens, causal_mask)

391 - else:

392 - raise AttributeError(f"Invalid attention type {str(attn_type)}")

393 -

394 -

395 -class ROCmFlashAttentionImpl(AttentionImpl):

396 - """

397 - If the input tensors contain prompt tokens, the layout is as follows:

398 - |<--------------- num_prompt_tokens -------------->|

400 -

401 - Otherwise, the layout is as follows:

402 - |<------------------ num_generation_tokens (M) ----------------->|

404 -

405 - Generation tokens can contain padding when cuda-graph is used.

406 - Currently, prompt tokens don't contain any padding.

407 -

408 - The prompts might have different lengths, while the generation tokens

409 - always have length 1.

410 -

411 - If chunked prefill is enabled, prefill tokens and decode tokens can be

412 - batched together in a flattened 1D query.

413 -

414 - |<----- num_prefill_tokens ---->|<------- num_decode_tokens ----------->|

415 - |<-prompt_0->|...|<-prompt_N-1->|<-generation_0->|...|<-generation_M-1->|

416 -

417 - Currently, cuda graph is disabled for chunked prefill, meaning there's no

418 - padding between prefill and decode tokens.

419 - """

420 -

421 - def __init__(

422 - self,

423 - num_heads: int,

424 - head_size: int,

425 - scale: float,

426 - num_kv_heads: int,

427 - alibi_slopes: Optional[List[float]],

428 - sliding_window: Optional[int],

429 - kv_cache_dtype: str,

430 - logits_soft_cap: Optional[float] = None,

431 - attn_type: str = AttentionType.DECODER,

432 - kv_sharing_target_layer_name: Optional[str] = None,

433 - use_irope: bool = False,

434 - ) -> None:

435 - if kv_sharing_target_layer_name is not None:

436 - raise NotImplementedError("KV sharing is not supported in V0 "

437 - "ROCM_FLASH backend.")

438 - if use_irope:

439 - logger.warning_once(

440 - "Using irope in ROCm Flash Attention is not supported yet, it "

441 - "will fail back to global attention for long context.")

442 - if use_irope:

443 - logger.warning(

444 - "Using irope in V0 is not supported yet, it will fall back "

445 - "to global attention for long context.")

446 - if logits_soft_cap is None:

447 - # In flash-attn, setting logits_soft_cap as 0 means no soft cap.

448 - self.logits_soft_cap = 0.0

449 - else:

450 - self.logits_soft_cap = logits_soft_cap

451 - self.attn_type = attn_type

452 - self.num_heads = num_heads

453 - self.head_size = head_size

454 - self.scale = float(scale)

455 - self.num_kv_heads = num_kv_heads

456 - if alibi_slopes is not None:

457 - alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)

458 - self.alibi_slopes = alibi_slopes

459 - self.sliding_window = ((sliding_window, sliding_window)

460 - if sliding_window is not None else (-1, -1))

461 - self.kv_cache_dtype = kv_cache_dtype

462 -

463 - self.num_queries_per_kv = self.num_heads // self.num_kv_heads

464 -

465 - self.paged_attn_module = _get_paged_attn_module()

466 - supported_head_sizes = self.paged_attn_module.get_supported_head_sizes(

467 - )

468 -

469 - if head_size not in supported_head_sizes:

470 - raise ValueError(

471 - f"Head size {head_size} is not supported by PagedAttention. "

472 - f"Supported head sizes are: {supported_head_sizes}.")

473 -

474 - self.use_naive_attn = False

475 - # NOTE: Allow for switching between Triton and CK. Defaulting to triton.

476 - self.use_triton_flash_attn = envs.VLLM_USE_TRITON_FLASH_ATTN

477 - if self.use_triton_flash_attn:

478 - if logits_soft_cap is not None:

479 - raise ValueError(

480 - "ROCm Triton FlashAttention does not support attention"

481 - " logits soft capping."

482 - " please try using the ROCm CK "

483 - "FA backend instead by setting the env var "

484 - "`VLLM_USE_TRITON_FLASH_ATTN=0`")

485 -

486 - from vllm.attention.ops.triton_flash_attention import ( # noqa: F401

487 - triton_attention)

488 - self.triton_attn_func = triton_attention

489 - logger.debug("Using Triton FA in ROCmBackend")

490 - if self.sliding_window != (-1, -1):

491 - logger.warning("ROCm Triton FA does not currently support "

492 - "sliding window attention. If using half "

493 - "precision, please try using the ROCm CK "

494 - "FA backend instead by setting the env var "

495 - "`VLLM_USE_TRITON_FLASH_ATTN=0`")

496 - else:

497 - # if not using triton, navi3x/navi21/navi10 do not use flash-attn

498 - # either

499 - if not current_platform.has_device_capability(90):

500 - self.use_naive_attn = True

501 - else:

502 - try:

503 - from flash_attn import flash_attn_varlen_func # noqa: F401

504 - self.fa_attn_func = flash_attn_varlen_func

505 - logger.debug("Using CK FA in ROCmBackend")

506 - except ModuleNotFoundError:

507 - self.use_naive_attn = True

508 -

509 - if self.use_naive_attn:

510 - if logits_soft_cap is not None:

511 - raise ValueError(

512 - "ROCm Naive FlashAttention does not support "

513 - "attention logits soft capping.")

514 -

515 - self.sdpa_attn_func = _sdpa_attention

516 - logger.debug("Using naive (SDPA) attention in ROCmBackend")

517 -

518 - self.aiter_kv_scales_initialized = False

519 - self.force_fp8_attention = (

520 - get_current_vllm_config() is not None

521 - and get_current_vllm_config().model_config.override_attention_dtype

522 - == "fp8")

523 -

524 - def repeat_kv(self, x: torch.Tensor, n_rep: int) -> torch.Tensor:

525 - """torch.repeat_interleave(x, dim=1, repeats=n_rep)"""

526 - tokens, n_kv_heads, head_dim = x.shape

527 - return (x[:, :,

528 - None, :].expand(tokens, n_kv_heads, n_rep,

529 - head_dim).reshape(tokens, n_kv_heads * n_rep,

530 - head_dim))

531 -

532 - def fused_output_quant_supported(self, quant_key: QuantKey):

533 - if self.use_triton_flash_attn:

534 - return quant_key == kFp8StaticTensorSym

535 -

536 - # Only supported in the Triton backend

537 - return False

538 -

539 - def forward(

540 - self,

541 - layer: AttentionLayer,

542 - query: torch.Tensor,

543 - key: torch.Tensor,

544 - value: torch.Tensor,

545 - kv_cache: torch.Tensor,

546 - attn_metadata: ROCmFlashAttentionMetadata,

547 - output: Optional[torch.Tensor] = None,

548 - output_scale: Optional[torch.Tensor] = None,

549 - output_block_scale: Optional[torch.Tensor] = None,

550 - ) -> torch.Tensor:

551 - """Forward pass with FlashAttention and PagedAttention.

552 -

553 - For decoder-only models: query, key and value must be non-None.

554 -

555 - For encoder/decoder models:

556 - * ROCmFlashAttentionImpl.forward() may be invoked for both self- and

557 - cross-attention layers.

558 - * For self-attention: query, key and value must be non-None.

559 - * For cross-attention:

560 - * Query must be non-None

561 - * During prefill, key and value must be non-None; key and value

562 - get cached for use during decode.

563 - * During decode, key and value may be None, since:

564 - (1) key and value tensors were cached during prefill, and

565 - (2) cross-attention key and value tensors do not grow during

566 - decode

567 -

568 - A note on how the attn_type (attention type enum) argument impacts

569 - attention forward() behavior:

570 -

571 - * DECODER: normal decoder-only behavior;

572 - use decoder self-attention block table

573 - * ENCODER: no KV caching; pass encoder sequence

574 - attributes (encoder_seq_lens/encoder_seq_lens_tensor/

575 - max_encoder_seq_len) to kernel, in lieu of decoder

576 - sequence attributes (seq_lens/seq_lens_tensor/max_seq_len)

577 - * ENCODER_DECODER: cross-attention behavior;

578 - use cross-attention block table for caching KVs derived

579 - from encoder hidden states; since KV sequence lengths

580 - will match encoder sequence lengths, pass encoder sequence

581 - attributes to kernel (encoder_seq_lens/encoder_seq_lens_tensor/

582 - max_encoder_seq_len)

583 - * ENCODER_ONLY: bidirectional attention with no KV caching;

584 - use prefill sequence attributes

585 -

586 - Args:

587 - layer: Attention layer instance.

588 - query: shape = [num_tokens, num_heads * head_size]

589 - key: shape = [num_tokens, num_kv_heads * head_size]

590 - value: shape = [num_tokens, num_kv_heads * head_size]

591 - kv_cache: KV cache tensor with shape

592 - [2, num_blocks, block_size * num_kv_heads * head_size].

593 - NOTE: kv_cache will be an empty tensor with shape [0]

594 - for profiling run.

595 - attn_metadata: Metadata for attention.

596 - output: Optional output tensor.

597 - output_scale: Optional output scale tensor.

598 - output_block_scale: Optional output block scale tensor.

599 - Returns:

600 - shape = [num_tokens, num_heads * head_size]

601 - """

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

603 -

604 - if output_scale is not None and not self.use_triton_flash_attn:

605 - raise NotImplementedError(

606 - "fused output quantization only supported for Triton"

607 - " implementation in ROCMFlashAttentionImpl for now")

608 -

609 - if output_block_scale is not None:

610 - raise NotImplementedError(

611 - "fused nvfp4 output quantization is not supported"

612 - " for ROCMFlashAttentionImpl")

613 -

614 - query = query.view(-1, self.num_heads, self.head_size)

615 - if key is not None:

616 - assert value is not None

617 - key = key.view(-1, self.num_kv_heads, self.head_size)

618 - value = value.view(-1, self.num_kv_heads, self.head_size)

619 - else:

620 - assert value is None

621 -

622 - paged_attn = self.paged_attn_module

623 -

624 - # Reshaping kv tensors is required for AITER paged attention kernel

625 - # because it works on a different tensor shape,

626 - # when the size of one element is one byte (int8/fp8 dtypes).

627 - # This reshaping is only required on the first forward call

628 - # and the kv cache must not be empty.

629 - if (is_rocm_aiter_paged_attn_enabled() and kv_cache.dtype.itemsize == 1

630 - and not self.aiter_kv_scales_initialized

631 - and kv_cache.shape != torch.Size([0])):

632 - num_blocks = kv_cache.shape[1]

633 - block_size = kv_cache.shape[2] // (self.num_kv_heads *

634 - self.head_size)

635 - k_scale = torch.empty((self.num_kv_heads, num_blocks * block_size),

636 - dtype=torch.float32,

637 - device=kv_cache.device)

638 - v_scale = torch.empty((self.num_kv_heads, num_blocks * block_size),

639 - dtype=torch.float32,

640 - device=kv_cache.device)

641 - self.aiter_kv_scales_initialized = True

642 - k_scale.fill_(layer._k_scale.item())

643 - v_scale.fill_(layer._v_scale.item())

644 - layer._k_scale = k_scale

645 - layer._v_scale = v_scale

646 -

647 - # Only update KV cache for decoder self-attention

648 - # and encoder-decoder cross-attention

649 - if self.attn_type not in [

650 - AttentionType.ENCODER, AttentionType.ENCODER_ONLY

651 - ] and kv_cache.numel() > 0:

652 - key_cache, value_cache = paged_attn.split_kv_cache(

653 - kv_cache, self.num_kv_heads, self.head_size)

654 -

655 - if key is not None and value is not None:

656 - # Reshape the input keys and values and store them in the

657 - # cache. If kv_cache is not provided, the new key and value

658 - # tensors are not cached. This happens during the initial

659 - # memory profiling run.

660 - paged_attn.write_to_paged_cache(

661 - key,

662 - value,

663 - key_cache,

664 - value_cache,

665 - attn_metadata.slot_mapping

666 - if self.attn_type != AttentionType.ENCODER_DECODER else

667 - attn_metadata.cross_slot_mapping,

668 - self.kv_cache_dtype,

669 - layer._k_scale,

670 - layer._v_scale,

671 - )

672 -

673 - if self.attn_type != AttentionType.ENCODER:

674 - num_prefill_tokens = attn_metadata.num_prefill_tokens

675 - elif self.attn_type == AttentionType.ENCODER_ONLY:

676 - # For encoder-only models, all tokens are processed in one go

677 - num_prefill_tokens = query.shape[0]

678 - else:

679 - assert attn_metadata.num_encoder_tokens is not None

680 - num_prefill_tokens = attn_metadata.num_encoder_tokens

681 -

682 - # Query for decode. KV is not needed because it is already cached.

683 - decode_query = query[num_prefill_tokens:]

684 - # QKV for prefill.

685 - query = query[:num_prefill_tokens]

686 -

687 - # For encoder-only and encoder models,

688 - # we process all tokens at once

689 - # For decoder and encoder-decoder,

690 - # we may need to limit key/value to prefill tokens

691 - if key is not None and value is not None \

692 - and self.attn_type not in [AttentionType.ENCODER_DECODER,

693 - AttentionType.ENCODER_ONLY]:

694 - key = key[:num_prefill_tokens]

695 - value = value[:num_prefill_tokens]

696 -

697 - if prefill_meta := attn_metadata.prefill_metadata:

698 - # Prompt run.

699 - # normal attention and DECODER

700 - if self.attn_type == AttentionType.DECODER and (

701 - kv_cache.numel() == 0 or prefill_meta.block_tables is None

702 - or prefill_meta.block_tables.numel() == 0):

703 - (query_seq_start_loc, query_max_seq_len, key_seq_start_loc,

704 - key_max_seq_len, seq_lens,

705 - causal_mask) = (prefill_meta.seq_start_loc,

706 - prefill_meta.max_prefill_seq_len,

707 - prefill_meta.seq_start_loc,

708 - prefill_meta.max_prefill_seq_len,

709 - attn_metadata.seq_lens, True)

710 - # prefix-enabled attention and ENCODER/ENCODER_DECODER

711 - else:

712 - (query_seq_start_loc, query_max_seq_len, key_seq_start_loc,

713 - key_max_seq_len, seq_lens,

714 - causal_mask) = _get_seq_len_block_table_args(

715 - prefill_meta, self.attn_type)

716 - # Prompt run.

717 - if kv_cache.numel() == 0 or prefill_meta.block_tables.numel() == 0:

718 - # triton attention

719 - # When block_tables are not filled, it means q and k are the

720 - # prompt, and they have the same length.

721 - attn_masks = None

722 - if self.use_triton_flash_attn:

723 - if self.alibi_slopes is not None:

724 - attn_masks = _make_alibi_bias(

725 - self.alibi_slopes,

726 - query.dtype,

727 - seq_lens,

728 - make_attn_mask=causal_mask) # type: ignore

729 -

730 - use_fp8_scales = (layer._q_scale and layer._k_scale

731 - and layer._v_scale and layer._prob_scale

732 - and (self.kv_cache_dtype == "fp8"

733 - or self.force_fp8_attention))

734 -

735 - full_scales = (

736 - layer._q_scale.item(), layer._k_scale.item(),

737 - layer._v_scale.item(),

738 - layer._prob_scale.item()) if use_fp8_scales else None

739 - self.triton_attn_func(

740 - query,

741 - key,

742 - value,

743 - output[:num_prefill_tokens],

744 - query_seq_start_loc,

745 - key_seq_start_loc,

746 - query_max_seq_len,

747 - key_max_seq_len,

748 - causal_mask,

749 - self.scale,

750 - attn_masks[0][None]

751 - if attn_masks is not None else None,

752 - full_scales,

753 - output_scale,

754 - )

755 - elif self.use_naive_attn:

756 - if self.num_kv_heads != self.num_heads:

757 - # Interleave for MQA workaround.

758 - key = self.repeat_kv(key, self.num_queries_per_kv)

759 - value = self.repeat_kv(value, self.num_queries_per_kv)

760 - if self.alibi_slopes is not None:

761 - attn_masks = _make_alibi_bias(

762 - self.alibi_slopes,

763 - query.dtype,

764 - attn_metadata.seq_lens,

765 - make_attn_mask=causal_mask) # type: ignore

766 - query = query.movedim(0, query.dim() - 2)

767 - key = key.movedim(0, key.dim() - 2)

768 - value = value.movedim(0, value.dim() - 2)

769 - # sdpa math backend attention

770 - self.sdpa_attn_func(

771 - query,

772 - key,

773 - value,

774 - output[:num_prefill_tokens],

775 - query_seq_start_loc,

776 - num_prefill_tokens,

777 - self.num_heads,

778 - self.head_size,

779 - self.scale,

780 - attn_masks,

781 - )

782 - else:

783 - # upstream FA does not support an output arg, copy

784 - output[:num_prefill_tokens] = self.fa_attn_func(

785 - q=query,

786 - k=key,

787 - v=value,

788 - cu_seqlens_q=query_seq_start_loc,

789 - cu_seqlens_k=key_seq_start_loc,

790 - max_seqlen_q=prefill_meta.max_prefill_seq_len,

791 - max_seqlen_k=key_max_seq_len,

792 - softmax_scale=self.scale,

793 - causal=causal_mask,

794 - window_size=self.sliding_window,

795 - alibi_slopes=self.alibi_slopes,

796 - softcap=self.logits_soft_cap,

797 - )

798 -

799 - else:

800 - # prefix-enabled attention -

801 - # not applicable for encoder-only models

802 - if self.attn_type != AttentionType.ENCODER_ONLY:

803 - output[:num_prefill_tokens] = paged_attn.forward_prefix(

804 - query,

805 - key,

806 - value,

807 - self.kv_cache_dtype,

808 - key_cache,

809 - value_cache,

810 - prefill_meta.block_tables,

811 - prefill_meta.query_start_loc,

812 - prefill_meta.seq_lens_tensor,

813 - prefill_meta.max_query_len,

814 - self.alibi_slopes,

815 - self.sliding_window[0],

816 - layer._k_scale,

817 - layer._v_scale,

818 - )

819 - # Skip decode phase for encoder-only models

820 - if (decode_meta := attn_metadata.decode_metadata) and (

821 - self.attn_type != AttentionType.ENCODER_ONLY):

822 - # Decoding run.

823 - # Whether to use rocm custom paged attention or not

824 - num_seqs, num_heads, head_size = decode_query.shape

825 - block_size = value_cache.shape[3]

826 - gqa_ratio = num_heads // self.num_kv_heads

827 - from vllm.platforms.rocm import use_rocm_custom_paged_attention

828 - use_custom = use_rocm_custom_paged_attention(

829 - decode_query.dtype, head_size, block_size, gqa_ratio,

830 - decode_meta.max_decode_seq_len, self.sliding_window,

831 - self.kv_cache_dtype, self.alibi_slopes)

832 -

833 - if use_custom:

834 - max_seq_len = (decode_meta.max_decode_seq_len if self.attn_type

835 - != AttentionType.ENCODER_DECODER else

836 - decode_meta.max_encoder_seq_len)

837 - assert max_seq_len is not None

838 - max_num_partitions = (

839 - (max_seq_len + _PARTITION_SIZE_ROCM - 1) //

840 - _PARTITION_SIZE_ROCM)

841 - assert _PARTITION_SIZE_ROCM % block_size == 0

842 - tmp_output = torch.empty(

843 - size=(num_seqs, num_heads, max_num_partitions, head_size),

844 - dtype=query.dtype,

845 - device=output.device,

846 - )

847 - exp_sums = torch.empty(

848 - size=(num_seqs, num_heads, max_num_partitions),

849 - dtype=torch.float32,

850 - device=output.device,

851 - )

852 - max_logits = torch.empty_like(exp_sums)

853 -

854 - query_start_loc = None

855 - ops.paged_attention_rocm(

856 - output[num_prefill_tokens:],

857 - exp_sums,

858 - max_logits,

859 - tmp_output,

860 - decode_query,

861 - key_cache,

862 - value_cache,

863 - self.num_kv_heads,

864 - self.scale,

865 - decode_meta.block_tables

866 - if self.attn_type != AttentionType.ENCODER_DECODER else

867 - decode_meta.cross_block_tables,

868 - decode_meta.seq_lens_tensor

869 - if self.attn_type != AttentionType.ENCODER_DECODER else

870 - decode_meta.encoder_seq_lens_tensor,

871 - query_start_loc,

872 - block_size,

873 - max_seq_len,

874 - self.alibi_slopes,

875 - self.kv_cache_dtype,

876 - layer._k_scale,

877 - layer._v_scale,

878 - output_scale,

879 - )

880 - else:

881 - # PagedAttention does not support fused quant, manually quantize

882 - if output_scale is None:

883 - out_pa = output[num_prefill_tokens:]

884 - else:

885 - out_pa = torch.empty_like(output[num_prefill_tokens:],

886 - dtype=query.dtype)

887 -

888 - out_pa[:] = paged_attn.forward_decode(

889 - decode_query,

890 - key_cache,

891 - value_cache,

892 - decode_meta.block_tables

893 - if self.attn_type != AttentionType.ENCODER_DECODER else

894 - decode_meta.cross_block_tables,

895 - decode_meta.seq_lens_tensor

896 - if self.attn_type != AttentionType.ENCODER_DECODER else

897 - decode_meta.encoder_seq_lens_tensor,

898 - decode_meta.max_decode_seq_len

899 - if self.attn_type != AttentionType.ENCODER_DECODER else

900 - decode_meta.max_encoder_seq_len,

901 - self.kv_cache_dtype,

902 - self.num_kv_heads,

903 - self.scale,

904 - self.alibi_slopes,

905 - layer._k_scale,

906 - layer._v_scale,

907 - )

908 -

909 - # Manually perform quantization

910 - if output_scale is not None:

911 - out_uq = out_pa.view(-1, self.num_heads * self.head_size)

912 - out_q = output.view(-1, self.num_heads * self.head_size)

913 - ops.scaled_fp8_quant(out_uq,

914 - output_scale,

915 - output=out_q[num_prefill_tokens:])

916 -

917 - # Reshape the output tensor.

918 - return output.view(-1, self.num_heads * self.head_size)

919 -

920 -

921 -def _sdpa_attention(

922 - query: torch.Tensor,

923 - key: torch.Tensor,

924 - value: torch.Tensor,

925 - output: torch.Tensor,

926 - seq_lens: torch.Tensor,

927 - num_tokens: int,

928 - num_heads: int,

929 - head_size: int,

930 - scale: float,

931 - attn_masks: Optional[List[torch.Tensor]] = None,

932 -) -> torch.Tensor:

933 - start = 0

934 - assert output.shape == (num_tokens, num_heads, head_size)

935 - assert output.dtype == query.dtype

936 - assert output.device == query.device

937 -

938 - for i, seq_len in enumerate(seq_lens):

939 - end = start + seq_len

940 - with torch.nn.attention.sdpa_kernel(

941 - torch.nn.attention.SDPBackend.MATH):

942 - sub_out = torch.nn.functional.scaled_dot_product_attention(

943 - query[:, start:end, :],

944 - key[:, start:end, :],

945 - value[:, start:end, :],

946 - dropout_p=0.0,

947 - is_causal=attn_masks is None,

948 - attn_mask=attn_masks[i] if attn_masks else None,

949 - scale=scale).movedim(query.dim() - 2, 0)

950 - output[start:end, :, :] = sub_out

951 - start = end

952 -

953 - return output

vllm/attention/backends/triton_mla.py deleted

-111

@@ -1,111 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -from typing import List, Optional, Type

5 -

6 -import torch

7 -

8 -from vllm.attention.backends.abstract import (AttentionType,

9 - is_quantized_kv_cache)

10 -from vllm.attention.backends.mla.common import (MLACommonBackend,

11 - MLACommonImpl,

12 - MLACommonMetadata)

13 -from vllm.attention.ops.triton_decode_attention import decode_attention_fwd

14 -

15 -

16 -class TritonMLABackend(MLACommonBackend):

17 -

18 - @staticmethod

19 - def get_name() -> str:

20 - return "TRITON_MLA"

21 -

22 - @staticmethod

23 - def get_impl_cls() -> Type["TritonMLAImpl"]:

24 - return TritonMLAImpl

25 -

26 -

27 -class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):

28 -

29 - def __init__(

30 - self,

31 - num_heads: int,

32 - head_size: int,

33 - scale: float,

34 - num_kv_heads: int,

35 - alibi_slopes: Optional[List[float]],

36 - sliding_window: Optional[int],

37 - kv_cache_dtype: str,

38 - logits_soft_cap: Optional[float],

39 - attn_type: str,

40 - kv_sharing_target_layer_name: Optional[str],

41 - # MLA Specific Arguments

42 - **mla_args) -> None:

43 - super().__init__(num_heads, head_size, scale, num_kv_heads,

44 - alibi_slopes, sliding_window, kv_cache_dtype,

45 - logits_soft_cap, attn_type,

46 - kv_sharing_target_layer_name, **mla_args)

47 -

48 - unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]

49 - if any(unsupported_features):

50 - raise NotImplementedError(

51 - "TritonMLAImpl does not support one of the following: "

52 - "alibi_slopes, sliding_window, logits_soft_cap")

53 -

54 - if attn_type != AttentionType.DECODER:

55 - raise NotImplementedError("Encoder self-attention and "

56 - "encoder/decoder cross-attention "

57 - "are not implemented for "

58 - "TritonMLAImpl")

59 -

60 - if is_quantized_kv_cache(self.kv_cache_dtype):

61 - raise NotImplementedError(

62 - "TritonMLA with FP8 KV cache not yet supported")

63 -

64 - def _forward_decode(

65 - self,

66 - q_nope: torch.Tensor,

67 - q_pe: torch.Tensor,

68 - kv_c_and_k_pe_cache: torch.Tensor,

69 - attn_metadata: MLACommonMetadata,

70 - ) -> torch.Tensor:

71 - assert kv_c_and_k_pe_cache.numel() > 0

72 -

73 - decode_meta = attn_metadata.decode_metadata

74 - assert decode_meta is not None

75 - B = q_nope.shape[0]

76 -

77 - q = torch.cat([q_nope, q_pe], dim=-1)

78 - o = torch.zeros(B,

79 - self.num_heads,

80 - self.kv_lora_rank,

81 - dtype=q.dtype,

82 - device=q.device)

83 -

84 - num_kv_splits = 4 # TODO: heuristic

85 -

86 - # TODO(lucas) Allocate ahead of time

87 - attn_logits = torch.empty(

88 - (

89 - B,

90 - self.num_heads,

91 - num_kv_splits,

92 - # NOTE(lucas) idk why the +1 is here but sglang has it so we

93 - # just mirror that

94 - self.kv_lora_rank + 1,

95 - ),

96 - dtype=torch.float32,

97 - device=q.device,

98 - )

99 -

100 - # Add a head dim of 1

101 - kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.unsqueeze(2)

102 - kv_c_cache = kv_c_and_k_pe_cache[..., :self.kv_lora_rank]

103 - PAGE_SIZE = kv_c_and_k_pe_cache.size(1)

104 -

105 - # Run MQA

106 - decode_attention_fwd(q, kv_c_and_k_pe_cache, kv_c_cache, o,

107 - decode_meta.block_tables,

108 - decode_meta.seq_lens_tensor, attn_logits,

109 - num_kv_splits, self.scale, PAGE_SIZE)

110 -

111 - return self._v_up_proj(o)

vllm/attention/backends/utils.py

+6 -8

@@ -338,10 +338,9 @@ class CommonAttentionState(AttentionState):

338 # The encoder decoder model works only with XFormers and

339 # Flash Attention backend. Assert the same.

340 assert self.runner.attn_backend.get_name() in \

341 - ["XFORMERS", "FLASH_ATTN", "ROCM_FLASH"], \

342 - f"Expected attn_backend name to be either 'XFORMERS'," \

343 - f"'ROCM_FLASH', or 'FLASH_ATTN', but " \

344 - f"got '{self.runner.attn_backend.get_name()}'"

341 + ["XFORMERS", "FLASH_ATTN"], \

342 + f"Expected attn_backend name to be either 'XFORMERS' or " \

343 + f"'FLASH_ATTN', but got '{self.runner.attn_backend.get_name()}'"

344 self._update_captured_metadata_for_enc_dec_model(

345 batch_size=batch_size, attn_metadata=attn_metadata)

346

@@ -360,10 +359,9 @@ class CommonAttentionState(AttentionState):

359 # The encoder decoder model works only with XFormers and

360 # Flash Attention backend. Assert the same.

361 assert self.runner.attn_backend.get_name() in \

363 - ["XFORMERS", "FLASH_ATTN", "ROCM_FLASH"], \

364 - f"Expected attn_backend name to be either 'XFORMERS'," \

365 - f"'ROCM_FLASH', or 'FLASH_ATTN', but " \

366 - f"got '{self.runner.attn_backend.get_name()}'"

362 + ["XFORMERS", "FLASH_ATTN"], \

363 + f"Expected attn_backend name to be either 'XFORMERS' or " \

364 + f"'FLASH_ATTN', but got '{self.runner.attn_backend.get_name()}'"

365 self._add_additional_input_buffers_for_enc_dec_model(

366 attn_metadata=attn_metadata, input_buffers=input_buffers)

367 return input_buffers

vllm/attention/backends/xformers.py deleted

-805

@@ -1,805 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -"""Attention layer with xFormers and PagedAttention."""

4 -from dataclasses import dataclass

5 -from typing import Dict, List, Optional, Tuple, Type

6 -

7 -import torch

8 -from xformers import ops as xops

9 -from xformers.ops.fmha.attn_bias import (AttentionBias,

10 - BlockDiagonalCausalMask,

11 - BlockDiagonalMask,

12 - LowerTriangularMaskWithTensorBias)

13 -

14 -from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,

15 - AttentionLayer,

16 - AttentionMetadata, AttentionType)

17 -from vllm.attention.backends.utils import (

18 - CommonAttentionState, CommonMetadataBuilder,

19 - get_num_prefill_decode_query_kv_tokens, get_seq_len_block_table_args,

20 - is_all_cross_attn_metadata_set, is_all_encoder_attn_metadata_set)

21 -from vllm.attention.ops.paged_attn import (PagedAttention,

22 - PagedAttentionMetadata)

23 -from vllm.logger import init_logger

24 -

25 -logger = init_logger(__name__)

26 -

27 -

28 -class XFormersBackend(AttentionBackend):

29 -

30 - @staticmethod

31 - def get_name() -> str:

32 - return "XFORMERS"

33 -

34 - @staticmethod

35 - def get_impl_cls() -> Type["XFormersImpl"]:

36 - return XFormersImpl

37 -

38 - @staticmethod

39 - def get_metadata_cls() -> Type["AttentionMetadata"]:

40 - return XFormersMetadata

41 -

42 - @staticmethod

43 - def get_builder_cls() -> Type["XFormersMetadataBuilder"]:

44 - return XFormersMetadataBuilder

45 -

46 - @staticmethod

47 - def get_state_cls() -> Type["CommonAttentionState"]:

48 - return CommonAttentionState

49 -

50 - @staticmethod

51 - def get_kv_cache_shape(

52 - num_blocks: int,

53 - block_size: int,

54 - num_kv_heads: int,

55 - head_size: int,

56 - ) -> Tuple[int, ...]:

57 - return PagedAttention.get_kv_cache_shape(num_blocks, block_size,

58 - num_kv_heads, head_size)

59 -

60 - @staticmethod

61 - def swap_blocks(

62 - src_kv_cache: torch.Tensor,

63 - dst_kv_cache: torch.Tensor,

64 - src_to_dst: Dict[int, int],

65 - ) -> None:

66 - PagedAttention.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dst)

67 -

68 - @staticmethod

69 - def copy_blocks(

70 - kv_caches: List[torch.Tensor],

71 - src_to_dists: torch.Tensor,

72 - ) -> None:

73 - PagedAttention.copy_blocks(kv_caches, src_to_dists)

74 -

75 -

76 -@dataclass

77 -class XFormersMetadata(AttentionMetadata, PagedAttentionMetadata):

78 - """Metadata for XFormersbackend.

79 -

80 - NOTE: Any python object stored here is not updated when it is

81 - cuda-graph replayed. If you have values that need to be changed

82 - dynamically, it should be stored in tensor. The tensor has to be

83 - updated from `CUDAGraphRunner.forward` API.

84 - """

85 -

86 - # |---------- N-1 iteration --------|

87 - # |---------------- N iteration ---------------------|

88 - # |- tokenA -|......................|-- newTokens ---|

89 - # |---------- context_len ----------|

90 - # |-------------------- seq_len ----------------------|

91 - # |-- query_len ---|

92 -

93 - # seq_lens stored as a tensor.

94 - seq_lens_tensor: Optional[torch.Tensor]

95 -

96 - # FIXME: It is for flash attn.

97 - # Maximum sequence length among prefill batch. 0 if there are decoding

98 - # requests only.

99 - max_prefill_seq_len: int

100 - # Maximum sequence length among decode batch. 0 if there are prefill

101 - # requests only.

102 - max_decode_seq_len: int

103 -

104 - # Whether or not if cuda graph is enabled.

105 - # Cuda-graph is currently enabled for decoding only.

106 - # TODO(woosuk): Move `use_cuda_graph` out since it's unrelated to attention.

107 - use_cuda_graph: bool

108 -

109 - # (batch_size,). The sequence length per sequence. Sequence length means

110 - # the computed tokens + new tokens None if it is a decoding.

111 - seq_lens: Optional[List[int]] = None

112 -

113 - # FIXME: It is for flash attn.

114 - # (batch_size + 1,). The cumulative sequence lengths of the sequences in

115 - # the batch, used to index into sequence. E.g., if the sequence length is

116 - # [4, 6], it is [0, 4, 10].

117 - seq_start_loc: Optional[torch.Tensor] = None

118 -

119 - # (batch_size,) A tensor of context lengths (tokens that are computed

120 - # so far).

121 - context_lens_tensor: Optional[torch.Tensor] = None

122 -

123 - # Maximum query length in the batch. None for decoding.

124 - max_query_len: Optional[int] = None

125 -

126 - # Max number of query tokens among request in the batch.

127 - max_decode_query_len: Optional[int] = None

128 -

129 - # (batch_size + 1,). The cumulative subquery lengths of the sequences in

130 - # the batch, used to index into subquery. E.g., if the subquery length

131 - # is [4, 6], it is [0, 4, 10].

132 - query_start_loc: Optional[torch.Tensor] = None

133 -

134 - # Self-attention prefill/decode metadata cache

135 - _cached_prefill_metadata: Optional["XFormersMetadata"] = None

136 - _cached_decode_metadata: Optional["XFormersMetadata"] = None

137 -

138 - # Begin encoder attn & enc/dec cross-attn fields...

139 -

140 - # Encoder sequence lengths representation

141 - encoder_seq_lens: Optional[List[int]] = None

142 - encoder_seq_lens_tensor: Optional[torch.Tensor] = None

143 - # FIXME: It is for flash attn.

144 - # (batch_size + 1,). The cumulative sequence lengths of the sequences in

145 - # the batch, used to index into sequence. E.g., if the sequence length is

146 - # [4, 6], it is [0, 4, 10].

147 - encoder_seq_start_loc: Optional[torch.Tensor] = None

148 -

149 - # Maximum sequence length among encoder sequences

150 - max_encoder_seq_len: Optional[int] = None

151 -

152 - # Number of tokens input to encoder

153 - num_encoder_tokens: Optional[int] = None

154 -

155 - # Cross-attention memory-mapping data structures: slot mapping

156 - # and block tables

157 - cross_slot_mapping: Optional[torch.Tensor] = None

158 - cross_block_tables: Optional[torch.Tensor] = None

159 -

160 - def __post_init__(self):

161 - # Set during the execution of the first attention op.

162 - # It is a list because it is needed to set per prompt

163 - # when alibi slopes is used. It is because of the limitation

164 - # from xformer API.

165 - # will not appear in the __repr__ and __init__

166 - self.attn_bias: Optional[List[AttentionBias]] = None

167 - self.encoder_attn_bias: Optional[List[AttentionBias]] = None

168 - self.cross_attn_bias: Optional[List[AttentionBias]] = None

169 -

170 - @property

171 - def is_all_encoder_attn_metadata_set(self):

172 - '''

173 - All attention metadata required for encoder attention is set.

174 - '''

175 - return is_all_encoder_attn_metadata_set(self)

176 -

177 - @property

178 - def is_all_cross_attn_metadata_set(self):

179 - '''

180 - All attention metadata required for enc/dec cross-attention is set.

181 -

182 - Superset of encoder attention required metadata.

183 - '''

184 - return is_all_cross_attn_metadata_set(self)

185 -

186 - @property

187 - def prefill_metadata(self) -> Optional["XFormersMetadata"]:

188 - if self.num_prefills == 0:

189 - return None

190 -

191 - if self._cached_prefill_metadata is not None:

192 - # Recover cached prefill-phase attention

193 - # metadata structure

194 - return self._cached_prefill_metadata

195 -

196 - assert ((self.seq_lens is not None)

197 - or (self.encoder_seq_lens is not None))

198 - assert ((self.seq_lens_tensor is not None)

199 - or (self.encoder_seq_lens_tensor is not None))

200 -

201 - # Compute some attn_metadata fields which default to None

202 - query_start_loc = (None if self.query_start_loc is None else

203 - self.query_start_loc[:self.num_prefills + 1])

204 - seq_start_loc = (None if self.seq_start_loc is None else

205 - self.seq_start_loc[:self.num_prefills + 1])

206 - slot_mapping = (None if self.slot_mapping is None else

207 - self.slot_mapping[:self.num_prefill_tokens])

208 - seq_lens = (None if self.seq_lens is None else

209 - self.seq_lens[:self.num_prefills])

210 - seq_lens_tensor = (None if self.seq_lens_tensor is None else

211 - self.seq_lens_tensor[:self.num_prefills])

212 - context_lens_tensor = (None if self.context_lens_tensor is None else

213 - self.context_lens_tensor[:self.num_prefills])

214 - block_tables = (None if self.block_tables is None else

215 - self.block_tables[:self.num_prefills])

216 -

217 - # Construct & cache prefill-phase attention metadata structure

218 - self._cached_prefill_metadata = XFormersMetadata(

219 - num_prefills=self.num_prefills,

220 - num_prefill_tokens=self.num_prefill_tokens,

221 - num_decode_tokens=0,

222 - slot_mapping=slot_mapping,

223 - multi_modal_placeholder_index_maps=self.

224 - multi_modal_placeholder_index_maps,

225 - enable_kv_scales_calculation=self.enable_kv_scales_calculation,

226 - seq_lens=seq_lens,

227 - seq_lens_tensor=seq_lens_tensor,

228 - max_query_len=self.max_query_len,

229 - max_prefill_seq_len=self.max_prefill_seq_len,

230 - max_decode_seq_len=0,

231 - query_start_loc=query_start_loc,

232 - seq_start_loc=seq_start_loc,

233 - context_lens_tensor=context_lens_tensor,

234 - block_tables=block_tables,

235 - use_cuda_graph=False,

236 - # Begin encoder & cross attn fields below...

237 - encoder_seq_lens=self.encoder_seq_lens,

238 - encoder_seq_lens_tensor=self.encoder_seq_lens_tensor,

239 - max_encoder_seq_len=self.max_encoder_seq_len,

240 - cross_slot_mapping=self.cross_slot_mapping,

241 - cross_block_tables=self.cross_block_tables)

242 - return self._cached_prefill_metadata

243 -

244 - @property

245 - def decode_metadata(self) -> Optional["XFormersMetadata"]:

246 - if self.num_decode_tokens == 0:

247 - return None

248 -

249 - if self._cached_decode_metadata is not None:

250 - # Recover cached decode-phase attention

251 - # metadata structure

252 - return self._cached_decode_metadata

253 - assert ((self.seq_lens_tensor is not None)

254 - or (self.encoder_seq_lens_tensor is not None))

255 -

256 - # Compute some attn_metadata fields which default to None

257 - slot_mapping = (None if self.slot_mapping is None else

258 - self.slot_mapping[self.num_prefill_tokens:])

259 - seq_lens_tensor = (None if self.seq_lens_tensor is None else

260 - self.seq_lens_tensor[self.num_prefills:])

261 - block_tables = (None if self.block_tables is None else

262 - self.block_tables[self.num_prefills:])

263 -

264 - # Construct & cache decode-phase attention metadata structure

265 - self._cached_decode_metadata = XFormersMetadata(

266 - num_prefills=0,

267 - num_prefill_tokens=0,

268 - num_decode_tokens=self.num_decode_tokens,

269 - slot_mapping=slot_mapping,

270 - multi_modal_placeholder_index_maps=None,

271 - enable_kv_scales_calculation=True,

272 - seq_lens_tensor=seq_lens_tensor,

273 - max_prefill_seq_len=0,

274 - max_decode_seq_len=self.max_decode_seq_len,

275 - block_tables=block_tables,

276 - use_cuda_graph=self.use_cuda_graph,

277 - # Begin encoder & cross attn fields below...

278 - encoder_seq_lens=self.encoder_seq_lens,

279 - encoder_seq_lens_tensor=self.encoder_seq_lens_tensor,

280 - max_encoder_seq_len=self.max_encoder_seq_len,

281 - cross_slot_mapping=self.cross_slot_mapping,

282 - cross_block_tables=self.cross_block_tables)

283 -

284 - # Batch may be composed of prefill|decodes, adjust query start indices

285 - # to refer to the start of decodes when the two are split apart.

286 - # E.g. in tokens:[3 prefills|6 decodes], query_start_loc=[3,9] => [0,6].

287 - if self._cached_decode_metadata.query_start_loc is not None:

288 - qs = self._cached_decode_metadata.query_start_loc

289 - self._cached_decode_metadata.query_start_loc = qs - qs[0]

290 - return self._cached_decode_metadata

291 -

292 -

293 -def _get_attn_bias(

294 - attn_metadata: XFormersMetadata,

295 - attn_type: str,

296 -) -> Optional[AttentionBias]:

297 - '''

298 - Extract appropriate attention bias from attention metadata

299 - according to attention type.

300 -

301 - Arguments:

302 -

303 - * attn_metadata: Attention metadata structure associated with attention

304 - * attn_type: encoder attention, decoder self-attention,

305 - encoder/decoder cross-attention

306 -

307 - Returns:

308 - * Appropriate attention bias value given the attention type

309 - '''

310 -

311 - if (attn_type == AttentionType.DECODER

312 - or attn_type == AttentionType.ENCODER_ONLY):

313 - return attn_metadata.attn_bias

314 - elif attn_type == AttentionType.ENCODER:

315 - return attn_metadata.encoder_attn_bias

316 - elif attn_type == AttentionType.ENCODER_DECODER:

317 - return attn_metadata.cross_attn_bias

318 - else:

319 - raise AttributeError(f"Invalid attention type {str(attn_type)}")

320 -

321 -

322 -def _set_attn_bias(

323 - attn_metadata: XFormersMetadata,

324 - attn_bias: List[Optional[AttentionBias]],

325 - attn_type: str,

326 -) -> None:

327 - '''

328 - Update appropriate attention bias field of attention metadata,

329 - according to attention type.

330 -

331 - Arguments:

332 -

333 - * attn_metadata: Attention metadata structure associated with attention

334 - * attn_bias: The desired attention bias value

335 - * attn_type: encoder attention, decoder self-attention,

336 - encoder/decoder cross-attention

337 - '''

338 -

339 - if (attn_type == AttentionType.DECODER

340 - or attn_type == AttentionType.ENCODER_ONLY):

341 - attn_metadata.attn_bias = attn_bias

342 - elif attn_type == AttentionType.ENCODER:

343 - attn_metadata.encoder_attn_bias = attn_bias

344 - elif attn_type == AttentionType.ENCODER_DECODER:

345 - attn_metadata.cross_attn_bias = attn_bias

346 - else:

347 - raise AttributeError(f"Invalid attention type {str(attn_type)}")

348 -

349 -

350 -class XFormersMetadataBuilder(CommonMetadataBuilder[XFormersMetadata]):

351 -

352 - _metadata_cls = XFormersMetadata

353 -

354 -

355 -class XFormersImpl(AttentionImpl[XFormersMetadata]):

356 - """

357 - If the input tensors contain prompt tokens, the layout is as follows:

358 - |<--------------- num_prefill_tokens ----------------->|

360 -

361 - Otherwise, the layout is as follows:

362 - |<----------------- num_decode_tokens ------------------>|

364 -

365 - Generation tokens can contain padding when cuda-graph is used.

366 - Currently, prompt tokens don't contain any padding.

367 -

368 - The prompts might have different lengths, while the generation tokens

369 - always have length 1.

370 -

371 - If chunked prefill is enabled, prefill tokens and decode tokens can be

372 - batched together in a flattened 1D query.

373 -

374 - |<----- num_prefill_tokens ---->|<------- num_decode_tokens --------->|

375 - |<-prefill_0->|...|<-prefill_N-1->|<--decode_0-->|...|<--decode_M-1-->|

376 -

377 - Currently, cuda graph is disabled for chunked prefill, meaning there's no

378 - padding between prefill and decode tokens.

379 - """

380 -

381 - def __init__(

382 - self,

383 - num_heads: int,

384 - head_size: int,

385 - scale: float,

386 - num_kv_heads: int,

387 - alibi_slopes: Optional[List[float]],

388 - sliding_window: Optional[int],

389 - kv_cache_dtype: str,

390 - logits_soft_cap: Optional[float] = None,

391 - attn_type: str = AttentionType.DECODER,

392 - kv_sharing_target_layer_name: Optional[str] = None,

393 - use_irope: bool = False,

394 - ) -> None:

395 - if kv_sharing_target_layer_name is not None:

396 - raise NotImplementedError("KV sharing is not supported in V0 "

397 - "XFORMERS backend.")

398 - if logits_soft_cap is not None:

399 - logger.warning_once("XFormers does not support logits soft cap. "

400 - "Outputs may be slightly off.")

401 - if use_irope:

402 - logger.warning_once(

403 - "Using irope in XFormers is not supported yet, it will fall"

404 - " back to global attention for long context.")

405 - self.num_heads = num_heads

406 - self.head_size = head_size

407 - self.scale = float(scale)

408 - self.num_kv_heads = num_kv_heads

409 - if alibi_slopes is not None:

410 - alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)

411 - self.alibi_slopes = alibi_slopes

412 - self.sliding_window = sliding_window

413 - self.kv_cache_dtype = kv_cache_dtype

414 -

415 - self.num_queries_per_kv = self.num_heads // self.num_kv_heads

416 -

417 - supported_head_sizes = PagedAttention.get_supported_head_sizes()

418 - if head_size not in supported_head_sizes:

419 - raise ValueError(

420 - f"Head size {head_size} is not supported by PagedAttention. "

421 - f"Supported head sizes are: {supported_head_sizes}.")

422 -

423 - self.attn_type = attn_type

424 -

425 - def forward(

426 - self,

427 - layer: AttentionLayer,

428 - query: torch.Tensor,

429 - key: Optional[torch.Tensor],

430 - value: Optional[torch.Tensor],

431 - kv_cache: torch.Tensor,

432 - attn_metadata: "XFormersMetadata",

433 - output: Optional[torch.Tensor] = None,

434 - output_scale: Optional[torch.Tensor] = None,

435 - output_block_scale: Optional[torch.Tensor] = None,

436 - ) -> torch.Tensor:

437 - """Forward pass with xFormers and PagedAttention.

438 -

439 - For decoder-only models: query, key and value must be non-None.

440 -

441 - For encoder/decoder models:

442 - * XFormersImpl.forward() may be invoked for both self- and cross-

443 - attention layers.

444 - * For self-attention: query, key and value must be non-None.

445 - * For cross-attention:

446 - * Query must be non-None

447 - * During prefill, key and value must be non-None; key and value

448 - get cached for use during decode.

449 - * During decode, key and value may be None, since:

450 - (1) key and value tensors were cached during prefill, and

451 - (2) cross-attention key and value tensors do not grow during

452 - decode

453 -

454 - A note on how the attn_type (attention type enum) argument impacts

455 - attention forward() behavior:

456 -

457 - * DECODER: normal decoder-only behavior;

458 - use decoder self-attention block table

459 - * ENCODER: no KV caching; pass encoder sequence

460 - attributes (encoder_seq_lens/encoder_seq_lens_tensor/

461 - max_encoder_seq_len) to kernel, in lieu of decoder

462 - sequence attributes (seq_lens/seq_lens_tensor/max_seq_len).

463 - Used for encoder branch of encoder-decoder models.

464 - * ENCODER_ONLY: no kv_caching, uses the normal attention

465 - attributes (seq_lens/seq_lens_tensor/max_seq_len).

466 - * ENCODER_DECODER: cross-attention behavior;

467 - use cross-attention block table for caching KVs derived

468 - from encoder hidden states; since KV sequence lengths

469 - will match encoder sequence lengths, pass encoder sequence

470 - attributes to kernel (encoder_seq_lens/encoder_seq_lens_tensor/

471 - max_encoder_seq_len)

472 -

473 - Args:

474 - layer: Attention layer instance.

475 - query: shape = [num_tokens, num_heads * head_size]

476 - key: shape = [num_tokens, num_kv_heads * head_size]

477 - value: shape = [num_tokens, num_kv_heads * head_size]

478 - kv_cache: KV cache tensor with shape

479 - [2, num_blocks, block_size * num_kv_heads * head_size].

480 - NOTE: kv_cache will be an empty tensor with shape [0]

481 - for profiling run.

482 - attn_metadata: Metadata for attention.

483 - output: Optional output tensor.

484 - output_scale: Optional output scale tensor.

485 - output_block_scale: Optional output block scale tensor.

486 - Returns:

487 - shape = [num_tokens, num_heads * head_size]

488 - """

489 - if output_scale is not None or output_block_scale is not None:

490 - raise NotImplementedError(

491 - "fused output quantization is not yet supported"

492 - " for XFormersImpl")

493 -

494 - attn_type = self.attn_type

495 - # Check that appropriate attention metadata attributes are

496 - # selected for the desired attention type

497 - if (attn_type == AttentionType.ENCODER

498 - and (not attn_metadata.is_all_encoder_attn_metadata_set)):

499 - raise AttributeError("Encoder attention requires setting "

500 - "encoder metadata attributes.")

501 -

502 - elif (attn_type == AttentionType.ENCODER_DECODER

503 - and (not attn_metadata.is_all_cross_attn_metadata_set)):

504 - raise AttributeError("Encoder/decoder cross-attention "

505 - "requires setting cross-attention "

506 - "metadata attributes.")

507 -

508 - query = query.view(-1, self.num_heads, self.head_size)

509 - if key is not None:

510 - assert value is not None

511 - key = key.view(-1, self.num_kv_heads, self.head_size)

512 - value = value.view(-1, self.num_kv_heads, self.head_size)

513 - else:

514 - assert value is None

515 -

516 - # Self-attention vs. cross-attention will impact

517 - # which KV cache memory-mapping & which

518 - # seqlen datastructures we utilize

519 -

520 - if (attn_type != AttentionType.ENCODER and kv_cache.numel() > 0):

521 - # KV-cache during decoder-self- or

522 - # encoder-decoder-cross-attention, but not

523 - # during encoder attention.

524 - #

525 - # Even if there are no new key/value pairs to cache,

526 - # we still need to break out key_cache and value_cache

527 - # i.e. for later use by paged attention

528 - key_cache, value_cache = PagedAttention.split_kv_cache(

529 - kv_cache, self.num_kv_heads, self.head_size)

530 -

531 - if (key is not None) and (value is not None):

532 -

533 - if attn_type == AttentionType.ENCODER_DECODER:

534 - # Update cross-attention KV cache (prefill-only)

535 - # During cross-attention decode, key & value will be None,

536 - # preventing this IF-statement branch from running

537 - updated_slot_mapping = attn_metadata.cross_slot_mapping

538 - else:

539 - # Update self-attention KV cache (prefill/decode)

540 - updated_slot_mapping = attn_metadata.slot_mapping

541 -

542 - # Reshape the input keys and values and store them in the cache.

543 - # If kv_cache is not provided, the new key and value tensors are

544 - # not cached. This happens during the initial memory

545 - # profiling run.

546 - PagedAttention.write_to_paged_cache(

547 - key, value, key_cache, value_cache, updated_slot_mapping,

548 - self.kv_cache_dtype, layer._k_scale, layer._v_scale)

549 - (num_prefill_query_tokens, num_prefill_kv_tokens,

550 - num_decode_query_tokens) = \

551 - get_num_prefill_decode_query_kv_tokens(attn_metadata, attn_type)

552 -

553 - output = torch.empty_like(query)

554 - # Query for decode. KV is not needed because it is already cached.

555 - decode_query = query[num_prefill_query_tokens:]

556 - # QKV for prefill.

557 - query = query[:num_prefill_query_tokens]

558 - if key is not None and value is not None:

559 - key = key[:num_prefill_kv_tokens]

560 - value = value[:num_prefill_kv_tokens]

561 -

562 - assert query.shape[0] == num_prefill_query_tokens

563 - assert decode_query.shape[0] == num_decode_query_tokens

564 -

565 - if prefill_meta := attn_metadata.prefill_metadata:

566 - # Prompt run.

567 - if kv_cache.numel() == 0 or prefill_meta.block_tables.numel() == 0:

568 - # normal attention.

569 - # block tables are empty if the prompt does not have a cached

570 - # prefix.

571 - out = self._run_memory_efficient_xformers_forward(

572 - query, key, value, prefill_meta, attn_type=attn_type)

573 - assert out.shape == output[:num_prefill_query_tokens].shape

574 - output[:num_prefill_query_tokens] = out

575 - else:

576 - assert attn_type != AttentionType.ENCODER_ONLY, (

577 - "Encoder-only models should not have prefix attention.")

578 -

579 - assert prefill_meta.query_start_loc is not None

580 - assert prefill_meta.max_query_len is not None

581 -

582 - # prefix-enabled attention

583 - # TODO(Hai) this triton kernel has regression issue (broke) to

584 - # deal with different data types between KV and FP8 KV cache,

585 - # to be addressed separately.

586 - out = PagedAttention.forward_prefix(

587 - query,

588 - key,

589 - value,

590 - self.kv_cache_dtype,

591 - key_cache,

592 - value_cache,

593 - prefill_meta.block_tables,

594 - prefill_meta.query_start_loc,

595 - prefill_meta.seq_lens_tensor,

596 - prefill_meta.max_query_len,

597 - self.alibi_slopes,

598 - self.sliding_window,

599 - layer._k_scale,

600 - layer._v_scale,

601 - )

602 - assert output[:num_prefill_query_tokens].shape == out.shape

603 - output[:num_prefill_query_tokens] = out

604 -

605 - if decode_meta := attn_metadata.decode_metadata:

606 - assert attn_type != AttentionType.ENCODER_ONLY, (

607 - "Encoder-only models should not have decode metadata.")

608 -

609 - (

610 - seq_lens_arg,

611 - max_seq_len_arg,

612 - block_tables_arg,

613 - ) = get_seq_len_block_table_args(decode_meta, False, attn_type)

614 -

615 - output[num_prefill_query_tokens:] = PagedAttention.forward_decode(

616 - decode_query,

617 - key_cache,

618 - value_cache,

619 - block_tables_arg,

620 - seq_lens_arg,

621 - max_seq_len_arg,

622 - self.kv_cache_dtype,

623 - self.num_kv_heads,

624 - self.scale,

625 - self.alibi_slopes,

626 - layer._k_scale,

627 - layer._v_scale,

628 - )

629 -

630 - # Reshape the output tensor.

631 - return output.view(-1, self.num_heads * self.head_size)

632 -

633 - def _run_memory_efficient_xformers_forward(

634 - self,

635 - query: torch.Tensor,

636 - key: torch.Tensor,

637 - value: torch.Tensor,

638 - attn_metadata: XFormersMetadata,

639 - attn_type: str = AttentionType.DECODER,

640 - ) -> torch.Tensor:

641 - """Attention for 1D query of multiple prompts. Multiple prompt

642 - tokens are flattened in to `query` input.

643 -

644 - See https://facebookresearch.github.io/xformers/components/ops.html

645 - for API spec.

646 -

647 - Args:

648 - query: shape = [num_prefill_tokens, num_heads, head_size]

649 - key: shape = [num_prefill_tokens, num_kv_heads, head_size]

650 - value: shape = [num_prefill_tokens, num_kv_heads, head_size]

651 - attn_metadata: Metadata for attention.

652 - attn_type: Select attention type, between encoder attention,

653 - decoder self-attention, or encoder/decoder cross-

654 - attention. Defaults to decoder self-attention,

655 - which is the vLLM default generally

656 - """

657 -

658 - original_query = query

659 - if self.num_kv_heads != self.num_heads:

660 - # GQA/MQA requires the shape [B, M, G, H, K].

661 - # Note that the output also has the same shape (which is different

662 - # from a spec from the doc).

663 - query = query.view(query.shape[0], self.num_kv_heads,

664 - self.num_queries_per_kv, query.shape[-1])

665 - key = key[:, :,

666 - None, :].expand(key.shape[0], self.num_kv_heads,

667 - self.num_queries_per_kv, key.shape[-1])

668 - value = value[:, :,

669 - None, :].expand(value.shape[0], self.num_kv_heads,

670 - self.num_queries_per_kv,

671 - value.shape[-1])

672 -

673 - # Set attention bias if not provided. This typically happens at

674 - # the very attention layer of every iteration.

675 - # FIXME(woosuk): This is a hack.

676 - attn_bias = _get_attn_bias(attn_metadata, attn_type)

677 - if attn_bias is None:

678 - if self.alibi_slopes is None:

679 -

680 - # Cross attention block of decoder branch of encoder-decoder

681 - # model uses seq_lens for dec / encoder_seq_lens for enc

682 - if (attn_type == AttentionType.ENCODER_DECODER):

683 - assert attn_metadata.seq_lens is not None

684 - assert attn_metadata.encoder_seq_lens is not None

685 -

686 - # Cross-attention mask is non-causal

687 - attn_bias = BlockDiagonalMask.from_seqlens(

688 - attn_metadata.seq_lens,

689 - attn_metadata.encoder_seq_lens,

690 - device=query.device)

691 -

692 - # Encoder branch of encoder-decoder model uses

693 - # attn_metadata.encoder_seq_lens

694 - elif attn_type == AttentionType.ENCODER:

695 -

696 - assert attn_metadata.encoder_seq_lens is not None

697 -

698 - # Encoder self-attention mask is non-causal

699 - attn_bias = BlockDiagonalMask.from_seqlens(

700 - attn_metadata.encoder_seq_lens, device=query.device)

701 -

702 - # Self-attention block of encoder-only model just

703 - # uses the seq_lens directly.

704 - elif attn_type == AttentionType.ENCODER_ONLY:

705 - assert attn_metadata.seq_lens is not None

706 -

707 - # Encoder self-attention mask is non-causal

708 - attn_bias = BlockDiagonalMask.from_seqlens(

709 - attn_metadata.seq_lens, device=query.device)

710 -

711 - # Self-attention block of decoder branch just

712 - # uses the seq_lens directly

713 - elif attn_type == AttentionType.DECODER:

714 - assert attn_metadata.seq_lens is not None

715 -

716 - # Decoder self-attention mask is causal

717 - attn_bias = BlockDiagonalCausalMask.from_seqlens(

718 - attn_metadata.seq_lens, device=query.device)

719 - else:

720 - raise ValueError("Unknown AttentionType: %s", attn_type)

721 -

722 - if self.sliding_window is not None:

723 - attn_bias = attn_bias.make_local_attention(

724 - self.sliding_window)

725 - attn_bias = [attn_bias]

726 - else:

727 - assert attn_type == AttentionType.DECODER

728 - assert attn_metadata.seq_lens is not None

729 - attn_bias = _make_alibi_bias(self.alibi_slopes,

730 - self.num_kv_heads, query.dtype,

731 - attn_metadata.seq_lens)

732 -

733 - _set_attn_bias(attn_metadata, attn_bias, attn_type)

734 -

735 - # No alibi slopes.

736 - # TODO(woosuk): Too many view operations. Let's try to reduce

737 - # them in the future for code readability.

738 - if self.alibi_slopes is None:

739 - # Add the batch dimension.

740 - query = query.unsqueeze(0)

741 - key = key.unsqueeze(0)

742 - value = value.unsqueeze(0)

743 - out = xops.memory_efficient_attention_forward(

744 - query,

745 - key,

746 - value,

747 - attn_bias=attn_bias[0],

748 - p=0.0,

749 - scale=self.scale)

750 - return out.view_as(original_query)

751 -

752 - # Attention with alibi slopes.

753 - # FIXME(woosuk): Because xformers does not support dynamic sequence

754 - # lengths with custom attention bias, we process each prompt one by

755 - # one. This is inefficient, especially when we have many short prompts.

756 - assert attn_metadata.seq_lens is not None

757 - output = torch.empty_like(original_query)

758 - start = 0

759 - for i, seq_len in enumerate(attn_metadata.seq_lens):

760 - end = start + seq_len

761 - out = xops.memory_efficient_attention_forward(

762 - query[None, start:end],

763 - key[None, start:end],

764 - value[None, start:end],

765 - attn_bias=attn_bias[i],

766 - p=0.0,

767 - scale=self.scale)

768 - # TODO(woosuk): Unnecessary copy. Optimize.

769 - output[start:end].copy_(out.view_as(original_query[start:end]))

770 - start += seq_len

771 - return output

772 -

773 -

774 -def _make_alibi_bias(

775 - alibi_slopes: torch.Tensor,

776 - num_kv_heads: int,

777 - dtype: torch.dtype,

778 - seq_lens: List[int],

779 -) -> List[AttentionBias]:

780 - attn_biases: List[AttentionBias] = []

781 - for seq_len in seq_lens:

782 - bias = torch.arange(seq_len, dtype=dtype)

783 - # NOTE(zhuohan): HF uses

784 - # `bias = bias[None, :].repeat(seq_len, 1)`

785 - # here. We find that both biases give the same results, but

786 - # the bias below more accurately follows the original ALiBi

787 - # paper.

788 - # Calculate a matrix where each element represents ith element- jth

789 - # element.

790 - bias = bias[None, :] - bias[:, None]

791 -

792 - padded_len = (seq_len + 7) // 8 * 8

793 - num_heads = alibi_slopes.shape[0]

794 - bias = torch.empty(

795 - 1, # batch size

796 - num_heads,

797 - seq_len,

798 - padded_len,

799 - device=alibi_slopes.device,

800 - dtype=dtype,

801 - )[:, :, :, :seq_len].copy_(bias)

802 - bias.mul_(alibi_slopes[:, None, None])

803 - attn_biases.append(LowerTriangularMaskWithTensorBias(bias))

804 -

805 - return attn_biases

vllm/config/model.py

+1 -6

@@ -32,8 +32,7 @@ from vllm.transformers_utils.config import (

32 from vllm.transformers_utils.runai_utils import (ObjectStorageModel,

33 is_runai_obj_uri)

34 from vllm.transformers_utils.utils import maybe_model_redirect

35 -from vllm.utils import (STR_DUAL_CHUNK_FLASH_ATTN_VAL, LayerBlockType,

36 - LazyLoader, common_broadcastable_dtype)

35 +from vllm.utils import LayerBlockType, LazyLoader, common_broadcastable_dtype

37 if TYPE_CHECKING:

38 from transformers import PretrainedConfig

@@ -1103,10 +1102,6 @@ class ModelConfig:

1102 self.hf_config.dual_chunk_attention_config[

1103 "sparse_attention_enabled"] = True

1104

1106 - if envs.VLLM_ATTENTION_BACKEND != STR_DUAL_CHUNK_FLASH_ATTN_VAL:

1107 - raise ValueError("please set VLLM_ATTENTION_BACKEND to "

1108 - f"{STR_DUAL_CHUNK_FLASH_ATTN_VAL}")

1109 -

1105 def verify_with_parallel_config(

1106 self,

1107 parallel_config: ParallelConfig,

vllm/distributed/kv_transfer/kv_connector/utils.py

+1 -1

@@ -44,7 +44,7 @@ class model_aware_kv_ops_helper:

44 # When VLLM_MLA_DISABLE=1, standard FA is used instead, leading

45 # to a kv_cache shape of [2, num_blks, blk_size,

46 # num_key_value_heads / tp, qk_nope_head_dim + qk_rope_head_dim].

47 - # For more details, see vllm/attention/backends/mla/common.py.

47 + # For more details, see vllm/v1/attention/backends/mla/common.py.

48 if self.is_deepseek_mla and self.use_mla_opt:

49 head_size = model_config.kv_lora_rank + \

50 model_config.qk_rope_head_dim

vllm/engine/arg_utils.py

+2 -13

@@ -44,8 +44,8 @@ from vllm.test_utils import MODEL_WEIGHTS_S3_BUCKET, MODELS_ON_S3

44 from vllm.transformers_utils.config import (get_model_path, is_interleaved,

45 maybe_override_with_speculators)

46 from vllm.transformers_utils.utils import check_gguf_file

47 -from vllm.utils import (STR_DUAL_CHUNK_FLASH_ATTN_VAL, FlexibleArgumentParser,

48 - GiB_bytes, get_ip, is_in_ray_actor)

47 +from vllm.utils import (FlexibleArgumentParser, GiB_bytes, get_ip,

48 + is_in_ray_actor)

49 from vllm.v1.sample.logits_processor import LogitsProcessor

51 # yapf: enable

@@ -1163,17 +1163,6 @@ class EngineArgs:

1163 self._set_default_args_v0(model_config)

1164 assert self.enable_chunked_prefill is not None

1165

1166 - if envs.VLLM_ATTENTION_BACKEND in [STR_DUAL_CHUNK_FLASH_ATTN_VAL]:

1167 - assert self.enforce_eager, (

1168 - "Cuda graph is not supported with DualChunkFlashAttention. "

1169 - "To run the model in eager mode, set 'enforce_eager=True' "

1170 - "or use '--enforce-eager' in the CLI.")

1171 - assert current_platform.is_cuda(), (

1172 - "DualChunkFlashAttention is only supported on CUDA platform.")

1173 - assert not use_v1, (

1174 - "DualChunkFlashAttention is not supported on V1 engine. "

1175 - "To run the model in V0 engine, try set 'VLLM_USE_V1=0'")

1176 -

1166 sliding_window: Optional[int] = None

1167 if not is_interleaved(model_config.hf_text_config):

1168 # Only set CacheConfig.sliding_window if the model is all sliding

vllm/envs.py

-1

@@ -529,7 +529,6 @@ environment_variables: dict[str, Callable[[], Any]] = {

529 # - "TORCH_SDPA": use torch.nn.MultiheadAttention

530 # - "FLASH_ATTN": use FlashAttention

531 # - "XFORMERS": use XFormers

532 - # - "ROCM_FLASH": use ROCmFlashAttention

532 # - "FLASHINFER": use flashinfer

533 # - "FLASHMLA": use FlashMLA

534 # - "FLASH_ATTN_MLA": use FlashAttention for MLA

vllm/model_executor/layers/mamba/mamba2_metadata.py

+12 -7

@@ -53,13 +53,18 @@ class Mamba2Metadata:

53 def get_platform_metadata_classes() -> tuple[type[AttentionMetadata], ...]:

54 """Returns the appropriate metadata classes for the current platform."""

55 if current_platform.is_rocm():

56 - from vllm.attention.backends.rocm_flash_attn import (

57 - ROCmFlashAttentionMetadata)

58 - return (ROCmFlashAttentionMetadata, PlaceholderAttentionMetadata)

59 - elif current_platform.is_cuda():

60 - from vllm.attention.backends.flash_attn import FlashAttentionMetadata

61 - from vllm.attention.backends.xformers import XFormersMetadata

62 - return (FlashAttentionMetadata, XFormersMetadata,

56 + from vllm.v1.attention.backends.rocm_aiter_fa import (

57 + AiterFlashAttentionMetadata)

58 + from vllm.v1.attention.backends.triton_attn import (

59 + TritonAttentionMetadata)

60 + return (AiterFlashAttentionMetadata, TritonAttentionMetadata,

61 + PlaceholderAttentionMetadata)

62 + if current_platform.is_cuda():

63 + from vllm.v1.attention.backends.flash_attn import (

64 + FlashAttentionMetadata)

65 + from vllm.v1.attention.backends.xformers import (

66 + XFormersAttentionMetadata)

67 + return (FlashAttentionMetadata, XFormersAttentionMetadata,

68 PlaceholderAttentionMetadata)

69 raise ValueError(

70 f"Unsupported platform for Mamba2: {current_platform.device_type}")

vllm/model_executor/models/deepseek_v2.py

+2 -1

@@ -478,7 +478,8 @@ class DeepseekV2MLAAttention(nn.Module):

478 Main reference: DeepseekV2 paper, and FlashInfer Implementation

479 (https://arxiv.org/abs/2405.04434 and https://github.com/flashinfer-ai/flashinfer/pull/551).

480

481 - For more info see MLACommonImpl in: vllm/attention/backends/mla/utils.py

481 + For more info see MLACommonImpl in:

482 + vllm/v1/attention/backends/mla/utils.py

483 """

484

485 def __init__(

vllm/platforms/cuda.py

+26 -113

@@ -226,8 +226,10 @@ class CudaPlatformBase(Platform):

226 kv_cache_dtype, block_size, use_v1, use_mla,

227 has_sink) -> str:

228 if use_mla:

229 - # TODO(lucas): refactor to be more concise

230 - # we should probably consider factoring out V1 here

229 + if not use_v1:

230 + raise RuntimeError(

231 + "MLA attention backends require the V1 engine. "

232 + "Set VLLM_USE_V1=1 to enable them.")

233

234 from vllm.attention.ops.flashmla import is_flashmla_supported

235 from vllm.attention.utils.fa_utils import flash_attn_supports_mla

@@ -246,35 +248,17 @@ class CudaPlatformBase(Platform):

248 use_triton = selected_backend == _Backend.TRITON_MLA or (

249 selected_backend is None)

250

249 - def _get_version(name, import_suffix) -> str:

250 - if use_v1:

251 - logger.info_once(f"Using {name} backend on V1 engine.")

252 - return f"vllm.v1.attention.backends.mla.{import_suffix}"

253 - else:

254 - logger.info_once(f"Using {name} backend.")

255 - return f"vllm.attention.backends.{import_suffix}"

256 -

251 if use_cutlassmla:

258 - if use_v1:

259 - logger.info_once("Using Cutlass MLA backend on V1 engine.")

260 - return ("vllm.v1.attention.backends.mla."

261 - "cutlass_mla.CutlassMLABackend")

262 - else:

263 - logger.warning(

264 - "Cutlass MLA backend is only supported on V1 engine")

252 + logger.info_once("Using Cutlass MLA backend on V1 engine.")

253 + return ("vllm.v1.attention.backends.mla."

254 + "cutlass_mla.CutlassMLABackend")

255 if use_flashinfermla:

266 - if use_v1:

267 - from vllm.v1.attention.backends.utils import (

268 - set_kv_cache_layout)

269 - set_kv_cache_layout("HND")

270 - logger.info_once(

271 - "Using FlashInfer MLA backend on V1 engine.")

272 - return ("vllm.v1.attention.backends.mla."

273 - "flashinfer_mla.FlashInferMLABackend")

274 - else:

275 - logger.warning(

276 - "FlashInfer MLA backend is only supported on V1 engine"

277 - )

256 + from vllm.v1.attention.backends.utils import (

257 + set_kv_cache_layout)

258 + set_kv_cache_layout("HND")

259 + logger.info_once("Using FlashInfer MLA backend on V1 engine.")

260 + return ("vllm.v1.attention.backends.mla."

261 + "flashinfer_mla.FlashInferMLABackend")

262 if use_flashmla:

263 if block_size != 64:

264 logger.warning(

@@ -282,20 +266,18 @@ class CudaPlatformBase(Platform):

266 " (currently only supports block size 64).",

267 block_size)

268 else:

285 - return _get_version("FlashMLA", "flashmla.FlashMLABackend")

286 - if use_flashattn:

287 - if use_v1:

288 - logger.info_once(

289 - "Using FlashAttention MLA backend on V1 engine.")

269 + logger.info_once("Using FlashMLA backend on V1 engine.")

270 return ("vllm.v1.attention.backends.mla."

291 - "flashattn_mla.FlashAttnMLABackend")

292 - else:

293 - logger.warning(

294 - "FlashAttention MLA backend is only supported on V1 "

295 - "engine.")

271 + "flashmla.FlashMLABackend")

272 + if use_flashattn:

273 + logger.info_once(

274 + "Using FlashAttention MLA backend on V1 engine.")

275 + return ("vllm.v1.attention.backends.mla."

276 + "flashattn_mla.FlashAttnMLABackend")

277 if use_triton:

297 - return _get_version("Triton MLA",

298 - "triton_mla.TritonMLABackend")

278 + logger.info_once("Using Triton MLA backend on V1 engine.")

279 + return ("vllm.v1.attention.backends.mla."

280 + "triton_mla.TritonMLABackend")

281 if use_v1:

282 FLASHINFER_V1 = "vllm.v1.attention.backends.flashinfer.FlashInferBackend" # noqa: E501

283 FLEX_ATTENTION_V1 = "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend" # noqa: E501

@@ -382,78 +364,9 @@ class CudaPlatformBase(Platform):

364 )

365 return FLEX_ATTENTION_V1

366

385 - # Backends for V0 engine

386 - if selected_backend == _Backend.XFORMERS:

387 - logger.info("Using XFormers backend.")

388 - return "vllm.attention.backends.xformers.XFormersBackend"

389 - elif selected_backend == _Backend.DUAL_CHUNK_FLASH_ATTN:

390 - logger.info("Using DualChunkFlashAttention backend.")

391 - return ("vllm.attention.backends.dual_chunk_flash_attn."

392 - "DualChunkFlashAttentionBackend")

393 - elif selected_backend == _Backend.DIFFERENTIAL_FLASH_ATTN:

394 - logger.info("Using DifferentialFlashAttention backend.")

395 - return ("vllm.attention.backends.differential_flash_attn."

396 - "DifferentialFlashAttentionBackend")

397 - elif selected_backend == _Backend.FLASH_ATTN:

398 - pass

399 - elif selected_backend:

400 - raise ValueError(

401 - f"Invalid attention backend for {cls.device_name}, "

402 - f"with use_v1: {use_v1} use_mla: {use_mla}")

403 -

404 - target_backend = _Backend.FLASH_ATTN

405 - if not cls.has_device_capability(80):

406 - # Volta and Turing NVIDIA GPUs.

407 - logger.info(

408 - "Cannot use FlashAttention-2 backend for Volta and Turing "

409 - "GPUs.")

410 - target_backend = _Backend.XFORMERS

411 - elif dtype not in (torch.float16, torch.bfloat16):

412 - logger.info(

413 - "Cannot use FlashAttention-2 backend for dtype other than "

414 - "torch.float16 or torch.bfloat16.")

415 - target_backend = _Backend.XFORMERS

416 - elif block_size % 16 != 0:

417 - logger.info(

418 - "Cannot use FlashAttention-2 backend for block size not "

419 - "divisible by 16.")

420 - target_backend = _Backend.XFORMERS

421 -

422 - # FlashAttn is valid for the model, checking if the package is

423 - # installed.

424 - if target_backend == _Backend.FLASH_ATTN:

425 - try:

426 - import vllm.vllm_flash_attn # noqa: F401

427 - from vllm.attention.backends.flash_attn import ( # noqa: F401

428 - FlashAttentionBackend, flash_attn_supports_fp8)

429 -

430 - supported_sizes = \

431 - FlashAttentionBackend.get_supported_head_sizes()

432 - if head_size not in supported_sizes:

433 - logger.info(

434 - "Cannot use FlashAttention-2 backend for head size %d.",

435 - head_size)

436 - target_backend = _Backend.XFORMERS

437 - fp8_kv_cache = (kv_cache_dtype is not None

438 - and kv_cache_dtype.startswith("fp8"))

439 - if (fp8_kv_cache and not flash_attn_supports_fp8()):

440 - logger.info(

441 - "Cannot use FlashAttention backend for FP8 KV cache.")

442 - target_backend = _Backend.XFORMERS

443 - except ImportError:

444 - logger.info(

445 - "Cannot use FlashAttention-2 backend because the "

446 - "vllm.vllm_flash_attn package is not found. "

447 - "Make sure that vllm_flash_attn was built and installed "

448 - "(on by default).")

449 - target_backend = _Backend.XFORMERS

450 -

451 - if target_backend == _Backend.XFORMERS:

452 - logger.info("Using XFormers backend.")

453 - return "vllm.attention.backends.xformers.XFormersBackend"

454 -

455 - logger.info("Using Flash Attention backend.")

456 - return "vllm.attention.backends.flash_attn.FlashAttentionBackend"

367 + raise RuntimeError(

368 + "V0 attention backends have been removed. Set VLLM_USE_V1=1 "

369 + "to select a supported backend.")

370

371 @classmethod

372 def get_punica_wrapper(cls) -> str:

vllm/platforms/rocm.py

+23 -38

@@ -191,6 +191,11 @@ class RocmPlatform(Platform):

191 kv_cache_dtype, block_size, use_v1, use_mla,

192 has_sink) -> str:

193 if use_mla:

194 + if not use_v1:

195 + raise RuntimeError(

196 + "MLA attention backends require the V1 engine. "

197 + "Set VLLM_USE_V1=1 to enable them.")

198 +

199 from vllm.v1.attention.backends.mla.rocm_aiter_mla import (

200 is_aiter_mla_enabled)

201

@@ -201,39 +206,24 @@ class RocmPlatform(Platform):

206

207 if selected_backend == _Backend.TRITON_MLA:

208 if block_size != 1:

204 - if use_v1:

205 - logger.info_once(

206 - "Using Triton MLA backend on V1 engine.")

207 - return ("vllm.v1.attention.backends.mla."

208 - "triton_mla.TritonMLABackend")

209 - else:

210 - logger.info("Using Triton MLA backend.")

211 - return "vllm.attention.backends.triton_mla.TritonMLABackend" # noqa: E501

212 - else:

213 - raise ValueError(

214 - f" The selected backend, {selected_backend.name},"

215 - f"does not support block size {block_size}.")

216 - elif selected_backend == _Backend.ROCM_AITER_MLA \

217 - or selected_backend == _Backend.ROCM_AITER_MLA_VLLM_V1:

209 + logger.info_once("Using Triton MLA backend on V1 engine.")

210 + return ("vllm.v1.attention.backends.mla."

211 + "triton_mla.TritonMLABackend")

212 + raise ValueError(

213 + f" The selected backend, {selected_backend.name},"

214 + f"does not support block size {block_size}.")

215 + if selected_backend in (_Backend.ROCM_AITER_MLA,

216 + _Backend.ROCM_AITER_MLA_VLLM_V1):

217 if block_size == 1:

219 - if use_v1:

220 - logger.info("Using AITER MLA backend on V1 engine.")

221 - return "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend" # noqa: E501

222 - else:

223 - logger.info("Using AITER MLA backend")

224 - return "vllm.attention.backends.rocm_aiter_mla.AiterMLABackend" # noqa: E501

225 - else:

226 - raise ValueError(

227 - f" The selected backend, {selected_backend.name},"

228 - f"does not support block size {block_size}."

229 - "(currently only supports block size 1)")

230 - else:

218 + logger.info("Using AITER MLA backend on V1 engine.")

219 + return "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend" # noqa: E501

220 raise ValueError(

221 f" The selected backend, {selected_backend.name},"

233 - f"is not MLA type while requested for MLA backend.")

234 -

235 - if selected_backend is None or selected_backend == _Backend.FLASH_ATTN:

236 - selected_backend = _Backend.ROCM_FLASH

222 + f"does not support block size {block_size}."

223 + "(currently only supports block size 1)")

224 + raise ValueError(

225 + f" The selected backend, {selected_backend.name},"

226 + f"is not MLA type while requested for MLA backend.")

227

228 if envs.VLLM_USE_V1:

229 if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA \

@@ -245,14 +235,9 @@ class RocmPlatform(Platform):

235 logger.info("Using Triton Attention backend on V1 engine.")

236 return ("vllm.v1.attention.backends."

237 "triton_attn.TritonAttentionBackend")

248 - if selected_backend == _Backend.ROCM_FLASH:

249 - if not cls.has_device_capability(90):

250 - # not Instinct series GPUs.

251 - logger.info("flash_attn is not supported on NAVI GPUs.")

252 - else:

253 - logger.info("%s is not supported in AMD GPUs.", selected_backend)

254 - logger.info("Using ROCmFlashAttention backend.")

255 - return "vllm.attention.backends.rocm_flash_attn.ROCmFlashAttentionBackend" # noqa: E501

238 + raise RuntimeError(

239 + "V0 attention backends have been removed. Set VLLM_USE_V1=1 "

240 + "to select a supported backend.")

241

242 @classmethod

243 def set_device(cls, device: torch.device) -> None:

vllm/utils/__init__.py

-2

@@ -157,10 +157,8 @@ STR_BACKEND_ENV_VAR: str = "VLLM_ATTENTION_BACKEND"

157 # register, corresponding to possible backends

158 STR_FLASHINFER_ATTN_VAL: str = "FLASHINFER"

159 STR_TORCH_SDPA_ATTN_VAL: str = "TORCH_SDPA"

160 -STR_ROCM_FLASH_ATTN_VAL: str = "ROCM_FLASH"

160 STR_XFORMERS_ATTN_VAL: str = "XFORMERS"

161 STR_FLASH_ATTN_VAL: str = "FLASH_ATTN"

163 -STR_DUAL_CHUNK_FLASH_ATTN_VAL: str = "DUAL_CHUNK_FLASH_ATTN"

162 STR_INVALID_VAL: str = "INVALID"

163

164 MB_bytes = 1_000_000