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

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

Enable V1 for Hybrid SSM/Attention Models (#20016)

Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com> Co-authored-by: Stanislaw Wozniak <stw@zurich.ibm.com> Co-authored-by: Tyler Michael Smith <tysmith@redhat.com> Co-authored-by: Chen Zhang <zhangch99@outlook.com>

Thomas Parnell committed Jul 4, 2025 at 19:46 UTC 2f35a022e648f83ea376ee7331deafaba36d7eac

14 files changed +401 -136

tests/models/language/generation/test_hybrid.py

+60 -10

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

4 import pytest

6 +from tests.models.registry import HF_EXAMPLE_MODELS

7 from tests.utils import multi_gpu_test

8 from vllm.engine.arg_utils import EngineArgs

9 from vllm.sampling_params import SamplingParams

@@ -19,31 +20,55 @@ pytestmark = pytest.mark.hybrid_model

20 SSM_MODELS = [

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

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

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

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

24 - # doesn't compare vLLM output with HF output.

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

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

24 ]

26 HYBRID_MODELS = [

27 "ai21labs/Jamba-tiny-dev",

31 - # NOTE: Currently the test failes due to HF transformers issue fixed in:

32 - # https://github.com/huggingface/transformers/pull/39033

33 - # We will enable vLLM test for Granite after next HF transformers release.

34 - # "ibm-granite/granite-4.0-tiny-preview",

28 # NOTE: Running Plamo2 in transformers implementation requires to install

29 # causal-conv1d package, which is not listed as a test dependency as it's

30 # not compatible with pip-compile.

31 "pfnet/plamo-2-1b",

32 "Zyphra/Zamba2-1.2B-instruct",

33 "hmellor/tiny-random-BambaForCausalLM",

34 + "ibm-ai-platform/Bamba-9B-v1",

35 + "nvidia/Nemotron-H-8B-Base-8K",

36 + "ibm-granite/granite-4.0-tiny-preview",

37 + "tiiuae/Falcon-H1-0.5B-Base",

38 +]

39 +

40 +HF_UNSUPPORTED_MODELS = [

41 + # The HF transformers implementation of

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

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

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

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

46 + # Note: I'm not seeing the same output from vLLM V0 vs. HF transformers

47 + # for Nemotron-H-8B; currently only compare vLLM V0 vs. vLLM V1

48 + "nvidia/Nemotron-H-8B-Base-8K",

49 + # NOTE: Currently the test fails due to HF transformers issue fixed in:

50 + # https://github.com/huggingface/transformers/pull/39033

51 + # We will enable vLLM test for Granite after next HF transformers release.

52 + "ibm-granite/granite-4.0-tiny-preview",

53 ]

55 V1_SUPPORTED_MODELS = [

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

57 + "ibm-ai-platform/Bamba-9B-v1",

58 + "Zyphra/Zamba2-1.2B-instruct",

59 + "nvidia/Nemotron-H-8B-Base-8K",

60 + "ibm-granite/granite-4.0-tiny-preview",

61 + "tiiuae/Falcon-H1-0.5B-Base",

62 ]

64 +ATTN_BLOCK_SIZES = {

65 + "ibm-ai-platform/Bamba-9B-v1": 528,

66 + "Zyphra/Zamba2-1.2B-instruct": 80,

67 + "nvidia/Nemotron-H-8B-Base-8K": 528,

68 + "ibm-granite/granite-4.0-tiny-preview": 400,

69 + "tiiuae/Falcon-H1-0.5B-Base": 800,

70 +}

71 +

72 # Avoid OOM

73 MAX_NUM_SEQS = 4

@@ -60,8 +85,16 @@ def test_models(

85 max_tokens: int,

86 num_logprobs: int,

87 ) -> None:

88 +

89 + try:

90 + model_info = HF_EXAMPLE_MODELS.find_hf_info(model)

91 + model_info.check_available_online(on_fail="skip")

92 + model_info.check_transformers_version(on_fail="skip")

93 + except ValueError:

94 + pass

95 +

96 with hf_runner(model) as hf_model:

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

97 + if model not in HF_UNSUPPORTED_MODELS:

98 hf_outputs = hf_model.generate_greedy_logprobs_limit(

99 example_prompts, max_tokens, num_logprobs)

100 else:

@@ -72,12 +105,21 @@ def test_models(

105 example_prompts, max_tokens, num_logprobs)

106

107 if model in V1_SUPPORTED_MODELS:

108 + if model in HYBRID_MODELS and model in ATTN_BLOCK_SIZES:

109 + block_size = ATTN_BLOCK_SIZES[model]

110 + else:

111 + block_size = 16

112 +

113 with monkeypatch.context() as m:

114 m.setenv("VLLM_USE_V1", "1")

115 + if model in HYBRID_MODELS:

116 + # required due to reorder_batch behaviour

117 + m.setenv("VLLM_ATTENTION_BACKEND", "FLASHINFER")

118 with vllm_runner(model,

119 max_num_seqs=MAX_NUM_SEQS,

120 enforce_eager=True,

80 - enable_prefix_caching=False) as vllm_model:

121 + enable_prefix_caching=False,

122 + block_size=block_size) as vllm_model:

123 vllm_v1_outputs = vllm_model.generate_greedy_logprobs(

124 example_prompts, max_tokens, num_logprobs)

125 else:

@@ -111,6 +153,14 @@ def test_batching(

153 max_tokens: int,

154 num_logprobs: int,

155 ) -> None:

156 +

157 + try:

158 + model_info = HF_EXAMPLE_MODELS.find_hf_info(model)

159 + model_info.check_available_online(on_fail="skip")

160 + model_info.check_transformers_version(on_fail="skip")

