| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import itertools |
| 4 | from dataclasses import dataclass, replace |
| 5 | from typing import Any |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm.config import VllmConfig |
| 10 | from vllm.utils.torch_utils import async_tensor_h2d |
| 11 | from vllm.v1.attention.backend import ( |
| 12 | AttentionBackend, |
| 13 | CommonAttentionMetadata, |
| 14 | ) |
| 15 | from vllm.v1.attention.backends.mamba_attn import ( |
| 16 | BaseMambaAttentionMetadata, |
| 17 | BaseMambaAttentionMetadataBuilder, |
| 18 | ) |
| 19 | from vllm.v1.kv_cache_interface import MambaSpec |
| 20 | |
| 21 | |
| 22 | def compute_varlen_chunk_metadata( |
| 23 | query_start_loc: torch.Tensor, |
| 24 | chunk_size: int, |
| 25 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 26 | """Build chunk-aligned, variable-length metadata used by Mamba2 SSD kernels. |
| 27 | |
| 28 | Given per-sequence cumulative token starts `query_start_loc` of shape [B+1] |
| 29 | and a physical `chunk_size`, returns three tensors on the same device: |
| 30 | - cu_chunk_seqlens: (nchunks+1,) int32 exclusive prefix-sum of |
| 31 | logical-chunk lengths (each logical chunk never crosses a sequence or |
| 32 | physical-chunk boundary). |
| 33 | - last_chunk_indices: (B,) int32 index of the last logical chunk |
| 34 | for each sequence (=-1 for empty sequences). |
| 35 | - seq_idx_chunks: (nchunks,) int32 sequence index for each logical |
| 36 | chunk in order. |
| 37 | |
| 38 | This is intentionally lightweight and CPU-side; it mirrors the metadata |
| 39 | produced by the V1 Mamba2 meta-data builder and is exported so tests |
| 40 | (and other callers) can avoid duplicating the logic. |
| 41 | """ |
| 42 | assert query_start_loc.ndim == 1, "query_start_loc must be 1-D [B+1]" |
| 43 | assert int(query_start_loc[0].item()) == 0, "query_start_loc[0] must be 0" |
| 44 | device = query_start_loc.device |
| 45 | |
| 46 | qsl64 = query_start_loc.to(torch.int64) |
| 47 | starts = qsl64[:-1].tolist() |
| 48 | ends = qsl64[1:].tolist() |
| 49 | total = int(qsl64[-1].item()) |
| 50 | |
| 51 | chunk_lens: list[int] = [] |
| 52 | seq_idx_chunks: list[int] = [] |
| 53 | last_chunk_indices: list[int] = [-1] * len(starts) |
| 54 | |
| 55 | for b, (s, e) in enumerate(zip(starts, ends)): |
| 56 | if e <= s: |
| 57 | # empty sequence |
| 58 | continue |
| 59 | pos = s |
| 60 | while pos < e: |
| 61 | # split at both sequence boundaries and physical chunk boundaries |
| 62 | room = chunk_size - (pos % chunk_size) |
| 63 | take = min(room, e - pos) |
| 64 | chunk_lens.append(int(take)) |
| 65 | seq_idx_chunks.append(b) |
| 66 | last_chunk_indices[b] = len(chunk_lens) - 1 |
| 67 | pos += take |
| 68 | |
| 69 | # Exclusive prefix sum over logical-chunk lengths |
| 70 | if chunk_lens: |
| 71 | cu_chunk_seqlens_list = [0] + list(itertools.accumulate(chunk_lens)) |
| 72 | # Final boundary must equal total tokens (check on host to avoid a sync) |
| 73 | assert cu_chunk_seqlens_list[-1] == total |
| 74 | else: |
| 75 | cu_chunk_seqlens_list = [0] |
| 76 | cu_chunk_seqlens = async_tensor_h2d( |
| 77 | cu_chunk_seqlens_list, dtype=torch.int32, device=device |
| 78 | ) |
| 79 | |
| 80 | # last_chunk_indices is empty when there are no sequences (len(starts) == 0). |
| 81 | last_chunk_indices_t = async_tensor_h2d( |
| 82 | last_chunk_indices, dtype=torch.int32, device=device |
| 83 | ) |
| 84 | seq_idx_chunks_t = async_tensor_h2d( |
| 85 | seq_idx_chunks, dtype=torch.int32, device=device |
| 86 | ) |
| 87 | return cu_chunk_seqlens, last_chunk_indices_t, seq_idx_chunks_t |
| 88 | |
| 89 | |
| 90 | class Mamba2AttentionBackend(AttentionBackend): |
| 91 | @staticmethod |
| 92 | def get_name() -> str: |
| 93 | return "MAMBA2_ATTN" |
| 94 | |
| 95 | @staticmethod |
| 96 | def get_builder_cls() -> type["Mamba2AttentionMetadataBuilder"]: |
| 97 | return Mamba2AttentionMetadataBuilder |
| 98 | |
| 99 | @classmethod |
| 100 | def is_ssm(cls) -> bool: |
| 101 | return True |
| 102 | |
| 103 | |
| 104 | @dataclass |
| 105 | class Mamba2AttentionMetadata(BaseMambaAttentionMetadata): |
| 106 | prep_initial_states: bool = False |
| 107 | chunk_size: int = 0 |
| 108 | |
| 109 | # Chunk-related metadata (only for prefill) |
| 110 | seq_idx_p: torch.Tensor | None = None |
| 111 | |
| 112 | |
| 113 | class Mamba2AttentionMetadataBuilder( |
| 114 | BaseMambaAttentionMetadataBuilder[Mamba2AttentionMetadata] |
| 115 | ): |
| 116 | metadata_cls = Mamba2AttentionMetadata |
| 117 | |
| 118 | def __init__( |
| 119 | self, |
| 120 | kv_cache_spec: MambaSpec, |
| 121 | layer_names: list[str], |
| 122 | vllm_config: VllmConfig, |
| 123 | device: torch.device, |
| 124 | ): |
| 125 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 126 | chunk_size = vllm_config.model_config.get_mamba_chunk_size() |
| 127 | assert chunk_size is not None, ( |
| 128 | "chunk_size needs to be set in the model config for Mamba2 models" |
| 129 | ) |
| 130 | self.chunk_size: int = chunk_size |
| 131 | |
| 132 | def build( |
| 133 | self, |
| 134 | common_prefix_len: int, |
| 135 | common_attn_metadata: CommonAttentionMetadata, |
| 136 | fast_build: bool = False, |
| 137 | **kwargs: Any, |
| 138 | ) -> Mamba2AttentionMetadata: |
| 139 | common = self._compute_common_metadata( |
| 140 | common_attn_metadata, |
| 141 | num_accepted_tokens=kwargs.get("num_accepted_tokens"), |
| 142 | num_decode_draft_tokens_cpu=kwargs.get("num_decode_draft_tokens_cpu"), |
| 143 | ) |
| 144 | |
| 145 | seq_idx_p = None |
| 146 | cu_chunk_seqlen_p = None |
| 147 | last_chunk_indices_p = None |
| 148 | prep_initial_states = False |
| 149 | |
| 150 | # Compute seq_idx for prefill only |
| 151 | if common.num_prefills > 0: |
| 152 | prep_initial_states = False |
| 153 | if common.has_initial_states_p is not None: |
| 154 | # Same condition as `has_initial_states_p`, but derived from CPU |
| 155 | # data so it needs no D2H. `seq_lens_cpu_upper_bound` is precise |
| 156 | # for prefill rows, which is all this slice covers. |
| 157 | num_computed_tokens_p_cpu, _ = self._prefill_cpu_metadata( |
| 158 | common_attn_metadata, |
| 159 | common.num_reqs, |
| 160 | common.num_prefills, |
| 161 | common.num_decode_tokens, |
| 162 | ) |
| 163 | prep_initial_states = bool((num_computed_tokens_p_cpu > 0).any()) |
| 164 | |
| 165 | cu_chunk_seqlen_p, seq_idx_p, last_chunk_indices_p = ( |
| 166 | self._build_chunk_metadata_tensors( |
| 167 | self.chunk_size, |
| 168 | common, |
| 169 | common_attn_metadata, |
| 170 | ) |
| 171 | ) |
| 172 | |
| 173 | return replace( |
| 174 | common, |
| 175 | prep_initial_states=prep_initial_states, |
| 176 | chunk_size=self.chunk_size, |
| 177 | seq_idx_p=seq_idx_p, |
| 178 | cu_chunk_seqlen_p=cu_chunk_seqlen_p, |
| 179 | last_chunk_indices_p=last_chunk_indices_p, |
| 180 | ) |