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

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, 23:30 UTC

[Model] Support Mamba2 (Codestral Mamba) (#9292)

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com> Co-authored-by: Yu Chin Fabian Lim <flim@sg.ibm.com>

Tyler Michael Smith committed Feb 17, 2025 at 07:17 UTC 1f69c4a892d35a362e9ffeaf74072aeecca2286b

9 files changed +376 -65

tests/models/decoder_only/language/test_mamba.py

+30 -13

@@ -4,6 +4,7 @@

4 Run `pytest tests/models/test_mamba.py`.

5 """

6 import pytest

7 +import torch

8 from transformers import AutoModelForCausalLM, AutoTokenizer

10 from vllm.engine.arg_utils import EngineArgs

@@ -11,7 +12,14 @@ from vllm.sampling_params import SamplingParams

13 from ...utils import check_outputs_equal

14 -MODELS = ["state-spaces/mamba-130m-hf", "tiiuae/falcon-mamba-tiny-dev"]

15 +MODELS = [

16 + "state-spaces/mamba-130m-hf",

17 + "tiiuae/falcon-mamba-tiny-dev",

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

19 + # Mamba2 is buggy for Codestral as it doesn't handle n_groups.

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

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

22 +]

25 # Use lower-level interfaces to create this greedy generator, as mamba will

@@ -21,6 +29,10 @@ def generate_greedy(model_name, example_prompts, max_tokens):

29 tokenizer = AutoTokenizer.from_pretrained(model_name)

30 model = AutoModelForCausalLM.from_pretrained(model_name)

32 + # Set the device (GPU if available, else CPU)

33 + device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

34 + model.to(device)

35 +

36 # Generate texts from the prompts

37 outputs = []

38 for prompt in example_prompts:

@@ -29,7 +41,9 @@ def generate_greedy(model_name, example_prompts, max_tokens):

41 input_ids = inputs["input_ids"].to(model.device)

43 # Generate text using the model's generate method directly

32 - generated_ids = model.generate(input_ids, max_new_tokens=max_tokens)

44 + generated_ids = model.generate(input_ids,

45 + max_new_tokens=max_tokens,

46 + do_sample=False)

47 generated_text = tokenizer.decode(generated_ids[0],

48 skip_special_tokens=True)

@@ -50,7 +64,8 @@ def test_models(

64 ) -> None:

65 hf_outputs = generate_greedy(model, example_prompts, max_tokens)

53 - with vllm_runner(model, dtype=dtype) as vllm_model:

67 + # Set max_num_seqs to keep Codestral from going OOM at fp32

68 + with vllm_runner(model, dtype=dtype, max_num_seqs=16) as vllm_model:

69 vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens)

71 # This test is for verifying whether the model's extra_repr

@@ -81,7 +96,7 @@ def test_batching(

96 ) -> None:

97 # To pass the small model tests, we need full precision.

98 for_loop_outputs = []

84 - with vllm_runner(model, dtype=dtype) as vllm_model:

99 + with vllm_runner(model, dtype=dtype, max_num_seqs=16) as vllm_model:

100 for prompt in example_prompts:

101 for_loop_outputs.append(

102 vllm_model.generate_greedy([prompt], max_tokens)[0])

@@ -165,20 +180,22 @@ def test_parallel_sampling(

180 max_tokens: int,

181 ) -> None:

182

168 - with vllm_runner(model, dtype=dtype) as vllm_model:

183 + # Numerical differences produce slightly different output for these

184 + if 'state-spaces' in model:

185 + example_prompts.pop(0)

186 + example_prompts.pop(0)

187 + example_prompts.pop(0)

188 +

189 + with vllm_runner(model, dtype=dtype, max_num_seqs=16) as vllm_model:

190 for_loop_outputs = []

191 for _ in range(10):

192 for_loop_outputs.append(

172 - # using example_prompts index 1 instead of 0 since with 0 the

173 - # logprobs get really close and the test doesn't pass

174 - vllm_model.generate_greedy([example_prompts[1]], max_tokens)

175 - [0])

193 + vllm_model.generate_greedy(example_prompts, max_tokens)[0])

194 sampling_params = SamplingParams(n=10,

195 temperature=0.001,

196 seed=0,

197 max_tokens=max_tokens)

180 - n_lt_1_outputs = vllm_model.generate([example_prompts[1]],

181 - sampling_params)

198 + n_lt_1_outputs = vllm_model.generate(example_prompts, sampling_params)

199 token_ids, texts = n_lt_1_outputs[0]

200 n_lt_1_outputs = [(token_id, text)

201 for token_id, text in zip(token_ids, texts)]

@@ -232,7 +249,7 @@ def test_models_preemption_recompute(

249 # Tests that outputs are identical with and w/o preemtions (recompute)

250 assert dtype == "float"

251

235 - with vllm_runner(model, dtype=dtype) as vllm_model:

252 + with vllm_runner(model, dtype=dtype, max_num_seqs=16) as vllm_model:

253 vllm_model.model.llm_engine.scheduler[

254 0].ENABLE_ARTIFICIAL_PREEMPT = True

255 preempt_vllm_outputs = vllm_model.generate_greedy(

@@ -283,7 +300,7 @@ def test_state_cleanup(

300 # This test is for verifying that the Mamba state is cleaned up between

301 # steps, If its not cleaned, an error would be expected.

302 try:

286 - with vllm_runner(model, dtype=dtype) as vllm_model:

303 + with vllm_runner(model, dtype=dtype, max_num_seqs=16) as vllm_model:

304 for _ in range(10):

305 vllm_model.generate_greedy([example_prompts[0]] * 100, 1)

306 except ValueError:

tests/models/registry.py

@@ -145,6 +145,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

145 "LLaMAForCausalLM": _HfExamplesInfo("decapoda-research/llama-7b-hf",

146 is_available_online=False),

147 "MambaForCausalLM": _HfExamplesInfo("state-spaces/mamba-130m-hf"),

148 + "Mamba2ForCausalLM": _HfExamplesInfo("mistralai/Mamba-Codestral-7B-v0.1",

149 + is_available_online=False),

150 "FalconMambaForCausalLM": _HfExamplesInfo("tiiuae/falcon-mamba-7b-instruct"), # noqa: E501

151 "MiniCPMForCausalLM": _HfExamplesInfo("openbmb/MiniCPM-2B-sft-bf16",

152 trust_remote_code=True),

vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py

+6 -2

@@ -293,7 +293,8 @@ def _chunk_scan_fwd_kernel(

293 dA_cs_m_boundary = tl.load(

294 dA_cumsum_ptr +

295 (pid_m * BLOCK_SIZE_M + c_off - 1) * stride_dA_cs_csize,

296 - mask=(pid_m * BLOCK_SIZE_M + c_off - 1) > -1,

296 + mask=(((pid_m * BLOCK_SIZE_M + c_off - 1) > -1)

297 + and ((pid_m * BLOCK_SIZE_M + c_off) < chunk_size)),

298 other=0.0).to(tl.float32)

299

300 if HAS_SEQ_IDX:

@@ -463,7 +464,10 @@ def _seq_idx_to_chunk_indices_offsets(seq_idx, chunk_size: int):

464 p += (s % chunk_size > 0)

465

466 # get the dimensions

466 - _s, _e = s // chunk_size + p, e // chunk_size + p + 1

467 + # - the + 1 for _e is to shift the boundary by one chunk

468 + # - this shifting is not needed if chunk_size divides e

469 + _s, _e = s // chunk_size + p, e // chunk_size + p + (e % chunk_size

470 + > 0)

471

472 # adjust inidces and offsets

473 chunk_indices[_s:_e] -= p

vllm/model_executor/models/bamba.py

+2 -19

@@ -440,23 +440,6 @@ class BambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

440 self.make_empty_intermediate_tensors = (

441 self.model.make_empty_intermediate_tensors)

442

443 - # follow jamba

444 - if self.scheduler_config is not None and \

445 - not self.model_config.enforce_eager:

446 - # for compilation

447 - if self.scheduler_config.max_num_seqs > \

448 - vllm_config.compilation_config.max_capture_size:

449 - self.max_batch_size = \

450 - vllm_config.compilation_config.max_capture_size

451 - else:

452 - self.max_batch_size = vllm_config.pad_for_cudagraph(

453 - self.scheduler_config.max_num_seqs)

454 - elif self.scheduler_config is not None:

455 - # for eager just take the scheduler_config if avail

456 - self.max_batch_size = self.scheduler_config.max_num_seqs

457 - else:

458 - self.max_batch_size = 8192 + 2

459 -

443 def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

444 return self.model.get_input_embeddings(input_ids)

445

@@ -474,8 +457,8 @@ class BambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

457 self.vllm_config.parallel_config, LayerBlockType.mamba)

458

459 self.mamba_cache = MambaCacheManager(

477 - self.lm_head.weight.dtype, num_mamba_layers,

478 - self.max_batch_size, *self._get_mamba_cache_shape())

460 + self.vllm_config, self.lm_head.weight.dtype, num_mamba_layers,

461 + *self._get_mamba_cache_shape())

462 mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)

463 hidden_states = self.model(input_ids, positions, kv_caches,

464 attn_metadata, mamba_cache_params,

vllm/model_executor/models/jamba.py

+2 -13

@@ -426,17 +426,6 @@ class JambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

426

427 self.make_empty_intermediate_tensors = (

428 self.model.make_empty_intermediate_tensors)

429 - if self.scheduler_config is not None and \

430 - not self.model_config.enforce_eager:

431 - if self.scheduler_config.max_num_seqs > \

432 - vllm_config.compilation_config.max_capture_size:

433 - self.max_batch_size = \

434 - vllm_config.compilation_config.max_capture_size

435 - else:

436 - self.max_batch_size = vllm_config.pad_for_cudagraph(

437 - self.scheduler_config.max_num_seqs)

438 - else:

439 - self.max_batch_size = 8192 + 2

429

430 def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

431 return self.model.get_input_embeddings(input_ids)

@@ -453,8 +442,8 @@ class JambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

442 num_mamba_layers = self.model_config.get_num_layers_by_block_type(

443 self.vllm_config.parallel_config, LayerBlockType.mamba)

444 self.mamba_cache = MambaCacheManager(

456 - self.lm_head.weight.dtype, num_mamba_layers,

457 - self.max_batch_size, *self._get_mamba_cache_shape())

445 + self.vllm_config, self.lm_head.weight.dtype, num_mamba_layers,

446 + *self._get_mamba_cache_shape())

447

448 mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)

449

vllm/model_executor/models/mamba.py

+3 -15

@@ -166,14 +166,13 @@ class MambaForCausalLM(nn.Module, HasInnerState, IsAttentionFree, SupportsPP):

166 config = vllm_config.model_config.hf_config

167 cache_config = vllm_config.cache_config

168 lora_config = vllm_config.lora_config

169 - scheduler_config = vllm_config.scheduler_config

169 + self.scheduler_config = vllm_config.scheduler_config

170 assert not cache_config.enable_prefix_caching, \

171 "Mamba does not support prefix caching"

172

173 super().__init__()

174 self.config = config

175 self.vllm_config = vllm_config

176 - self.scheduler_config = scheduler_config

176 self.model_config = vllm_config.model_config

177 self.backbone = MambaModel(vllm_config=vllm_config,

178 prefix=maybe_prefix(prefix, "backbone"))

@@ -202,17 +201,6 @@ class MambaForCausalLM(nn.Module, HasInnerState, IsAttentionFree, SupportsPP):

201

202 self.make_empty_intermediate_tensors = (

203 self.backbone.make_empty_intermediate_tensors)

205 - if self.scheduler_config is not None and \

206 - not self.model_config.enforce_eager:

207 - if self.scheduler_config.max_num_seqs > \

208 - vllm_config.compilation_config.max_capture_size:

209 - self.max_batch_size = \

210 - vllm_config.compilation_config.max_capture_size

211 - else:

212 - self.max_batch_size = vllm_config.pad_for_cudagraph(

213 - self.scheduler_config.max_num_seqs)

214 - else:

215 - self.max_batch_size = 8192 + 2

204

205 def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

206 return self.backbone.get_input_embeddings(input_ids)

@@ -229,8 +217,8 @@ class MambaForCausalLM(nn.Module, HasInnerState, IsAttentionFree, SupportsPP):

217 num_mamba_layers = self.model_config.get_num_layers_by_block_type(

218 self.vllm_config.parallel_config, LayerBlockType.mamba)

219 self.mamba_cache = MambaCacheManager(

232 - self.lm_head.weight.dtype, num_mamba_layers,

233 - self.max_batch_size, *self._get_mamba_cache_shape())

220 + self.vllm_config, self.lm_head.weight.dtype, num_mamba_layers,

221 + *self._get_mamba_cache_shape())

222

223 mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)

224

vllm/model_executor/models/mamba2.py new

+320

@@ -0,0 +1,320 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +"""PyTorch MAMBA2 model."""

3 +from typing import Iterable, List, Optional, Set, Tuple

4 +

5 +import torch

6 +from torch import nn

7 +from transformers import MambaConfig

8 +

9 +from vllm.attention.backends.abstract import AttentionMetadata

10 +from vllm.config import VllmConfig

11 +from vllm.distributed import divide, get_tensor_model_parallel_world_size

12 +from vllm.distributed.parallel_state import get_pp_group

13 +from vllm.model_executor.layers.layernorm import RMSNorm

14 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

15 +from vllm.model_executor.layers.mamba.mamba_mixer2 import (

16 + MambaMixer2, extra_groups_for_head_shards)

17 +from vllm.model_executor.layers.quantization.base_config import (

18 + QuantizationConfig)

19 +from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler

20 +from vllm.model_executor.layers.vocab_parallel_embedding import (

21 + DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)

22 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader

23 +from vllm.model_executor.models.interfaces import (HasInnerState,

24 + IsAttentionFree)

25 +from vllm.model_executor.models.mamba_cache import (MambaCacheManager,

26 + MambaCacheParams)

27 +from vllm.model_executor.sampling_metadata import SamplingMetadata

28 +from vllm.sequence import IntermediateTensors

29 +from vllm.utils import LayerBlockType

30 +

31 +from .utils import (is_pp_missing_parameter,

32 + make_empty_intermediate_tensors_factory, make_layers,

33 + maybe_prefix)

34 +

35 +KVCache = Tuple[torch.Tensor, torch.Tensor]

36 +

37 +

38 +class Mamba2DecoderLayer(nn.Module):

39 +

40 + def __init__(self,

41 + config: MambaConfig,

42 + quant_config: Optional[QuantizationConfig] = None) -> None:

43 + super().__init__()

44 + self.config = config

45 + self.mixer = MambaMixer2(hidden_size=config.hidden_size,

46 + ssm_state_size=config.state_size,

47 + conv_kernel_size=config.conv_kernel,

48 + intermediate_size=getattr(

49 + config, "intermediate_size",

50 + config.expand * config.hidden_size),

51 + use_conv_bias=config.use_conv_bias,

52 + use_bias=config.use_bias,

53 + n_groups=config.n_groups,

54 + num_heads=config.num_heads,

55 + head_dim=config.head_dim,

56 + rms_norm_eps=config.layer_norm_epsilon,

57 + activation=config.hidden_act,

58 + chunk_size=config.chunk_size,

59 + quant_config=quant_config)

60 +

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

62 +

63 + def forward(

64 + self,

65 + hidden_states: torch.Tensor,

66 + attn_metadata: AttentionMetadata,

67 + residual: Optional[torch.Tensor],

68 + mamba_cache_params: MambaCacheParams,

69 + sequence_idx: Optional[torch.Tensor],

70 + **kwargs,

71 + ):

72 + if residual is None:

73 + residual = hidden_states

74 + hidden_states = self.norm(hidden_states)

75 + else:

76 + hidden_states, residual = self.norm(hidden_states, residual)

77 +

78 + hidden_states = self.mixer(hidden_states, attn_metadata,

79 + mamba_cache_params, sequence_idx)

80 + return hidden_states, residual

81 +

82 +

83 +class Mamba2Model(nn.Module):

84 +

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

86 + super().__init__()

87 +

88 + config = vllm_config.model_config.hf_config

89 + quant_config = vllm_config.quant_config

90 + lora_config = vllm_config.lora_config

91 + is_lora_enabled = bool(lora_config)

92 + assert not is_lora_enabled

93 +

94 + self.config = config

95 + lora_vocab = ((lora_config.lora_extra_vocab_size *

96 + (lora_config.max_loras or 1)) if lora_config else 0)

97 + self.vocab_size = config.vocab_size + lora_vocab

98 + self.org_vocab_size = config.vocab_size

99 +

100 + self.embeddings = VocabParallelEmbedding(

101 + self.vocab_size,

102 + config.hidden_size,

103 + org_num_embeddings=config.vocab_size,

104 + )

105 +

106 + self.start_layer, self.end_layer, self.layers = make_layers(

107 + config.num_hidden_layers,

108 + lambda prefix: Mamba2DecoderLayer(config,

109 + quant_config=quant_config),

110 + prefix=f"{prefix}.layers")

111 +

112 + self.norm_f = RMSNorm(config.hidden_size,

113 + eps=config.layer_norm_epsilon)

114 + self.make_empty_intermediate_tensors = (

115 + make_empty_intermediate_tensors_factory(

116 + ["hidden_states", "residual"], config.hidden_size))

117 +

118 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

119 + return self.embeddings(input_ids)

120 +

121 + def forward(

122 + self,

123 + input_ids: torch.Tensor,

124 + positions: torch.Tensor,

125 + attn_metadata: AttentionMetadata,

126 + mamba_cache_params: MambaCacheParams,

127 + intermediate_tensors: Optional[IntermediateTensors] = None,

128 + inputs_embeds: Optional[torch.Tensor] = None,

129 + ) -> torch.Tensor:

130 + if get_pp_group().is_first_rank:

131 + if inputs_embeds is not None:

132 + hidden_states = inputs_embeds

133 + else:

134 + hidden_states = self.get_input_embeddings(input_ids)

135 + residual = None

136 + else:

137 + assert intermediate_tensors is not None

138 + hidden_states = intermediate_tensors["hidden_states"]

139 + residual = intermediate_tensors["residual"]

140 +

141 + # pass a sequence index tensor, that is required for

142 + # proper continuous batching computation including

143 + # chunked prefill

144 + seq_idx = None

145 + if attn_metadata.num_prefills > 0:

146 + seq_idx = torch.zeros_like(input_ids, dtype=torch.int32)

147 + for i, (srt, end) in enumerate(

148 + zip(

149 + attn_metadata.query_start_loc,

150 + attn_metadata.query_start_loc[1:],

151 + )):

152 + seq_idx[srt:end] = i

153 + seq_idx.unsqueeze_(0)

154 +

155 + for i in range(len(self.layers)):

156 + layer = self.layers[i]

157 +

158 + hidden_states, residual = layer(

159 + positions=positions,

160 + hidden_states=hidden_states,

161 + attn_metadata=attn_metadata,

162 + residual=residual,

163 + mamba_cache_params=mamba_cache_params.at_layer_idx(

164 + i - self.start_layer),

165 + sequence_idx=seq_idx)

166 +

167 + if not get_pp_group().is_last_rank:

168 + return IntermediateTensors({

169 + "hidden_states": hidden_states,

170 + "residual": residual

171 + })

172 +

173 + hidden_states, _ = self.norm_f(hidden_states, residual)

174 +

175 + return hidden_states

176 +

177 +

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

179 +

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

181 + config = vllm_config.model_config.hf_config

182 + cache_config = vllm_config.cache_config

183 + lora_config = vllm_config.lora_config

184 + scheduler_config = vllm_config.scheduler_config

185 + assert not cache_config.enable_prefix_caching, \

186 + "Mamba does not support prefix caching"

187 +

188 + super().__init__()

189 + self.config = config

190 + self.vllm_config = vllm_config

191 + self.scheduler_config = scheduler_config

192 + self.model_config = vllm_config.model_config

193 + self.backbone = Mamba2Model(vllm_config=vllm_config,

194 + prefix=maybe_prefix(prefix, "backbone"))

195 + self.unpadded_vocab_size = config.vocab_size

196 + if lora_config:

197 + self.unpadded_vocab_size += lora_config.lora_extra_vocab_size

198 +

199 + self.lm_head = ParallelLMHead(

200 + self.unpadded_vocab_size,

201 + config.hidden_size,

202 + org_num_embeddings=config.vocab_size,

203 + padding_size=DEFAULT_VOCAB_PADDING_SIZE

204 + # We need bigger padding if using lora for kernel

205 + # compatibility

206 + if not lora_config else lora_config.lora_vocab_padding_size,

207 + )

208 + if config.tie_word_embeddings:

209 + self.lm_head = self.lm_head.tie_weights(self.backbone.embeddings)

210 +

211 + # Used to track and store by the Mamba cache between steps.

212 + self.mamba_cache: Optional[MambaCacheManager] = None

213 +

214 + self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,

215 + config.vocab_size)

216 + self.sampler = get_sampler()

217 +

218 + self.make_empty_intermediate_tensors = (

219 + self.backbone.make_empty_intermediate_tensors)

220 +

221 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

222 + return self.backbone.get_input_embeddings(input_ids)

223 +

224 + def forward(self,

225 + input_ids: torch.Tensor,

226 + positions: torch.Tensor,

227 + kv_caches: List[KVCache],

228 + attn_metadata: AttentionMetadata,

229 + intermediate_tensors: Optional[IntermediateTensors] = None,

230 + inputs_embeds: Optional[torch.Tensor] = None,

231 + **kwargs):

232 + if self.mamba_cache is None:

233 + num_mamba_layers = self.model_config.get_num_layers_by_block_type(

234 + self.vllm_config.parallel_config, LayerBlockType.mamba)

235 + self.mamba_cache = MambaCacheManager(

236 + self.vllm_config, self.lm_head.weight.dtype, num_mamba_layers,

237 + *self._get_mamba_cache_shape())

238 +

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

240 +

241 + hidden_states = self.backbone(input_ids, positions, attn_metadata,

242 + mamba_cache_params, intermediate_tensors,

243 + inputs_embeds)

244 +

245 + return hidden_states

246 +

247 + def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):

248 + return self.mamba_cache.copy_inputs_before_cuda_graphs(

249 + input_buffers, **kwargs)

250 +

251 + def get_seqlen_agnostic_capture_inputs(self, batch_size: int):

252 + return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)

253 +

254 + def _get_mamba_cache_shape(

255 + self) -> Tuple[Tuple[int, int], Tuple[int, int]]:

256 + world_size = get_tensor_model_parallel_world_size()

257 +

258 + conv_state_shape, temporal_state_shape = None, None

259 +

260 + intermediate_size = getattr(

261 + self.config, "intermediate_size",

262 + self.config.expand * self.config.hidden_size)

263 +

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

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

266 + n_groups = (

267 + self.config.n_groups +

268 + extra_groups_for_head_shards(self.config.n_groups, world_size))

269 +

270 + # - heads and n_groups are TP-ed

271 + conv_dim = (intermediate_size + 2 * n_groups * self.config.state_size)

272 + conv_state_shape = (

273 + divide(conv_dim, world_size),

274 + self.config.conv_kernel - 1,

275 + )

276 +

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

278 + # - they are typically small

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

280 + temporal_state_shape = (

281 + divide(self.config.num_heads, world_size),

282 + self.config.head_dim,

283 + self.config.state_size,

284 + )

285 + return conv_state_shape, temporal_state_shape

286 +

287 + def compute_logits(self, hidden_states: torch.Tensor,

288 + sampling_metadata: SamplingMetadata) -> torch.Tensor:

289 + logits = self.logits_processor(self.lm_head, hidden_states,

290 + sampling_metadata)

291 + return logits

292 +

293 + def sample(

294 + self,

295 + logits: Optional[torch.Tensor],

296 + sampling_metadata: SamplingMetadata,

297 + ) -> Optional[SamplerOutput]:

298 + next_tokens = self.sampler(logits, sampling_metadata)

299 + return next_tokens

300 +

301 + def load_weights(self, weights: Iterable[Tuple[str,

302 + torch.Tensor]]) -> Set[str]:

303 + params_dict = dict(self.named_parameters())

304 + loaded_params: Set[str] = set()

305 + for name, loaded_weight in weights:

306 + if "A_log" in name:

307 + name = name.replace("A_log", "A")

308 +

309 + # Skip loading extra bias for GPTQ models.

310 + if name.endswith(".bias") and name not in params_dict:

311 + continue

312 + if is_pp_missing_parameter(name, self):

313 + continue

314 +

315 + param = params_dict[name]

316 + weight_loader = getattr(param, "weight_loader",

317 + default_weight_loader)

318 + weight_loader(param, loaded_weight)

319 + loaded_params.add(name)

320 + return loaded_params

vllm/model_executor/models/mamba_cache.py

+10 -3

@@ -1,11 +1,12 @@

1 # SPDX-License-Identifier: Apache-2.0

3 from dataclasses import dataclass

4 -from typing import Dict, List

4 +from typing import Dict, List, Tuple

6 import torch

8 from vllm.attention.backends.utils import PAD_SLOT_ID

9 +from vllm.config import VllmConfig

12 @dataclass

@@ -22,8 +23,14 @@ class MambaCacheParams:

24 class MambaCacheManager:

25 - def __init__(self, dtype, num_mamba_layers, max_batch_size,

26 - conv_state_shape, temporal_state_shape):

26 + def __init__(self, vllm_config: VllmConfig, dtype: torch.dtype,

27 + num_mamba_layers: int, conv_state_shape: Tuple[int, int],

28 + temporal_state_shape: Tuple[int, int]):

29 +

30 + # Determine max batch size to set size of MambaCache

31 + max_batch_size = vllm_config.scheduler_config.max_num_seqs

32 + if not vllm_config.model_config.enforce_eager:

33 + max_batch_size = vllm_config.pad_for_cudagraph(max_batch_size)

35 conv_state = torch.empty(size=(num_mamba_layers, max_batch_size) +

36 conv_state_shape,

vllm/model_executor/models/registry.py

@@ -71,6 +71,7 @@ _TEXT_GENERATION_MODELS = {

71 "LLaMAForCausalLM": ("llama", "LlamaForCausalLM"),

72 "MambaForCausalLM": ("mamba", "MambaForCausalLM"),

73 "FalconMambaForCausalLM": ("mamba", "MambaForCausalLM"),

74 + "Mamba2ForCausalLM": ("mamba2", "Mamba2ForCausalLM"),

75 "MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"),

76 "MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"),

77 "MistralForCausalLM": ("llama", "LlamaForCausalLM"),