161 + except ValueError:

162 + pass

163 +

164 for_loop_outputs = []

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

166 for prompt in example_prompts:

tests/models/registry.py

+1 -1

@@ -169,7 +169,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

169 "ExaoneForCausalLM": _HfExamplesInfo("LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct"), # noqa: E501

170 "Fairseq2LlamaForCausalLM": _HfExamplesInfo("mgleize/fairseq2-dummy-Llama-3.2-1B"), # noqa: E501

171 "FalconForCausalLM": _HfExamplesInfo("tiiuae/falcon-7b"),

172 - "FalconH1ForCausalLM":_HfExamplesInfo("tiiuae/Falcon-H1-1.5B-Instruct",

172 + "FalconH1ForCausalLM":_HfExamplesInfo("tiiuae/Falcon-H1-0.5B-Base",

173 min_transformers_version="4.53"),

174 "GemmaForCausalLM": _HfExamplesInfo("google/gemma-1.1-2b-it"),

175 "Gemma2ForCausalLM": _HfExamplesInfo("google/gemma-2-9b"),

tests/v1/test_oracle.py

-1

@@ -13,7 +13,6 @@ UNSUPPORTED_MODELS_V1 = [

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

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

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

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

16 "BAAI/bge-m3", # embedding

17 ]

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

+1 -1

@@ -108,7 +108,7 @@ def _selective_scan_update_kernel(

108 # is the same as the batch id.

109 if HAS_STATE_BATCH_INDICES:

110 state_batch_indices_ptr += pid_b

111 - state_batch_idx = tl.load(state_batch_indices_ptr)

111 + state_batch_idx = tl.load(state_batch_indices_ptr).to(tl.int64)

112 state_ptr += (state_batch_idx * stride_state_batch +

113 pid_h * stride_state_head)

114 else:

vllm/model_executor/models/bamba.py

+30 -15

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

9 from torch import nn

10 from transformers import BambaConfig

12 +from vllm import envs

13 from vllm.attention.layer import Attention

14 from vllm.config import CacheConfig, VllmConfig

15 from vllm.distributed import divide, get_tensor_model_parallel_world_size

@@ -36,7 +37,7 @@ from vllm.sequence import IntermediateTensors

37 from vllm.utils import LayerBlockType

39 from .interfaces import (HasInnerState, IsHybrid, SupportsLoRA, SupportsPP,

39 - SupportsQuant, SupportsV0Only)

40 + SupportsQuant)

41 from .utils import (AutoWeightsLoader, is_pp_missing_parameter,

42 make_empty_intermediate_tensors_factory, make_layers,

43 maybe_prefix)

@@ -97,7 +98,9 @@ class BambaMixerDecoderLayer(nn.Module):

98 head_dim=config.mamba_d_head,

99 rms_norm_eps=config.rms_norm_eps,

100 activation=config.hidden_act,

100 - quant_config=quant_config)

101 + quant_config=quant_config,

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

103 + chunk_size=config.mamba_chunk_size)

104

105 self.feed_forward = BambaMLP(config, quant_config=quant_config)

106 self.input_layernorm = RMSNorm(config.hidden_size,

@@ -313,10 +316,14 @@ class BambaModel(nn.Module):

316

317 attn_metadata = get_forward_context().attn_metadata

318

316 - mamba2_metadata = prepare_mamba2_metadata(

317 - chunk_size=self.config.mamba_chunk_size,

318 - attn_metadata=attn_metadata,

319 - )

319 + if not envs.VLLM_USE_V1:

320 + mamba2_metadata = prepare_mamba2_metadata(

321 + chunk_size=self.config.mamba_chunk_size,

322 + attn_metadata=attn_metadata,

323 + )

324 + else:

325 + # v1 get mamba2_metadata from forward_context

326 + mamba2_metadata = None

327

328 if get_pp_group().is_first_rank:

329 if inputs_embeds is not None:

@@ -337,7 +344,8 @@ class BambaModel(nn.Module):

344 num_attn += 1

345

346 layer_mamba_cache_params = None

340 - if isinstance(layer, BambaMixerDecoderLayer):

347 + if isinstance(layer,

348 + BambaMixerDecoderLayer) and mamba_cache_params:

349 layer_mamba_cache_params = mamba_cache_params.at_layer_idx(

350 i - num_attn)

351

@@ -411,7 +419,7 @@ class BambaModel(nn.Module):

419

420

421 class BambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

414 - IsHybrid, SupportsV0Only, SupportsQuant):

422 + IsHybrid, SupportsQuant):

