| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from dataclasses import dataclass |
| 4 | from typing import Any |
| 5 | |
| 6 | import numpy as np |
| 7 | import torch |
| 8 | import torch.nn as nn |
| 9 | |
| 10 | from vllm.config import VllmConfig |
| 11 | from vllm.config.compilation import CUDAGraphMode |
| 12 | from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFuncsByType |
| 13 | from vllm.triton_utils import tl, triton |
| 14 | from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder |
| 15 | from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder |
| 16 | from vllm.v1.attention.backends.short_conv_attn import ( |
| 17 | PleShortConvAttentionMetadataBuilder, |
| 18 | ShortConvAttentionMetadataBuilder, |
| 19 | ) |
| 20 | from vllm.v1.core.sched.output import NewRequestData |
| 21 | from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec |
| 22 | from vllm.v1.utils import CpuGpuBuffer |
| 23 | from vllm.v1.worker.gpu.attn_utils import build_attn_metadata |
| 24 | from vllm.v1.worker.gpu.input_batch import InputBatch |
| 25 | from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache |
| 26 | from vllm.v1.worker.gpu.model_states.default import DefaultModelState |
| 27 | from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata |
| 28 | from vllm.v1.worker.gpu.model_states.recoverssm import RecoverSSMState |
| 29 | from vllm.v1.worker.mamba_utils import ( |
| 30 | MambaSpecDecodeGPUContext, |
| 31 | get_mamba_group_ids, |
| 32 | get_mamba_groups, |
| 33 | preprocess_mamba_align_fused_kernel, |
| 34 | validate_mamba_state_copy_funcs, |
| 35 | ) |
| 36 | from vllm.v1.worker.utils import AttentionGroup |
| 37 | |
| 38 | |
| 39 | @dataclass |
| 40 | class MambaHybridAttnMetadata(ModelSpecificAttnMetadata): |
| 41 | is_prefilling: torch.Tensor |
| 42 | num_accepted_tokens: torch.Tensor | None = None |
| 43 | num_decode_draft_tokens_cpu: torch.Tensor | None = None |
| 44 | |
| 45 | def get_extra_common_attn_kwargs( |
| 46 | self, kv_cache_group_id: int, num_reqs: int |
| 47 | ) -> dict[str, Any]: |
| 48 | return {"is_prefilling": self.is_prefilling[:num_reqs]} |
| 49 | |
| 50 | def get_extra_attn_kwargs( |
| 51 | self, attn_metadata_builder: Any, num_reqs: int |
| 52 | ) -> dict[str, Any]: |
| 53 | if not isinstance( |
| 54 | attn_metadata_builder, |
| 55 | ( |
| 56 | Mamba2AttentionMetadataBuilder, |
| 57 | GDNAttentionMetadataBuilder, |
| 58 | ShortConvAttentionMetadataBuilder, |
| 59 | PleShortConvAttentionMetadataBuilder, |
| 60 | ), |
| 61 | ): |
| 62 | return {} |
| 63 | return { |
| 64 | "num_accepted_tokens": None |
| 65 | if self.num_accepted_tokens is None |
| 66 | else self.num_accepted_tokens[:num_reqs], |
| 67 | "num_decode_draft_tokens_cpu": None |
| 68 | if self.num_decode_draft_tokens_cpu is None |
| 69 | else self.num_decode_draft_tokens_cpu[:num_reqs], |
| 70 | } |
| 71 | |
| 72 | |
| 73 | class MambaHybridModelState(DefaultModelState): |
| 74 | """Model state for hybrid attention + Mamba / linear-attention models.""" |
| 75 | |
| 76 | def __init__( |
| 77 | self, |
| 78 | vllm_config: VllmConfig, |
| 79 | model: nn.Module, |
| 80 | encoder_cache: EncoderCache | None, |
| 81 | device: torch.device, |
| 82 | ) -> None: |
| 83 | super().__init__(vllm_config, model, encoder_cache, device) |
| 84 | self.cache_config = vllm_config.cache_config |
| 85 | self.num_accepted_tokens_gpu = torch.ones( |
| 86 | self.max_num_reqs, dtype=torch.int32, device=self.device |
| 87 | ) |
| 88 | # Pre-copy "align" prefix-cache state (V2). The migration of each |
| 89 | # request's mamba state across block boundaries runs as a fused GPU |
| 90 | # kernel reusing the postprocess copy machinery, so the per-step src |
| 91 | # columns and the running state_idx are kept GPU-resident. |
| 92 | self._align_mode = self.cache_config.mamba_cache_mode == "align" |
| 93 | self.recoverssm = ( |
| 94 | RecoverSSMState() if self.cache_config.use_kda_recoverssm else None |
| 95 | ) |
| 96 | if self._align_mode: |
| 97 | self._mamba_state_idx_gpu = torch.zeros( |
| 98 | self.max_num_reqs, dtype=torch.int32, device=self.device |
| 99 | ) |
| 100 | self._mamba_src_col_gpu = torch.full( |
| 101 | (self.max_num_reqs,), -1, dtype=torch.int32, device=self.device |
| 102 | ) |
| 103 | self._mamba_src_off_gpu = torch.zeros( |
| 104 | self.max_num_reqs, dtype=torch.int32, device=self.device |
| 105 | ) |
| 106 | self._mamba_ctx: MambaSpecDecodeGPUContext | None = None |
| 107 | self._mamba_group_ids: list[int] = [] |
| 108 | self._mamba_spec: MambaSpec | None = None |
| 109 | self._mamba_state_copy_funcs: MambaStateCopyFuncsByType | None = None |
| 110 | |
| 111 | def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: |
| 112 | super().add_request(req_index, new_req_data) |
| 113 | # Must reset the speculative acceptance count in this idx which could be stale. |
| 114 | self.num_accepted_tokens_gpu[req_index].fill_(1) |
| 115 | if self._align_mode: |
| 116 | # Seed the running state block from the resumed/prefilled position. |
| 117 | self._mamba_state_idx_gpu[req_index].fill_( |
| 118 | (new_req_data.num_computed_tokens - 1) // self.cache_config.block_size |
| 119 | ) |
| 120 | |
| 121 | def _get_mamba_group_info( |
| 122 | self, kv_cache_config: KVCacheConfig |
| 123 | ) -> tuple[list[int], MambaSpec]: |
| 124 | if self._mamba_spec is None: |
| 125 | mamba_groups = get_mamba_groups(kv_cache_config) |
| 126 | mamba_spec = next(iter(mamba_groups)) |
| 127 | assert all( |
| 128 | spec.block_size == mamba_spec.block_size |
| 129 | and spec.num_speculative_blocks == mamba_spec.num_speculative_blocks |
| 130 | and spec.mamba_cache_mode == mamba_spec.mamba_cache_mode |
| 131 | for spec in mamba_groups |
| 132 | ), "all mamba groups must share cache scheduling parameters" |
| 133 | self._mamba_group_ids = get_mamba_group_ids(mamba_groups) |
| 134 | self._mamba_spec = mamba_spec |
| 135 | return self._mamba_group_ids, self._mamba_spec |
| 136 | |
| 137 | def _ensure_align_ctx( |
| 138 | self, |
| 139 | kv_cache_config: KVCacheConfig, |
| 140 | mamba_group_ids: list[int], |
| 141 | block_tables: tuple[torch.Tensor, ...], |
| 142 | ) -> MambaSpecDecodeGPUContext: |
| 143 | if self._mamba_state_copy_funcs is None: |
| 144 | mamba_groups = get_mamba_groups(kv_cache_config) |
| 145 | mamba_types = {spec.mamba_type for spec in mamba_groups} |
| 146 | copy_funcs = self.model.get_mamba_state_copy_funcs(mamba_types) |
| 147 | validate_mamba_state_copy_funcs(mamba_groups, copy_funcs) |
| 148 | self._mamba_state_copy_funcs = copy_funcs |
| 149 | copy_funcs = self._mamba_state_copy_funcs |
| 150 | if self._mamba_ctx is None: |
| 151 | # Both SD and DS conv layouts support a >0 spec-decode shift: the |
| 152 | # fused pre-copy kernel (``_copy_mamba_state_block``) applies the |
| 153 | # ``token_bias = num_accepted - 1`` window shift per conv layout |
| 154 | # (SD: contiguous slice; DS: per-dim-row strided slice), matching |
| 155 | # the V1 ``get_conv_copy_spec`` semantics. |
| 156 | self._mamba_ctx = MambaSpecDecodeGPUContext.create( |
| 157 | max_num_reqs=self.max_num_reqs, |
| 158 | kv_cache_config=kv_cache_config, |
| 159 | copy_funcs=copy_funcs, |
| 160 | device=self.device, |
| 161 | make_buffer=lambda n, dtype: CpuGpuBuffer( |
| 162 | n, dtype=dtype, device=self.device |
| 163 | ), |
| 164 | ) |
| 165 | ctx = self._mamba_ctx |
| 166 | if not ctx.is_initialized: |
| 167 | forward_context = self.vllm_config.compilation_config.static_forward_context |
| 168 | # block_tables are batch-order slices of the persistent |
| 169 | # input_block_tables (stable data_ptr), so the metadata is captured |
| 170 | # once here and reused across steps. |
| 171 | ctx.initialize_from_forward_context( |
| 172 | kv_cache_config, |
| 173 | forward_context, |
| 174 | copy_funcs, |
| 175 | [block_tables[gid] for gid in mamba_group_ids], |
| 176 | ) |
| 177 | return ctx |
| 178 | |
| 179 | def preprocess_state( |
| 180 | self, |
| 181 | input_batch: InputBatch, |
| 182 | block_tables: tuple[torch.Tensor, ...], |
| 183 | kv_cache_config: KVCacheConfig, |
| 184 | num_computed_tokens: torch.Tensor, |
| 185 | ) -> None: |
| 186 | """Migrate each request's mamba state across block boundaries before the |
| 187 | forward (V1 align semantics, done on GPU). Runs on real batches only |
| 188 | (dummy DP/profiling runs skip preprocess_state), and before |
| 189 | ``prepare_attn`` gathers ``num_accepted_tokens``, so the boundary reset |
| 190 | is visible to the forward kernels. |
| 191 | """ |
| 192 | if not self._align_mode: |
| 193 | return |
| 194 | num_reqs = input_batch.num_reqs |
| 195 | if num_reqs == 0: |
| 196 | return |
| 197 | mamba_group_ids, mamba_spec = self._get_mamba_group_info(kv_cache_config) |
| 198 | ctx = self._ensure_align_ctx(kv_cache_config, mamba_group_ids, block_tables) |
| 199 | |
| 200 | # The state-advance + pre-copy kernels run every step; they fast-exit per |
| 201 | # request when src_col < 0 or src_col == dst_col, so no copy happens on |
| 202 | # steps that don't cross a block boundary. (Skipping the launch entirely |
| 203 | # would need a V1-style async-D2H of the actual num_computed, since |
| 204 | # num_computed_tokens_np is an optimistic mirror under async scheduling; |
| 205 | # the launch cost is ~0.3% of TPOT, so the GPU fast-exit suffices.) |
| 206 | block = 256 |
| 207 | grid = (triton.cdiv(num_reqs, block),) |
| 208 | preprocess_mamba_align_fused_kernel[grid]( |
| 209 | input_batch.idx_mapping, |
| 210 | self._mamba_state_idx_gpu, |
| 211 | num_computed_tokens, |
| 212 | input_batch.query_start_loc, |
| 213 | self.num_accepted_tokens_gpu, |
| 214 | self._mamba_src_col_gpu, |
| 215 | self._mamba_src_off_gpu, |
| 216 | num_reqs, |
| 217 | BLOCK_SIZE=block, |
| 218 | MAMBA_BLOCK_SIZE=mamba_spec.block_size, |
| 219 | ) |
| 220 | ctx.run_fused_precopy( |
| 221 | num_reqs, |
| 222 | self._mamba_state_idx_gpu, |
| 223 | self._mamba_src_col_gpu, |
| 224 | self._mamba_src_off_gpu, |
| 225 | input_batch.idx_mapping, |
| 226 | ) |
| 227 | |
| 228 | def prepare_attn( |
| 229 | self, |
| 230 | input_batch: InputBatch, |
| 231 | cudagraph_mode: CUDAGraphMode, |
| 232 | block_tables: tuple[torch.Tensor, ...], |
| 233 | slot_mappings: torch.Tensor, |
| 234 | attn_groups: list[list[AttentionGroup]], |
| 235 | kv_cache_config: KVCacheConfig, |
| 236 | for_capture: bool = False, |
| 237 | ubatch_idx: int = 0, |
| 238 | model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None, |
| 239 | ) -> dict[str, Any]: |
| 240 | assert ubatch_idx == 0, "DBO is not supported" |
| 241 | assert model_specific_attn_metadata is None |
| 242 | if cudagraph_mode == CUDAGraphMode.FULL: |
| 243 | num_reqs = input_batch.num_reqs_after_padding |
| 244 | num_tokens = input_batch.num_tokens_after_padding |
| 245 | else: |
| 246 | num_reqs = input_batch.num_reqs |
| 247 | num_tokens = input_batch.num_tokens |
| 248 | query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np) |
| 249 | # Prefer the promised bound: a capture dummy's measured max is its even |
| 250 | # split, not the length the graph must replay. |
| 251 | max_query_len = input_batch.max_query_len |
| 252 | if max_query_len is None: |
| 253 | max_query_len = input_batch.num_scheduled_tokens.max().item() |
| 254 | seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound |
| 255 | if for_capture: |
| 256 | # Capture with worst-case max_seq_len so the graph is valid at any replay. |
| 257 | max_seq_len = self.max_model_len |
| 258 | else: |
| 259 | max_seq_len = seq_lens_cpu_upper_bound[:num_reqs].max().item() |
| 260 | |
| 261 | is_prefilling_np = input_batch.is_prefilling_np |
| 262 | if input_batch.prefill_runs_as_decode_np is not None: |
| 263 | # A prompt tail the scheduler padded with placeholder drafts must run |
| 264 | # as a spec-decode row: the prefill kernels can't roll them back. |
| 265 | is_prefilling_np = is_prefilling_np & ~input_batch.prefill_runs_as_decode_np |
| 266 | is_prefilling = torch.zeros(num_reqs, dtype=torch.bool, device="cpu") |
| 267 | is_prefilling[: input_batch.num_reqs] = torch.from_numpy(is_prefilling_np) |
| 268 | # During CUDAGraph capture, num_decode_draft_tokens_cpu and num_accepted_tokens |
| 269 | # are created by attn_metadata_builder.build_for_cudagraph_capture, so we only |
| 270 | # compute them during actual (non-capture) forward execution. |
| 271 | num_accepted_tokens = None |
| 272 | num_decode_draft_tokens_cpu = None |
| 273 | if not for_capture and self.vllm_config.num_speculative_tokens > 0: |
| 274 | num_accepted_tokens = self.num_accepted_tokens_gpu.new_ones(num_reqs) |
| 275 | num_accepted_tokens[: input_batch.num_reqs] = self.num_accepted_tokens_gpu[ |
| 276 | input_batch.idx_mapping |
| 277 | ] |
| 278 | |
| 279 | # GDN uses >= 0 to select spec-decode rows, so non-decode rows |
| 280 | # need the -1 sentinel rather than a raw zero draft count. |
| 281 | num_decode_draft_tokens_np = np.full(num_reqs, -1, dtype=np.int32) |
| 282 | num_draft_tokens_per_req = input_batch.num_draft_tokens_per_req |
| 283 | if num_draft_tokens_per_req is not None: |
| 284 | # Test request state, not num_scheduled_tokens == draft_count+1: |
| 285 | # adaptive rewrites num_scheduled_tokens to an even split, so that |
| 286 | # equality rarely holds and would demote every verify row to decode. |
| 287 | is_decode = ~is_prefilling_np & (input_batch.num_scheduled_tokens > 0) |
| 288 | spec_decode_mask = (num_draft_tokens_per_req > 0) & is_decode |
| 289 | num_decode_draft_tokens_np[: input_batch.num_reqs] = np.where( |
| 290 | spec_decode_mask, num_draft_tokens_per_req, -1 |
| 291 | ) |
| 292 | num_decode_draft_tokens_cpu = torch.from_numpy(num_decode_draft_tokens_np) |
| 293 | |
| 294 | if self._align_mode: |
| 295 | mamba_group_ids, _ = self._get_mamba_group_info(kv_cache_config) |
| 296 | aligned_index_builders = [] |
| 297 | for group_idx, group_id in enumerate(mamba_group_ids): |
| 298 | for group in attn_groups[group_id]: |
| 299 | builder = group.get_metadata_builder(0) |
| 300 | if hasattr(builder, "mamba_aligned_state_indices"): |
| 301 | aligned_index_builders.append((group_idx, builder)) |
| 302 | if aligned_index_builders: |
| 303 | ctx = self._ensure_align_ctx( |
| 304 | kv_cache_config, mamba_group_ids, block_tables |
| 305 | ) |
| 306 | all_group_indices = ctx.compute_aligned_state_indices( |
| 307 | input_batch.seq_lens, num_reqs |
| 308 | ) |
| 309 | for group_idx, builder in aligned_index_builders: |
| 310 | builder.mamba_aligned_state_indices = all_group_indices[group_idx] |
| 311 | |
| 312 | mamba_attn_metadata = MambaHybridAttnMetadata( |
| 313 | is_prefilling=is_prefilling, |
| 314 | num_accepted_tokens=num_accepted_tokens, |
| 315 | num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, |
| 316 | ) |
| 317 | attn_metadata = build_attn_metadata( |
| 318 | attn_groups=attn_groups, |
| 319 | num_reqs=num_reqs, |
| 320 | num_tokens=num_tokens, |
| 321 | query_start_loc_gpu=input_batch.query_start_loc, |
| 322 | query_start_loc_cpu=query_start_loc_cpu, |
| 323 | max_query_len=max_query_len, |
| 324 | seq_lens=input_batch.seq_lens, |
| 325 | max_seq_len=max_seq_len, |
| 326 | block_tables=block_tables, |
| 327 | slot_mappings=slot_mappings, |
| 328 | kv_cache_config=kv_cache_config, |
| 329 | seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, |
| 330 | dcp_local_seq_lens=input_batch.dcp_local_seq_lens, |
| 331 | positions=input_batch.positions, |
| 332 | model_specific_attn_metadata=mamba_attn_metadata, |
| 333 | for_cudagraph_capture=for_capture, |
| 334 | rswa_prefix_lens=input_batch.prompt_lens, |
| 335 | ) |
| 336 | if self.recoverssm is not None: |
| 337 | self.recoverssm.record_step( |
| 338 | attn_metadata, attn_groups, for_capture=for_capture |
| 339 | ) |
| 340 | return attn_metadata |
| 341 | |
| 342 | def postprocess_state( |
| 343 | self, |
| 344 | idx_mapping: torch.Tensor, |
| 345 | num_sampled: torch.Tensor | int, |
| 346 | num_computed_tokens: torch.Tensor | None = None, |
| 347 | ) -> None: |
| 348 | # Chunked prefill does not sample a token, so num_sampled can be 0. |
| 349 | # Mamba treats num_accepted_tokens=1 as the neutral non-spec value. |
| 350 | num_reqs = idx_mapping.shape[0] |
| 351 | if num_reqs: |
| 352 | if not isinstance(num_sampled, int): |
| 353 | # idx_mapping may contain -1 sentinels (filtered rows) under PP; the |
| 354 | # kernel skips them rather than scattering with a host-side gather. |
| 355 | _scatter_num_accepted_kernel[(num_reqs,)]( |
| 356 | idx_mapping, |
| 357 | num_sampled, |
| 358 | self.num_accepted_tokens_gpu, |
| 359 | ) |
| 360 | else: |
| 361 | # Fill with single value. |
| 362 | _fill_num_accepted_kernel[(num_reqs,)]( |
| 363 | idx_mapping, |
| 364 | self.num_accepted_tokens_gpu, |
| 365 | max(num_sampled, 1), |
| 366 | ) |
| 367 | |
| 368 | if self.recoverssm is not None: |
| 369 | self.recoverssm.commit_step( |
| 370 | num_sampled, |
| 371 | idx_mapping, |
| 372 | state_indices=(self._mamba_state_idx_gpu if self._align_mode else None), |
| 373 | num_accepted_tokens=self.num_accepted_tokens_gpu, |
| 374 | ) |
| 375 | |
| 376 | if not num_reqs: |
| 377 | return |
| 378 | |
| 379 | # Align: save the running state to the block-aligned position when |
| 380 | # spec-decode acceptance leaves the sequence non-block-aligned (mirrors |
| 381 | # the V1 align postprocess). num_computed_tokens already holds the |
| 382 | # post-step advanced count. |
| 383 | if ( |
| 384 | self._align_mode |
| 385 | and num_computed_tokens is not None |
| 386 | and self._mamba_ctx is not None |
| 387 | ): |
| 388 | self._mamba_ctx.run_fused_postprocess_align( |
| 389 | num_reqs, |
| 390 | self.num_accepted_tokens_gpu, |
| 391 | self._mamba_state_idx_gpu, |
| 392 | num_computed_tokens, |
| 393 | idx_mapping, |
| 394 | ) |
| 395 | |
| 396 | |
| 397 | @triton.jit |
| 398 | def _scatter_num_accepted_kernel( |
| 399 | idx_mapping_ptr, # [num_reqs] batch_idx -> req_state_idx (-1 to skip) |
| 400 | num_sampled_ptr, # [num_reqs] |
| 401 | num_accepted_ptr, # [max_num_reqs] |
| 402 | ): |
| 403 | row = tl.program_id(0) |
| 404 | req_state_idx = tl.load(idx_mapping_ptr + row) |
| 405 | if req_state_idx < 0: |
| 406 | return |
| 407 | num_sampled = tl.load(num_sampled_ptr + row) |
| 408 | tl.store(num_accepted_ptr + req_state_idx, tl.maximum(num_sampled, 1)) |
| 409 | |
| 410 | |
| 411 | @triton.jit |
| 412 | def _fill_num_accepted_kernel( |
| 413 | idx_mapping_ptr, # [num_reqs] batch_idx -> req_state_idx (-1 to skip) |
| 414 | num_accepted_ptr, # [max_num_reqs] |
| 415 | num_sampled, |
| 416 | ): |
| 417 | row = tl.program_id(0) |
| 418 | req_state_idx = tl.load(idx_mapping_ptr + row) |
| 419 | if req_state_idx < 0: |
| 420 | return |
| 421 | tl.store(num_accepted_ptr + req_state_idx, num_sampled) |