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

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, 01:00 UTC

[V1] [Hybrid] Enable piecewise CUDA Graph for mamba layers (#21194)

Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>

Thomas Parnell committed Jul 19, 2025 at 21:27 UTC 881e3cbe3b3cef5d6fc50ca0c19e30a9dd11c452

10 files changed +100 -31

tests/models/language/generation/test_hybrid.py

-1

@@ -104,7 +104,6 @@ def test_models(

104 m.setenv("VLLM_ATTENTION_BACKEND", "FLASHINFER")

105 with vllm_runner(model,

106 max_num_seqs=MAX_NUM_SEQS,

107 - enforce_eager=True,

107 enable_prefix_caching=False) as vllm_model:

108 vllm_v1_outputs = vllm_model.generate_greedy_logprobs(

109 example_prompts, max_tokens, num_logprobs)

vllm/config.py

@@ -4312,6 +4312,7 @@ class CompilationConfig:

4312 self.splitting_ops = [] if self.full_cuda_graph else [

4313 "vllm.unified_attention",

4314 "vllm.unified_attention_with_output",

4315 + "vllm.mamba_mixer2",

4316 ]

4317

4318

vllm/model_executor/layers/mamba/mamba_mixer2.py

+66 -9

@@ -13,7 +13,7 @@ from vllm.distributed import (divide, get_tensor_model_parallel_rank,

13 get_tensor_model_parallel_world_size,

14 tensor_model_parallel_all_gather,

15 tensor_model_parallel_all_reduce)

16 -from vllm.forward_context import get_forward_context

16 +from vllm.forward_context import ForwardContext, get_forward_context

17 from vllm.model_executor.custom_op import CustomOp

18 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

19 RowParallelLinear)

@@ -33,6 +33,8 @@ from vllm.model_executor.model_loader.weight_utils import (

33 LoaderFunction, composed_weight_loader, sharded_weight_loader)

34 from vllm.model_executor.models.mamba_cache import MambaCacheParams

35 from vllm.model_executor.utils import set_weight_attrs

36 +from vllm.platforms import current_platform

37 +from vllm.utils import direct_register_custom_op

38 from vllm.v1.attention.backends.mamba_attn import Mamba2AttentionMetadata

40 # Added by the IBM Team, 2024

@@ -424,14 +426,36 @@ class MambaMixer2(MambaBase, CustomOp):

426 def forward_native(

427 self,

428 hidden_states: torch.Tensor,

427 - conv_state: torch.Tensor,

428 - ssm_state: torch.Tensor,

429 + output: torch.Tensor,

430 + mamba_cache_params: MambaCacheParams,

431 + mamba2_metadata: Mamba2Metadata,

432 + mup_vector: Optional[torch.Tensor] = None,

433 ):

434 pass

435

436 + def forward(

437 + self,

438 + hidden_states: torch.Tensor,

439 + output: torch.Tensor,

440 + mamba_cache_params: MambaCacheParams,

441 + mamba2_metadata: Mamba2Metadata,

442 + mup_vector: Optional[torch.Tensor] = None,

443 + ):

444 + if not envs.VLLM_USE_V1:

445 + CustomOp.forward(self, hidden_states, output, mamba_cache_params,

446 + mamba2_metadata, mup_vector)

447 + else:

448 + torch.ops.vllm.mamba_mixer2(

449 + hidden_states,

450 + output,

451 + self.prefix,

452 + mup_vector,

453 + )

454 +

455 def forward_cuda(

456 self,

457 hidden_states: torch.Tensor,

458 + output: torch.Tensor,

459 mamba_cache_params: MambaCacheParams,

460 mamba2_metadata: Mamba2Metadata,

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

@@ -517,6 +541,7 @@ class MambaMixer2(MambaBase, CustomOp):

541 num_prefill_tokens = attn_metadata.num_prefill_tokens # token count

542 has_prefill = num_prefills > 0

543 has_decode = num_decodes > 0

544 + num_actual_tokens = num_prefill_tokens + num_decodes

545

546 # NOTE: V0 put prefill before decode, v1 puts decode before prefill

547 # Separate prefill and decode by splitting varlen input

@@ -524,18 +549,18 @@ class MambaMixer2(MambaBase, CustomOp):

549 # NOTE: V0 put prefill before decode, v1 puts decode before prefill

550 if envs.VLLM_USE_V1:

551 hidden_states_B_C_d, hidden_states_B_C_p = torch.split(

527 - hidden_states_B_C,

552 + hidden_states_B_C[:num_actual_tokens],

553 [num_decodes, num_prefill_tokens],

554 dim=0,

555 )

556 dt_d, dt_p = torch.split(

532 - dt,

557 + dt[:num_actual_tokens],

558 [num_decodes, num_prefill_tokens],

559 dim=0,

560 )

561 # Split along batch dimension

562 state_indices_tensor_d, state_indices_tensor_p = torch.split(

538 - state_indices_tensor,

563 + state_indices_tensor[:num_actual_tokens],

564 [num_decodes, num_prefills],

565 dim=0,

566 )

@@ -696,11 +721,10 @@ class MambaMixer2(MambaBase, CustomOp):

721 # GatedRMSNorm internally applying SiLU to the gate

722 # SiLU is applied internally before normalization, unlike standard

723 # norm usage

699 - hidden_states = self.norm(hidden_states, gate)

724 + hidden_states = self.norm(hidden_states, gate[:num_actual_tokens])

725

726 # 5. Final linear projection

702 - out, _ = self.out_proj(hidden_states)

703 - return out

727 + output[:num_actual_tokens], _ = self.out_proj(hidden_states)

728

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

730 return get_mamba_state_shape(

@@ -712,3 +736,36 @@ class MambaMixer2(MambaBase, CustomOp):

736 state_size=self.ssm_state_size,

737 conv_kernel=self.conv_kernel_size,

738 )

739 +

740 +

741 +def mamba_mixer2(

742 + hidden_states: torch.Tensor,

743 + output: torch.Tensor,

744 + layer_name: str,

745 + mup_vector: Optional[torch.Tensor] = None,

746 +) -> None:

747 + forward_context: ForwardContext = get_forward_context()

748 + self = forward_context.no_compile_layers[layer_name]

749 + self.forward_cuda(hidden_states=hidden_states,

750 + output=output,

751 + mamba_cache_params=None,

752 + mamba2_metadata=None,

753 + mup_vector=mup_vector)

754 +

755 +

756 +def mamba_mixer2_fake(

757 + hidden_states: torch.Tensor,

758 + output: torch.Tensor,

759 + layer_name: str,

760 + mup_vector: Optional[torch.Tensor] = None,

761 +) -> None:

762 + return

763 +

764 +

765 +direct_register_custom_op(

766 + op_name="mamba_mixer2",

767 + op_func=mamba_mixer2,

768 + mutates_args=["output"],

769 + fake_impl=mamba_mixer2_fake,

770 + dispatch_key=current_platform.dispatch_key,

771 +)

vllm/model_executor/models/bamba.py

+6 -5

@@ -11,6 +11,7 @@ from transformers import BambaConfig

12 from vllm import envs

13 from vllm.attention.layer import Attention

14 +from vllm.compilation.decorators import support_torch_compile

15 from vllm.config import CacheConfig, VllmConfig

16 from vllm.distributed import get_tensor_model_parallel_world_size

17 from vllm.distributed.parallel_state import get_pp_group

@@ -122,11 +123,10 @@ class BambaMixerDecoderLayer(nn.Module):

123 hidden_states, residual = self.input_layernorm(

124 hidden_states, residual)

125

125 - hidden_states = self.mamba(hidden_states, mamba_cache_params,

126 - mamba2_metadata)

126 + output = torch.empty_like(hidden_states)

127 + self.mamba(hidden_states, output, mamba_cache_params, mamba2_metadata)

128 # Fully Connected

128 - hidden_states, residual = self.pre_ff_layernorm(

129 - hidden_states, residual)

129 + hidden_states, residual = self.pre_ff_layernorm(output, residual)

130 hidden_states = self.feed_forward(hidden_states)

131 return hidden_states, residual

132

@@ -169,7 +169,7 @@ class BambaAttentionDecoderLayer(nn.Module):

169 self.max_position_embeddings = max_position_embeddings

170

171 if hasattr(config, "partial_rotary_factor"):

172 - rotary_dim = self.head_dim * config.partial_rotary_factor

172 + rotary_dim = int(self.head_dim * config.partial_rotary_factor)

173 elif hasattr(config, "attn_rotary_emb"):

174 rotary_dim = config.attn_rotary_emb # for backward compatibility

175 else:

@@ -258,6 +258,7 @@ ALL_DECODER_LAYER_TYPES = {

258 }

259

260

261 +@support_torch_compile

262 class BambaModel(nn.Module):

263

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

vllm/model_executor/models/falcon_h1.py

+6 -2

@@ -10,6 +10,7 @@ from transformers import FalconH1Config

11 from vllm import envs

12 from vllm.attention.layer import Attention

13 +from vllm.compilation.decorators import support_torch_compile

14 from vllm.config import CacheConfig, VllmConfig

15 from vllm.distributed import get_tensor_model_parallel_world_size

16 from vllm.distributed.parallel_state import get_pp_group

@@ -179,13 +180,15 @@ class FalconH1SSMDecoderLayer(nn.Module):

180 mamba2_metadata: Mamba2Metadata,

181 **kwargs,

182 ):

182 - hidden_states = self.mamba(

183 + output = torch.empty_like(hidden_states)

184 + self.mamba(

185 hidden_states,

186 + output,

187 mamba_cache_params,

188 mamba2_metadata=mamba2_metadata,

189 mup_vector=self.mup_vector,

190 )

188 - return hidden_states, residual

191 + return output, residual

192

193

194 class FalconH1AttentionDecoderLayer(nn.Module):

@@ -398,6 +401,7 @@ class FalconH1ParallelHybrid(nn.Module):

401 return hidden_states

402

403

404 +@support_torch_compile

405 class FalconH1Model(nn.Module):

406

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

vllm/model_executor/models/granitemoehybrid.py

+5 -3

@@ -11,6 +11,7 @@ from transformers import GraniteMoeHybridConfig

12 from vllm import envs

13 from vllm.attention.layer import Attention

14 +from vllm.compilation.decorators import support_torch_compile

15 from vllm.config import CacheConfig, VllmConfig

16 from vllm.distributed import get_tensor_model_parallel_world_size

17 from vllm.distributed.parallel_state import get_pp_group

@@ -104,9 +105,9 @@ class GraniteMoeHybridMambaDecoderLayer(nn.Module):

105 ):

106 residual = hidden_states

107 hidden_states = self.input_layernorm(hidden_states)

107 - hidden_states = self.mamba(hidden_states, mamba_cache_params,

108 - mamba2_metadata)

109 - hidden_states = residual + hidden_states * self.residual_multiplier

108 + output = torch.empty_like(hidden_states)

109 + self.mamba(hidden_states, output, mamba_cache_params, mamba2_metadata)

110 + hidden_states = residual + output * self.residual_multiplier

111

112 residual = hidden_states

113 hidden_states = self.post_attention_layernorm(hidden_states)

@@ -307,6 +308,7 @@ ALL_DECODER_LAYER_TYPES = {

308 }

309

310

311 +@support_torch_compile

312 class GraniteMoeHybridModel(nn.Module):

313

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

vllm/model_executor/models/mamba2.py

+5 -3

@@ -10,6 +10,7 @@ from transformers import MambaConfig

11 from vllm import envs

12 from vllm.attention.backends.abstract import AttentionMetadata

13 +from vllm.compilation.decorators import support_torch_compile

14 from vllm.config import VllmConfig

15 from vllm.distributed.parallel_state import get_pp_group

16 from vllm.forward_context import get_forward_context

@@ -79,11 +80,12 @@ class Mamba2DecoderLayer(nn.Module):

80 else:

81 hidden_states, residual = self.norm(hidden_states, residual)

82 - hidden_states = self.mixer(hidden_states, mamba_cache_params,

83 - mamba2_metadata)

84 - return hidden_states, residual

83 + output = torch.empty_like(hidden_states)

84 + self.mixer(hidden_states, output, mamba_cache_params, mamba2_metadata)

85 + return output, residual

88 +@support_torch_compile

89 class Mamba2Model(nn.Module):

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

vllm/model_executor/models/nemotron_h.py

+5 -3

@@ -25,6 +25,7 @@ from torch import nn

26 from vllm import envs

27 from vllm.attention.layer import Attention

28 +from vllm.compilation.decorators import support_torch_compile

29 from vllm.config import CacheConfig, VllmConfig

30 from vllm.distributed import get_tensor_model_parallel_world_size

31 from vllm.distributed.parallel_state import get_pp_group

@@ -172,9 +173,9 @@ class NemotronHMambaDecoderLayer(nn.Module):

173 else:

174 hidden_states, residual = self.norm(hidden_states, residual)

175

175 - hidden_states = self.mixer(hidden_states, mamba_cache_params,

176 - mamba2_metadata)

177 - return hidden_states, residual

176 + output = torch.empty_like(hidden_states)

177 + self.mixer(hidden_states, output, mamba_cache_params, mamba2_metadata)

178 + return output, residual

179

180

181 class NemotronHAttention(nn.Module):

@@ -292,6 +293,7 @@ ALL_DECODER_LAYER_TYPES = {

293 }

294

295

296 +@support_torch_compile

297 class NemotronHModel(nn.Module):

298

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

vllm/model_executor/models/zamba2.py

+6 -2

@@ -17,6 +17,7 @@ from transformers import Zamba2Config

18 from vllm import envs

19 from vllm.attention.layer import Attention

20 +from vllm.compilation.decorators import support_torch_compile

21 from vllm.config import CacheConfig, VllmConfig

22 from vllm.distributed import get_tensor_model_parallel_world_size

23 from vllm.forward_context import get_forward_context

@@ -548,14 +549,16 @@ class Zamba2MambaDecoderLayer(nn.Module):

549 hidden_states = self.input_layernorm(hidden_states)

550

551 # Process through Mamba mixer

551 - hidden_states = self.mamba(

552 + output = torch.empty_like(hidden_states)

553 + self.mamba(

554 hidden_states,

555 + output,

556 mamba_cache_params=mamba_cache_params,

557 mamba2_metadata=mamba2_metadata,

558 )

559

560 # residual connection after mamba

558 - hidden_states = residual + hidden_states

561 + hidden_states = residual + output

562

563 return hidden_states

564

@@ -646,6 +649,7 @@ class Zamba2HybridLayer(nn.Module):

649 return layer_outputs

650

651

652 +@support_torch_compile

653 class Zamba2Model(nn.Module):

654 """Core Zamba2 model combining transformer and Mamba architectures.

655

vllm/v1/worker/gpu_model_runner.py

-3

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

2753 if self.vllm_config.speculative_config is not None:

2754 raise NotImplementedError(

2755 "Mamba with speculative decoding is not supported yet.")

2756 - if not self.vllm_config.model_config.enforce_eager:

2757 - raise NotImplementedError(

2758 - "Mamba with cuda graph is not supported yet.")

2756 if self.vllm_config.cache_config.enable_prefix_caching:

2757 raise NotImplementedError(

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