423 packed_modules_mapping = {

424 "qkv_proj": [

425 "q_proj",

@@ -475,15 +483,22 @@ class BambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

483 intermediate_tensors: Optional[IntermediateTensors] = None,

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

485 **kwargs):

478 - if self.mamba_cache is None:

486

480 - num_mamba_layers = self.model_config.get_num_layers_by_block_type(

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

487 + mamba_cache_params = None

488 + if not envs.VLLM_USE_V1:

489 + if self.mamba_cache is None:

490 + num_mamba_layers = \

491 + self.model_config.get_num_layers_by_block_type(

492 + self.vllm_config.parallel_config,

493 + LayerBlockType.mamba

494 + )

495 +

496 + self.mamba_cache = MambaCacheManager(

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

498 + num_mamba_layers, *self._get_mamba_cache_shape())

499 +

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

501

483 - self.mamba_cache = MambaCacheManager(

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

485 - *self._get_mamba_cache_shape())

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

502 hidden_states = self.model(input_ids, positions, mamba_cache_params,

503 intermediate_tensors, inputs_embeds)

504

vllm/model_executor/models/falcon_h1.py

+40 -17

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

8 from torch import nn

9 from transformers import FalconH1Config

11 +from vllm import envs

12 from vllm.attention.layer import Attention

13 from vllm.config import CacheConfig, VllmConfig

14 from vllm.distributed import divide, get_tensor_model_parallel_world_size

@@ -33,8 +34,7 @@ from vllm.model_executor.models.mamba_cache import (MambaCacheManager,

34 from vllm.model_executor.sampling_metadata import SamplingMetadata

35 from vllm.sequence import IntermediateTensors

36 -from .interfaces import (HasInnerState, IsHybrid, SupportsLoRA, SupportsPP,

37 - SupportsV0Only)

37 +from .interfaces import HasInnerState, IsHybrid, SupportsLoRA, SupportsPP

38 from .utils import (PPMissingLayer, is_pp_missing_parameter,

39 make_empty_intermediate_tensors_factory, make_layers,

40 maybe_prefix)

@@ -85,6 +85,7 @@ class FalconH1SSMDecoderLayer(nn.Module):

85 config: FalconH1Config,

86 cache_config: Optional[CacheConfig] = None,

87 quant_config: Optional[QuantizationConfig] = None,

88 + prefix: str = "",

89 ) -> None:

90 super().__init__()

91 self.config = config

@@ -107,6 +108,8 @@ class FalconH1SSMDecoderLayer(nn.Module):

108 activation=config.hidden_act,

109 quant_config=quant_config,

110 use_rms_norm=config.mamba_rms_norm,

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

112 + chunk_size=config.mamba_chunk_size,

113 )

114 # n_groups is overridden later by `MambaMixer2`

115 self.groups_time_state_size = self.mamba.n_groups * config.mamba_d_state

@@ -316,6 +319,7 @@ class FalconH1ParallelHybrid(nn.Module):

319 prefix: str = "",

320 ) -> None:

321 super().__init__()

322 +

323 # Instantiate the attention branch

324 self.self_attn = FalconH1AttentionDecoderLayer(

325 config=config,

@@ -323,11 +327,18 @@ class FalconH1ParallelHybrid(nn.Module):

327 quant_config=quant_config,

328 prefix=prefix,

329 )

330 +

331 + # In V1 all attention/ssm layers must have

332 + # different index in prefix

333 + ssm_layer_idx = config.num_hidden_layers + layer_idx

334 + ssm_prefix = prefix.split(".")[0] + f".{ssm_layer_idx}"

335 +

336 # Instantiate the SSM branch

337 self.mamba = FalconH1SSMDecoderLayer(

338 config=config,

339 cache_config=cache_config,

340 quant_config=quant_config,

341 + prefix=ssm_prefix,

342 )

343 self.ssm_out_multiplier = config.ssm_out_multiplier

344 self.ssm_in_multiplier = config.ssm_in_multiplier

@@ -452,10 +463,16 @@ class FalconH1Model(nn.Module):

463 # proper continuous batching computation including

464 # chunked prefill

465 attn_metadata = get_forward_context().attn_metadata

455 - mamba2_metadata = prepare_mamba2_metadata(

456 - chunk_size=self.config.mamba_chunk_size,

457 - attn_metadata=attn_metadata,

458 - )

466 +

467 + if not envs.VLLM_USE_V1:

468 + mamba2_metadata = prepare_mamba2_metadata(

469 + chunk_size=self.config.mamba_chunk_size,

470 + attn_metadata=attn_metadata,

471 + )

472 + else:

473 + # v1 get mamba2_metadata from forward_context

474 + mamba2_metadata = None

475 +

476 if get_pp_group().is_first_rank:

477 if inputs_embeds is not None:

478 hidden_states = inputs_embeds * self.embedding_multiplier

@@ -468,7 +485,9 @@ class FalconH1Model(nn.Module):

485

486 for i in range(self.start_layer, self.end_layer):

487 layer = self.layers[i]

471 - layer_mamba_cache_params = mamba_cache_params.at_layer_idx(i)

488 + layer_mamba_cache_params = None

489 + if mamba_cache_params:

490 + layer_mamba_cache_params = mamba_cache_params.at_layer_idx(i)

491 hidden_states = layer(

492 positions=positions,

493 hidden_states=hidden_states,

@@ -484,7 +503,7 @@ class FalconH1Model(nn.Module):

503

504

505 class FalconH1ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

487 - IsHybrid, SupportsV0Only):

506 + IsHybrid):

507 packed_modules_mapping = {

508 "qkv_proj": ["q_proj", "k_proj", "v_proj"],

509 "gate_up_proj": ["gate_proj", "up_proj"],

@@ -558,15 +577,19 @@ class FalconH1ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

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

578 **kwargs,

579 ):

561 - if self.mamba_cache is None:

562 - self.mamba_cache = MambaCacheManager(

563 - self.vllm_config,

564 - self.lm_head.weight.dtype

565 - if hasattr(self.lm_head, 'weight') else torch.bfloat16,

566 - self.config.num_hidden_layers,

567 - *self._get_mamba_cache_shape(),

568 - )

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

580 +

581 + mamba_cache_params = None

582 + if not envs.VLLM_USE_V1:

583 + if self.mamba_cache is None:

584 + self.mamba_cache = MambaCacheManager(

585 + self.vllm_config,

586 + self.lm_head.weight.dtype if hasattr(

587 + self.lm_head, 'weight') else torch.bfloat16,

588 + self.config.num_hidden_layers,

589 + *self._get_mamba_cache_shape(),

590 + )

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

592 +

