| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | from collections.abc import Sequence |
| 6 | |
| 7 | import torch |
| 8 | from torch import nn |
| 9 | |
| 10 | from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config |
| 11 | from vllm.config.mamba import MambaBackendEnum |
| 12 | from vllm.distributed import ( |
| 13 | divide, |
| 14 | get_tensor_model_parallel_rank, |
| 15 | get_tensor_model_parallel_world_size, |
| 16 | tensor_model_parallel_all_gather, |
| 17 | tensor_model_parallel_all_reduce, |
| 18 | ) |
| 19 | from vllm.forward_context import ForwardContext, get_forward_context |
| 20 | from vllm.logger import init_logger |
| 21 | from vllm.model_executor.custom_op import CustomOp, PluggableLayer |
| 22 | from vllm.model_executor.layers.attention import Attention |
| 23 | from vllm.model_executor.layers.linear import ( |
| 24 | ColumnParallelLinear, |
| 25 | MergedColumnParallelLinear, |
| 26 | RowParallelLinear, |
| 27 | ) |
| 28 | from vllm.model_executor.layers.mamba.abstract import MambaBase |
| 29 | from vllm.model_executor.layers.mamba.mamba_utils import ( |
| 30 | MambaStateDtypeCalculator, |
| 31 | MambaStateShapeCalculator, |
| 32 | is_conv_state_dim_first, |
| 33 | ) |
| 34 | from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( |
| 35 | causal_conv1d_fn, |
| 36 | causal_conv1d_update, |
| 37 | ) |
| 38 | from vllm.model_executor.layers.mamba.ops.layernorm_gated import rms_norm_gated |
| 39 | from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_output_only import ( # noqa: E501 |
| 40 | selective_state_update_replayssm_output_only, |
| 41 | ) |
| 42 | from vllm.model_executor.layers.mamba.ops.ssd_combined import ( |
| 43 | mamba_chunk_scan_combined_varlen, |
| 44 | ) |
| 45 | from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( |
| 46 | commit_replayssm_ring_trackers, |
| 47 | reset_replayssm_ring_trackers, |
| 48 | selective_state_update, |
| 49 | selective_state_update_replayssm_flashinfer, |
| 50 | ) |
| 51 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 52 | from vllm.model_executor.model_loader.weight_utils import ( |
| 53 | LoaderFunction, |
| 54 | composed_weight_loader, |
| 55 | sharded_weight_loader, |
| 56 | ) |
| 57 | from vllm.model_executor.parameter import BasevLLMParameter |
| 58 | from vllm.model_executor.utils import set_weight_attrs |
| 59 | from vllm.platforms import current_platform |
| 60 | from vllm.utils.torch_utils import ( |
| 61 | LayerNameType, |
| 62 | _encode_layer_name, |
| 63 | _resolve_layer_name, |
| 64 | direct_register_custom_op, |
| 65 | ) |
| 66 | from vllm.v1.attention.backend import AttentionMetadata |
| 67 | from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata |
| 68 | from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum |
| 69 | from vllm.v1.kv_cache_interface import KVCacheGroupSpec |
| 70 | |
| 71 | logger = init_logger(__name__) |
| 72 | |
| 73 | |
| 74 | def _view_mtp_decode_tensor( |
| 75 | tensor: torch.Tensor, decode_batch: int, spec_query_len: int |
| 76 | ) -> torch.Tensor: |
| 77 | return tensor.view(decode_batch, spec_query_len, *tensor.shape[1:]) |
| 78 | |
| 79 | |
| 80 | # Added by the IBM Team, 2024 |
| 81 | |
| 82 | |
| 83 | # Adapted from transformers.models.mamba2.modeling_mamba2.MambaRMSNormGated |
| 84 | # --8<-- [start:mixer2_gated_rms_norm] |
| 85 | @CustomOp.register("mixer2_gated_rms_norm") |
| 86 | class Mixer2RMSNormGated(CustomOp): |
| 87 | # --8<-- [end:mixer2_gated_rms_norm] |
| 88 | |
| 89 | def __init__( |
| 90 | self, |
| 91 | full_hidden_size: int, |
| 92 | full_n_groups: int, |
| 93 | use_rms_norm: bool = True, |
| 94 | eps: float = 1e-6, |
| 95 | ): |
| 96 | super().__init__() |
| 97 | self.tp_size = get_tensor_model_parallel_world_size() |
| 98 | self.tp_rank = get_tensor_model_parallel_rank() |
| 99 | self.full_hidden_size = full_hidden_size |
| 100 | self.group_size = full_hidden_size // full_n_groups |
| 101 | self.per_rank_hidden_size = full_hidden_size // self.tp_size |
| 102 | self.n_groups = full_hidden_size // self.group_size |
| 103 | |
| 104 | self.variance_epsilon = eps |
| 105 | self.use_rms_norm = use_rms_norm |
| 106 | if self.use_rms_norm: |
| 107 | # Register norm weight only if we're actually applying RMSNorm |
| 108 | self.weight = nn.Parameter(torch.ones(self.per_rank_hidden_size)) |
| 109 | set_weight_attrs(self.weight, {"weight_loader": sharded_weight_loader(0)}) |
| 110 | else: |
| 111 | # Avoid checkpoint mismatch by skipping unused parameter |
| 112 | self.register_parameter("weight", None) |
| 113 | assert self.full_hidden_size % self.tp_size == 0, ( |
| 114 | "Tensor parallel world size must divide hidden size." |
| 115 | ) |
| 116 | |
| 117 | def forward_native( |
| 118 | self, |
| 119 | x: torch.Tensor, |
| 120 | gate: torch.Tensor, |
| 121 | ): |
| 122 | # Three tensor-parallel cases: |
| 123 | # 1. n_groups is 1 |
| 124 | # In this case we parallelize along the reduction dim. |
| 125 | # Each rank computes a local sum of squares followed by AllReduce |
| 126 | # 2. tp_size divides n_groups |
| 127 | # Each rank only reduces within its local group(s). |
| 128 | # No collective ops necessary. |
| 129 | # 3. The general case can be pretty complicated so we AllGather |
| 130 | # the input and then redundantly compute the RMSNorm. |
| 131 | input_dtype = x.dtype |
| 132 | x = x * nn.functional.silu(gate.to(torch.float32)) |
| 133 | if not self.use_rms_norm: |
| 134 | return x.to(input_dtype) |
| 135 | |
| 136 | if self.n_groups == 1: |
| 137 | if self.tp_size > 1: |
| 138 | # Compute local sum and then reduce to obtain global sum |
| 139 | local_sums = x.pow(2).sum(dim=-1, keepdim=True) |
| 140 | global_sums = tensor_model_parallel_all_reduce(local_sums) |
| 141 | # Calculate the variance |
| 142 | count = self.tp_size * x.shape[-1] |
| 143 | variance = global_sums / count |
| 144 | |
| 145 | else: |
| 146 | variance = x.pow(2).mean(-1, keepdim=True) |
| 147 | x = x * torch.rsqrt(variance + self.variance_epsilon) |
| 148 | else: |
| 149 | redundant_tp: bool = self.n_groups % self.tp_size != 0 |
| 150 | if redundant_tp: |
| 151 | # To handle the general case, redundantly apply the variance |
| 152 | x = tensor_model_parallel_all_gather(x, -1) |
| 153 | |
| 154 | *prefix_dims, hidden_dim = x.shape |
| 155 | group_count = hidden_dim // self.group_size |
| 156 | x_grouped = x.view(*prefix_dims, group_count, self.group_size) |
| 157 | variance = x_grouped.pow(2).mean(-1, keepdim=True) |
| 158 | x_grouped = x_grouped * torch.rsqrt(variance + self.variance_epsilon) |
| 159 | x = x_grouped.view(*prefix_dims, hidden_dim) |
| 160 | |
| 161 | if redundant_tp: |
| 162 | start = self.per_rank_hidden_size * self.tp_rank |
| 163 | end = start + self.per_rank_hidden_size |
| 164 | x = x[..., start:end] |
| 165 | |
| 166 | return self.weight * x.to(input_dtype) |
| 167 | |
| 168 | def forward_cuda( |
| 169 | self, |
| 170 | x: torch.Tensor, |
| 171 | gate: torch.Tensor, |
| 172 | ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: |
| 173 | input_dtype = x.dtype |
| 174 | if not self.use_rms_norm: |
| 175 | # Keep gate in float32 for numerical stability during silu |
| 176 | return x * nn.functional.silu(gate.to(torch.float32)).to(input_dtype) |
| 177 | |
| 178 | if ((self.n_groups % self.tp_size) != 0) or self.n_groups != 1: |
| 179 | return self.forward_native(x, gate) |
| 180 | |
| 181 | return rms_norm_gated( |
| 182 | x, |
| 183 | self.weight.data, |
| 184 | bias=None, |
| 185 | z=gate, |
| 186 | eps=self.variance_epsilon, |
| 187 | norm_before_gate=False, |
| 188 | ) |
| 189 | |
| 190 | def forward_xpu( |
| 191 | self, |
| 192 | x: torch.Tensor, |
| 193 | gate: torch.Tensor, |
| 194 | ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: |
| 195 | return self.forward_cuda(x, gate) |
| 196 | |
| 197 | |
| 198 | def mamba_v2_sharded_weight_loader( |
| 199 | shard_spec: list[tuple[int, int, float]], |
| 200 | tp_size: int, |
| 201 | tp_rank: int, |
| 202 | ) -> LoaderFunction: |
| 203 | """Create a weight loader for mamba v2. This ensures that the projections |
| 204 | are correctly sharded so that they can be split into x, B, C. It also |
| 205 | ensures that all the groups corresponding to a head shard is placed |
| 206 | together with it. |
| 207 | """ |
| 208 | |
| 209 | def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: |
| 210 | # - track boundary of (sharded) param, and loaded_weight, respectively |
| 211 | boundary, loaded_boundary = 0, 0 |
| 212 | |
| 213 | # - iterate over the shard specs |
| 214 | for full_dim, extra, duplicate_groups in shard_spec: |
| 215 | # - full dim is the model dim (before TP). |
| 216 | # - extra > 0, means there is expected overall increase |
| 217 | # of dimensions. This is so because of replication. |
| 218 | # - ratio is used map the tp_rank to the actual shard |
| 219 | # rank. This is useful when there is replication of |
| 220 | # groups to accompany head shards. |
| 221 | |
| 222 | # - size of the loaded shard |
| 223 | shard_size = full_dim // tp_size |
| 224 | |
| 225 | # - compute the rank into the loaded shard. |
| 226 | # - if there is replication, different TP shards will |
| 227 | # take from the same rank. |
| 228 | # NOTE: currently we only support duplication |
| 229 | # in the case where num_groups == 1 |
| 230 | rank = 0 if duplicate_groups else tp_rank |
| 231 | |
| 232 | # - leftmost boundary index into loaded weight. |
| 233 | loaded_skip = rank * shard_size |
| 234 | loaded_start_idx = loaded_boundary + loaded_skip |
| 235 | |
| 236 | # - take these many dims from the loaded weight. |
| 237 | take = min(shard_size, full_dim - extra - loaded_skip) |
| 238 | |
| 239 | # - always shard on dim 0 |
| 240 | # - the ignore is for a mundane mypy error as it does not |
| 241 | # seem to handle slices well. |
| 242 | # https://github.com/python/mypy/issues/2410 |
| 243 | target_slice = param.data[boundary : (boundary + take), ...] |
| 244 | param.data[ |
| 245 | boundary : (boundary + take), ... # type: ignore[misc] |
| 246 | ] = loaded_weight[ |
| 247 | loaded_start_idx : ( |
| 248 | loaded_start_idx + take |
| 249 | ) # type: ignore[misc] |
| 250 | ].view_as(target_slice) |
| 251 | |
| 252 | # move indexing boundaries |
| 253 | boundary += shard_size |
| 254 | loaded_boundary += full_dim - extra |
| 255 | |
| 256 | return loader |
| 257 | |
| 258 | |
| 259 | # Adapted from transformers.models.mamba.modeling_mamba.MambaMixer |
| 260 | # --8<-- [start:mamba_mixer2] |
| 261 | @PluggableLayer.register("mamba_mixer2") |
| 262 | class MambaMixer2(MambaBase, PluggableLayer): |
| 263 | """Compute ∆, A, B, C, and D the state space parameters and compute |
| 264 | the `contextualized_states`. A, D are input independent |
| 265 | (see Mamba paper [1] Section 3.5.2 "Interpretation of A" |
| 266 | for why A isn't selective) ∆, B, C are input-dependent |
| 267 | (this is a key difference between Mamba and the linear time |
| 268 | invariant S4, and is why Mamba is called |
| 269 | **selective** state spaces) |
| 270 | """ |
| 271 | |
| 272 | # --8<-- [end:mamba_mixer2] |
| 273 | |
| 274 | def __init__( |
| 275 | self, |
| 276 | hidden_size: int, |
| 277 | ssm_state_size: int, |
| 278 | conv_kernel_size: int, |
| 279 | intermediate_size: int, |
| 280 | use_conv_bias: bool, |
| 281 | use_bias: bool, |
| 282 | n_groups: int = 1, |
| 283 | num_heads: int = 128, |
| 284 | head_dim: int = 64, |
| 285 | rms_norm_eps: float = 1e-5, |
| 286 | activation: str = "silu", |
| 287 | use_rms_norm: bool = True, |
| 288 | model_config: ModelConfig | None = None, |
| 289 | cache_config: CacheConfig | None = None, |
| 290 | quant_config: QuantizationConfig | None = None, |
| 291 | prefix: str = "", |
| 292 | ): |
| 293 | super().__init__() |
| 294 | |
| 295 | # For TP, the sharding plan is as follows: |
| 296 | # - for the conv modules, since |
| 297 | # conv_dim = intermediate_size * 2 * n_groups * ssm_state_size, |
| 298 | # we shard intermediate_size and n_groups |
| 299 | # - since intermediate_size = n_heads * head_dim, sharding on |
| 300 | # intermediate_size is achieved by sharding on n_heads. |
| 301 | # - IF, world_size divides groups, then sharding |
| 302 | # (n_groups / world_size, n_heads / world_size) |
| 303 | # also maintains the invariant n_heads % n_groups == 0 |
| 304 | # - HOWEVER IF, world_size DOES NOT divide groups, then we need |
| 305 | # to allocate extra space in the shard, such that groups |
| 306 | # may be replicated to follow the head shard. |
| 307 | # - NOTE: currently for the world size DOES NOT divide groups |
| 308 | # case, we only support the case when n_groups == 1 |
| 309 | self.tp_size = get_tensor_model_parallel_world_size() |
| 310 | tp_rank = get_tensor_model_parallel_rank() |
| 311 | |
| 312 | assert num_heads % self.tp_size == 0, ( |
| 313 | "Tensor parallel world size must divide num heads." |
| 314 | ) |
| 315 | |
| 316 | assert (n_groups % self.tp_size) == 0 or n_groups == 1, ( |
| 317 | "If tensor parallel world size does not divide num_groups, " |
| 318 | "then num_groups must equal 1." |
| 319 | ) |
| 320 | |
| 321 | self.ssm_state_size = ssm_state_size |
| 322 | self.conv_kernel_size = conv_kernel_size |
| 323 | self.activation = activation |
| 324 | |
| 325 | self.intermediate_size = intermediate_size |
| 326 | self.head_dim = head_dim |
| 327 | self.num_heads = num_heads |
| 328 | |
| 329 | self.n_groups = n_groups |
| 330 | if n_groups % self.tp_size != 0: |
| 331 | # - for TP we shard conv_dim by sharding on n_groups, |
| 332 | # - but if n_groups cannot divide tp_size, we need to |
| 333 | # extend some extra groups |
| 334 | groups = MambaStateShapeCalculator.extra_groups_for_head_shards( |
| 335 | n_groups, self.tp_size |
| 336 | ) |
| 337 | self.n_groups = n_groups + groups |
| 338 | |
| 339 | self.groups_ssm_state_size = self.n_groups * self.ssm_state_size |
| 340 | self.conv_dim = intermediate_size + 2 * self.groups_ssm_state_size |
| 341 | |
| 342 | if n_groups % self.tp_size == 0: |
| 343 | self.conv1d = MergedColumnParallelLinear( |
| 344 | input_size=conv_kernel_size, |
| 345 | output_sizes=[ |
| 346 | intermediate_size, |
| 347 | self.groups_ssm_state_size, |
| 348 | self.groups_ssm_state_size, |
| 349 | ], |
| 350 | bias=use_conv_bias, |
| 351 | quant_config=None, |
| 352 | prefix=f"{prefix}.conv1d", |
| 353 | ) |
| 354 | |
| 355 | self.in_proj = MergedColumnParallelLinear( |
| 356 | input_size=hidden_size, |
| 357 | output_sizes=[ |
| 358 | intermediate_size, |
| 359 | intermediate_size, |
| 360 | self.groups_ssm_state_size, |
| 361 | self.groups_ssm_state_size, |
| 362 | self.num_heads, |
| 363 | ], |
| 364 | bias=use_bias, |
| 365 | quant_config=quant_config, |
| 366 | prefix=f"{prefix}.in_proj", |
| 367 | ) |
| 368 | else: |
| 369 | # This is the n_groups == 1 case, |
| 370 | # where we need to duplicate groups if TP>1. |
| 371 | |
| 372 | self.conv1d = ColumnParallelLinear( |
| 373 | input_size=conv_kernel_size, |
| 374 | output_size=self.conv_dim, |
| 375 | bias=use_conv_bias, |
| 376 | quant_config=None, |
| 377 | prefix=f"{prefix}.conv1d", |
| 378 | ) |
| 379 | |
| 380 | self.in_proj = ColumnParallelLinear( |
| 381 | input_size=hidden_size, |
| 382 | output_size=intermediate_size + self.conv_dim + self.num_heads, |
| 383 | bias=use_bias, |
| 384 | quant_config=quant_config, |
| 385 | prefix=f"{prefix}.in_proj", |
| 386 | ) |
| 387 | |
| 388 | # - because in_proj is a concatenation of 3 weights, we |
| 389 | # need to interleave them before sharding |
| 390 | # - use the custom weight loader mamba_v2_sharded_weight_loader |
| 391 | # for conv1d.bias, covn1d.weight and in_proj.weight |
| 392 | # - need to set these settings, to assign the groups |
| 393 | # to the head shards |
| 394 | group_shard_settings = ( |
| 395 | self.groups_ssm_state_size, # expected model size |
| 396 | (self.n_groups - n_groups) * self.ssm_state_size, # extra dims assigned |
| 397 | n_groups == 1, # if there was only one group |
| 398 | ) |
| 399 | intermediate_settings = (intermediate_size, 0, False) |
| 400 | head_settings = (self.num_heads, 0, False) |
| 401 | |
| 402 | # - the weight already has a "weight_loader" attribute |
| 403 | # which set_weight_attrs will raise if we do not |
| 404 | # delete before trying to override it |
| 405 | # - ditto for the other two weights below |
| 406 | delattr(self.conv1d.bias, "weight_loader") |
| 407 | set_weight_attrs( |
| 408 | self.conv1d.bias, |
| 409 | { |
| 410 | "weight_loader": mamba_v2_sharded_weight_loader( |
| 411 | [ |
| 412 | intermediate_settings, |
| 413 | group_shard_settings, |
| 414 | group_shard_settings, |
| 415 | ], |
| 416 | self.tp_size, |
| 417 | tp_rank, |
| 418 | ) |
| 419 | }, |
| 420 | ) |
| 421 | |
| 422 | delattr(self.conv1d.weight, "weight_loader") |
| 423 | set_weight_attrs( |
| 424 | self.conv1d.weight, |
| 425 | { |
| 426 | "weight_loader": mamba_v2_sharded_weight_loader( |
| 427 | [ |
| 428 | intermediate_settings, |
| 429 | group_shard_settings, |
| 430 | group_shard_settings, |
| 431 | ], |
| 432 | self.tp_size, |
| 433 | tp_rank, |
| 434 | ) |
| 435 | }, |
| 436 | ) |
| 437 | |
| 438 | # Create the custom weight loader for Mamba sharding with group |
| 439 | # replication. This handles the interleaved projections correctly. |
| 440 | mamba_loader = mamba_v2_sharded_weight_loader( |
| 441 | [ |
| 442 | intermediate_settings, # for gate |
| 443 | intermediate_settings, |
| 444 | group_shard_settings, |
| 445 | group_shard_settings, |
| 446 | head_settings, # for dt |
| 447 | ], |
| 448 | self.tp_size, |
| 449 | tp_rank, |
| 450 | ) |
| 451 | |
| 452 | # Apply the custom weight loader to in_proj.weight and its scales. |
| 453 | # Works for both non-quantized (Parameter) and quantized |
| 454 | # (ModelWeightParameter which extends BasevLLMParameter) |
| 455 | for attr_name in ("weight", "weight_scale", "input_scale"): |
| 456 | param = getattr(self.in_proj, attr_name, None) |
| 457 | if param is not None: |
| 458 | if isinstance(param, BasevLLMParameter): |
| 459 | param.weight_loader = mamba_loader |
| 460 | else: |
| 461 | if hasattr(param, "weight_loader"): |
| 462 | delattr(param, "weight_loader") |
| 463 | set_weight_attrs(param, {"weight_loader": mamba_loader}) |
| 464 | |
| 465 | # unsqueeze to fit conv1d weights shape into the linear weights shape. |
| 466 | # Can't do this in `weight_loader` since it already exists in |
| 467 | # `ColumnParallelLinear` and `MergedColumnParallelLinear`, |
| 468 | # and `set_weight_attrs` doesn't allow to override it |
| 469 | self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1) |
| 470 | conv_weights = self.conv1d.weight.view( |
| 471 | self.conv1d.weight.size(0), self.conv1d.weight.size(2) |
| 472 | ) |
| 473 | self.register_buffer("conv_weights", conv_weights, persistent=False) |
| 474 | |
| 475 | # - these are TPed by heads to reduce the size of the |
| 476 | # temporal shape |
| 477 | self.A = nn.Parameter( |
| 478 | torch.empty( |
| 479 | divide(num_heads, self.tp_size), |
| 480 | dtype=torch.float32, |
| 481 | ) |
| 482 | ) |
| 483 | self.D = nn.Parameter(torch.ones(num_heads // self.tp_size)) |
| 484 | self.dt_bias = nn.Parameter(torch.ones(num_heads // self.tp_size)) |
| 485 | self.use_rms_norm = use_rms_norm |
| 486 | |
| 487 | set_weight_attrs(self.D, {"weight_loader": sharded_weight_loader(0)}) |
| 488 | a_weight_loader = composed_weight_loader( |
| 489 | sharded_weight_loader(0), lambda x: -torch.exp(x.float()) |
| 490 | ) |
| 491 | set_weight_attrs(self.A, {"weight_loader": a_weight_loader}) |
| 492 | set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)}) |
| 493 | |
| 494 | self.out_proj = RowParallelLinear( |
| 495 | intermediate_size, |
| 496 | hidden_size, |
| 497 | bias=use_bias, |
| 498 | input_is_parallel=True, |
| 499 | quant_config=quant_config, |
| 500 | prefix=f"{prefix}.out_proj", |
| 501 | ) |
| 502 | |
| 503 | self.norm = Mixer2RMSNormGated( |
| 504 | intermediate_size, n_groups, self.use_rms_norm, eps=rms_norm_eps |
| 505 | ) |
| 506 | |
| 507 | self._ssd_kernels_warmed_up = False |
| 508 | |
| 509 | # - get hidden_states, B and C after depthwise convolution. |
| 510 | self.split_hidden_states_B_C_fn = lambda hidden_states_B_C: torch.split( |
| 511 | hidden_states_B_C, |
| 512 | [ |
| 513 | self.intermediate_size // self.tp_size, |
| 514 | self.groups_ssm_state_size // self.tp_size, |
| 515 | self.groups_ssm_state_size // self.tp_size, |
| 516 | ], |
| 517 | dim=-1, |
| 518 | ) |
| 519 | |
| 520 | vllm_config = get_current_vllm_config() |
| 521 | compilation_config = vllm_config.compilation_config |
| 522 | if prefix in compilation_config.static_forward_context: |
| 523 | raise ValueError(f"Duplicate layer name: {prefix}") |
| 524 | compilation_config.static_forward_context[prefix] = self |
| 525 | |
| 526 | self.model_config = model_config |
| 527 | self.cache_config = cache_config |
| 528 | self.prefix = prefix |
| 529 | self.use_replayssm = ( |
| 530 | cache_config.use_replayssm if cache_config is not None else False |
| 531 | ) |
| 532 | self.replayssm_buffer_len = ( |
| 533 | cache_config.replayssm_buffer_len |
| 534 | if cache_config is not None and cache_config.use_replayssm |
| 535 | else None |
| 536 | ) |
| 537 | self.mamba_config = vllm_config.mamba_config |
| 538 | if self.use_replayssm and self.num_heads % self.tp_size != 0: |
| 539 | raise ValueError( |
| 540 | "--use-replayssm requires tensor-parallel heads to divide evenly" |
| 541 | ) |
| 542 | # ReplaySSM appends x/dt/B rings to (conv_state, ssm_state). |
| 543 | _n_state = 5 if self.use_replayssm else 2 |
| 544 | self.kv_cache = tuple(torch.tensor([]) for _ in range(_n_state)) |
| 545 | self._replayssm_ring_start = torch.empty(0, dtype=torch.int32) |
| 546 | self._replayssm_prev_num_accepted = torch.empty(0, dtype=torch.int32) |
| 547 | self._replayssm_prev_query_len = torch.empty(0, dtype=torch.int32) |
| 548 | self._commits_replayssm_trackers = True |
| 549 | self._updates_replayssm_trackers = True |
| 550 | |
| 551 | self.num_spec = vllm_config.num_speculative_tokens |
| 552 | |
| 553 | # Pre-compute sizes for forward pass |
| 554 | self.tped_intermediate_size = self.intermediate_size // self.tp_size |
| 555 | self.tped_conv_size = self.conv_dim // self.tp_size |
| 556 | self.tped_dt_size = self.num_heads // self.tp_size |
| 557 | |
| 558 | self.split_hidden_states_B_C_fn = lambda hidden_states_B_C: torch.split( |
| 559 | hidden_states_B_C, |
| 560 | [ |
| 561 | self.tped_intermediate_size, |
| 562 | self.groups_ssm_state_size // self.tp_size, |
| 563 | self.groups_ssm_state_size // self.tp_size, |
| 564 | ], |
| 565 | dim=-1, |
| 566 | ) |
| 567 | |
| 568 | # Check if running on Blackwell (SM100+) for kernel tuning |
| 569 | self.is_blackwell = current_platform.is_device_capability_family(100) |
| 570 | |
| 571 | def forward( |
| 572 | self, |
| 573 | hidden_states: torch.Tensor, |
| 574 | mup_vector: torch.Tensor | None = None, |
| 575 | ): |
| 576 | # 1. Gated MLP's linear projection |
| 577 | projected_states, _ = self.in_proj(hidden_states) |
| 578 | if mup_vector is not None: |
| 579 | projected_states = projected_states * mup_vector |
| 580 | |
| 581 | # 2. Prepare inputs for conv + SSM |
| 582 | ssm_output = torch.empty( |
| 583 | [ |
| 584 | hidden_states.shape[0], |
| 585 | (self.num_heads // self.tp_size) * self.head_dim, |
| 586 | ], |
| 587 | dtype=hidden_states.dtype, |
| 588 | device=hidden_states.device, |
| 589 | ) |
| 590 | |
| 591 | # 3. conv + SSM |
| 592 | # (split `projected_states` into hidden_states_B_C, dt in the custom op to |
| 593 | # ensure it is not treated as an intermediate tensor by torch compile) |
| 594 | torch.ops.vllm.mamba_mixer2( |
| 595 | projected_states, |
| 596 | ssm_output, |
| 597 | _encode_layer_name(self.prefix), |
| 598 | ) |
| 599 | |
| 600 | # 4. gated MLP |
| 601 | # GatedRMSNorm internally applying SiLU to the gate |
| 602 | # SiLU is applied internally before normalization, unlike standard |
| 603 | # norm usage |
| 604 | gate = projected_states[..., : self.tped_intermediate_size] |
| 605 | hidden_states = self.norm(ssm_output, gate) |
| 606 | |
| 607 | # 5. Final linear projection |
| 608 | output, _ = self.out_proj(hidden_states) |
| 609 | |
| 610 | return output |
| 611 | |
| 612 | def _warmup_ssd_kernels(self, projected_states: torch.Tensor) -> None: |
| 613 | """Run a minimal SSD forward pass to trigger Triton autotuning |
| 614 | while GPU memory is still plentiful (before SSM cache allocation). |
| 615 | """ |
| 616 | if self._ssd_kernels_warmed_up: |
| 617 | return |
| 618 | self._ssd_kernels_warmed_up = True |
| 619 | logger.info_once("Warming up Mamba2 SSD Triton kernels...") |
| 620 | |
| 621 | device = projected_states.device |
| 622 | dtype = projected_states.dtype |
| 623 | |
| 624 | nheads = self.num_heads // self.tp_size |
| 625 | ngroups = self.n_groups // self.tp_size |
| 626 | headdim = self.head_dim |
| 627 | dstate = self.ssm_state_size |
| 628 | |
| 629 | if self.model_config is None: |
| 630 | return |
| 631 | chunk_size = self.model_config.get_mamba_chunk_size() |
| 632 | |
| 633 | # Triton's autotuner includes tensor dtypes in its cache key, |
| 634 | # so state_dtype must match what real inference uses. |
| 635 | ssm_state_dtype = self.get_state_dtype()[1] |
| 636 | |
| 637 | # SSD kernel autotune keys depend on dtype and head dimensions, |
| 638 | # not on sequence length or batch size, so a single shape suffices. |
| 639 | seqlen = chunk_size |
| 640 | batch = 1 |
| 641 | nchunks = seqlen // chunk_size # = 1 |
| 642 | |
| 643 | x = torch.randn(seqlen, nheads, headdim, device=device, dtype=dtype) |
| 644 | dt = torch.randn(seqlen, nheads, device=device, dtype=dtype) |
| 645 | B = torch.randn(seqlen, ngroups, dstate, device=device, dtype=dtype) |
| 646 | C = torch.randn(seqlen, ngroups, dstate, device=device, dtype=dtype) |
| 647 | cu_seqlens = torch.tensor([0, seqlen], device=device, dtype=torch.int32) |
| 648 | cu_chunk_seqlens = torch.tensor( |
| 649 | [i * chunk_size for i in range(nchunks + 1)], |
| 650 | device=device, |
| 651 | dtype=torch.int32, |
| 652 | ) |
| 653 | last_chunk_indices = torch.tensor( |
| 654 | [nchunks - 1], device=device, dtype=torch.int32 |
| 655 | ) |
| 656 | seq_idx = torch.zeros(nchunks, device=device, dtype=torch.int32) |
| 657 | out = torch.empty(seqlen, nheads, headdim, device=device, dtype=dtype) |
| 658 | |
| 659 | # Two kernels (_state_passing_fwd, _chunk_scan_fwd) use |
| 660 | # HAS_INITSTATES as a constexpr, producing separate compiled |
| 661 | # binaries. Warm up both code paths so neither triggers |
| 662 | # JIT compilation during inference. |
| 663 | for use_initial_states in (False, True): |
| 664 | initial_states = ( |
| 665 | torch.randn( |
| 666 | batch, |
| 667 | nheads, |
| 668 | headdim, |
| 669 | dstate, |
| 670 | device=device, |
| 671 | dtype=ssm_state_dtype, |
| 672 | ) |
| 673 | if use_initial_states |
| 674 | else None |
| 675 | ) |
| 676 | try: |
| 677 | mamba_chunk_scan_combined_varlen( |
| 678 | x=x, |
| 679 | dt=dt, |
| 680 | A=self.A, |
| 681 | B=B, |
| 682 | C=C, |
| 683 | chunk_size=chunk_size, |
| 684 | cu_seqlens=cu_seqlens, |
| 685 | cu_chunk_seqlens=cu_chunk_seqlens, |
| 686 | last_chunk_indices=last_chunk_indices, |
| 687 | seq_idx=seq_idx, |
| 688 | out=out, |
| 689 | D=self.D, |
| 690 | z=None, |
| 691 | dt_bias=self.dt_bias, |
| 692 | initial_states=initial_states, |
| 693 | dt_softplus=True, |
| 694 | dt_limit=(0.0, float("inf")), |
| 695 | state_dtype=ssm_state_dtype, |
| 696 | ) |
| 697 | except Exception: |
| 698 | logger.warning( |
| 699 | "Mamba2 SSD kernel warmup failed for layer %s " |
| 700 | "(initial_states=%s). First inference may experience " |
| 701 | "latency spike or OOM due to autotuner.", |
| 702 | self.prefix, |
| 703 | use_initial_states, |
| 704 | exc_info=True, |
| 705 | ) |
| 706 | |
| 707 | logger.debug("Mamba2 SSD kernel warmup completed for layer %s", self.prefix) |
| 708 | torch.accelerator.empty_cache() |
| 709 | |
| 710 | def conv_ssm_forward( |
| 711 | self, |
| 712 | projected_states: torch.Tensor, |
| 713 | output: torch.Tensor, |
| 714 | ): |
| 715 | hidden_states_B_C, dt = torch.split( |
| 716 | projected_states[..., self.tped_intermediate_size :], |
| 717 | [self.tped_conv_size, self.tped_dt_size], |
| 718 | dim=-1, |
| 719 | ) |
| 720 | |
| 721 | forward_context = get_forward_context() |
| 722 | # attn_metadata contains metadata necessary for the mamba2 triton |
| 723 | # kernels to operate in continuous batching and in chunked prefill |
| 724 | # modes; they are computed at top-level model forward since they |
| 725 | # stay the same and reused for all mamba layers in the same iteration |
| 726 | attn_metadata_raw = forward_context.attn_metadata |
| 727 | |
| 728 | assert self.cache_config is not None |
| 729 | use_spec_decode = self.num_spec > 0 |
| 730 | ring_start = prev_num_accepted = prev_query_len = None |
| 731 | |
| 732 | attn_metadata: AttentionMetadata | None = None |
| 733 | if attn_metadata_raw is not None: |
| 734 | assert isinstance(attn_metadata_raw, dict) |
| 735 | attn_metadata = attn_metadata_raw.get(self.prefix) |
| 736 | if attn_metadata is not None: |
| 737 | assert isinstance(attn_metadata, Mamba2AttentionMetadata) |
| 738 | # conv_state must be (..., dim, width-1) for the conv kernels. |
| 739 | # DS layout stores it that way directly; SD layout needs a |
| 740 | # transpose (which keeps dim contiguous via stride tricks). |
| 741 | conv_state = ( |
| 742 | self.kv_cache[0] |
| 743 | if is_conv_state_dim_first() |
| 744 | else self.kv_cache[0].transpose(-1, -2) |
| 745 | ) |
| 746 | ssm_state = self.kv_cache[1] |
| 747 | if self.use_replayssm: |
| 748 | x_cache, dt_cache, B_cache = self.kv_cache[2:5] |
| 749 | if self.mamba_config.backend == MambaBackendEnum.FLASHINFER: |
| 750 | ring_start = self._replayssm_ring_start |
| 751 | prev_num_accepted = self._replayssm_prev_num_accepted |
| 752 | prev_query_len = self._replayssm_prev_query_len |
| 753 | else: |
| 754 | x_cache = dt_cache = B_cache = None |
| 755 | has_initial_states_p = attn_metadata.has_initial_states_p |
| 756 | prep_initial_states = attn_metadata.prep_initial_states |
| 757 | chunk_size = attn_metadata.chunk_size |
| 758 | seq_idx_p = attn_metadata.seq_idx_p |
| 759 | query_start_loc_p = attn_metadata.query_start_loc_p |
| 760 | cu_chunk_seqlen_p = attn_metadata.cu_chunk_seqlen_p |
| 761 | last_chunk_indices_p = attn_metadata.last_chunk_indices_p |
| 762 | state_indices_tensor_p = attn_metadata.state_indices_tensor_p |
| 763 | state_indices_tensor_d = attn_metadata.state_indices_tensor_d |
| 764 | replayssm_state_indices_d = attn_metadata.replayssm_state_indices_d |
| 765 | num_accepted_tokens = attn_metadata.num_accepted_tokens |
| 766 | query_start_loc_d = attn_metadata.query_start_loc_d |
| 767 | num_decodes = attn_metadata.num_decodes |
| 768 | num_decode_tokens = attn_metadata.num_decode_tokens |
| 769 | |
| 770 | if attn_metadata is None: |
| 771 | # V1 profile run -- warm up SSD kernels so that autotuning |
| 772 | # completes before SSM cache allocation. |
| 773 | self._warmup_ssd_kernels(projected_states) |
| 774 | hidden_states_B_C = ( |
| 775 | hidden_states_B_C.transpose(0, 1).clone().transpose(0, 1) |
| 776 | ).contiguous() |
| 777 | hidden_states, _B, _C = self.split_hidden_states_B_C_fn(hidden_states_B_C) |
| 778 | return hidden_states |
| 779 | |
| 780 | num_prefills = attn_metadata.num_prefills |
| 781 | num_prefill_tokens = attn_metadata.num_prefill_tokens |
| 782 | has_prefill = num_prefills > 0 |
| 783 | has_decode = num_decodes > 0 |
| 784 | num_actual_tokens = num_prefill_tokens + num_decode_tokens |
| 785 | |
| 786 | # Split along token dimension |
| 787 | hidden_states_B_C_d, hidden_states_B_C_p = torch.split( |
| 788 | hidden_states_B_C[:num_actual_tokens], |
| 789 | [num_decode_tokens, num_prefill_tokens], |
| 790 | dim=0, |
| 791 | ) |
| 792 | dt_d, dt_p = torch.split( |
| 793 | dt[:num_actual_tokens], |
| 794 | [num_decode_tokens, num_prefill_tokens], |
| 795 | dim=0, |
| 796 | ) |
| 797 | |
| 798 | preallocated_ssm_out_d, preallocated_ssm_out_p = torch.split( |
| 799 | output[:num_actual_tokens], |
| 800 | [num_decode_tokens, num_prefill_tokens], |
| 801 | dim=0, |
| 802 | ) |
| 803 | |
| 804 | # Process prefill requests |
| 805 | if has_prefill: |
| 806 | # 2. Convolution sequence transformation |
| 807 | # - It will read the initial states for every sequence, |
| 808 | # that has "has_initial_states_p" == True, |
| 809 | # from "cache_indices", using "state_indices_tensor_p". |
| 810 | # - It updates the "conv_state" cache in positions pointed |
| 811 | # to by "state_indices_tensor_p". |
| 812 | # In particular, it will always write the state at the |
| 813 | # sequence end. |
| 814 | x = hidden_states_B_C_p.transpose( |
| 815 | 0, 1 |
| 816 | ) # this is the form that causal-conv see |
| 817 | hidden_states_B_C_p = causal_conv1d_fn( |
| 818 | x, |
| 819 | self.conv_weights, |
| 820 | self.conv1d.bias, |
| 821 | activation=self.activation, |
| 822 | conv_states=conv_state, |
| 823 | has_initial_state=has_initial_states_p, |
| 824 | cache_indices=state_indices_tensor_p, |
| 825 | metadata=attn_metadata, |
| 826 | query_start_loc=query_start_loc_p, |
| 827 | ).transpose(0, 1)[:num_prefill_tokens] |
| 828 | |
| 829 | hidden_states_p, B_p, C_p = self.split_hidden_states_B_C_fn( |
| 830 | hidden_states_B_C_p |
| 831 | ) |
| 832 | |
| 833 | # 3. State Space Model sequence transformation |
| 834 | initial_states = None |
| 835 | if has_initial_states_p is not None and prep_initial_states: |
| 836 | assert state_indices_tensor_p is not None |
| 837 | initial_states = torch.where( |
| 838 | has_initial_states_p[:, None, None, None], |
| 839 | ssm_state[state_indices_tensor_p], |
| 840 | 0, |
| 841 | ) |
| 842 | |
| 843 | # NOTE: final output is an in-place update of out tensor |
| 844 | assert preallocated_ssm_out_p is not None |
| 845 | varlen_states = mamba_chunk_scan_combined_varlen( |
| 846 | hidden_states_p.view( |
| 847 | num_prefill_tokens, self.num_heads // self.tp_size, self.head_dim |
| 848 | ), |
| 849 | dt_p, |
| 850 | self.A, |
| 851 | B_p.view(num_prefill_tokens, self.n_groups // self.tp_size, -1), |
| 852 | C_p.view(num_prefill_tokens, self.n_groups // self.tp_size, -1), |
| 853 | chunk_size=chunk_size, |
| 854 | D=self.D, |
| 855 | z=None, |
| 856 | dt_bias=self.dt_bias, |
| 857 | seq_idx=seq_idx_p, |
| 858 | cu_seqlens=query_start_loc_p, |
| 859 | cu_chunk_seqlens=cu_chunk_seqlen_p, |
| 860 | last_chunk_indices=last_chunk_indices_p, |
| 861 | initial_states=initial_states, |
| 862 | dt_softplus=True, |
| 863 | dt_limit=(0.0, float("inf")), |
| 864 | out=preallocated_ssm_out_p.view(num_prefill_tokens, -1, self.head_dim), |
| 865 | state_dtype=ssm_state.dtype, |
| 866 | ) |
| 867 | |
| 868 | # update ssm states |
| 869 | # - varlen state is a (num_prefills, nheads, headdim, dstate) |
| 870 | # tensor |
| 871 | assert state_indices_tensor_p is not None |
| 872 | ssm_state[state_indices_tensor_p] = varlen_states |
| 873 | if ring_start is not None and self._updates_replayssm_trackers: |
| 874 | assert prev_num_accepted is not None |
| 875 | reset_replayssm_ring_trackers( |
| 876 | ring_start, |
| 877 | prev_num_accepted, |
| 878 | prev_query_len, |
| 879 | state_indices_tensor_p, |
| 880 | ) |
| 881 | |
| 882 | # Process decode requests |
| 883 | if has_decode: |
| 884 | assert state_indices_tensor_d is not None |
| 885 | if ( |
| 886 | self.use_replayssm |
| 887 | and self.mamba_config.backend == MambaBackendEnum.FLASHINFER |
| 888 | and use_spec_decode |
| 889 | and self._commits_replayssm_trackers |
| 890 | ): |
| 891 | assert ring_start is not None |
| 892 | assert prev_num_accepted is not None |
| 893 | assert prev_query_len is not None |
| 894 | assert replayssm_state_indices_d is not None |
| 895 | assert num_accepted_tokens is not None |
| 896 | assert query_start_loc_d is not None |
| 897 | assert x_cache is not None |
| 898 | assert self.replayssm_buffer_len is not None |
| 899 | # The previous MTP query's accepted prefix is known only now. |
| 900 | # Commit it once, before this cache group's layers evaluate the |
| 901 | # current speculative query with the shared ring positions. |
| 902 | commit_replayssm_ring_trackers( |
| 903 | ring_start, |
| 904 | prev_num_accepted, |
| 905 | prev_query_len, |
| 906 | replayssm_state_indices_d, |
| 907 | num_accepted_tokens, |
| 908 | query_start_loc_d, |
| 909 | logical_window=self.replayssm_buffer_len, |
| 910 | ring_buffer_len=x_cache.size(2), |
| 911 | ) |
| 912 | |
| 913 | # 2. Convolution sequence transformation |
| 914 | hidden_states_B_C_d = causal_conv1d_update( |
| 915 | hidden_states_B_C_d, |
| 916 | conv_state, |
| 917 | self.conv_weights, |
| 918 | self.conv1d.bias, |
| 919 | self.activation, |
| 920 | conv_state_indices=state_indices_tensor_d, |
| 921 | num_accepted_tokens=num_accepted_tokens, |
| 922 | query_start_loc=query_start_loc_d, |
| 923 | # ReplaySSM keeps one physical state block while a speculative |
| 924 | # decode call still processes the full target + draft window. |
| 925 | max_query_len=( |
| 926 | 1 + self.num_spec |
| 927 | if self.use_replayssm and use_spec_decode |
| 928 | else state_indices_tensor_d.size(-1) |
| 929 | ), |
| 930 | ) |
| 931 | |
| 932 | hidden_states_d, B_d, C_d = self.split_hidden_states_B_C_fn( |
| 933 | hidden_states_B_C_d |
| 934 | ) |
| 935 | |
| 936 | # 3. State Space Model sequence transformation |
| 937 | n_groups = self.n_groups // self.tp_size |
| 938 | A_d = ( |
| 939 | self.A[:, None, ...][:, :, None] |
| 940 | .expand(-1, self.head_dim, self.ssm_state_size) |
| 941 | .to(dtype=torch.float32) |
| 942 | ) |
| 943 | dt_d = dt_d[:, :, None].expand(-1, -1, self.head_dim) |
| 944 | dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim) |
| 945 | D_d = self.D[:, None, ...].expand(-1, self.head_dim) |
| 946 | B_d = B_d.view(-1, n_groups, B_d.shape[1] // n_groups) |
| 947 | C_d = C_d.view(-1, n_groups, C_d.shape[1] // n_groups) |
| 948 | hidden_states_d = hidden_states_d.view( |
| 949 | -1, self.num_heads // self.tp_size, self.head_dim |
| 950 | ) |
| 951 | |
| 952 | assert preallocated_ssm_out_d is not None |
| 953 | # - the hidden is reshaped into (bs, num_heads, head_dim) |
| 954 | # - mamba_cache_params.ssm_state's slots will be selected |
| 955 | # using state_indices_tensor_d |
| 956 | # NOTE: final output is an in-place update of out tensor |
| 957 | preallocated_ssm_out_d = preallocated_ssm_out_d.view( |
| 958 | num_decode_tokens, -1, self.head_dim |
| 959 | ) |
| 960 | if self.use_replayssm: |
| 961 | assert self.replayssm_buffer_len is not None |
| 962 | if self.mamba_config.backend == MambaBackendEnum.FLASHINFER: |
| 963 | assert ring_start is not None |
| 964 | assert prev_num_accepted is not None |
| 965 | assert prev_query_len is not None |
| 966 | assert attn_metadata.replayssm_scratch is not None |
| 967 | fi_cu_seqlens = query_start_loc_d |
| 968 | fi_max_seqlen = None |
| 969 | if use_spec_decode: |
| 970 | spec_query_len = 1 + self.num_spec |
| 971 | fi_max_seqlen = spec_query_len |
| 972 | assert replayssm_state_indices_d is not None |
| 973 | decode_batch = replayssm_state_indices_d.size(0) |
| 974 | if num_decode_tokens == decode_batch * spec_query_len: |
| 975 | hidden_states_d = _view_mtp_decode_tensor( |
| 976 | hidden_states_d, decode_batch, spec_query_len |
| 977 | ) |
| 978 | dt_d = _view_mtp_decode_tensor( |
| 979 | dt_d, decode_batch, spec_query_len |
| 980 | ) |
| 981 | B_d = _view_mtp_decode_tensor( |
| 982 | B_d, decode_batch, spec_query_len |
| 983 | ) |
| 984 | C_d = _view_mtp_decode_tensor( |
| 985 | C_d, decode_batch, spec_query_len |
| 986 | ) |
| 987 | preallocated_ssm_out_d = _view_mtp_decode_tensor( |
| 988 | preallocated_ssm_out_d, |
| 989 | decode_batch, |
| 990 | spec_query_len, |
| 991 | ) |
| 992 | fi_cu_seqlens = None |
| 993 | fi_max_seqlen = None |
| 994 | selective_state_update_replayssm_flashinfer( |
| 995 | ssm_state, |
| 996 | hidden_states_d, |
| 997 | dt_d, |
| 998 | A_d, |
| 999 | B_d, |
| 1000 | C_d, |
| 1001 | preallocated_ssm_out_d, |
| 1002 | x_cache, |
| 1003 | B_cache, |
| 1004 | dt_cache, |
| 1005 | ring_start, |
| 1006 | prev_num_accepted, |
| 1007 | prev_query_len, |
| 1008 | logical_window=self.replayssm_buffer_len, |
| 1009 | D=D_d, |
| 1010 | dt_bias=dt_bias, |
| 1011 | dt_softplus=True, |
| 1012 | state_batch_indices=replayssm_state_indices_d, |
| 1013 | scratch=attn_metadata.replayssm_scratch, |
| 1014 | update_trackers=( |
| 1015 | self._updates_replayssm_trackers and not use_spec_decode |
| 1016 | ), |
| 1017 | enable_stochastic_rounding=( |
| 1018 | self.mamba_config.enable_stochastic_rounding |
| 1019 | ), |
| 1020 | stochastic_rounding_philox_rounds=( |
| 1021 | self.mamba_config.stochastic_rounding_philox_rounds |
| 1022 | ), |
| 1023 | cu_seqlens=fi_cu_seqlens, |
| 1024 | max_seqlen=fi_max_seqlen, |
| 1025 | enable_pdl=False, |
| 1026 | ) |
| 1027 | else: |
| 1028 | selective_state_update_replayssm_output_only( |
| 1029 | ssm_state, |
| 1030 | hidden_states_d, |
| 1031 | dt_d, |
| 1032 | A_d, |
| 1033 | B_d, |
| 1034 | C_d, |
| 1035 | D_d, |
| 1036 | dt_bias, |
| 1037 | dt_softplus=True, |
| 1038 | x_cache=x_cache, |
| 1039 | dt_cache=dt_cache, |
| 1040 | B_cache=B_cache, |
| 1041 | bc_pre=attn_metadata.bc_pre_scratch, |
| 1042 | write_pos=attn_metadata.write_pos_d, |
| 1043 | is_flush=attn_metadata.is_flush_d, |
| 1044 | max_cache_len=self.replayssm_buffer_len, |
| 1045 | state_batch_indices=state_indices_tensor_d, |
| 1046 | out=preallocated_ssm_out_d, |
| 1047 | enable_stochastic_rounding=( |
| 1048 | self.mamba_config.enable_stochastic_rounding |
| 1049 | ), |
| 1050 | cache_philox_rounds=( |
| 1051 | self.mamba_config.stochastic_rounding_philox_rounds |
| 1052 | ), |
| 1053 | ) |
| 1054 | else: |
| 1055 | selective_state_update( |
| 1056 | ssm_state, |
| 1057 | hidden_states_d, |
| 1058 | dt_d, |
| 1059 | A_d, |
| 1060 | B_d, |
| 1061 | C_d, |
| 1062 | D_d, |
| 1063 | dt_bias, |
| 1064 | dt_softplus=True, |
| 1065 | state_batch_indices=state_indices_tensor_d, |
| 1066 | dst_state_batch_indices=state_indices_tensor_d, |
| 1067 | out=preallocated_ssm_out_d, |
| 1068 | num_accepted_tokens=num_accepted_tokens, |
| 1069 | cu_seqlens=query_start_loc_d, |
| 1070 | is_blackwell=self.is_blackwell, |
| 1071 | ) |
| 1072 | |
| 1073 | def get_state_dtype(self) -> tuple[torch.dtype, ...]: |
| 1074 | assert self.model_config is not None |
| 1075 | assert self.cache_config is not None |
| 1076 | base_dtype = MambaStateDtypeCalculator.mamba2_state_dtype( |
| 1077 | self.model_config.dtype, |
| 1078 | self.cache_config.mamba_cache_dtype, |
| 1079 | self.cache_config.mamba_ssm_cache_dtype, |
| 1080 | ) |
| 1081 | if self.use_replayssm: |
| 1082 | return MambaStateDtypeCalculator.append_replayssm_ring( |
| 1083 | base_dtype, self.model_config.dtype |
| 1084 | ) |
| 1085 | return base_dtype |
| 1086 | |
| 1087 | def get_state_shape(self) -> tuple[tuple[int, ...], ...]: |
| 1088 | tp_world_size = get_tensor_model_parallel_world_size() |
| 1089 | base_shape = MambaStateShapeCalculator.mamba2_state_shape( |
| 1090 | intermediate_size=self.intermediate_size, |
| 1091 | tp_world_size=tp_world_size, |
| 1092 | n_groups=self.n_groups, |
| 1093 | num_heads=self.num_heads, |
| 1094 | head_dim=self.head_dim, |
| 1095 | state_size=self.ssm_state_size, |
| 1096 | conv_kernel=self.conv_kernel_size, |
| 1097 | num_spec=self.num_spec, |
| 1098 | ) |
| 1099 | if self.use_replayssm: |
| 1100 | assert self.replayssm_buffer_len is not None |
| 1101 | return MambaStateShapeCalculator.append_replayssm_ring( |
| 1102 | base_shapes=base_shape, |
| 1103 | n_groups=self.n_groups, |
| 1104 | tp_world_size=tp_world_size, |
| 1105 | logical_window=self.replayssm_buffer_len, |
| 1106 | backend=self.mamba_config.backend, |
| 1107 | num_speculative_tokens=self.num_spec, |
| 1108 | ) |
| 1109 | return base_shape |
| 1110 | |
| 1111 | @property |
| 1112 | def mamba_type(self) -> MambaAttentionBackendEnum: |
| 1113 | return MambaAttentionBackendEnum.MAMBA2 |
| 1114 | |
| 1115 | |
| 1116 | def share_replayssm_ring_trackers( |
| 1117 | ordered_layer_names: list[str], |
| 1118 | forward_context: dict[str, Attention], |
| 1119 | kv_cache_groups: Sequence[KVCacheGroupSpec] | None = None, |
| 1120 | ) -> None: |
| 1121 | """Share ring cursors within each cache-slot index namespace. |
| 1122 | |
| 1123 | Layers backed by one KV-cache group use the same physical block indices and |
| 1124 | can therefore share cursors. Different KV-cache groups may assign different |
| 1125 | block indices to the same request and must keep separate cursor tensors. For |
| 1126 | speculative decode, the first local layer commits the preceding acceptance; |
| 1127 | for standard decode, the final local layer advances after all layers consume |
| 1128 | the previous values. |
| 1129 | """ |
| 1130 | replayssm_mixers: dict[str, MambaMixer2] = {} |
| 1131 | for layer_name in ordered_layer_names: |
| 1132 | layer = forward_context[layer_name] |
| 1133 | if ( |
| 1134 | isinstance(layer, MambaMixer2) |
| 1135 | and layer.use_replayssm |
| 1136 | and layer.mamba_config.backend == MambaBackendEnum.FLASHINFER |
| 1137 | ): |
| 1138 | replayssm_mixers[layer_name] = layer |
| 1139 | |
| 1140 | layer_to_group: dict[str, int] = {} |
| 1141 | if kv_cache_groups: |
| 1142 | for group_idx, group in enumerate(kv_cache_groups): |
| 1143 | for layer_name in group.layer_names: |
| 1144 | layer_to_group[layer_name] = group_idx |
| 1145 | |
| 1146 | groups_by_namespace: dict[int | str, list[str]] = {} |
| 1147 | for layer_name in ordered_layer_names: |
| 1148 | if layer_name not in replayssm_mixers: |
| 1149 | continue |
| 1150 | namespace = layer_to_group.get(layer_name, layer_name) |
| 1151 | groups_by_namespace.setdefault(namespace, []).append(layer_name) |
| 1152 | |
| 1153 | for group_layer_names in groups_by_namespace.values(): |
| 1154 | first_layer_name = group_layer_names[0] |
| 1155 | last_layer_name = group_layer_names[-1] |
| 1156 | |
| 1157 | first_mixer = replayssm_mixers[group_layer_names[0]] |
| 1158 | first_state = first_mixer.kv_cache[1] |
| 1159 | num_blocks, device = first_state.shape[0], first_state.device |
| 1160 | for layer_name in group_layer_names: |
| 1161 | state = replayssm_mixers[layer_name].kv_cache[1] |
| 1162 | if (state.shape[0], state.device) != (num_blocks, device): |
| 1163 | raise ValueError( |
| 1164 | "ReplaySSM layers in one cache group must share cache capacity" |
| 1165 | ) |
| 1166 | |
| 1167 | ring_start = torch.zeros(num_blocks, dtype=torch.int32, device=device) |
| 1168 | prev_num_accepted = torch.zeros_like(ring_start) |
| 1169 | prev_query_len = torch.zeros_like(ring_start) |
| 1170 | for layer_name in group_layer_names: |
| 1171 | mixer = replayssm_mixers[layer_name] |
| 1172 | mixer._replayssm_ring_start = ring_start |
| 1173 | mixer._replayssm_prev_num_accepted = prev_num_accepted |
| 1174 | mixer._replayssm_prev_query_len = prev_query_len |
| 1175 | mixer._commits_replayssm_trackers = False |
| 1176 | mixer._updates_replayssm_trackers = False |
| 1177 | |
| 1178 | replayssm_mixers[first_layer_name]._commits_replayssm_trackers = True |
| 1179 | replayssm_mixers[last_layer_name]._updates_replayssm_trackers = True |
| 1180 | |
| 1181 | |
| 1182 | def mamba_mixer2( |
| 1183 | projected_states: torch.Tensor, |
| 1184 | output: torch.Tensor, |
| 1185 | layer_name: LayerNameType, |
| 1186 | ) -> None: |
| 1187 | layer_name = _resolve_layer_name(layer_name) |
| 1188 | forward_context: ForwardContext = get_forward_context() |
| 1189 | self = forward_context.no_compile_layers[layer_name] |
| 1190 | self.conv_ssm_forward(projected_states=projected_states, output=output) |
| 1191 | |
| 1192 | |
| 1193 | direct_register_custom_op( |
| 1194 | op_name="mamba_mixer2", |
| 1195 | op_func=mamba_mixer2, |
| 1196 | mutates_args=["output"], |
| 1197 | ) |