| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from dataclasses import dataclass |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | from vllm.config import CompilationConfig, VllmConfig |
| 8 | from vllm.v1.attention.backend import ( |
| 9 | AttentionBackend, |
| 10 | AttentionCGSupport, |
| 11 | AttentionMetadataBuilder, |
| 12 | CommonAttentionMetadata, |
| 13 | ) |
| 14 | from vllm.v1.attention.backends.utils import ( |
| 15 | PAD_SLOT_ID, |
| 16 | mamba_get_block_table_tensor, |
| 17 | split_decodes_and_prefills, |
| 18 | ) |
| 19 | from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec |
| 20 | |
| 21 | |
| 22 | class LinearAttentionBackend(AttentionBackend): |
| 23 | @staticmethod |
| 24 | def get_name() -> str: |
| 25 | return "LINEAR_ATTN" |
| 26 | |
| 27 | @staticmethod |
| 28 | def get_builder_cls() -> type["LinearAttentionMetadataBuilder"]: |
| 29 | return LinearAttentionMetadataBuilder |
| 30 | |
| 31 | @classmethod |
| 32 | def is_ssm(cls) -> bool: |
| 33 | return True |
| 34 | |
| 35 | |
| 36 | @dataclass |
| 37 | class LinearAttentionMetadata: |
| 38 | num_prefills: int |
| 39 | num_prefill_tokens: int |
| 40 | num_decodes: int |
| 41 | num_decode_tokens: int |
| 42 | query_start_loc: torch.Tensor |
| 43 | seq_lens: torch.Tensor |
| 44 | |
| 45 | state_indices_tensor: torch.Tensor # shape: [batch,] |
| 46 | |
| 47 | |
| 48 | class LinearAttentionMetadataBuilder(AttentionMetadataBuilder[LinearAttentionMetadata]): |
| 49 | kv_cache_spec: MambaSpec |
| 50 | reorder_batch_threshold: int = 1 |
| 51 | |
| 52 | _cudagraph_support = AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE |
| 53 | |
| 54 | def __init__( |
| 55 | self, |
| 56 | kv_cache_spec: MambaSpec, |
| 57 | layer_names: list[str], |
| 58 | vllm_config: VllmConfig, |
| 59 | device: torch.device, |
| 60 | ): |
| 61 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 62 | |
| 63 | def build( |
| 64 | self, |
| 65 | common_prefix_len: int, |
| 66 | common_attn_metadata: CommonAttentionMetadata, |
| 67 | fast_build: bool = False, |
| 68 | ) -> LinearAttentionMetadata: |
| 69 | query_start_loc = common_attn_metadata.query_start_loc |
| 70 | seq_lens = common_attn_metadata.seq_lens |
| 71 | |
| 72 | state_indices_tensor = mamba_get_block_table_tensor( |
| 73 | common_attn_metadata.block_table_tensor, |
| 74 | common_attn_metadata.seq_lens, |
| 75 | self.kv_cache_spec, |
| 76 | self.vllm_config.cache_config.mamba_cache_mode, |
| 77 | )[:, 0] |
| 78 | |
| 79 | num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( |
| 80 | split_decodes_and_prefills( |
| 81 | common_attn_metadata, decode_threshold=self.reorder_batch_threshold |
| 82 | ) |
| 83 | ) |
| 84 | |
| 85 | attn_metadata = LinearAttentionMetadata( |
| 86 | num_prefills=num_prefills, |
| 87 | num_prefill_tokens=num_prefill_tokens, |
| 88 | num_decodes=num_decodes, |
| 89 | num_decode_tokens=num_decode_tokens, |
| 90 | query_start_loc=query_start_loc, |
| 91 | seq_lens=seq_lens, |
| 92 | state_indices_tensor=state_indices_tensor, |
| 93 | ) |
| 94 | return attn_metadata |
| 95 | |
| 96 | |
| 97 | class BailingLinearAttentionBackend(LinearAttentionBackend): |
| 98 | @staticmethod |
| 99 | def get_name() -> str: |
| 100 | return "BAILING_LINEAR_ATTN" |
| 101 | |
| 102 | @staticmethod |
| 103 | def get_builder_cls() -> type["BailingLinearAttentionMetadataBuilder"]: |
| 104 | return BailingLinearAttentionMetadataBuilder |
| 105 | |
| 106 | |
| 107 | @dataclass |
| 108 | class BailingLinearAttentionMetadata(LinearAttentionMetadata): |
| 109 | state_indices_tensor_d: torch.Tensor | None = None |
| 110 | state_indices_tensor_p: torch.Tensor | None = None |
| 111 | num_accepted_tokens: torch.Tensor | None = None |
| 112 | query_start_loc_d: torch.Tensor | None = None |
| 113 | |
| 114 | |
| 115 | class BailingLinearAttentionMetadataBuilder(LinearAttentionMetadataBuilder): |
| 116 | supports_spec_decode_metadata = True |
| 117 | supports_update_block_table: bool = False |
| 118 | |
| 119 | @classmethod |
| 120 | def get_cudagraph_support( |
| 121 | cls, |
| 122 | vllm_config: VllmConfig, |
| 123 | kv_cache_spec: KVCacheSpec, |
| 124 | ) -> AttentionCGSupport: |
| 125 | return AttentionCGSupport.UNIFORM_BATCH |
| 126 | |
| 127 | def __init__( |
| 128 | self, |
| 129 | kv_cache_spec: MambaSpec, |
| 130 | layer_names: list[str], |
| 131 | vllm_config: VllmConfig, |
| 132 | device: torch.device, |
| 133 | ): |
| 134 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 135 | self.compilation_config: CompilationConfig = vllm_config.compilation_config |
| 136 | self.num_spec_tokens: int = vllm_config.num_speculative_tokens |
| 137 | self.use_spec_decode: bool = self.num_spec_tokens > 0 |
| 138 | self.decode_cudagraph_max_bs: int = vllm_config.scheduler_config.max_num_seqs |
| 139 | if self.compilation_config.max_cudagraph_capture_size is not None: |
| 140 | self.decode_cudagraph_max_bs = min( |
| 141 | self.decode_cudagraph_max_bs, |
| 142 | self.compilation_config.max_cudagraph_capture_size, |
| 143 | ) |
| 144 | self.decode_state_indices_tensor: torch.Tensor = torch.empty( |
| 145 | (self.decode_cudagraph_max_bs, 1 + self.num_spec_tokens), |
| 146 | dtype=torch.int32, |
| 147 | device=device, |
| 148 | ) |
| 149 | self.decode_legacy_state_indices_tensor: torch.Tensor = torch.empty( |
| 150 | (self.decode_cudagraph_max_bs,), |
| 151 | dtype=torch.int32, |
| 152 | device=device, |
| 153 | ) |
| 154 | self.decode_query_start_loc: torch.Tensor = torch.empty( |
| 155 | (self.decode_cudagraph_max_bs + 1,), |
| 156 | dtype=torch.int32, |
| 157 | device=device, |
| 158 | ) |
| 159 | self.decode_num_accepted_tokens: torch.Tensor = torch.empty( |
| 160 | (self.decode_cudagraph_max_bs,), |
| 161 | dtype=torch.int32, |
| 162 | device=device, |
| 163 | ) |
| 164 | self._init_reorder_batch_threshold(1, self.use_spec_decode) |
| 165 | |
| 166 | def build_for_cudagraph_capture( |
| 167 | self, |
| 168 | common_attn_metadata: CommonAttentionMetadata, |
| 169 | ) -> BailingLinearAttentionMetadata: |
| 170 | num_accepted_tokens = None |
| 171 | if self.use_spec_decode: |
| 172 | assert common_attn_metadata.max_query_len <= 1 + self.num_spec_tokens, ( |
| 173 | "Bailing linear attention only supports speculative decoding " |
| 174 | "with query length <= 1 + number of speculative tokens." |
| 175 | ) |
| 176 | num_accepted_tokens = torch.diff(common_attn_metadata.query_start_loc) |
| 177 | return self.build( |
| 178 | common_prefix_len=0, |
| 179 | common_attn_metadata=common_attn_metadata, |
| 180 | num_accepted_tokens=num_accepted_tokens, |
| 181 | ) |
| 182 | |
| 183 | def build( # type: ignore[override] |
| 184 | self, |
| 185 | common_prefix_len: int, |
| 186 | common_attn_metadata: CommonAttentionMetadata, |
| 187 | fast_build: bool = False, |
| 188 | *, |
| 189 | num_accepted_tokens: torch.Tensor | None = None, |
| 190 | num_decode_draft_tokens_cpu: torch.Tensor | None = None, |
| 191 | ) -> BailingLinearAttentionMetadata: |
| 192 | query_start_loc = common_attn_metadata.query_start_loc |
| 193 | seq_lens = common_attn_metadata.seq_lens |
| 194 | num_reqs = common_attn_metadata.num_reqs |
| 195 | use_spec_decode = self.use_spec_decode and num_accepted_tokens is not None |
| 196 | |
| 197 | state_indices_tensor = mamba_get_block_table_tensor( |
| 198 | common_attn_metadata.block_table_tensor, |
| 199 | common_attn_metadata.seq_lens, |
| 200 | self.kv_cache_spec, |
| 201 | self.vllm_config.cache_config.mamba_cache_mode, |
| 202 | ) |
| 203 | if state_indices_tensor.dim() == 1: |
| 204 | state_indices_tensor = state_indices_tensor.unsqueeze(-1) |
| 205 | |
| 206 | decode_threshold = self.reorder_batch_threshold if use_spec_decode else 1 |
| 207 | num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( |
| 208 | split_decodes_and_prefills( |
| 209 | common_attn_metadata, |
| 210 | decode_threshold=decode_threshold, |
| 211 | ) |
| 212 | ) |
| 213 | state_indices_tensor_d, state_indices_tensor_p = torch.split( |
| 214 | state_indices_tensor, |
| 215 | [num_decodes, num_prefills], |
| 216 | dim=0, |
| 217 | ) |
| 218 | state_indices_tensor_p = state_indices_tensor_p[:, 0] |
| 219 | |
| 220 | query_start_loc_d = None |
| 221 | if use_spec_decode: |
| 222 | assert num_accepted_tokens is not None |
| 223 | state_indices_tensor_d = state_indices_tensor_d[ |
| 224 | :, : 1 + self.num_spec_tokens |
| 225 | ] |
| 226 | query_start_loc_d = query_start_loc[: num_decodes + 1] |
| 227 | num_accepted_tokens = num_accepted_tokens[:num_decodes] |
| 228 | else: |
| 229 | state_indices_tensor_d = state_indices_tensor_d[:, 0] |
| 230 | num_accepted_tokens = None |
| 231 | |
| 232 | legacy_state_indices_tensor = state_indices_tensor[:, 0] |
| 233 | cudagraph_mode = self.compilation_config.cudagraph_mode |
| 234 | use_full_cudagraph = ( |
| 235 | cudagraph_mode is not None and cudagraph_mode.has_full_cudagraphs() |
| 236 | ) |
| 237 | if ( |
| 238 | num_prefills == 0 |
| 239 | and num_decodes <= self.decode_cudagraph_max_bs |
| 240 | and use_full_cudagraph |
| 241 | ): |
| 242 | padded_bs = num_reqs |
| 243 | is_padded_decode = seq_lens[:num_decodes] == 0 |
| 244 | if state_indices_tensor_d.dim() > 1: |
| 245 | state_indices_tensor_d = torch.where( |
| 246 | is_padded_decode.unsqueeze(1), |
| 247 | torch.full_like(state_indices_tensor_d, PAD_SLOT_ID), |
| 248 | state_indices_tensor_d, |
| 249 | ) |
| 250 | self.decode_state_indices_tensor[:num_decodes].copy_( |
| 251 | state_indices_tensor_d, |
| 252 | non_blocking=True, |
| 253 | ) |
| 254 | state_indices_tensor_d = self.decode_state_indices_tensor[:padded_bs] |
| 255 | state_indices_tensor_d[num_decodes:] = PAD_SLOT_ID |
| 256 | |
| 257 | self.decode_legacy_state_indices_tensor[:num_decodes].copy_( |
| 258 | torch.where( |
| 259 | is_padded_decode, |
| 260 | torch.full_like( |
| 261 | legacy_state_indices_tensor[:num_decodes], |
| 262 | PAD_SLOT_ID, |
| 263 | ), |
| 264 | legacy_state_indices_tensor[:num_decodes], |
| 265 | ), |
| 266 | non_blocking=True, |
| 267 | ) |
| 268 | legacy_state_indices_tensor = self.decode_legacy_state_indices_tensor[ |
| 269 | :padded_bs |
| 270 | ] |
| 271 | legacy_state_indices_tensor[num_decodes:] = PAD_SLOT_ID |
| 272 | if state_indices_tensor_d.dim() == 1: |
| 273 | state_indices_tensor_d = legacy_state_indices_tensor |
| 274 | |
| 275 | if use_spec_decode and num_accepted_tokens is not None: |
| 276 | assert query_start_loc_d is not None |
| 277 | self.decode_query_start_loc[: num_decodes + 1].copy_( |
| 278 | query_start_loc_d, |
| 279 | non_blocking=True, |
| 280 | ) |
| 281 | decode_num_query_tokens = query_start_loc_d[-1] |
| 282 | query_start_loc_d = self.decode_query_start_loc[: padded_bs + 1] |
| 283 | query_start_loc_d[num_decodes + 1 :] = decode_num_query_tokens |
| 284 | |
| 285 | self.decode_num_accepted_tokens[:num_decodes].copy_( |
| 286 | num_accepted_tokens, |
| 287 | non_blocking=True, |
| 288 | ) |
| 289 | num_accepted_tokens = self.decode_num_accepted_tokens[:padded_bs] |
| 290 | num_accepted_tokens[num_decodes:] = 1 |
| 291 | |
| 292 | return BailingLinearAttentionMetadata( |
| 293 | num_prefills=num_prefills, |
| 294 | num_prefill_tokens=num_prefill_tokens, |
| 295 | num_decodes=num_decodes, |
| 296 | num_decode_tokens=num_decode_tokens, |
| 297 | query_start_loc=query_start_loc, |
| 298 | seq_lens=seq_lens, |
| 299 | state_indices_tensor=legacy_state_indices_tensor, |
| 300 | state_indices_tensor_d=state_indices_tensor_d, |
| 301 | state_indices_tensor_p=state_indices_tensor_p, |
| 302 | num_accepted_tokens=num_accepted_tokens, |
| 303 | query_start_loc_d=query_start_loc_d, |
| 304 | ) |