593 hidden_states = self.model(

594 input_ids,

595 positions,

vllm/model_executor/models/granitemoehybrid.py

+32 -17

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

9 from torch import nn

10 from transformers import GraniteMoeHybridConfig

12 +from vllm import envs

13 from vllm.attention.layer import Attention

14 from vllm.config import CacheConfig, VllmConfig

15 from vllm.distributed import divide, get_tensor_model_parallel_world_size

@@ -36,7 +37,7 @@ from vllm.utils import LayerBlockType

37 from .granitemoe import GraniteMoeMoE

38 from .granitemoeshared import GraniteMoeSharedMLP

39 from .interfaces import (HasInnerState, IsHybrid, SupportsLoRA, SupportsPP,

39 - SupportsQuant, SupportsV0Only)

40 + SupportsQuant)

41 from .utils import (AutoWeightsLoader, is_pp_missing_parameter,

42 make_empty_intermediate_tensors_factory, make_layers,

43 maybe_prefix)

@@ -67,7 +68,9 @@ class GraniteMoeHybridMambaDecoderLayer(nn.Module):

68 head_dim=config.mamba_d_head,

69 rms_norm_eps=config.rms_norm_eps,

70 activation=config.hidden_act,

70 - quant_config=quant_config)

71 + quant_config=quant_config,

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

73 + chunk_size=config.mamba_chunk_size)

75 self.block_sparse_moe = None

76 if getattr(config, "num_local_experts", 0) > 0:

@@ -361,10 +364,15 @@ class GraniteMoeHybridModel(nn.Module):

364 ) -> torch.Tensor:

365

366 attn_metadata = get_forward_context().attn_metadata

364 - mamba2_metadata = prepare_mamba2_metadata(

365 - chunk_size=self.config.mamba_chunk_size,

366 - attn_metadata=attn_metadata,

367 - )

367 +

368 + if not envs.VLLM_USE_V1:

369 + mamba2_metadata = prepare_mamba2_metadata(

370 + chunk_size=self.config.mamba_chunk_size,

371 + attn_metadata=attn_metadata,

372 + )

373 + else:

374 + # v1 get mamba2_metadata from forward_context

375 + mamba2_metadata = None

376

377 if get_pp_group().is_first_rank:

378 if inputs_embeds is not None:

@@ -386,7 +394,9 @@ class GraniteMoeHybridModel(nn.Module):

394 num_attn += 1

395

396 layer_mamba_cache_params = None

389 - if isinstance(layer, GraniteMoeHybridMambaDecoderLayer):

397 + if isinstance(

398 + layer,

399 + GraniteMoeHybridMambaDecoderLayer) and mamba_cache_params:

400 layer_mamba_cache_params = mamba_cache_params.at_layer_idx(

401 i - num_attn)

402

@@ -501,8 +511,7 @@ class GraniteMoeHybridModel(nn.Module):

511

512

513 class GraniteMoeHybridForCausalLM(nn.Module, HasInnerState, SupportsLoRA,

504 - SupportsPP, IsHybrid, SupportsV0Only,

505 - SupportsQuant):

514 + SupportsPP, IsHybrid, SupportsQuant):

515 packed_modules_mapping = {

516 "qkv_proj": [

517 "q_proj",

@@ -571,14 +580,20 @@ class GraniteMoeHybridForCausalLM(nn.Module, HasInnerState, SupportsLoRA,

580 intermediate_tensors: Optional[IntermediateTensors] = None,

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

582 **kwargs):

574 - if self.mamba_cache is None:

575 - num_mamba_layers = self.model_config.get_num_layers_by_block_type(

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

577 - self.mamba_cache = MambaCacheManager(

578 - self.vllm_config, self.model_config.dtype, num_mamba_layers,

579 - *self._get_mamba_cache_shape())

580 -

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

583 +

584 + mamba_cache_params = None

585 + if not envs.VLLM_USE_V1:

586 + if self.mamba_cache is None:

587 + num_mamba_layers = (

588 + self.model_config.get_num_layers_by_block_type(

589 + self.vllm_config.parallel_config,

590 + LayerBlockType.mamba))

591 + self.mamba_cache = MambaCacheManager(

592 + self.vllm_config, self.model_config.dtype,

593 + num_mamba_layers, *self._get_mamba_cache_shape())

594 +

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

596 +

597 hidden_states = self.model(input_ids, positions, mamba_cache_params,

598 intermediate_tensors, inputs_embeds)

599

vllm/model_executor/models/nemotron_h.py

+30 -15

@@ -23,6 +23,7 @@ from typing import Optional

23 import torch

24 from torch import nn

26 +from vllm import envs

27 from vllm.attention.layer import Attention

28 from vllm.config import CacheConfig, VllmConfig

29 from vllm.distributed import divide, get_tensor_model_parallel_world_size

@@ -44,8 +45,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

45 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

46 from vllm.model_executor.models.interfaces import (HasInnerState, IsHybrid,

47 SupportsLoRA, SupportsPP,

47 - SupportsQuant,

48 - SupportsV0Only)

48 + SupportsQuant)

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

50 MambaCacheParams)

51 from vllm.model_executor.models.utils import (

@@ -153,6 +153,8 @@ class NemotronHMambaDecoderLayer(nn.Module):

153 rms_norm_eps=config.rms_norm_eps,

154 activation=config.mamba_hidden_act,

155 quant_config=quant_config,

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

157 + chunk_size=config.chunk_size,

158 )

159

160 self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

@@ -348,10 +350,14 @@ class NemotronHModel(nn.Module):

350

351 attn_metadata = get_forward_context().attn_metadata

352

351 - mamba2_metadata = prepare_mamba2_metadata(

352 - chunk_size=self.config.chunk_size,

353 - attn_metadata=attn_metadata,

354 - )

353 + if not envs.VLLM_USE_V1:

354 + mamba2_metadata = prepare_mamba2_metadata(

355 + chunk_size=self.config.chunk_size,

356 + attn_metadata=attn_metadata,

357 + )

358 + else:

359 + # v1 get mamba2_metadata from forward_context

360 + mamba2_metadata = None

361

362 if get_pp_group().is_first_rank:

363 if inputs_embeds is not None:

@@ -369,7 +375,8 @@ class NemotronHModel(nn.Module):

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

376 layer = self.layers[i]

377 layer_mamba_cache_params = None

372 - if isinstance(layer, NemotronHMambaDecoderLayer):

378 + if isinstance(layer,

379 + NemotronHMambaDecoderLayer) and mamba_cache_params:

380 layer_mamba_cache_params = mamba_cache_params.at_layer_idx(

381 i - num_non_mamba_layers)

382 else:

@@ -437,7 +444,7 @@ class NemotronHModel(nn.Module):

444

445

446 class NemotronHForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

440 - IsHybrid, SupportsV0Only, SupportsQuant):

