@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a89209b78d

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 04:45 UTC

[v1] Support mamba2 (#19327)

Signed-off-by: Chen Zhang <zhangch99@outlook.com>

Chen Zhang committed Jun 19, 2025 at 04:34 UTC a89209b78de0aa3d47a684a34861336d3ba7c0fe

9 files changed +583 -121

tests/models/language/generation/test_hybrid.py

+42 -11

@@ -17,9 +17,10 @@ SSM_MODELS = [

17 "state-spaces/mamba-130m-hf",

18 "tiiuae/falcon-mamba-tiny-dev",

19 # TODO: Compare to a Mamba2 model. The HF transformers implementation of

20 - # Mamba2 is buggy for Codestral as it doesn't handle n_groups.

20 + # Mamba2 is buggy for Codestral as it doesn't handle n_groups, so the test

21 + # doesn't compare vLLM output with HF output.

22 # See https://github.com/huggingface/transformers/pull/35943

22 - # "mistralai/Mamba-Codestral-7B-v0.1",

23 + "mistralai/Mamba-Codestral-7B-v0.1",

24 ]

26 HYBRID_MODELS = [

@@ -35,6 +36,10 @@ HYBRID_MODELS = [

36 "hmellor/tiny-random-BambaForCausalLM",

37 ]

39 +V1_SUPPORTED_MODELS = [

40 + "mistralai/Mamba-Codestral-7B-v0.1",

41 +]

42 +

43 # Avoid OOM

44 MAX_NUM_SEQS = 4

@@ -46,24 +51,50 @@ def test_models(

51 hf_runner,

52 vllm_runner,

53 example_prompts,

54 + monkeypatch,

55 model: str,

56 max_tokens: int,

57 num_logprobs: int,

58 ) -> None:

59 with hf_runner(model) as hf_model:

54 - hf_outputs = hf_model.generate_greedy_logprobs_limit(

55 - example_prompts, max_tokens, num_logprobs)

60 + if model != "mistralai/Mamba-Codestral-7B-v0.1":

61 + hf_outputs = hf_model.generate_greedy_logprobs_limit(

62 + example_prompts, max_tokens, num_logprobs)

63 + else:

64 + hf_outputs = None

66 with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:

58 - vllm_outputs = vllm_model.generate_greedy_logprobs(

67 + vllm_v0_outputs = vllm_model.generate_greedy_logprobs(

68 example_prompts, max_tokens, num_logprobs)

61 - check_logprobs_close(

62 - outputs_0_lst=hf_outputs,

63 - outputs_1_lst=vllm_outputs,

64 - name_0="hf",

65 - name_1="vllm",

66 - )

70 + if model in V1_SUPPORTED_MODELS:

71 + with monkeypatch.context() as m:

72 + m.setenv("VLLM_USE_V1", "1")

73 + with vllm_runner(model,

74 + max_num_seqs=MAX_NUM_SEQS,

75 + enforce_eager=True,

76 + enable_prefix_caching=False) as vllm_model:

77 + vllm_v1_outputs = vllm_model.generate_greedy_logprobs(

78 + example_prompts, max_tokens, num_logprobs)

79 + else:

80 + vllm_v1_outputs = None

81 +

82 + if hf_outputs is not None:

83 + check_logprobs_close(

84 + outputs_0_lst=hf_outputs,

85 + outputs_1_lst=vllm_v0_outputs,

86 + name_0="hf",

87 + name_1="vllm-v0",

88 + )

89 +

90 + if model in V1_SUPPORTED_MODELS:

91 + ref_outputs = hf_outputs if hf_outputs is not None else vllm_v0_outputs

92 + check_logprobs_close(

93 + outputs_0_lst=ref_outputs,

94 + outputs_1_lst=vllm_v1_outputs,

95 + name_0="hf" if hf_outputs is not None else "vllm-v0",

96 + name_1="vllm-v1",

97 + )

100 @pytest.mark.parametrize("model", SSM_MODELS + HYBRID_MODELS)

tests/v1/test_oracle.py

+1 -1

@@ -12,7 +12,7 @@ from vllm.engine.async_llm_engine import AsyncLLMEngine

12 UNSUPPORTED_MODELS_V1 = [

13 "openai/whisper-large-v3", # transcription

14 "facebook/bart-large-cnn", # encoder decoder

15 - "mistralai/Mamba-Codestral-7B-v0.1", # mamba

15 + "state-spaces/mamba-130m-hf", # mamba1

16 "hmellor/tiny-random-BambaForCausalLM", # hybrid

17 "BAAI/bge-m3", # embedding

18 ]

vllm/engine/arg_utils.py

+6 -1

@@ -1355,12 +1355,17 @@ class EngineArgs:

1355 recommend_to_remove=False)

1356 return False

1357

1358 - # No Mamba or Encoder-Decoder so far.

1358 + # No Encoder-Decoder, not all Mamba so far.

1359 if not model_config.is_v1_compatible:

1360 _raise_or_fallback(feature_name=model_config.architectures,

1361 recommend_to_remove=False)

1362 return False

1363

1364 + # V1 mamba models are unoptimized.

1365 + if model_config.has_inner_state and _warn_or_fallback(

1366 + feature_name="Mamba"):

1367 + return False

1368 +

1369 # No Concurrent Partial Prefills so far.

1370 if (self.max_num_partial_prefills

1371 != SchedulerConfig.max_num_partial_prefills

vllm/model_executor/layers/mamba/mamba_mixer2.py

+175 -60

@@ -6,7 +6,9 @@ from typing import Optional, Union

6 import torch

7 from torch import nn

9 +from vllm import envs

10 from vllm.attention.backends.abstract import AttentionMetadata

11 +from vllm.config import get_current_vllm_config

12 from vllm.distributed import (divide, get_tensor_model_parallel_rank,

13 get_tensor_model_parallel_world_size,

14 tensor_model_parallel_all_gather,

@@ -27,6 +29,7 @@ from vllm.model_executor.model_loader.weight_utils import (

29 LoaderFunction, composed_weight_loader, sharded_weight_loader)

30 from vllm.model_executor.models.mamba_cache import MambaCacheParams

31 from vllm.model_executor.utils import set_weight_attrs

32 +from vllm.v1.attention.backends.mamba_attn import Mamba2AttentionMetadata

34 # Added by the IBM Team, 2024

@@ -227,20 +230,22 @@ class MambaMixer2(CustomOp):

230 """

231

232 def __init__(

230 - self,

231 - hidden_size: int,

232 - ssm_state_size: int,

233 - conv_kernel_size: int,

234 - intermediate_size: int,

235 - use_conv_bias: bool,

236 - use_bias: bool,

237 - n_groups: int = 1,

238 - num_heads: int = 128,

239 - head_dim: int = 64,

240 - rms_norm_eps: float = 1e-5,

241 - activation: str = "silu",

242 - use_rms_norm: bool = True,

243 - quant_config: Optional[QuantizationConfig] = None,

233 + self,

234 + hidden_size: int,

235 + ssm_state_size: int,

236 + conv_kernel_size: int,

237 + intermediate_size: int,

238 + use_conv_bias: bool,

239 + use_bias: bool,

240 + n_groups: int = 1,

241 + num_heads: int = 128,

242 + head_dim: int = 64,

243 + rms_norm_eps: float = 1e-5,

244 + activation: str = "silu",

245 + use_rms_norm: bool = True,

246 + quant_config: Optional[QuantizationConfig] = None,

247 + prefix: str = "",

248 + chunk_size: int = -1, # the chunk size used by v1

249 ):

250 super().__init__()

251

@@ -273,6 +278,7 @@ class MambaMixer2(CustomOp):

278 ), "Tensor parallel currently not supported for quantized models."

279

280 self.ssm_state_size = ssm_state_size

281 + self.conv_kernel_size = conv_kernel_size

282 self.activation = activation

283

284 self.intermediate_size = intermediate_size

@@ -411,6 +417,22 @@ class MambaMixer2(CustomOp):

417 self.use_rms_norm,

418 eps=rms_norm_eps)

419

420 + if envs.VLLM_USE_V1:

421 + compilation_config = get_current_vllm_config().compilation_config

422 + if prefix in compilation_config.static_forward_context:

423 + raise ValueError(f"Duplicate layer name: {prefix}")

424 + compilation_config.static_forward_context[prefix] = self

425 + # The outer list is for v0 PP virtual engine. Though this code path

426 + # only runs for v1, we have to do this to unify with the interface

427 + # of Attention + v0 PP.

428 + # The inner tuple is (conv_state, ssm_state)

429 + self.kv_cache = [(torch.tensor([]), torch.tensor([]))]

430 + assert chunk_size != -1, "chunk_size must be set for v1"

431 +

432 + # NOTE: chunk_size may be -1 for models without v1 support

433 + self.chunk_size = chunk_size

434 + self.prefix = prefix

435 +

436 def forward_native(

437 self,

438 hidden_states: torch.Tensor,

@@ -426,17 +448,37 @@ class MambaMixer2(CustomOp):

448 mamba2_metadata: Mamba2Metadata,

449 mup_vector: Optional[torch.Tensor] = None,

450 ):

451 + forward_context = get_forward_context()

452 # mamba2_metadata contains metadata necessary for the mamba2 triton

453 # kernels to operate in continuous batching and in chunked prefill

454 # modes; they are computed at top-level model forward since they

455 # stay the same and reused for all mamba layers in the same iteration

433 - attn_metadata: AttentionMetadata = get_forward_context().attn_metadata

434 -

435 - num_prefills = attn_metadata.num_prefills # request count

436 - num_decodes = attn_metadata.num_decode_tokens # token count (=request)

437 - num_prefill_tokens = attn_metadata.num_prefill_tokens # token count

438 - has_prefill = num_prefills > 0

439 - has_decode = num_decodes > 0

456 + attn_metadata: AttentionMetadata = forward_context.attn_metadata

457 + if envs.VLLM_USE_V1:

458 + if attn_metadata is not None:

459 + assert isinstance(attn_metadata, dict)

460 + attn_metadata = attn_metadata[self.prefix]

461 + assert isinstance(attn_metadata, Mamba2AttentionMetadata)

462 + self_kv_cache = self.kv_cache[forward_context.virtual_engine]

463 + conv_state = self_kv_cache[0]

464 + ssm_state = self_kv_cache[1]

465 + state_indices_tensor = attn_metadata.state_indices_tensor

466 + has_initial_states_p = attn_metadata.has_initial_states

467 + prep_initial_states = attn_metadata.prep_initial_states

468 + chunk_size = attn_metadata.chunk_size

469 + seq_idx_p = attn_metadata.seq_idx

470 + chunk_indices_p = attn_metadata.chunk_indices

471 + chunk_offsets_p = attn_metadata.chunk_offsets

472 + else:

473 + conv_state = mamba_cache_params.conv_state

474 + ssm_state = mamba_cache_params.ssm_state

475 + state_indices_tensor = mamba_cache_params.state_indices_tensor

476 + has_initial_states_p = mamba2_metadata.has_initial_states

477 + prep_initial_states = mamba2_metadata.prep_initial_states

478 + chunk_size = mamba2_metadata.chunk_size

479 + seq_idx_p = mamba2_metadata.seq_idx

480 + chunk_indices_p = mamba2_metadata.chunk_indices

481 + chunk_offsets_p = mamba2_metadata.chunk_offsets

482

483 groups_time_state_size = self.n_groups * self.ssm_state_size

484

@@ -459,27 +501,6 @@ class MambaMixer2(CustomOp):

501 conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0),

502 self.conv1d.weight.size(2))

503

462 - # Separate prefill and decode by splitting varlen input

463 - # Split along token dimension

464 - hidden_states_B_C_p, hidden_states_B_C_d = torch.split(

465 - hidden_states_B_C,

466 - [num_prefill_tokens, num_decodes],

467 - dim=0,

468 - )

469 - dt_p, dt_d = torch.split(

470 - dt,

471 - [num_prefill_tokens, num_decodes],

472 - dim=0,

473 - )

474 - # Split along batch dimension

475 - state_indices_tensor_p, state_indices_tensor_d = torch.split(

476 - mamba_cache_params.state_indices_tensor,

477 - [num_prefills, num_decodes],

478 - dim=0,

479 - )

480 - query_start_loc_p = (attn_metadata.query_start_loc[:num_prefills + 1]

481 - if has_prefill else None)

482 -

504 # - get hidden_states, B and C after depthwise convolution.

505 split_hidden_states_B_C_fn = lambda hidden_states_B_C: torch.split(

506 hidden_states_B_C,

@@ -491,20 +512,80 @@ class MambaMixer2(CustomOp):

512 dim=-1,

513 )

514

515 + if envs.VLLM_USE_V1 and attn_metadata is None:

516 + # V1 profile run

517 + hidden_states_B_C = (hidden_states_B_C.transpose(

518 + 0, 1).clone().transpose(0, 1)).contiguous()

519 + hidden_states, _B, _C = split_hidden_states_B_C_fn(

520 + hidden_states_B_C)

521 + hidden_states = self.norm(hidden_states, gate)

522 + out, _ = self.out_proj(hidden_states)

523 + return out

524 +

525 + num_prefills = attn_metadata.num_prefills # request count

526 + num_decodes = attn_metadata.num_decode_tokens # token count (=request)

527 + num_prefill_tokens = attn_metadata.num_prefill_tokens # token count

528 + has_prefill = num_prefills > 0

529 + has_decode = num_decodes > 0

530 +

531 + # NOTE: V0 put prefill before decode, v1 puts decode before prefill

532 + # Separate prefill and decode by splitting varlen input

533 + # Split along token dimension

534 + if envs.VLLM_USE_V1:

535 + hidden_states_B_C_d, hidden_states_B_C_p = torch.split(

536 + hidden_states_B_C,

537 + [num_decodes, num_prefill_tokens],

538 + dim=0,

539 + )

540 + dt_d, dt_p = torch.split(

541 + dt,

542 + [num_decodes, num_prefill_tokens],

543 + dim=0,

544 + )

545 + # Split along batch dimension

546 + state_indices_tensor_d, state_indices_tensor_p = torch.split(

547 + state_indices_tensor,

548 + [num_decodes, num_prefills],

549 + dim=0,

550 + )

551 + query_start_loc_p = (

552 + attn_metadata.query_start_loc[-num_prefills - 1:] -

553 + num_decodes if has_prefill else None)

554 + else:

555 + hidden_states_B_C_p, hidden_states_B_C_d = torch.split(

556 + hidden_states_B_C,

557 + [num_prefill_tokens, num_decodes],

558 + dim=0,

559 + )

560 + dt_p, dt_d = torch.split(

561 + dt,

562 + [num_prefill_tokens, num_decodes],

563 + dim=0,

564 + )

565 + # Split along batch dimension

566 + state_indices_tensor_p, state_indices_tensor_d = torch.split(

567 + state_indices_tensor,

568 + [num_prefills, num_decodes],

569 + dim=0,

570 + )

571 + query_start_loc_p = (attn_metadata.query_start_loc[:num_prefills +

572 + 1]

573 + if has_prefill else None)

574 +

575 ssd_output_list = []

576

577 # Process prefill requests

578 if has_prefill:

579 # 2. Convolution sequence transformation

580 # - "cache_indices" updates the conv_state cache in positions

500 - # pointed to by "mamba_cache_params.state_indices_tensor"

581 + # pointed to by "state_indices_tensor"

582 hidden_states_B_C_p = causal_conv1d_fn(

583 hidden_states_B_C_p.transpose(0, 1),

584 conv_weights,

585 self.conv1d.bias,

586 activation=self.activation,

506 - conv_states=mamba_cache_params.conv_state,

507 - has_initial_state=mamba2_metadata.has_initial_states,

587 + conv_states=conv_state,

588 + has_initial_state=has_initial_states_p,

589 cache_indices=state_indices_tensor_p,

590 query_start_loc=query_start_loc_p).transpose(

591 0, 1)[:num_prefill_tokens]

@@ -516,12 +597,11 @@ class MambaMixer2(CustomOp):

597

598 # 3. State Space Model sequence transformation

599 initial_states = None

519 - if (mamba2_metadata.has_initial_states is not None

520 - and mamba2_metadata.prep_initial_states):

600 + if (has_initial_states_p is not None and prep_initial_states):

601 # making a copy of the states

602 initial_states = torch.where(

523 - mamba2_metadata.has_initial_states[:, None, None, None],

524 - mamba_cache_params.ssm_state[state_indices_tensor_p], 0)

603 + has_initial_states_p[:, None, None, None],

604 + ssm_state[state_indices_tensor_p], 0)

605

606 scan_output, varlen_state = mamba_chunk_scan_combined(

607 hidden_states_p.view(1, num_prefill_tokens,

@@ -533,14 +613,14 @@ class MambaMixer2(CustomOp):

613 -1),

614 C_p.view(1, num_prefill_tokens, self.n_groups // self.tp_size,

615 -1),

536 - chunk_size=mamba2_metadata.chunk_size,

616 + chunk_size=chunk_size,

617 D=self.D,

618 z=None,

619 dt_bias=self.dt_bias,

540 - seq_idx=mamba2_metadata.seq_idx,

541 - chunk_indices=mamba2_metadata.chunk_indices,

542 - chunk_offsets=mamba2_metadata.chunk_offsets,

543 - cu_seqlens=attn_metadata.query_start_loc[:num_prefills + 1],

620 + seq_idx=seq_idx_p,

621 + chunk_indices=chunk_indices_p,

622 + chunk_offsets=chunk_offsets_p,

623 + cu_seqlens=query_start_loc_p,

624 initial_states=initial_states,

625 return_varlen_states=True,

626 return_final_states=False,

@@ -550,7 +630,7 @@ class MambaMixer2(CustomOp):

630

631 # update ssm states

632 # - varlen state is a (num_prefills, nheads, headdim, dstate) tensor

553 - mamba_cache_params.ssm_state[state_indices_tensor_p] = varlen_state

633 + ssm_state[state_indices_tensor_p] = varlen_state

634

635 # - reshape

636 ssd_output_list.append(scan_output.view(num_prefill_tokens, -1))

@@ -560,7 +640,7 @@ class MambaMixer2(CustomOp):

640 # 2. Convolution sequence transformation

641 hidden_states_B_C_d = causal_conv1d_update(

642 hidden_states_B_C_d,

563 - mamba_cache_params.conv_state,

643 + conv_state,

644 conv_weights,

645 self.conv1d.bias,

646 self.activation,

@@ -586,7 +666,7 @@ class MambaMixer2(CustomOp):

666 # using state_indices_tensor_d

667

668 hidden_states_d = selective_state_update(

589 - mamba_cache_params.ssm_state,

669 + ssm_state,

670 hidden_states_d,

671 dt_d,

672 A_d,

@@ -598,9 +678,16 @@ class MambaMixer2(CustomOp):

678 dt_softplus=True,

679 state_batch_indices=state_indices_tensor_d,

680 )

601 - ssd_output_list.append(

602 - hidden_states_d.view(-1, (self.num_heads // self.tp_size) *

603 - self.head_dim))

681 +

682 + if envs.VLLM_USE_V1:

683 + ssd_output_list.insert(

684 + 0,

685 + hidden_states_d.view(-1, (self.num_heads // self.tp_size) *

686 + self.head_dim))

687 + else:

688 + ssd_output_list.append(

689 + hidden_states_d.view(-1, (self.num_heads // self.tp_size) *

690 + self.head_dim))

691

692 # Merge prefill and decode outputs before passing to gated MLP

693 hidden_states = torch.vstack(ssd_output_list)

@@ -614,3 +701,31 @@ class MambaMixer2(CustomOp):

701 # 5. Final linear projection

702 out, _ = self.out_proj(hidden_states)

703 return out

704 +

705 + def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:

706 + world_size = get_tensor_model_parallel_world_size()

707 +

708 + conv_state_shape, temporal_state_shape = None, None

709 +

710 + # if n_groups is not divisible by world_size, need to extend the shards

711 + # to ensure all groups needed by a head is sharded along with it

712 + n_groups = (self.n_groups +

713 + extra_groups_for_head_shards(self.n_groups, world_size))

714 +

715 + # - heads and n_groups are TP-ed

716 + conv_dim = (self.intermediate_size +

717 + 2 * n_groups * self.ssm_state_size)

718 + conv_state_shape = (

719 + divide(conv_dim, world_size),

720 + self.conv_kernel_size - 1,

721 + )

722 +

723 + # These are not TP-ed as they depend on A, dt_bias, D

724 + # - they are typically small

725 + # e.g., (h_heads, d_head, d_state) = (128, 64, 128)

726 + temporal_state_shape = (

727 + divide(self.num_heads, world_size),

728 + self.head_dim,

729 + self.ssm_state_size,

730 + )

731 + return conv_state_shape, temporal_state_shape

vllm/model_executor/models/mamba2.py

+33 -21

@@ -8,6 +8,7 @@ import torch

8 from torch import nn

9 from transformers import MambaConfig

11 +from vllm import envs

12 from vllm.attention.backends.abstract import AttentionMetadata

13 from vllm.config import VllmConfig

14 from vllm.distributed import divide, get_tensor_model_parallel_world_size

@@ -25,8 +26,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

26 DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)

27 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

28 from vllm.model_executor.models.interfaces import (HasInnerState,

28 - IsAttentionFree,

29 - SupportsV0Only)

29 + IsAttentionFree)

30 from vllm.model_executor.models.mamba_cache import (MambaCacheManager,

31 MambaCacheParams)

32 from vllm.model_executor.sampling_metadata import SamplingMetadata

@@ -44,7 +44,8 @@ class Mamba2DecoderLayer(nn.Module):

45 def __init__(self,

46 config: MambaConfig,

47 - quant_config: Optional[QuantizationConfig] = None) -> None:

47 + quant_config: Optional[QuantizationConfig] = None,

48 + prefix: str = "") -> None:

49 super().__init__()

50 self.config = config

51 self.mixer = MambaMixer2(hidden_size=config.hidden_size,

@@ -60,7 +61,9 @@ class Mamba2DecoderLayer(nn.Module):

61 head_dim=config.head_dim,

62 rms_norm_eps=config.layer_norm_epsilon,

63 activation=config.hidden_act,

63 - quant_config=quant_config)

64 + quant_config=quant_config,

65 + prefix=f"{prefix}.mixer",

66 + chunk_size=config.chunk_size)

68 self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)

@@ -108,8 +111,8 @@ class Mamba2Model(nn.Module):

111

112 self.start_layer, self.end_layer, self.layers = make_layers(

113 config.num_hidden_layers,

111 - lambda prefix: Mamba2DecoderLayer(config,

112 - quant_config=quant_config),

114 + lambda prefix: Mamba2DecoderLayer(

115 + config, quant_config=quant_config, prefix=prefix),

116 prefix=f"{prefix}.layers")

117

118 self.norm_f = RMSNorm(config.hidden_size,

@@ -142,10 +145,14 @@ class Mamba2Model(nn.Module):

145

146 attn_metadata: AttentionMetadata = get_forward_context().attn_metadata

147

145 - mamba2_metadata = prepare_mamba2_metadata(

146 - chunk_size=self.config.chunk_size,

147 - attn_metadata=attn_metadata,

148 - )

148 + if not envs.VLLM_USE_V1:

149 + mamba2_metadata = prepare_mamba2_metadata(

150 + chunk_size=self.config.chunk_size,

151 + attn_metadata=attn_metadata,

152 + )

153 + else:

154 + # v1 get mamba2_metadata from forward_context

155 + mamba2_metadata = None

156

157 for i in range(len(self.layers)):

158 layer = self.layers[i]

@@ -155,7 +162,7 @@ class Mamba2Model(nn.Module):

162 hidden_states=hidden_states,

163 residual=residual,

164 mamba_cache_params=mamba_cache_params.at_layer_idx(

158 - i - self.start_layer),

165 + i - self.start_layer) if mamba_cache_params else None,

166 mamba2_metadata=mamba2_metadata)

167

168 if not get_pp_group().is_last_rank:

@@ -190,8 +197,7 @@ class Mamba2Model(nn.Module):

197 return loaded_params

198

199

193 -class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree,

194 - SupportsV0Only):

200 +class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree):

201

202 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

203 config = vllm_config.model_config.hf_config

@@ -242,14 +248,20 @@ class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree,

248 intermediate_tensors: Optional[IntermediateTensors] = None,

249 inputs_embeds: Optional[torch.Tensor] = None,

250 **kwargs):

245 - if self.mamba_cache is None:

246 - num_mamba_layers = self.model_config.get_num_layers_by_block_type(

247 - self.vllm_config.parallel_config, LayerBlockType.mamba)

248 - self.mamba_cache = MambaCacheManager(

249 - self.vllm_config, self.lm_head.weight.dtype, num_mamba_layers,

250 - *self._get_mamba_cache_shape())

251 -

252 - mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)

251 + if not envs.VLLM_USE_V1:

252 + if self.mamba_cache is None:

253 + num_mamba_layers = (

254 + self.model_config.get_num_layers_by_block_type(

255 + self.vllm_config.parallel_config,

256 + LayerBlockType.mamba))

257 + self.mamba_cache = MambaCacheManager(

258 + self.vllm_config, self.lm_head.weight.dtype,

259 + num_mamba_layers, *self._get_mamba_cache_shape())

260 +

261 + mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)

262 + else:

263 + # NOTE: mamba_cache_params is not needed for v1

264 + mamba_cache_params = None

265

266 hidden_states = self.backbone(input_ids, positions, mamba_cache_params,

267 intermediate_tensors, inputs_embeds)

vllm/v1/attention/backends/mamba_attn.py new

+192

@@ -0,0 +1,192 @@

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 TYPE_CHECKING

5 +

6 +import torch

7 +

8 +from vllm.attention.backends.abstract import AttentionBackend

9 +from vllm.config import VllmConfig, get_layers_from_vllm_config

10 +from vllm.model_executor.layers.mamba.mamba2_metadata import (

11 + _query_start_loc_to_chunk_indices_offsets)

12 +from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,

13 + CommonAttentionMetadata)

14 +from vllm.v1.kv_cache_interface import MambaSpec

15 +from vllm.v1.worker.block_table import BlockTable

16 +

17 +if TYPE_CHECKING:

18 + from vllm.v1.core.sched.output import SchedulerOutput

19 + from vllm.v1.worker.gpu_input_batch import InputBatch

20 + from vllm.v1.worker.gpu_model_runner import GPUModelRunner

21 +

22 +

23 +def get_mamba2_chunk_size(vllm_config: VllmConfig) -> int:

24 + from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2

25 + layers = get_layers_from_vllm_config(vllm_config, MambaMixer2)

26 + chunk_sizes = set(layer.chunk_size for layer in layers.values())

27 + assert len(

28 + chunk_sizes) == 1, "All Mamba2 layers must have the same chunk size"

29 + return chunk_sizes.pop()

30 +

31 +

32 +class Mamba2AttentionBackend(AttentionBackend):

33 +

34 + @staticmethod

35 + def get_builder_cls() -> type["Mamba2AttentionMetadataBuilder"]:

36 + return Mamba2AttentionMetadataBuilder

37 +

38 +

39 +@dataclass

40 +class Mamba2AttentionMetadata:

41 + num_prefills: int

42 + num_prefill_tokens: int

43 + num_decodes: int

44 + num_decode_tokens: int

45 + query_start_loc: torch.Tensor

46 + seq_lens: torch.Tensor

47 +

48 + has_initial_states: torch.Tensor

49 + prep_initial_states: bool

50 + chunk_size: int

51 + seq_idx: torch.Tensor

52 + chunk_indices: torch.Tensor

53 + chunk_offsets: torch.Tensor

54 +

55 + state_indices_tensor: torch.Tensor # shape: [batch,]

56 +

57 +

58 +class Mamba2AttentionMetadataBuilder(

59 + AttentionMetadataBuilder[Mamba2AttentionMetadata]):

60 +

61 + def __init__(self, runner: "GPUModelRunner", kv_cache_spec: MambaSpec,

62 + block_table: BlockTable):

63 + self.runner = runner

64 + self.kv_cache_spec = kv_cache_spec

65 + self.block_table = block_table

66 + self.chunk_size = get_mamba2_chunk_size(runner.vllm_config)

67 +

68 + def reorder_batch(self, input_batch: "InputBatch",

69 + scheduler_output: "SchedulerOutput") -> bool:

70 + # NOTE (Chen): Copied from MLACommonMetadataBuilder and

71 + # FlashInferMetadataBuilder. Should be refactored later to avoid code

72 + # duplication of these 3 functions.

73 + # We now want to reorder the batch so that the "decode" requests are and

74 + # the front and the "prefill" requests are at the using the least amount

75 + # swaps possible. (NOTE for now we loosely use "decode" to mean requests

76 + # where attention is likely memory-bound and "prefill" to mean requests

77 + # where attention is likely compute-bound, TODO(lucas): figure out a

78 + # better naming here)

79 + decodes = []

80 + prefills = []

81 + num_decode_tokens = 0

82 + num_prefill_tokens = 0

83 +

84 + for i, req_id in enumerate(input_batch.req_ids):

85 + num_tokens = scheduler_output.num_scheduled_tokens[req_id]

86 + # for now treat 1 scheduled token as "decode" even if its not,

87 + # we should update this to something like < 8 in the future but

88 + # currently the decode run only supports num_tokens = 1

89 + if num_tokens == 1:

90 + decodes.append(i)

91 + num_decode_tokens += num_tokens

92 + else:

93 + prefills.append(i)

94 + num_prefill_tokens += num_tokens

95 +

96 + # We hope that this is fairly minimal since decodes

97 + # should be around for a number of iterations so hopefully they are

98 + # relatively stationary (and new request are generally appended to the

99 + # persistent batch so already should be at the back)

100 + # To achieve this we loop over the decodes in descending order and

101 + # the prefills in ascending order. We swap decodes from the "back"

102 + # i.e. past where the last decode should be in the reodorered with

103 + # prefills from the front of the batch.

104 + # `decodes` and `prefills` are already in ascending order just based on

105 + # the above loop

106 + num_decodes = len(decodes)

107 + num_prefills = len(prefills)

108 + modified_batch = False

109 +

110 + for i in range(1, min(num_decodes, num_prefills) + 1):

111 + # If the decode is at the "back" of the batch, i, we can swap it

112 + # with the prefill closest to the front of the batch

113 + decode_idx = decodes[num_decodes - i]

114 + if decode_idx < num_decodes:

115 + break

116 +

117 + input_batch.swap_states(prefills[i - 1], decode_idx)

118 + modified_batch = True

119 +

120 + # Save for next `build` call

121 + # TODO(lucas): this is a bit of a hack, we should probably have a

122 + # better way of doing this

123 + self._num_decodes = num_decodes

124 + self._num_prefills = num_prefills

125 + self._num_decode_tokens = num_decode_tokens

126 + self._num_prefill_tokens = num_prefill_tokens

127 +

128 + return modified_batch

129 +

130 + def build(self, common_prefix_len: int,

131 + common_attn_metadata: CommonAttentionMetadata):

132 + num_reqs = common_attn_metadata.num_reqs

133 + query_start_loc = common_attn_metadata.query_start_loc

134 + seq_lens = common_attn_metadata.seq_lens

135 +

136 + seq_idx = None

137 + chunk_indices, chunk_offsets = None, None

138 + # Need flags to indicate if there are initial states

139 + # currently we really only support the FlashAttention backend

140 + has_initial_states = None

141 + prep_initial_states = False

142 +

143 + state_indices_tensor = self.block_table.block_table[:num_reqs, 0]

144 +

145 + # Compute seq_idx, chunk_indices and chunk_offsets for prefill only

146 + if self._num_prefills > 0:

147 + #[batch,]

148 + has_initial_states_cpu = (

149 + self.runner.input_batch.

150 + num_computed_tokens_cpu_tensor[num_reqs -

151 + self._num_prefills:num_reqs]

152 + > 0)

153 + prep_initial_states = torch.any(has_initial_states_cpu).item()

154 + has_initial_states = has_initial_states_cpu.to(

155 + query_start_loc.device)

156 +

157 + query_start_loc_p = common_attn_metadata.query_start_loc[

158 + -self._num_prefills - 1:] - self._num_decode_tokens

159 +

160 + seq_idx = torch.repeat_interleave(

161 + torch.arange(self._num_prefills,

162 + dtype=torch.int32,

163 + device=query_start_loc_p.device),

164 + query_start_loc_p.diff(),

165 + output_size=self._num_prefill_tokens)

166 + seq_idx.unsqueeze_(0)

167 +

168 + # We compute metadata for chunked prefill once at the top level

169 + # model forward and reuse them in mamba layers. If not needed,

170 + # they will be ignored inside mamba kernels.

171 + if prep_initial_states:

172 + chunk_indices, chunk_offsets = (

173 + _query_start_loc_to_chunk_indices_offsets(

174 + query_start_loc_p, self.chunk_size,

175 + self._num_prefill_tokens))

176 +

177 + attn_metadata = Mamba2AttentionMetadata(

178 + num_prefills=self._num_prefills,

179 + num_prefill_tokens=self._num_prefill_tokens,

180 + num_decodes=self._num_decodes,

181 + num_decode_tokens=self._num_decode_tokens,

182 + query_start_loc=query_start_loc,

183 + seq_lens=seq_lens,

184 + has_initial_states=has_initial_states,

185 + prep_initial_states=prep_initial_states,

186 + chunk_size=self.chunk_size,

187 + seq_idx=seq_idx,

188 + chunk_indices=chunk_indices,

189 + chunk_offsets=chunk_offsets,

190 + state_indices_tensor=state_indices_tensor,

191 + )

192 + return attn_metadata

vllm/v1/core/single_type_kv_cache_manager.py

+42 -1

@@ -8,7 +8,7 @@ from vllm.utils import cdiv

8 from vllm.v1.core.block_pool import BlockPool

9 from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock

10 from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheSpec,

11 - SlidingWindowSpec)

11 + MambaSpec, SlidingWindowSpec)

12 from vllm.v1.request import Request

@@ -52,6 +52,7 @@ class SingleTypeKVCacheManager(ABC):

53 self.caching_hash_fn = caching_hash_fn

54 self.kv_cache_group_id = kv_cache_group_id

55 + self._null_block = block_pool.null_block

57 def get_num_blocks_to_allocate(

58 self, request_id: str, num_tokens: int,

@@ -390,9 +391,49 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

391 return 0

392

393

394 +class MambaManager(SingleTypeKVCacheManager):

395 +

396 + @classmethod

397 + def find_longest_cache_hit(

398 + cls,

399 + block_hashes: list[BlockHash],

400 + max_length: int,

401 + kv_cache_group_ids: list[int],

402 + block_pool: BlockPool,

403 + kv_cache_spec: KVCacheSpec,

404 + use_eagle: bool,

405 + ) -> tuple[list[KVCacheBlock], ...]:

406 + assert isinstance(

407 + kv_cache_spec,

408 + MambaSpec), ("MambaManager can only be used for mamba groups")

409 + # Prefix caching is not supported for mamba now. Always return empty

410 + # list.

411 + computed_blocks: tuple[list[KVCacheBlock], ...] = tuple(

412 + [] for _ in range(len(kv_cache_group_ids)))

413 + return computed_blocks

414 +

415 + def remove_skipped_blocks(self, request_id: str,

416 + num_computed_tokens: int) -> None:

417 + # Each request will always have 1 block at this moment, so no need to

418 + # remove blocks.

419 + pass

420 +

421 + def get_num_common_prefix_blocks(self, request_id: str,

422 + num_running_requests: int) -> int:

423 + return 0

424 +

425 + def allocate_new_blocks(self, request_id: str,

426 + num_tokens: int) -> list[KVCacheBlock]:

427 + new_blocks = super().allocate_new_blocks(request_id, num_tokens)

428 + assert len(self.req_to_blocks[request_id]) == 1, (

429 + "MambaManager should only allocate 1 block for each request.")

430 + return new_blocks

431 +

432 +

433 spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {

434 FullAttentionSpec: FullAttentionManager,

435 SlidingWindowSpec: SlidingWindowManager,

436 + MambaSpec: MambaManager,

437 }

438

439

vllm/v1/kv_cache_interface.py

+24

@@ -3,6 +3,7 @@

4 import copy

5 from dataclasses import dataclass

6 +from math import prod

7 from typing import Optional

9 import torch

@@ -154,6 +155,29 @@ class SlidingWindowSpec(AttentionSpec):

155 return (cdiv(num_tokens, self.block_size) + 1) * self.page_size_bytes

156

157

158 +@dataclass

159 +class MambaSpec(KVCacheSpec):

160 + shapes: tuple[tuple[int, ...], ...]

161 + dtype: torch.dtype

162 +

163 + def __post_init__(self):

164 + self.num_elements = sum(prod(shape) for shape in self.shapes)

165 +

166 + @property

167 + def type_id(self) -> str:

168 + return f"mamba_{self.shapes}_{self.dtype}"

169 +

170 + @property

171 + def page_size_bytes(self) -> int:

172 + return self.num_elements * get_dtype_size(self.dtype)

173 +

174 + def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:

175 + # We allocate 1 block for each request now, so max_memory_usage_bytes is

176 + # the same as page_size_bytes.

177 + # Need to update this when supporting prefix caching.

178 + return self.page_size_bytes

179 +

180 +

181 @dataclass

182 class KVCacheTensor:

183 """

vllm/v1/worker/gpu_model_runner.py

+68 -26

@@ -29,6 +29,7 @@ from vllm.distributed.parallel_state import (

29 from vllm.forward_context import (DPMetadata, get_forward_context,

30 set_forward_context)

31 from vllm.logger import init_logger

32 +from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2

33 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding

34 from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader

35 from vllm.multimodal import MULTIMODAL_REGISTRY

@@ -38,12 +39,14 @@ from vllm.sampling_params import SamplingType

39 from vllm.sequence import IntermediateTensors

40 from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,

41 GiB_bytes, LazyLoader, async_tensor_h2d, cdiv,

41 - check_use_alibi, is_pin_memory_available)

42 + check_use_alibi, get_dtype_size,

43 + is_pin_memory_available)

44 +from vllm.v1.attention.backends.mamba_attn import Mamba2AttentionBackend

45 from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,

46 CommonAttentionMetadata)

47 from vllm.v1.core.encoder_cache_manager import compute_encoder_budget

48 from vllm.v1.kv_cache_interface import (AttentionSpec, FullAttentionSpec,

46 - KVCacheConfig, KVCacheSpec,

49 + KVCacheConfig, KVCacheSpec, MambaSpec,

50 SlidingWindowSpec)

51 from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, LogprobsTensors,

52 ModelRunnerOutput)

@@ -2093,28 +2096,31 @@ class GPUModelRunner(LoRAModelRunnerMixin):

2096 for i, kv_cache_group_spec in enumerate(

2097 kv_cache_config.kv_cache_groups):

2098 kv_cache_spec = kv_cache_group_spec.kv_cache_spec

2096 - if not isinstance(kv_cache_spec, AttentionSpec):

2097 - raise NotImplementedError(

2098 - "Only AttentionSpec is supported for now.")

2099 - attn_backend_i = get_attn_backend(

2100 - kv_cache_spec.head_size,

2101 - self.dtype,

2102 - kv_cache_spec.dtype,

2103 - kv_cache_spec.block_size,

2104 - self.model_config.is_attention_free,

2105 - use_mla=kv_cache_spec.use_mla,

2106 - )

2107 - if attn_backend_i is None:

2108 - error_msg = (

2109 - f"Error with get_attn_backend: {kv_cache_spec.head_size=}, "

2110 - f"{self.dtype=}, {kv_cache_spec.dtype=}, "

2111 - f"{kv_cache_spec.block_size=}, "

2112 - f"{self.model_config.is_attention_free=}, "

2113 - f"{kv_cache_spec.use_mla=}")

2114 - logger.error(error_msg)

2115 - raise NotImplementedError(

2116 - "Non-Attention backend is not supported by V1 "

2117 - "GPUModelRunner.")

2099 + if isinstance(kv_cache_spec, AttentionSpec):

2100 + attn_backend_i = get_attn_backend(

2101 + kv_cache_spec.head_size,

2102 + self.dtype,

2103 + kv_cache_spec.dtype,

2104 + kv_cache_spec.block_size,

2105 + self.model_config.is_attention_free,

2106 + use_mla=kv_cache_spec.use_mla,

2107 + )

2108 + if attn_backend_i is None:

2109 + error_msg = (f"Error with get_attn_backend: "

2110 + f"{kv_cache_spec.head_size=}, "

2111 + f"{self.dtype=}, {kv_cache_spec.dtype=}, "

2112 + f"{kv_cache_spec.block_size=}, "

2113 + f"{self.model_config.is_attention_free=}, "

2114 + f"{kv_cache_spec.use_mla=}")

2115 + logger.error(error_msg)

2116 + raise NotImplementedError(

2117 + "Non-Attention backend is not supported by V1 "

2118 + "GPUModelRunner.")

2119 + elif isinstance(kv_cache_spec, MambaSpec):

2120 + attn_backend_i = Mamba2AttentionBackend

2121 + else:

2122 + raise ValueError(

2123 + f"Unknown KV cache spec type: {type(kv_cache_spec)}")

2124

2125 block_table_i = self.input_batch.block_table[i]

2126 attn_metadata_builder_i = attn_backend_i.get_builder_cls()(

@@ -2242,6 +2248,22 @@ class GPUModelRunner(LoRAModelRunnerMixin):

2248 kv_caches[layer_name] = kv_cache_raw_tensors[

2249 layer_name].view(dtype).view(kv_cache_shape).permute(

2250 *inv_order)

2251 + elif isinstance(kv_cache_spec, MambaSpec):

2252 + raw_tensor = kv_cache_raw_tensors[layer_name]

2253 + dtype = kv_cache_spec.dtype

2254 + state_tensors = []

2255 + start_pos = 0

2256 + for shape in kv_cache_spec.shapes:

2257 + target_shape = (num_blocks, *shape)

2258 + size_in_bytes = np.prod(shape) * get_dtype_size(

2259 + dtype) * num_blocks

2260 + tensor = raw_tensor[start_pos:start_pos +

2261 + size_in_bytes]

2262 + tensor = tensor.view(dtype).view(target_shape)

2263 + state_tensors.append(tensor)

2264 + start_pos += size_in_bytes

2265 + assert start_pos == raw_tensor.numel()

2266 + kv_caches[layer_name] = tuple(state_tensors)

2267 else:

2268 raise NotImplementedError

2269 return kv_caches

@@ -2307,11 +2329,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):

2329 format. Layers that do not need KV cache are not included.

2330 """

2331

2310 - layers = get_layers_from_vllm_config(self.vllm_config, Attention)

2332 block_size = self.vllm_config.cache_config.block_size

2333 use_mla = self.vllm_config.model_config.use_mla

2334 kv_cache_spec: dict[str, KVCacheSpec] = {}

2314 - for layer_name, attn_module in layers.items():

2335 + attn_layers = get_layers_from_vllm_config(self.vllm_config, Attention)

2336 + for layer_name, attn_module in attn_layers.items():

2337 if (kv_tgt_layer :=

2338 attn_module.kv_sharing_target_layer_name) is not None:

2339 # The layer doesn't need its own KV cache and will use that of

@@ -2351,4 +2373,24 @@ class GPUModelRunner(LoRAModelRunnerMixin):

2373 raise ValueError(

2374 f"Unknown attention type: {attn_module.attn_type}")

2375

2376 + mamba_layers = get_layers_from_vllm_config(self.vllm_config,

2377 + MambaMixer2)

2378 + if len(mamba_layers) > 0:

2379 + if self.vllm_config.speculative_config is not None:

2380 + raise NotImplementedError(

2381 + "Mamba with speculative decoding is not supported yet.")

2382 + if not self.vllm_config.model_config.enforce_eager:

2383 + raise NotImplementedError(

2384 + "Mamba with cuda graph is not supported yet.")

2385 + if self.vllm_config.cache_config.enable_prefix_caching:

2386 + raise NotImplementedError(

2387 + "Prefix caching is not supported for Mamba yet.")

2388 + max_model_len = self.vllm_config.model_config.max_model_len

2389 + # Set block_size to max_model_len, so that mamba model will always

2390 + # have only one block in the KV cache.

2391 + for layer_name, mamba_module in mamba_layers.items():

2392 + kv_cache_spec[layer_name] = MambaSpec(

2393 + shapes=mamba_module.get_state_shape(),

2394 + dtype=self.kv_cache_dtype,

2395 + block_size=max_model_len)

2396 return kv_cache_spec