| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import functools |
| 5 | from collections.abc import Callable |
| 6 | from dataclasses import dataclass |
| 7 | from typing import Literal, TypeAlias |
| 8 | |
| 9 | import torch |
| 10 | |
| 11 | import vllm.envs as envs |
| 12 | from vllm.config.cache import MambaDType |
| 13 | from vllm.config.mamba import MambaBackendEnum |
| 14 | from vllm.config.model import ModelDType |
| 15 | from vllm.distributed import divide |
| 16 | from vllm.logger import init_logger |
| 17 | from vllm.utils.torch_utils import ( |
| 18 | STR_DTYPE_TO_TORCH_DTYPE, |
| 19 | get_kv_cache_torch_dtype, |
| 20 | ) |
| 21 | from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum |
| 22 | |
| 23 | logger = init_logger(__name__) |
| 24 | |
| 25 | ConvStateLayoutType = Literal["SD", "DS"] |
| 26 | |
| 27 | |
| 28 | @functools.lru_cache |
| 29 | def get_conv_state_layout() -> ConvStateLayoutType: |
| 30 | """Return the SSM conv state layout. |
| 31 | |
| 32 | SD = (state_len, dim) — dim is the innermost contiguous dimension. |
| 33 | DS = (dim, state_len) — TP-sharded dim is on dim-1 (like HND for KV |
| 34 | cache), consistent with SSM temporal state layout. |
| 35 | """ |
| 36 | layout: ConvStateLayoutType | None = envs.VLLM_SSM_CONV_STATE_LAYOUT |
| 37 | if layout is not None: |
| 38 | logger.info_once( |
| 39 | "VLLM_SSM_CONV_STATE_LAYOUT env detected. " |
| 40 | "Setting SSM conv state layout to %s.", |
| 41 | layout, |
| 42 | ) |
| 43 | return layout |
| 44 | |
| 45 | return "SD" |
| 46 | |
| 47 | |
| 48 | def is_conv_state_dim_first() -> bool: |
| 49 | """True when the conv state is stored as (dim, state_len) per block.""" |
| 50 | return get_conv_state_layout() == "DS" |
| 51 | |
| 52 | |
| 53 | class MambaStateDtypeCalculator: |
| 54 | @classmethod |
| 55 | def linear_attention_state_dtype( |
| 56 | cls, |
| 57 | model_dtype: ModelDType | torch.dtype, |
| 58 | mamba_cache_dtype: MambaDType, |
| 59 | ) -> tuple[torch.dtype, ...]: |
| 60 | state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype) |
| 61 | return (state_dtype,) |
| 62 | |
| 63 | @classmethod |
| 64 | def mamba1_state_dtype( |
| 65 | cls, |
| 66 | model_dtype: ModelDType | torch.dtype, |
| 67 | mamba_cache_dtype: MambaDType, |
| 68 | mamba_ssm_cache_dtype: MambaDType, |
| 69 | ) -> tuple[torch.dtype, ...]: |
| 70 | return cls._mamba_state_dtype( |
| 71 | model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype |
| 72 | ) |
| 73 | |
| 74 | @classmethod |
| 75 | def mamba2_state_dtype( |
| 76 | cls, |
| 77 | model_dtype: ModelDType | torch.dtype, |
| 78 | mamba_cache_dtype: MambaDType, |
| 79 | mamba_ssm_cache_dtype: MambaDType, |
| 80 | ) -> tuple[torch.dtype, ...]: |
| 81 | return cls._mamba_state_dtype( |
| 82 | model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype |
| 83 | ) |
| 84 | |
| 85 | @classmethod |
| 86 | def append_replayssm_ring( |
| 87 | cls, |
| 88 | base_dtypes: tuple[torch.dtype, ...], |
| 89 | model_dtype: ModelDType | torch.dtype, |
| 90 | ) -> tuple[torch.dtype, ...]: |
| 91 | """Append the ReplaySSM ring dtypes to a base ``(conv, ssm)`` tuple: |
| 92 | ``(x_cache, dt_cache, B_cache)`` = ``(activation, fp32, activation)``. |
| 93 | """ |
| 94 | activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype) |
| 95 | return (*base_dtypes, activation_dtype, torch.float32, activation_dtype) |
| 96 | |
| 97 | @classmethod |
| 98 | def _mamba_state_dtype( |
| 99 | cls, |
| 100 | model_dtype: ModelDType | torch.dtype, |
| 101 | mamba_cache_dtype: MambaDType, |
| 102 | mamba_ssm_cache_dtype: MambaDType, |
| 103 | ) -> tuple[torch.dtype, ...]: |
| 104 | conv_state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype) |
| 105 | if mamba_ssm_cache_dtype == "auto": |
| 106 | temporal_state_dtype = conv_state_dtype |
| 107 | else: |
| 108 | temporal_state_dtype = STR_DTYPE_TO_TORCH_DTYPE[mamba_ssm_cache_dtype] |
| 109 | |
| 110 | return (conv_state_dtype, temporal_state_dtype) |
| 111 | |
| 112 | @classmethod |
| 113 | def short_conv_state_dtype( |
| 114 | cls, |
| 115 | model_dtype: ModelDType | torch.dtype, |
| 116 | mamba_cache_dtype: MambaDType, |
| 117 | ) -> tuple[torch.dtype, ...]: |
| 118 | conv_state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype) |
| 119 | return (conv_state_dtype,) |
| 120 | |
| 121 | @classmethod |
| 122 | def gated_delta_net_state_dtype( |
| 123 | cls, |
| 124 | model_dtype: ModelDType | torch.dtype, |
| 125 | mamba_cache_dtype: MambaDType, |
| 126 | mamba_ssm_cache_dtype: MambaDType = "auto", |
| 127 | ) -> tuple[torch.dtype, torch.dtype]: |
| 128 | return cls._mamba_state_dtype( |
| 129 | model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype |
| 130 | ) |
| 131 | |
| 132 | @classmethod |
| 133 | def kda_state_dtype( |
| 134 | cls, |
| 135 | model_dtype: ModelDType | torch.dtype, |
| 136 | mamba_cache_dtype: MambaDType, |
| 137 | mamba_ssm_cache_dtype: MambaDType = "auto", |
| 138 | ) -> tuple[torch.dtype, torch.dtype]: |
| 139 | conv_state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype) |
| 140 | if mamba_ssm_cache_dtype == "auto": |
| 141 | recurrent_state_dtype = torch.float32 |
| 142 | elif mamba_ssm_cache_dtype in ("float32", "bfloat16"): |
| 143 | recurrent_state_dtype = STR_DTYPE_TO_TORCH_DTYPE[mamba_ssm_cache_dtype] |
| 144 | else: |
| 145 | raise ValueError( |
| 146 | "KDA recurrent state supports only auto, float32, and bfloat16; " |
| 147 | f"got {mamba_ssm_cache_dtype}" |
| 148 | ) |
| 149 | return (conv_state_dtype, recurrent_state_dtype) |
| 150 | |
| 151 | @classmethod |
| 152 | def append_kda_recoverssm_record( |
| 153 | cls, |
| 154 | base_dtypes: tuple[torch.dtype, ...], |
| 155 | model_dtype: ModelDType | torch.dtype, |
| 156 | ) -> tuple[torch.dtype, ...]: |
| 157 | activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype) |
| 158 | return (*base_dtypes, torch.float32, activation_dtype) |
| 159 | |
| 160 | |
| 161 | class MambaStateShapeCalculator: |
| 162 | @classmethod |
| 163 | def linear_attention_state_shape( |
| 164 | cls, |
| 165 | num_heads: int, |
| 166 | tp_size: int, |
| 167 | head_dim: int, |
| 168 | ) -> tuple[tuple[int, int, int]]: |
| 169 | state_shape = (num_heads // tp_size, head_dim, head_dim) |
| 170 | return (state_shape,) |
| 171 | |
| 172 | @staticmethod |
| 173 | def _orient_conv_shape(dim: int, state_len: int) -> tuple[int, int]: |
| 174 | """Return (dim, state_len) for DS layout, (state_len, dim) for SD.""" |
| 175 | if is_conv_state_dim_first(): |
| 176 | return (dim, state_len) |
| 177 | return (state_len, dim) |
| 178 | |
| 179 | @classmethod |
| 180 | def mamba1_state_shape( |
| 181 | cls, |
| 182 | tp_world_size: int, |
| 183 | intermediate_size: int, |
| 184 | state_size: int, |
| 185 | conv_kernel: int, |
| 186 | ) -> tuple[tuple[int, int], tuple[int, int]]: |
| 187 | conv_dim = divide(intermediate_size, tp_world_size) |
| 188 | conv_state_shape = cls._orient_conv_shape(conv_dim, conv_kernel - 1) |
| 189 | |
| 190 | temporal_state_shape = (divide(intermediate_size, tp_world_size), state_size) |
| 191 | |
| 192 | return conv_state_shape, temporal_state_shape |
| 193 | |
| 194 | @classmethod |
| 195 | def mamba2_state_shape( |
| 196 | cls, |
| 197 | tp_world_size: int, |
| 198 | intermediate_size: int, |
| 199 | n_groups: int, |
| 200 | num_heads: int, |
| 201 | head_dim: int, |
| 202 | state_size: int, |
| 203 | conv_kernel: int, |
| 204 | num_spec: int = 0, |
| 205 | ) -> tuple[tuple[int, int], tuple[int, int, int]]: |
| 206 | # if n_groups is not divisible by world_size, need to extend the shards |
| 207 | # to ensure all groups needed by a head is sharded along with it |
| 208 | n_groups = n_groups + cls.extra_groups_for_head_shards(n_groups, tp_world_size) |
| 209 | # heads and n_groups are TP-ed |
| 210 | conv_dim = intermediate_size + 2 * n_groups * state_size |
| 211 | |
| 212 | conv_state_shape = cls._orient_conv_shape( |
| 213 | divide(conv_dim, tp_world_size), conv_kernel - 1 + num_spec |
| 214 | ) |
| 215 | |
| 216 | # These are not TP-ed as they depend on A, dt_bias, D |
| 217 | # - they are typically small |
| 218 | # e.g., (h_heads, head_dim, state_size) = (128, 64, 128) |
| 219 | temporal_state_shape = (divide(num_heads, tp_world_size), head_dim, state_size) |
| 220 | return conv_state_shape, temporal_state_shape |
| 221 | |
| 222 | @classmethod |
| 223 | def append_replayssm_ring( |
| 224 | cls, |
| 225 | base_shapes: tuple[tuple[int, int], tuple[int, int, int]], |
| 226 | n_groups: int, |
| 227 | tp_world_size: int, |
| 228 | logical_window: int, |
| 229 | backend: MambaBackendEnum, |
| 230 | num_speculative_tokens: int = 0, |
| 231 | ) -> tuple[ |
| 232 | tuple[int, int], |
| 233 | tuple[int, int, int], |
| 234 | tuple[int, int, int], |
| 235 | tuple[int, int], |
| 236 | tuple[int, int, int], |
| 237 | ]: |
| 238 | """Append the physical ReplaySSM ring shapes. |
| 239 | |
| 240 | ``base_shapes[1]`` is ``(nheads // tp, head_dim, state_size)``; |
| 241 | B_cache uses the un-extended ``n_groups``. |
| 242 | """ |
| 243 | ring_buffer_len = logical_window |
| 244 | if backend == MambaBackendEnum.FLASHINFER: |
| 245 | # FlashInfer keeps the live window and current verify window together. |
| 246 | ring_buffer_len += 1 + num_speculative_tokens |
| 247 | local_nheads, head_dim, state_size = base_shapes[1] |
| 248 | local_ngroups = divide(n_groups, tp_world_size) |
| 249 | return ( |
| 250 | *base_shapes, |
| 251 | (local_nheads, ring_buffer_len, head_dim), |
| 252 | (local_nheads, ring_buffer_len), |
| 253 | (local_ngroups, ring_buffer_len, state_size), |
| 254 | ) |
| 255 | |
| 256 | @classmethod |
| 257 | def short_conv_state_shape( |
| 258 | cls, |
| 259 | tp_world_size: int, |
| 260 | intermediate_size: int, |
| 261 | conv_kernel: int, |
| 262 | num_spec: int = 0, |
| 263 | ) -> tuple[tuple[int, int]]: |
| 264 | conv_dim = divide(intermediate_size, tp_world_size) |
| 265 | conv_state_shape = cls._orient_conv_shape(conv_dim, conv_kernel - 1 + num_spec) |
| 266 | return (conv_state_shape,) |
| 267 | |
| 268 | @classmethod |
| 269 | def extra_groups_for_head_shards(cls, ngroups: int, tp_size: int): |
| 270 | """Compute the increase in group numbers to account for |
| 271 | replication in order to accompany the head shards.""" |
| 272 | # in the case ngoups % tp_size == 0, this will be zero |
| 273 | if ngroups % tp_size == 0: |
| 274 | return 0 |
| 275 | |
| 276 | # for n_groups == 1, this is exactly tp_size - n_groups |
| 277 | return tp_size - ngroups |
| 278 | |
| 279 | @classmethod |
| 280 | def gated_delta_net_state_shape( |
| 281 | cls, |
| 282 | tp_world_size: int, |
| 283 | num_k_heads: int, |
| 284 | num_v_heads: int, |
| 285 | head_k_dim: int, |
| 286 | head_v_dim: int, |
| 287 | conv_kernel_size: int, |
| 288 | num_spec: int = 0, |
| 289 | ): |
| 290 | conv_dim = head_k_dim * num_k_heads * 2 + head_v_dim * num_v_heads |
| 291 | conv_state_shape = cls._orient_conv_shape( |
| 292 | divide(conv_dim, tp_world_size), |
| 293 | conv_kernel_size - 1 + num_spec, |
| 294 | ) |
| 295 | |
| 296 | temporal_state_shape = ( |
| 297 | divide(num_v_heads, tp_world_size), |
| 298 | head_v_dim, |
| 299 | head_k_dim, |
| 300 | ) |
| 301 | return conv_state_shape, temporal_state_shape |
| 302 | |
| 303 | @classmethod |
| 304 | def kda_state_shape( |
| 305 | cls, |
| 306 | tp_world_size: int, |
| 307 | num_heads: int, |
| 308 | head_dim: int, |
| 309 | num_k_heads: int | None = None, |
| 310 | head_k_dim: int | None = None, |
| 311 | conv_kernel_size: int = 4, |
| 312 | num_spec: int = 0, |
| 313 | ) -> tuple[tuple[int, int], tuple[int, int, int]]: |
| 314 | if num_k_heads is None: |
| 315 | num_k_heads = num_heads |
| 316 | if head_k_dim is None: |
| 317 | head_k_dim = head_dim |
| 318 | |
| 319 | proj_size = num_heads * head_dim |
| 320 | proj_k_size = num_k_heads * head_k_dim |
| 321 | |
| 322 | conv_dim = proj_size + 2 * proj_k_size |
| 323 | conv_state_shape = cls._orient_conv_shape( |
| 324 | divide(conv_dim, tp_world_size), conv_kernel_size - 1 + num_spec |
| 325 | ) |
| 326 | recurrent_state_shape = (divide(num_heads, tp_world_size), head_dim, head_dim) |
| 327 | return (conv_state_shape, recurrent_state_shape) |
| 328 | |
| 329 | @classmethod |
| 330 | def append_kda_recoverssm_record( |
| 331 | cls, |
| 332 | base_shapes: tuple[tuple[int, int], tuple[int, int, int]], |
| 333 | num_heads: int, |
| 334 | head_dim: int, |
| 335 | tp_world_size: int, |
| 336 | spec_query_len: int, |
| 337 | ) -> tuple[ |
| 338 | tuple[int, int], |
| 339 | tuple[int, int, int], |
| 340 | tuple[int, int, int], |
| 341 | tuple[int, int, int], |
| 342 | ]: |
| 343 | local_num_heads = divide(num_heads, tp_world_size) |
| 344 | return ( |
| 345 | *base_shapes, |
| 346 | (local_num_heads, spec_query_len, head_dim), |
| 347 | (local_num_heads, spec_query_len, 2 * head_dim), |
| 348 | ) |
| 349 | |
| 350 | |
| 351 | @dataclass |
| 352 | class MambaCopySpec: |
| 353 | """Data class specifying the memory-copy parameters for Mamba states used for |
| 354 | prefix caching in align mode. |
| 355 | |
| 356 | Attributes: |
| 357 | start_addr (int): Starting address for the memory copy operation. |
| 358 | num_elements (int): Number of elements to copy from the starting address. |
| 359 | |
| 360 | """ |
| 361 | |
| 362 | start_addr: int |
| 363 | num_elements: int |
| 364 | |
| 365 | |
| 366 | MambaStateCopyFunc: TypeAlias = Callable[ |
| 367 | [torch.Tensor, list[int], int, int], MambaCopySpec |
| 368 | ] |
| 369 | """ |
| 370 | Type alias for a function that computes a MambaCopySpec for copying state slices. |
| 371 | Parameters: |
| 372 | state: torch.Tensor - the Mamba state tensor (e.g., conv or temporal states). |
| 373 | block_ids: list[int] - the list of block indices for the state to copy. |
| 374 | cur_block_idx: int - current block index within `block_ids` to copy from. |
| 375 | num_accepted_tokens: int - number of accepted tokens used to compute the copy offset. |
| 376 | Range: 1 .. 1 + num_speculative_tokens (inclusive). |
| 377 | """ |
| 378 | MambaStateCopyFuncs: TypeAlias = tuple[MambaStateCopyFunc, ...] |
| 379 | MambaStateCopyFuncsByType: TypeAlias = dict[ |
| 380 | MambaAttentionBackendEnum, MambaStateCopyFuncs |
| 381 | ] |
| 382 | |
| 383 | |
| 384 | def get_conv_copy_spec( |
| 385 | state: torch.Tensor, |
| 386 | block_ids: list[int], |
| 387 | cur_block_idx: int, |
| 388 | num_accepted_tokens: int, |
| 389 | ) -> MambaCopySpec: |
| 390 | """Return a MambaCopySpec for copying a convolutional state slice. |
| 391 | |
| 392 | Works for both SD layout ``(num_blocks, state_len, dim)`` and |
| 393 | DS layout ``(num_blocks, dim, state_len)``. |
| 394 | """ |
| 395 | src_block_id = block_ids[cur_block_idx] |
| 396 | offset = num_accepted_tokens - 1 |
| 397 | if is_conv_state_dim_first(): |
| 398 | # DS offset > 0 is handled by the fused postprocess kernel. |
| 399 | assert offset == 0, ( |
| 400 | "DS conv state with num_accepted_tokens > 1 must be handled by " |
| 401 | "the fused postprocess kernel, not get_conv_copy_spec" |
| 402 | ) |
| 403 | src_state = state[src_block_id] |
| 404 | else: |
| 405 | # SD layout: (num_blocks, state_len, dim), with dim contiguous. |
| 406 | src_state = state[src_block_id, offset:] |
| 407 | return MambaCopySpec( |
| 408 | start_addr=src_state.data_ptr(), num_elements=src_state.numel() |
| 409 | ) |
| 410 | |
| 411 | |
| 412 | def get_temporal_copy_spec( |
| 413 | state: torch.Tensor, |
| 414 | block_ids: list[int], |
| 415 | cur_block_idx: int, |
| 416 | num_accepted_tokens: int, |
| 417 | ) -> MambaCopySpec: |
| 418 | """Return a MambaCopySpec for copying a temporal state slice.""" |
| 419 | src_block_id = block_ids[cur_block_idx + num_accepted_tokens - 1] |
| 420 | src_state = state[src_block_id] |
| 421 | return MambaCopySpec( |
| 422 | start_addr=src_state.data_ptr(), num_elements=src_state.numel() |
| 423 | ) |
| 424 | |
| 425 | |
| 426 | class MambaStateCopyFuncCalculator: |
| 427 | @classmethod |
| 428 | def linear_attention_state_copy_func(cls): |
| 429 | return (get_temporal_copy_spec,) |
| 430 | |
| 431 | @classmethod |
| 432 | def mamba1_state_copy_func(cls): |
| 433 | return (get_conv_copy_spec, get_temporal_copy_spec) |
| 434 | |
| 435 | @classmethod |
| 436 | def mamba2_state_copy_func(cls): |
| 437 | return get_conv_copy_spec, get_temporal_copy_spec |
| 438 | |
| 439 | @classmethod |
| 440 | def short_conv_state_copy_func(cls): |
| 441 | return (get_conv_copy_spec,) |
| 442 | |
| 443 | @classmethod |
| 444 | def gated_delta_net_state_copy_func(cls): |
| 445 | return (get_conv_copy_spec, get_temporal_copy_spec) |
| 446 | |
| 447 | @classmethod |
| 448 | def kda_state_copy_func(cls): |
| 449 | return (get_conv_copy_spec, get_temporal_copy_spec) |