447 + IsHybrid, SupportsQuant):

448 packed_modules_mapping = {

449 "qkv_proj": [

450 "q_proj",

@@ -499,15 +506,23 @@ class NemotronHForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

506 intermediate_tensors: Optional[IntermediateTensors] = None,

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

508 **kwargs):

502 - if self.mamba_cache is None:

509

504 - num_mamba_layers = self.model_config.get_num_layers_by_block_type(

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

510 + mamba_cache_params = None

511 + if not envs.VLLM_USE_V1:

512 + if self.mamba_cache is None:

513 +

514 + num_mamba_layers = \

515 + self.model_config.get_num_layers_by_block_type(

516 + self.vllm_config.parallel_config,

517 + LayerBlockType.mamba

518 + )

519 +

520 + self.mamba_cache = MambaCacheManager(

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

522 + num_mamba_layers, *self._get_mamba_cache_shape())

523 +

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

525

507 - self.mamba_cache = MambaCacheManager(

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

509 - *self._get_mamba_cache_shape())

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

526 hidden_states = self.model(input_ids, positions, mamba_cache_params,

527 intermediate_tensors, inputs_embeds)

528

vllm/model_executor/models/zamba2.py

+64 -37

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

15 from torch import nn

16 from transformers import Zamba2Config

18 +from vllm import envs

19 from vllm.attention.layer import Attention

20 from vllm.config import CacheConfig, VllmConfig

21 from vllm.distributed import divide, get_tensor_model_parallel_world_size

@@ -41,7 +42,7 @@ from vllm.model_executor.models.mamba_cache import (MambaCacheManager,

42 from vllm.model_executor.sampling_metadata import SamplingMetadata

43 from vllm.sequence import IntermediateTensors

44 -from .interfaces import HasInnerState, IsHybrid, SupportsV0Only

45 +from .interfaces import HasInnerState, IsHybrid

46 from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix

@@ -58,6 +59,7 @@ class Zamba2LoRA(nn.Module):

59 rank: int,

60 output_dim: Union[int, list[int]],

61 quant_config: Optional[QuantizationConfig] = None,

62 + prefix: str = "",

63 ):

64 """Initialize the attention layer.

@@ -283,6 +285,7 @@ class Zamba2MLP(nn.Module):

285 bare_block_idx: int,

286 num_hybrid_layers: dict[int, int],

287 quant_config: Optional[QuantizationConfig] = None,

288 + prefix: str = "",

289 ) -> None:

290 """Initialize the MLP layer.

291

@@ -471,11 +474,10 @@ class Zamba2MambaDecoderLayer(nn.Module):

474 computation depending on configuration.

475 """

476

474 - def __init__(

475 - self,

476 - config: Zamba2Config,

477 - quant_config: Optional[QuantizationConfig] = None,

478 - ) -> None:

477 + def __init__(self,

478 + config: Zamba2Config,

479 + quant_config: Optional[QuantizationConfig] = None,

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

481 """Initialize the Mamba decoder layer.

482

483 Args:

@@ -486,20 +488,21 @@ class Zamba2MambaDecoderLayer(nn.Module):

488

489 # Initialize Mamba mixer with expanded intermediate size

490 intermediate_size = config.mamba_expand * config.hidden_size

489 - self.mamba = MambaMixer2(

490 - hidden_size=config.hidden_size,

491 - ssm_state_size=config.mamba_d_state,

492 - conv_kernel_size=config.mamba_d_conv,

493 - intermediate_size=intermediate_size,

494 - use_conv_bias=config.use_conv_bias,

495 - use_bias=config.add_bias_linear,

496 - n_groups=config.mamba_ngroups,

497 - num_heads=config.n_mamba_heads,

498 - head_dim=intermediate_size // config.n_mamba_heads,

499 - rms_norm_eps=config.rms_norm_eps,

500 - activation="silu",

501 - quant_config=quant_config,

502 - )

491 + self.mamba = MambaMixer2(hidden_size=config.hidden_size,

492 + ssm_state_size=config.mamba_d_state,

493 + conv_kernel_size=config.mamba_d_conv,

494 + intermediate_size=intermediate_size,

495 + use_conv_bias=config.use_conv_bias,

496 + use_bias=config.add_bias_linear,

497 + n_groups=config.mamba_ngroups,

498 + num_heads=config.n_mamba_heads,

499 + head_dim=intermediate_size //

500 + config.n_mamba_heads,

501 + rms_norm_eps=config.rms_norm_eps,

502 + activation="silu",

503 + quant_config=quant_config,

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

505 + chunk_size=config.chunk_size)

506

507 # Input normalization

508 self.input_layernorm = RMSNorm(config.hidden_size,

@@ -573,6 +576,7 @@ class Zamba2HybridLayer(nn.Module):

576 config: Zamba2Config,

577 block_idx: int,

578 quant_config: Optional[QuantizationConfig] = None,

579 + prefix: str = "",

580 ) -> None:

581 """Initialize the hybrid layer.

582

@@ -589,7 +593,8 @@ class Zamba2HybridLayer(nn.Module):

593 bias=False,

594 quant_config=quant_config)

595 self.mamba_decoder = Zamba2MambaDecoderLayer(config,

592 - quant_config=quant_config)

596 + quant_config=quant_config,

597 + prefix=prefix)

598

599 def forward(

600 self,

@@ -699,14 +704,23 @@ class Zamba2Model(nn.Module):

704 # Initialize layers according to block type configuration

705 layers = []

706 for layer_idx, layer_type in enumerate(config.layers_block_type):

707 + # tdoublep: avoid layers getting same index

708 + # somewhat hacky but correct (I think)

709 + prefix = str(len(layer2block_map) + layer_idx)

710 if layer_type == "hybrid":

711 block = next(blocks)

712 block_idx = layer2block_map[layer_idx]

713 layers.append(

706 - Zamba2HybridLayer(block, config, block_idx, quant_config))

714 + Zamba2HybridLayer(block,

715 + config,

716 + block_idx,

717 + quant_config,

718 + prefix=prefix))

719 else:

720 layers.append(

709 - Zamba2MambaDecoderLayer(config, quant_config=quant_config))

721 + Zamba2MambaDecoderLayer(config,

722 + quant_config=quant_config,

723 + prefix=prefix))

724 self.layers = nn.ModuleList(layers)

725

726 # Final layer normalization

@@ -751,19 +765,30 @@ class Zamba2Model(nn.Module):

765

766 attn_metadata = get_forward_context().attn_metadata

767

754 - mamba2_metadata = prepare_mamba2_metadata(

755 - chunk_size=self.config.chunk_size,

756 - attn_metadata=attn_metadata,

757 - )

768 + if not envs.VLLM_USE_V1:

769 + mamba2_metadata = prepare_mamba2_metadata(

770 + chunk_size=self.config.chunk_size,

771 + attn_metadata=attn_metadata,

772 + )

773 + else:

774 + # v1 get mamba2_metadata from forward_context

775 + mamba2_metadata = None

776

777 # Process through layers

778 original_hidden_states = torch.clone(hidden_states)

779 for layer_idx, layer in enumerate(self.layers):

780 +

781 + layer_mamba_cache_params = None

782 + if (isinstance(layer, (Zamba2HybridLayer, Zamba2MambaDecoderLayer))

783 + and mamba_cache_params):

784 + layer_mamba_cache_params = mamba_cache_params.at_layer_idx(

785 + layer_idx)

786 +

787 layer_outputs = layer(

788 hidden_states,

789 original_hidden_states=original_hidden_states,

790 positions=positions,

766 - mamba_cache_params=mamba_cache_params.at_layer_idx(layer_idx),

791 + mamba_cache_params=layer_mamba_cache_params,

792 mamba2_metadata=mamba2_metadata,

793 )

794 hidden_states = layer_outputs

@@ -803,7 +828,7 @@ class Zamba2Model(nn.Module):

828 return loaded_params

829

830

806 -class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsV0Only):

831 +class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid):

832 """Zamba2 model with causal language modeling head.

833

834 This class wraps the core Zamba2 model and adds:

@@ -897,14 +922,16 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsV0Only):

922 Output hidden states

923 """

924 # Initialize Mamba cache if needed

900 - if self.mamba_cache is None:

901 - num_mamba_layers = self.config.num_hidden_layers

902 - self.mamba_cache = MambaCacheManager(

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

904 - *self._get_mamba_cache_shape())

905 -

906 - # Get cache parameters for current run

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

925 + mamba_cache_params = None

926 + if not envs.VLLM_USE_V1:

927 + if self.mamba_cache is None:

928 + num_mamba_layers = self.config.num_hidden_layers

929 + self.mamba_cache = MambaCacheManager(

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

931 + num_mamba_layers, *self._get_mamba_cache_shape())

932 +

933 + # Get cache parameters for current run

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

935

936 # Forward pass through model

937 hidden_states = self.model(

vllm/v1/core/kv_cache_coordinator.py

+8 -3

@@ -27,6 +27,7 @@ class KVCacheCoordinator(ABC):

27 ):

28 self.kv_cache_config = kv_cache_config

29 self.max_model_len = max_model_len

30 + self.enable_caching = enable_caching

32 self.block_pool = BlockPool(kv_cache_config.num_blocks, enable_caching,

33 enable_kv_cache_events)

@@ -268,9 +269,13 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

269

270 self.full_attention_block_size = self.full_attention_spec.block_size

271 self.other_block_size = self.other_spec.block_size

271 - assert self.other_block_size % self.full_attention_block_size == 0, (

272 - "KVCacheCoordinator assumes the block_size of full attention "

273 - "layers is divisible by other layers now.")

272 +

273 + if self.enable_caching:

274 + # this requirement is only needed for the prefix caching logic

275 + divisible = self.other_block_size % self.full_attention_block_size

276 + assert divisible == 0, (

277 + "KVCacheCoordinator assumes the block_size of full "

278 + "attention layers is divisible by other layers now.")

279

280 if max(self.full_attention_group_ids) < min(self.other_group_ids):

281 self.full_attn_first = True

vllm/v1/core/kv_cache_manager.py

+10 -6

@@ -84,12 +84,15 @@ class KVCacheManager:

84 self.log_stats = log_stats

85 # FIXME: make prefix cache stats conditional on log_stats

86 self.prefix_cache_stats = PrefixCacheStats() if log_stats else None

87 - assert len(

88 - set(g.kv_cache_spec.block_size

89 - for g in kv_cache_config.kv_cache_groups)

90 - ) == 1, "Only one block size is supported for now"

91 - self.block_size = kv_cache_config.kv_cache_groups[

92 - 0].kv_cache_spec.block_size

87 +

88 + self.block_size: Optional[int] = None

89 + if self.enable_caching:

90 + assert len(

91 + set(g.kv_cache_spec.block_size

92 + for g in kv_cache_config.kv_cache_groups)

93 + ) == 1, "Only one block size is supported for now"

94 + self.block_size = kv_cache_config.kv_cache_groups[

95 + 0].kv_cache_spec.block_size

97 self.coordinator = get_kv_cache_coordinator(

98 kv_cache_config=kv_cache_config,

@@ -154,6 +157,7 @@ class KVCacheManager:

157 # if the scheduler has tried to schedule the request before.

158 block_hashes = self.req_to_block_hashes[request.request_id]

159 if not block_hashes:

160 + assert self.block_size is not None

161 block_hashes = hash_request_tokens(self.caching_hash_fn,

162 self.block_size, request)

163 self.req_to_block_hashes[request.request_id] = block_hashes

vllm/v1/core/kv_cache_utils.py

+4 -2

@@ -864,9 +864,11 @@ def _get_kv_cache_config_uniform_page_size(

864 kv_cache_groups=kv_cache_groups,

865 )

866

867 + min_block_size = min(

868 + [group.kv_cache_spec.block_size for group in kv_cache_groups])

869 +

870 # Print the KV cache size and maximum concurrency.

868 - num_tokens = num_blocks // len(

869 - grouped_layers) * vllm_config.cache_config.block_size

871 + num_tokens = num_blocks // len(grouped_layers) * min_block_size

872 num_tokens_str = f"{num_tokens:,}"

873 logger.info("GPU KV cache size: %s tokens", num_tokens_str)

874 max_model_len_str = f"{vllm_config.model_config.max_model_len:,}"

vllm/v1/kv_cache_interface.py

+6 -1

@@ -159,6 +159,7 @@ class SlidingWindowSpec(AttentionSpec):

159 class MambaSpec(KVCacheSpec):

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

161 dtype: torch.dtype

162 + page_size_padded: Optional[int] = None

163

164 def __post_init__(self):

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

@@ -169,7 +170,11 @@ class MambaSpec(KVCacheSpec):

170

171 @property

172 def page_size_bytes(self) -> int:

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

173 + page_size = self.num_elements * get_dtype_size(self.dtype)

174 + if self.page_size_padded is not None:

175 + assert self.page_size_padded >= page_size

176 + return self.page_size_padded

177 + return page_size

178

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

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

vllm/v1/worker/gpu_model_runner.py

+115 -10

@@ -334,6 +334,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):

334 # the same order of requests. We ensure this by only allowing the first

335 # group to reorder the batch and asserting that all other groups do not

336 # reorder the batch.

337 + # TODO(tdoublep): make this more flexible so that any group can

338 + # re-order the batch (not only the first).

339 + # TODO(tdoublep): verify this during engine init instead of at runtime

340 for i in range(1, len(self.kv_cache_config.kv_cache_groups)):

341 batch_reordered = self.attn_metadata_builders[i].reorder_batch(

342 self.input_batch, scheduler_output)

@@ -2449,6 +2452,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

2452 corresponding memory buffer for KV cache.

2453 """

2454 kv_caches: dict[str, torch.Tensor] = {}

2455 + has_attn, has_mamba = False, False

2456 for i, kv_cache_group_spec in enumerate(

2457 kv_cache_config.kv_cache_groups):

2458 kv_cache_spec = kv_cache_group_spec.kv_cache_spec

@@ -2458,6 +2462,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

2462 num_blocks = (raw_tensor.numel() //

2463 kv_cache_spec.page_size_bytes)

2464 if isinstance(kv_cache_spec, AttentionSpec):

2465 + has_attn = True

2466 kv_cache_shape = self.attn_backends[i].get_kv_cache_shape(

2467 num_blocks, kv_cache_spec.block_size,

2468 kv_cache_spec.num_kv_heads, kv_cache_spec.head_size)

@@ -2486,25 +2491,67 @@ class GPUModelRunner(LoRAModelRunnerMixin):

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

2492 *inv_order)

