| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import NamedTuple |
| 5 | |
| 6 | import torch |
| 7 | from torch import nn |
| 8 | from torch.nn.parameter import Parameter |
| 9 | |
| 10 | from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config |
| 11 | from vllm.distributed.parallel_state import ( |
| 12 | get_tensor_model_parallel_rank, |
| 13 | get_tensor_model_parallel_world_size, |
| 14 | ) |
| 15 | from vllm.forward_context import ForwardContext, get_forward_context |
| 16 | from vllm.model_executor.custom_op import PluggableLayer |
| 17 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 18 | from vllm.model_executor.layers.linear import ( |
| 19 | ColumnParallelLinear, |
| 20 | MergedColumnParallelLinear, |
| 21 | RowParallelLinear, |
| 22 | ) |
| 23 | from vllm.model_executor.layers.mamba.abstract import MambaBase |
| 24 | from vllm.model_executor.layers.mamba.mamba_utils import ( |
| 25 | MambaStateDtypeCalculator, |
| 26 | MambaStateShapeCalculator, |
| 27 | is_conv_state_dim_first, |
| 28 | ) |
| 29 | from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( |
| 30 | causal_conv1d_fn, |
| 31 | causal_conv1d_update, |
| 32 | ) |
| 33 | from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_scan_fn |
| 34 | from vllm.model_executor.layers.mamba.ops.ssu_dispatch import selective_state_update |
| 35 | from vllm.model_executor.utils import set_weight_attrs |
| 36 | from vllm.platforms import current_platform |
| 37 | from vllm.utils.torch_utils import ( |
| 38 | LayerNameType, |
| 39 | _encode_layer_name, |
| 40 | _resolve_layer_name, |
| 41 | direct_register_custom_op, |
| 42 | ) |
| 43 | from vllm.v1.attention.backend import AttentionMetadata |
| 44 | from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionMetadata |
| 45 | from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum |
| 46 | |
| 47 | |
| 48 | # Adapted from transformers.models.mamba.modeling_mamba.MambaMixer |
| 49 | # --8<-- [start:mamba_mixer] |
| 50 | @PluggableLayer.register("mamba_mixer") |
| 51 | class MambaMixer(MambaBase, PluggableLayer): |
| 52 | """Compute ∆, A, B, C, and D the state space parameters and compute |
| 53 | the `contextualized_states`. A, D are input independent |
| 54 | (see Mamba paper [1] Section 3.5.2 "Interpretation of A" |
| 55 | for why A isn't selective) ∆, B, C are input-dependent |
| 56 | (this is a key difference between Mamba and the linear time |
| 57 | invariant S4, and is why Mamba is called |
| 58 | **selective** state spaces) |
| 59 | """ |
| 60 | |
| 61 | # --8<-- [end:mamba_mixer] |
| 62 | |
| 63 | def __init__( |
| 64 | self, |
| 65 | hidden_size: int, |
| 66 | ssm_state_size: int, |
| 67 | conv_kernel_size: int, |
| 68 | intermediate_size: int, |
| 69 | time_step_rank: int, |
| 70 | use_conv_bias: bool, |
| 71 | use_bias: bool, |
| 72 | use_rms_norm: bool, |
| 73 | rms_norm_has_weight: bool = True, |
| 74 | rms_norm_eps: float = 1e-5, |
| 75 | activation="silu", |
| 76 | is_lora_enabled: bool = False, |
| 77 | model_config: ModelConfig | None = None, |
| 78 | cache_config: CacheConfig | None = None, |
| 79 | prefix: str = "", |
| 80 | ): |
| 81 | super().__init__() |
| 82 | self.time_step_rank = time_step_rank |
| 83 | self.ssm_state_size = ssm_state_size |
| 84 | self.use_rms_norm = use_rms_norm |
| 85 | self.activation = activation |
| 86 | self.is_lora_enabled = is_lora_enabled |
| 87 | self.conv_kernel_size = conv_kernel_size |
| 88 | self.intermediate_size = intermediate_size |
| 89 | |
| 90 | self.conv1d = ColumnParallelLinear( |
| 91 | input_size=conv_kernel_size, |
| 92 | output_size=intermediate_size, |
| 93 | bias=use_conv_bias, |
| 94 | prefix=f"{prefix}.conv1d", |
| 95 | ) |
| 96 | # unsqueeze to fit conv1d weights shape into the linear weights shape. |
| 97 | # Can't do this in `weight_loader` since it already exists in |
| 98 | # `ColumnParallelLinear` and `set_weight_attrs` |
| 99 | # doesn't allow to override it |
| 100 | self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1) |
| 101 | |
| 102 | self.in_proj = MergedColumnParallelLinear( |
| 103 | hidden_size, |
| 104 | [intermediate_size] * 2, |
| 105 | bias=use_bias, |
| 106 | prefix=f"{prefix}.in_proj", |
| 107 | ) |
| 108 | |
| 109 | # selective projection used to make dt, B and C input dependent |
| 110 | self.x_proj = RowParallelLinear( |
| 111 | intermediate_size, |
| 112 | time_step_rank + ssm_state_size * 2, |
| 113 | bias=False, |
| 114 | prefix=f"{prefix}.x_proj", |
| 115 | ) |
| 116 | # time step projection (discretization) - |
| 117 | # In the forward we need to apply dt_proj without the bias, |
| 118 | # as the bias is added in the selective scan kernel. |
| 119 | self.dt_proj = ColumnParallelLinear( |
| 120 | time_step_rank, |
| 121 | intermediate_size, |
| 122 | bias=True, |
| 123 | skip_bias_add=True, |
| 124 | prefix=f"{prefix}.dt_proj", |
| 125 | ) |
| 126 | |
| 127 | def weight_loader(param: Parameter, loaded_weight: torch.Tensor): |
| 128 | tp_rank = get_tensor_model_parallel_rank() |
| 129 | tp_size = get_tensor_model_parallel_world_size() |
| 130 | param.data.copy_( |
| 131 | loaded_weight.data.split(loaded_weight.shape[0] // tp_size, dim=0)[ |
| 132 | tp_rank |
| 133 | ] |
| 134 | ) |
| 135 | |
| 136 | def A_weight_loader(param: Parameter, loaded_weight: torch.Tensor): |
| 137 | weight_loader(param, -torch.exp(loaded_weight.float())) |
| 138 | |
| 139 | tp_size = get_tensor_model_parallel_world_size() |
| 140 | self.A = nn.Parameter( |
| 141 | torch.empty( |
| 142 | intermediate_size // tp_size, |
| 143 | ssm_state_size, |
| 144 | dtype=torch.float32, |
| 145 | ) |
| 146 | ) |
| 147 | self.D = nn.Parameter(torch.ones(intermediate_size // tp_size)) |
| 148 | |
| 149 | set_weight_attrs(self.D, {"weight_loader": weight_loader}) |
| 150 | set_weight_attrs(self.A, {"weight_loader": A_weight_loader}) |
| 151 | |
| 152 | self.out_proj = RowParallelLinear( |
| 153 | intermediate_size, |
| 154 | hidden_size, |
| 155 | bias=use_bias, |
| 156 | input_is_parallel=True, |
| 157 | prefix=f"{prefix}.out_proj", |
| 158 | ) |
| 159 | |
| 160 | self.dt_layernorm = ( |
| 161 | RMSNorm( |
| 162 | time_step_rank, |
| 163 | eps=rms_norm_eps, |
| 164 | has_weight=rms_norm_has_weight, |
| 165 | ) |
| 166 | if use_rms_norm |
| 167 | else None |
| 168 | ) |
| 169 | |
| 170 | self.b_layernorm = ( |
| 171 | RMSNorm( |
| 172 | ssm_state_size, |
| 173 | eps=rms_norm_eps, |
| 174 | has_weight=rms_norm_has_weight, |
| 175 | ) |
| 176 | if use_rms_norm |
| 177 | else None |
| 178 | ) |
| 179 | |
| 180 | self.c_layernorm = ( |
| 181 | RMSNorm( |
| 182 | ssm_state_size, |
| 183 | eps=rms_norm_eps, |
| 184 | has_weight=rms_norm_has_weight, |
| 185 | ) |
| 186 | if use_rms_norm |
| 187 | else None |
| 188 | ) |
| 189 | |
| 190 | compilation_config = get_current_vllm_config().compilation_config |
| 191 | if prefix in compilation_config.static_forward_context: |
| 192 | raise ValueError(f"Duplicate layer name: {prefix}") |
| 193 | compilation_config.static_forward_context[prefix] = self |
| 194 | # The inner tuple is (conv_state, ssm_state) |
| 195 | self.kv_cache = (torch.tensor([]), torch.tensor([])) |
| 196 | |
| 197 | self.model_config = model_config |
| 198 | self.cache_config = cache_config |
| 199 | self.prefix = prefix |
| 200 | |
| 201 | def _ssm_transform( |
| 202 | self, x: torch.Tensor |
| 203 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 204 | # LoRA kernel requires contiguous tensor. |
| 205 | # ROCm: Non-contiguous tensors cause incorrect GEMM |
| 206 | # results when batch > 1. |
| 207 | if self.is_lora_enabled or current_platform.is_rocm(): |
| 208 | x = x.contiguous() |
| 209 | ssm_params = self.x_proj(x)[0] |
| 210 | time_step, B, C = torch.split( |
| 211 | ssm_params, |
| 212 | [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], |
| 213 | dim=-1, |
| 214 | ) |
| 215 | if self.use_rms_norm: |
| 216 | assert self.dt_layernorm is not None |
| 217 | assert self.b_layernorm is not None |
| 218 | assert self.c_layernorm is not None |
| 219 | time_step = self.dt_layernorm(time_step.contiguous()) |
| 220 | B = self.b_layernorm(B.contiguous()) |
| 221 | C = self.c_layernorm(C.contiguous()) |
| 222 | |
| 223 | # ROCm: tensor from split is non-contiguous, causing incorrect |
| 224 | # GEMM results in dt_proj. |
| 225 | if current_platform.is_rocm(): |
| 226 | time_step = time_step.contiguous() |
| 227 | |
| 228 | discrete_time_step = self.dt_proj(time_step)[0].transpose(-2, -1) |
| 229 | return discrete_time_step, B, C |
| 230 | |
| 231 | def forward(self, hidden_states: torch.Tensor, output: torch.Tensor): |
| 232 | torch.ops.vllm.mamba_mixer( |
| 233 | hidden_states, |
| 234 | output, |
| 235 | _encode_layer_name(self.prefix), |
| 236 | ) |
| 237 | |
| 238 | def forward_impl(self, hidden_states: torch.Tensor, output: torch.Tensor): |
| 239 | """Run the Mamba-1 SSM pipeline. |
| 240 | |
| 241 | Steps |
| 242 | ----- |
| 243 | 1. Apply the gated-MLP linear projection to the raw input. |
| 244 | 2. Pass the projected sequence through the convolutional mixing layer. |
| 245 | 3. Feed the result into the State-Space Model (SSM) blocks. |
| 246 | 4. Perform the recurrence y ← SSM(A, B, C, Δ)(x) |
| 247 | to produce contextual representations. |
| 248 | 5. Project the contextualised sequence back |
| 249 | to the output embedding dimension. |
| 250 | |
| 251 | Batch handling |
| 252 | -------------- |
| 253 | Prefill and decode tokens are processed by dedicated CUDA |
| 254 | kernels for both the convolutional (conv1d) and SSM stages. |
| 255 | In the case of a mixed batch (containing both prefill and |
| 256 | decode tokens), both sets of kernels are executed independently |
| 257 | and their outputs are concatenated before the final output projection. |
| 258 | """ |
| 259 | forward_context: ForwardContext = get_forward_context() |
| 260 | attn_metadata_raw = forward_context.attn_metadata |
| 261 | |
| 262 | assert self.cache_config is not None |
| 263 | |
| 264 | attn_metadata: AttentionMetadata | None = None |
| 265 | if attn_metadata_raw is not None: |
| 266 | assert isinstance(attn_metadata_raw, dict) |
| 267 | attn_metadata = attn_metadata_raw.get(self.prefix) |
| 268 | if attn_metadata is not None: |
| 269 | assert isinstance(attn_metadata, Mamba1AttentionMetadata) |
| 270 | query_start_loc_p = attn_metadata.query_start_loc_p |
| 271 | state_indices_tensor_p = attn_metadata.state_indices_tensor_p |
| 272 | state_indices_tensor_d = attn_metadata.state_indices_tensor_d |
| 273 | conv_state = ( |
| 274 | self.kv_cache[0] |
| 275 | if is_conv_state_dim_first() |
| 276 | else self.kv_cache[0].transpose(-1, -2) |
| 277 | ) |
| 278 | ssm_state = self.kv_cache[1] |
| 279 | has_initial_states_p = attn_metadata.has_initial_states_p |
| 280 | cu_chunk_seqlen_p = attn_metadata.cu_chunk_seqlen_p |
| 281 | last_chunk_indices_p = attn_metadata.last_chunk_indices_p |
| 282 | |
| 283 | # 1. Gated MLP's linear projection |
| 284 | projected_states = self.in_proj(hidden_states)[0].transpose(-2, -1) |
| 285 | hidden_states_BC, gate = projected_states.chunk(2, dim=-2) |
| 286 | |
| 287 | conv_weights = self.conv1d.weight.view( |
| 288 | self.conv1d.weight.size(0), self.conv1d.weight.size(2) |
| 289 | ) |
| 290 | |
| 291 | if attn_metadata is None: |
| 292 | # V1 profile run |
| 293 | hidden_states_BC = hidden_states_BC.contiguous() |
| 294 | return self.out_proj(hidden_states_BC.transpose(-2, -1))[0] |
| 295 | |
| 296 | num_prefill_tokens = attn_metadata.num_prefill_tokens # token count |
| 297 | num_decode_tokens = attn_metadata.num_decode_tokens |
| 298 | has_prefill = num_prefill_tokens > 0 |
| 299 | has_decode = num_decode_tokens > 0 |
| 300 | num_actual_tokens = num_prefill_tokens + num_decode_tokens |
| 301 | |
| 302 | prefill_decode_split = split_batch_to_prefill_and_decode( |
| 303 | hidden_states_BC, |
| 304 | gate, |
| 305 | num_prefill_tokens, |
| 306 | num_decode_tokens, |
| 307 | ) |
| 308 | hidden_states_BC_p = prefill_decode_split.hidden_states_BC_p |
| 309 | hidden_states_BC_d = prefill_decode_split.hidden_states_BC_d |
| 310 | gate_p = prefill_decode_split.gate_p |
| 311 | gate_d = prefill_decode_split.gate_d |
| 312 | |
| 313 | ssm_outputs = [] |
| 314 | |
| 315 | if has_prefill: |
| 316 | # 2. Convolution sequence transformation |
| 317 | conv_out_p = causal_conv1d_fn( |
| 318 | hidden_states_BC_p, |
| 319 | conv_weights, |
| 320 | self.conv1d.bias, |
| 321 | activation=self.activation, |
| 322 | conv_states=conv_state, |
| 323 | has_initial_state=has_initial_states_p, |
| 324 | cache_indices=state_indices_tensor_p, |
| 325 | query_start_loc=query_start_loc_p, |
| 326 | metadata=attn_metadata, |
| 327 | ) |
| 328 | # 3. State Space Model sequence transformations. |
| 329 | discrete_time_step_p, B_p, C_p = self._ssm_transform( |
| 330 | conv_out_p.transpose(-2, -1) |
| 331 | ) |
| 332 | time_proj_bias = self._time_proj_bias() |
| 333 | |
| 334 | # 4. Perform the recurrence y ← SSM(A, B, C, Δ)(x) |
| 335 | scan_out_p = selective_scan_fn( |
| 336 | conv_out_p, |
| 337 | ssm_state, |
| 338 | discrete_time_step_p, |
| 339 | self.A, |
| 340 | B_p.transpose(-2, -1), |
| 341 | C_p.transpose(-2, -1), |
| 342 | self.D.float(), |
| 343 | gate_p, |
| 344 | time_proj_bias, |
| 345 | delta_softplus=True, |
| 346 | cache_indices=state_indices_tensor_p, |
| 347 | has_initial_state=has_initial_states_p, |
| 348 | query_start_loc=query_start_loc_p, |
| 349 | cu_chunk_seqlen=cu_chunk_seqlen_p, |
| 350 | last_chunk_indices=last_chunk_indices_p, |
| 351 | ) |
| 352 | ssm_outputs.append(scan_out_p) |
| 353 | |
| 354 | if has_decode: |
| 355 | # state_indices_tensor_d is assigned when attn_metadata is not None, |
| 356 | # and has_decode is only True when attn_metadata is not None |
| 357 | assert state_indices_tensor_d is not None |
| 358 | # 2. Convolution sequence transformation |
| 359 | conv_out_d = causal_conv1d_update( |
| 360 | hidden_states_BC_d.transpose(0, 1), |
| 361 | conv_state, |
| 362 | conv_weights, |
| 363 | self.conv1d.bias, |
| 364 | self.activation, |
| 365 | conv_state_indices=state_indices_tensor_d, |
| 366 | ).transpose(0, 1) |
| 367 | |
| 368 | # 3. State Space Model sequence transformation. |
| 369 | discrete_time_step_d, B_d, C_d = self._ssm_transform( |
| 370 | conv_out_d.transpose(-2, -1) |
| 371 | ) |
| 372 | time_proj_bias = self._time_proj_bias() |
| 373 | |
| 374 | # 4. Perform the recurrence y ← SSM(A, B, C, Δ)(x) |
| 375 | scan_outputs_d = torch.empty_like(hidden_states_BC_d.transpose(0, 1)) |
| 376 | selective_state_update( |
| 377 | ssm_state, |
| 378 | conv_out_d.transpose(0, 1), |
| 379 | discrete_time_step_d.transpose(0, 1), |
| 380 | self.A, |
| 381 | B_d, |
| 382 | C_d, |
| 383 | self.D, |
| 384 | time_proj_bias, |
| 385 | z=gate_d.transpose(0, 1), |
| 386 | dt_softplus=True, |
| 387 | state_batch_indices=state_indices_tensor_d, |
| 388 | dst_state_batch_indices=state_indices_tensor_d, |
| 389 | out=scan_outputs_d, |
| 390 | ) |
| 391 | scan_outputs_d = scan_outputs_d.transpose(0, 1) |
| 392 | |
| 393 | ssm_outputs.insert(0, scan_outputs_d) |
| 394 | |
| 395 | scan_outputs_combined = ( |
| 396 | ssm_outputs[0] if len(ssm_outputs) == 1 else torch.cat(ssm_outputs, dim=-1) |
| 397 | ) |
| 398 | |
| 399 | # 5. Final output projection |
| 400 | if self.is_lora_enabled: # Lora kernel requires contiguous tensor. |
| 401 | scan_outputs_combined = scan_outputs_combined.transpose(-2, -1).contiguous() |
| 402 | out = self.out_proj(scan_outputs_combined)[0] |
| 403 | else: |
| 404 | out = self.out_proj(scan_outputs_combined.transpose(-2, -1))[0] |
| 405 | |
| 406 | output[:num_actual_tokens] = out |
| 407 | |
| 408 | def get_state_dtype(self) -> tuple[torch.dtype]: |
| 409 | assert self.model_config is not None |
| 410 | assert self.cache_config is not None |
| 411 | return MambaStateDtypeCalculator.mamba1_state_dtype( |
| 412 | self.model_config.dtype, |
| 413 | self.cache_config.mamba_cache_dtype, |
| 414 | self.cache_config.mamba_ssm_cache_dtype, |
| 415 | ) |
| 416 | |
| 417 | def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]: |
| 418 | return MambaStateShapeCalculator.mamba1_state_shape( |
| 419 | tp_world_size=get_tensor_model_parallel_world_size(), |
| 420 | intermediate_size=self.intermediate_size, |
| 421 | state_size=self.ssm_state_size, |
| 422 | conv_kernel=self.conv_kernel_size, |
| 423 | ) |
| 424 | |
| 425 | @property |
| 426 | def mamba_type(self) -> MambaAttentionBackendEnum: |
| 427 | return MambaAttentionBackendEnum.MAMBA1 |
| 428 | |
| 429 | def _time_proj_bias(self) -> torch.Tensor | None: |
| 430 | if hasattr(self.dt_proj, "bias") and self.dt_proj.bias is not None: |
| 431 | return self.dt_proj.bias.float() |
| 432 | return None |
| 433 | |
| 434 | |
| 435 | class PrefillDecodeSplit(NamedTuple): |
| 436 | hidden_states_BC_p: torch.Tensor |
| 437 | hidden_states_BC_d: torch.Tensor |
| 438 | gate_p: torch.Tensor |
| 439 | gate_d: torch.Tensor |
| 440 | |
| 441 | |
| 442 | def split_batch_to_prefill_and_decode( |
| 443 | hidden_states_BC: torch.Tensor, |
| 444 | gate: torch.Tensor, |
| 445 | num_prefill_tokens: int, |
| 446 | num_decode_tokens: int, |
| 447 | ) -> PrefillDecodeSplit: |
| 448 | num_actual_tokens = num_prefill_tokens + num_decode_tokens |
| 449 | |
| 450 | # In v1, decode tokens come first, then prefill tokens. |
| 451 | hidden_states_BC_d, hidden_states_BC_p = torch.split( |
| 452 | hidden_states_BC[..., :num_actual_tokens], |
| 453 | [num_decode_tokens, num_prefill_tokens], |
| 454 | dim=-1, |
| 455 | ) |
| 456 | gate_d, gate_p = torch.split( |
| 457 | gate[..., :num_actual_tokens], [num_decode_tokens, num_prefill_tokens], dim=-1 |
| 458 | ) |
| 459 | |
| 460 | return PrefillDecodeSplit( |
| 461 | hidden_states_BC_p=hidden_states_BC_p, |
| 462 | hidden_states_BC_d=hidden_states_BC_d, |
| 463 | gate_p=gate_p, |
| 464 | gate_d=gate_d, |
| 465 | ) |
| 466 | |
| 467 | |
| 468 | def mamba_mixer( |
| 469 | hidden_states: torch.Tensor, |
| 470 | output: torch.Tensor, |
| 471 | layer_name: LayerNameType, |
| 472 | ) -> None: |
| 473 | layer_name = _resolve_layer_name(layer_name) |
| 474 | forward_context: ForwardContext = get_forward_context() |
| 475 | self = forward_context.no_compile_layers[layer_name] |
| 476 | self.forward_impl(hidden_states=hidden_states, output=output) |
| 477 | |
| 478 | |
| 479 | direct_register_custom_op( |
| 480 | op_name="mamba_mixer", |
| 481 | op_func=mamba_mixer, |
| 482 | mutates_args=["output"], |
| 483 | ) |