2493 elif isinstance(kv_cache_spec, MambaSpec):

2494 + has_mamba = True

2495 raw_tensor = kv_cache_raw_tensors[layer_name]

2496 dtype = kv_cache_spec.dtype

2497 + num_element_per_page = (kv_cache_spec.page_size_bytes //

2498 + get_dtype_size(dtype))

2499 state_tensors = []

2492 - start_pos = 0

2500 + storage_offset = 0

2501 for shape in kv_cache_spec.shapes:

2502 target_shape = (num_blocks, *shape)

2495 - size_in_bytes = np.prod(shape) * get_dtype_size(

2496 - dtype) * num_blocks

2497 - tensor = raw_tensor[start_pos:start_pos +

2498 - size_in_bytes]

2499 - tensor = tensor.view(dtype).view(target_shape)

2503 + stride = torch.empty(target_shape).stride()

2504 + target_stride = (num_element_per_page, *stride[1:])

2505 + tensor = torch.as_strided(

2506 + raw_tensor.view(dtype),

2507 + size=target_shape,

2508 + stride=target_stride,

2509 + storage_offset=storage_offset,

2510 + )

2511 state_tensors.append(tensor)

2501 - start_pos += size_in_bytes

2502 - assert start_pos == raw_tensor.numel()

2503 - kv_caches[layer_name] = tuple(state_tensors)

2512 + storage_offset += stride[0]

2513 +

2514 + kv_caches[layer_name] = state_tensors

2515 else:

2516 raise NotImplementedError

2517 +

2518 + if has_attn and has_mamba:

2519 + self._verify_hybrid_attention_mamba_layout(kv_cache_config,

2520 + kv_cache_raw_tensors)

2521 +

2522 return kv_caches

2523

2524 + def _verify_hybrid_attention_mamba_layout(

2525 + self, kv_cache_config: KVCacheConfig,

2526 + kv_cache_raw_tensors: dict[str, torch.Tensor]) -> None:

2527 + """

2528 + Verify that the KV cache memory layout is compatible for

2529 + models with both attention and mamba KV cache groups.

2530 +

2531 + Args:

2532 + kv_cache_config: The KV cache config

2533 + kv_cache_raw_tensors: The KV cache buffer of each layer.

2534 + """

2535 +

2536 + for i, kv_cache_group_spec in enumerate(

2537 + kv_cache_config.kv_cache_groups):

2538 + kv_cache_spec = kv_cache_group_spec.kv_cache_spec

2539 + for layer_name in kv_cache_group_spec.layer_names:

2540 + raw_tensor = kv_cache_raw_tensors[layer_name]

2541 + num_blocks = (raw_tensor.numel() //

2542 + kv_cache_spec.page_size_bytes)

2543 + if isinstance(kv_cache_spec, AttentionSpec):

2544 + kv_cache_shape = self.attn_backends[i].get_kv_cache_shape(

2545 + num_blocks, kv_cache_spec.block_size,

2546 + kv_cache_spec.num_kv_heads, kv_cache_spec.head_size)

2547 + if kv_cache_shape[0] != num_blocks or kv_cache_shape[

2548 + 1] != 2:

2549 + raise ValueError(

2550 + "Hybrid models in V1 require an attention "

2551 + "backend with kv_cache_shape="

2552 + "(num_blocks, 2, ...). Please try setting "

2553 + "VLLM_ATTENTION_BACKEND=FLASHINFER")

2554 +

2555 def initialize_kv_cache_tensors(

2556 self, kv_cache_config: KVCacheConfig) -> dict[str, torch.Tensor]:

2557 """

@@ -2623,11 +2670,69 @@ class GPUModelRunner(LoRAModelRunnerMixin):

2670 raise NotImplementedError(

2671 "Prefix caching is not supported for Mamba yet.")

2672 max_model_len = self.vllm_config.model_config.max_model_len

2673 +

2674 + page_size_padded = self._maybe_pad_mamba_page_size(

2675 + attn_layers, mamba_layers, kv_cache_spec, max_model_len,

2676 + block_size)

2677 +

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

2679 # have only one block in the KV cache.

2680 for layer_name, mamba_module in mamba_layers.items():

2681 kv_cache_spec[layer_name] = MambaSpec(

2682 shapes=mamba_module.get_state_shape(),

2683 dtype=self.kv_cache_dtype,

2632 - block_size=max_model_len)

2684 + block_size=max_model_len,

2685 + page_size_padded=page_size_padded)

2686 +

2687 return kv_cache_spec

2688 +

2689 + def _maybe_pad_mamba_page_size(

2690 + self,

2691 + attn_layers: dict[str, Attention],

2692 + mamba_layers: dict[str, MambaMixer2],

2693 + kv_cache_spec: dict[str, KVCacheSpec],

2694 + max_model_len: int,

2695 + block_size: int,

2696 + ) -> Optional[int]:

2697 + """

2698 + Ensure that page size of attention KV cache groups is greater than or

2699 + equal to the mamba KV cache groups. If not, we suggest to the user

2700 + how to set the attention block size to ensure that it is.

2701 +

2702 + If the attention page size is strictly greater than the mamba page size,

2703 + we pad the mamba page size to make them equal.

2704 +

2705 + Args:

2706 + attn_layers: Attention layers

2707 + mamba_layers: Mamba layers

2708 + kv_cache_spec: KV cache spec (populated with attention layers)

2709 +

2710 + Returns:

2711 + Optional[int]: Mamba page size with padding (None if no padding).

2712 + """

2713 +

2714 + if len(attn_layers) == 0:

2715 + return None

2716 +

2717 + attn_layer_name = next(iter(attn_layers))

2718 + attn_page_size = kv_cache_spec[attn_layer_name].page_size_bytes

2719 + mamba_layer_name = next(iter(mamba_layers))

2720 + mamba_page_size = MambaSpec(

2721 + shapes=mamba_layers[mamba_layer_name].get_state_shape(),

2722 + dtype=self.kv_cache_dtype,

2723 + block_size=max_model_len).page_size_bytes

2724 + if attn_page_size < mamba_page_size:

2725 + # attention page size (for 16 tokens)

2726 + attn_page_size_16 = 16 * attn_page_size // block_size

2727 + # some attention backends (e.g. FA) only support setting

2728 + # block size to multiple of 16, so let's suggest a value

2729 + # that would work (note: FA is currently not compatible

2730 + # with mamba layers, use FlashInfer instead).

2731 + suggest_attn_block_size = 16 * cdiv(mamba_page_size,

2732 + attn_page_size_16)

2733 + raise ValueError(

2734 + "Attention block size should be increased to at least "

2735 + f"{suggest_attn_block_size} in order to match "

2736 + "the mamba page size")

2737 +

2738 + return attn_page_size