47e605092b · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -19,6 +19,7 @@

19 """Gemma 4 model implementation for vLLM."""

20

21 from collections.abc import Iterable

22 +from dataclasses import replace

23 from itertools import islice

24

25 import regex as re

@@ -32,6 +33,7 @@ from vllm.distributed import (

33 get_tensor_model_parallel_rank,

34 get_tensor_model_parallel_world_size,

35 )

36 +from vllm.forward_context import get_forward_context

37 from vllm.logger import init_logger

38 from vllm.model_executor.layers.activation import GeluAndMul

39 from vllm.model_executor.layers.attention import Attention

@@ -56,6 +58,7 @@ from vllm.model_executor.model_loader.weight_utils import (

58 maybe_remap_kv_scale_name,

59 )

60 from vllm.sequence import IntermediateTensors

61 +from vllm.v1.attention.backends.utils import KVSharingFastPrefillMetadata

62

63 from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP

64 from .utils import (

@@ -636,7 +639,205 @@ class Gemma4DecoderLayer(nn.Module):

639 return hidden_states, None

640

641

639 -@support_torch_compile

642 +def _run_decoder_layers(

643 + decoder_layers: list[Gemma4DecoderLayer],

644 + layer_idx_start: int,

645 + positions: torch.Tensor,

646 + hidden_states: torch.Tensor,

647 + per_layer_inputs: torch.Tensor | None = None,

648 + **kwargs,

649 +) -> torch.Tensor:

650 + """Run a slice of decoder layers with PLE extraction."""

651 + residual = None

652 + for idx, layer in enumerate(decoder_layers):

653 + layer_idx = idx + layer_idx_start

654 + layer_per_input = (

655 + per_layer_inputs[:, layer_idx, :] if per_layer_inputs is not None else None

656 + )

657 + hidden_states, residual = layer(

658 + positions,

659 + hidden_states,

660 + residual,

661 + per_layer_input=layer_per_input,

662 + **kwargs,

663 + )

664 + return hidden_states

665 +

666 +

667 +@support_torch_compile(

668 + enable_if=lambda vllm_config: vllm_config.cache_config.kv_sharing_fast_prefill

669 +)

670 +class Gemma4SelfDecoderLayers(nn.Module):

671 + """Compiled wrapper: embedding + non-KV-shared layers (YOCO first half).

672 +

673 + Owns the embedding and PLE modules so they are inside the compiled

674 + graph. Gemma4Model delegates embedding methods here.

675 + """

676 +

677 + def __init__(

678 + self,

679 + *,

680 + vllm_config: VllmConfig,

681 + prefix: str = "",

682 + decoder_layers: list[Gemma4DecoderLayer],

683 + layer_idx_start: int,

684 + embed_tokens: VocabParallelEmbedding,

685 + normalizer: torch.Tensor,

686 + embed_tokens_per_layer: VocabParallelEmbedding | None,

687 + embed_scale_per_layer: torch.Tensor | None,

688 + per_layer_model_projection: ColumnParallelLinear | None,

689 + per_layer_projection_norm: RMSNorm | None,

690 + per_layer_input_scale: torch.Tensor | None,

691 + per_layer_projection_scale: torch.Tensor | None,

692 + ):

693 + super().__init__()

694 + self.decoder_layers = decoder_layers

695 + self.layer_idx_start = layer_idx_start

696 +

697 + config = _get_text_config(vllm_config.model_config.hf_config)

698 + self.config = config

699 + self.hidden_size_per_layer_input = getattr(

700 + config, "hidden_size_per_layer_input", 0

701 + )

702 + self.vocab_size_per_layer_input = getattr(

703 + config, "vocab_size_per_layer_input", config.vocab_size

704 + )

705 +

706 + # Shared references to modules owned by Gemma4Model — must be

707 + # inside this nn.Module so torch.compile captures them.

708 + self.embed_tokens = embed_tokens

709 + self.normalizer = normalizer

710 + self.embed_tokens_per_layer = embed_tokens_per_layer

711 + self.embed_scale_per_layer = embed_scale_per_layer

712 + self.per_layer_model_projection = per_layer_model_projection

713 + self.per_layer_projection_norm = per_layer_projection_norm

714 + self.per_layer_input_scale = per_layer_input_scale

715 + self.per_layer_projection_scale = per_layer_projection_scale

716 +

717 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

718 + return self.embed_tokens(input_ids) * self.normalizer

719 +

720 + def get_per_layer_inputs(self, input_ids: torch.Tensor) -> torch.Tensor | None:

721 + """Get per-layer embeddings from embed_tokens_per_layer.

722 +

723 + Returns:

724 + Per-layer embeddings (num_tokens, num_layers,

725 + hidden_size_per_layer_input)

726 + """

727 + if self.embed_tokens_per_layer is None:

728 + return None

729 + per_layer_inputs_mask = torch.logical_and(

730 + input_ids >= 0,

731 + input_ids < self.vocab_size_per_layer_input,

732 + )

733 + per_layer_inputs_tokens = torch.where(

734 + per_layer_inputs_mask, input_ids, torch.zeros_like(input_ids)

735 + )

736 + per_layer_embeds = self.embed_tokens_per_layer(per_layer_inputs_tokens)

737 + per_layer_embeds = per_layer_embeds * self.embed_scale_per_layer

738 + return per_layer_embeds.reshape(

739 + *input_ids.shape,

740 + self.config.num_hidden_layers,

741 + self.hidden_size_per_layer_input,

742 + )

743 +

744 + def project_per_layer_inputs(

745 + self,

746 + inputs_embeds: torch.Tensor,

747 + per_layer_inputs: torch.Tensor | None,

748 + ) -> torch.Tensor | None:

749 + """Project inputs_embeds and combine with per_layer_inputs.

750 +

751 + Steps:

752 + 1. Project inputs_embeds: hidden_size → total_ple_dim

753 + 2. Scale by hidden_size^{-0.5}

754 + 3. Reshape to (num_tokens, num_layers, per_layer_dim)

755 + 4. Normalize with per_layer_projection_norm

756 + 5. Combine: (projection + per_layer_inputs) * 1/sqrt(2)

757 + """

758 + if self.per_layer_model_projection is None:

759 + return None

760 + per_layer_projection = self.per_layer_model_projection(inputs_embeds)

761 + per_layer_projection = per_layer_projection * self.per_layer_projection_scale

762 + per_layer_projection = per_layer_projection.reshape(

763 + *inputs_embeds.shape[:-1],

764 + self.config.num_hidden_layers,

765 + self.hidden_size_per_layer_input,

766 + )

767 + per_layer_projection = self.per_layer_projection_norm(per_layer_projection)

768 + if per_layer_inputs is None:

769 + return per_layer_projection

770 + return (per_layer_projection + per_layer_inputs) * self.per_layer_input_scale

771 +

772 + def forward(

773 + self,

774 + input_ids: torch.Tensor | None,

775 + positions: torch.Tensor,

776 + inputs_embeds: torch.Tensor | None = None,

777 + per_layer_inputs: torch.Tensor | None = None,

778 + **kwargs,

779 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

780 + if inputs_embeds is not None:

781 + hidden_states = inputs_embeds

782 + per_layer_inputs = self.project_per_layer_inputs(

783 + hidden_states, per_layer_inputs

784 + )

785 + else:

786 + hidden_states = self.embed_input_ids(input_ids)

787 + per_layer_embeds = self.get_per_layer_inputs(input_ids)

788 + per_layer_inputs = self.project_per_layer_inputs(

789 + hidden_states, per_layer_embeds

790 + )

791 +

792 + hidden_states = _run_decoder_layers(

793 + self.decoder_layers,

794 + self.layer_idx_start,

795 + positions,

796 + hidden_states,

797 + per_layer_inputs,

798 + **kwargs,

799 + )

800 + return hidden_states, per_layer_inputs

801 +

802 +

803 +@support_torch_compile(

804 + enable_if=lambda vllm_config: vllm_config.cache_config.kv_sharing_fast_prefill

805 +)

806 +class Gemma4CrossDecoderLayers(nn.Module):

807 + """Cross-decoder layers (YOCO second half, KV-shared)."""

808 +

809 + def __init__(

810 + self,

811 + *,

812 + vllm_config: VllmConfig,

813 + prefix: str = "",

814 + decoder_layers: list[Gemma4DecoderLayer],

815 + layer_idx_start: int,

816 + ):

817 + super().__init__()

818 + self.decoder_layers = decoder_layers

819 + self.layer_idx_start = layer_idx_start

820 +

821 + def forward(

822 + self,

823 + positions: torch.Tensor,

824 + hidden_states: torch.Tensor,

825 + per_layer_inputs: torch.Tensor | None = None,

826 + **kwargs,

827 + ) -> torch.Tensor:

828 + return _run_decoder_layers(

829 + self.decoder_layers,

830 + self.layer_idx_start,

831 + positions,

832 + hidden_states,

833 + per_layer_inputs,

834 + **kwargs,

835 + )

836 +

837 +

838 +@support_torch_compile(

839 + enable_if=lambda vllm_config: not vllm_config.cache_config.kv_sharing_fast_prefill

840 +)

841 class Gemma4Model(nn.Module):

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

843 super().__init__()

@@ -740,6 +941,75 @@ class Gemma4Model(nn.Module):

941 torch.tensor(config.hidden_size**0.5),

942 persistent=False,

943 )

944 +

945 + # --- You Only Cache Once (YOCO) split for fast prefill ---

946 + first_kv_shared_layer_idx = config.num_hidden_layers - getattr(

947 + config, "num_kv_shared_layers", 0

948 + )

949 +

950 + from vllm.compilation.backends import set_model_tag

951 +

952 + # Layers 0..(K-1) are self-decoder layers in YOCO

953 + with set_model_tag("self_decoder"):

954 + self.self_decoder = Gemma4SelfDecoderLayers(

955 + vllm_config=vllm_config,

956 + prefix=f"{prefix}.self_decoder",

957 + decoder_layers=self.layers[:first_kv_shared_layer_idx],

958 + layer_idx_start=0,

959 + embed_tokens=self.embed_tokens,

960 + normalizer=self.normalizer,

961 + embed_tokens_per_layer=getattr(self, "embed_tokens_per_layer", None),

962 + embed_scale_per_layer=getattr(self, "embed_scale_per_layer", None),

963 + per_layer_model_projection=getattr(

964 + self, "per_layer_model_projection", None

965 + ),

966 + per_layer_projection_norm=getattr(

967 + self, "per_layer_projection_norm", None

968 + ),

969 + per_layer_input_scale=getattr(self, "per_layer_input_scale", None),

970 + per_layer_projection_scale=getattr(

971 + self, "per_layer_projection_scale", None

972 + ),

973 + )

974 + # Layers K..(N-1) are cross-decoder layers in YOCO

975 + with set_model_tag("cross_decoder"):

976 + self.cross_decoder = Gemma4CrossDecoderLayers(

977 + vllm_config=vllm_config,

978 + prefix=f"{prefix}.cross_decoder",

979 + decoder_layers=self.layers[first_kv_shared_layer_idx:],

980 + layer_idx_start=first_kv_shared_layer_idx,

981 + )

982 +

983 + self.fast_prefill_enabled = cache_config.kv_sharing_fast_prefill

984 +

985 + if self.fast_prefill_enabled:

986 + # Allocate static buffers for CUDAGraph

987 + max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens

988 + device = next(self.parameters()).device

989 + self.positions = torch.zeros(

990 + max_num_tokens, dtype=torch.int64, device=device

991 + )

992 + self.hidden_states = torch.zeros(

993 + (max_num_tokens, config.hidden_size),

994 + dtype=self.embed_tokens.weight.dtype,

995 + device=device,

996 + )

997 + if (

998 + self.hidden_size_per_layer_input

999 + and self.hidden_size_per_layer_input > 0

1000 + ):

1001 + self.per_layer_inputs = torch.zeros(

1002 + (

1003 + max_num_tokens,

1004 + config.num_hidden_layers,

1005 + self.hidden_size_per_layer_input,

1006 + ),

1007 + dtype=self.embed_tokens.weight.dtype,

1008 + device=device,

1009 + )

1010 + else:

1011 + self.per_layer_inputs = None

1012 +

1013 # Custom factory that includes per_layer_inputs for PLE-enabled PP.

1014 # per_layer_inputs has shape (batch, num_layers, per_layer_dim),

1015 # which differs from the standard (batch, hidden_size) shape,

@@ -776,47 +1046,22 @@ class Gemma4Model(nn.Module):

1046 self.make_empty_intermediate_tensors = _make_empty_intermediate_tensors

1047

1048 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

779 - return self.embed_tokens(input_ids) * self.normalizer

1049 + return self.self_decoder.embed_input_ids(input_ids)

1050

781 - def get_per_layer_inputs(self, input_ids: torch.Tensor) -> torch.Tensor:

1051 + def get_per_layer_inputs(self, input_ids: torch.Tensor) -> torch.Tensor | None:

1052 """Get per-layer embeddings from embed_tokens_per_layer.

1053

1054 Returns:

1055 Per-layer embeddings (num_tokens, num_layers,

1056 hidden_size_per_layer_input)

1057 """

788 - if self.embed_tokens_per_layer is None:

789 - return None

790 -

791 - # Handle out-of-vocab tokens for PLE (vocab_size_per_layer_input may

792 - # be smaller than the main vocab_size).

793 - per_layer_inputs_mask = torch.logical_and(

794 - input_ids >= 0,

795 - input_ids < self.vocab_size_per_layer_input,

796 - )

797 - per_layer_inputs_tokens = torch.where(

798 - per_layer_inputs_mask, input_ids, torch.zeros_like(input_ids)

799 - )

800 -

801 - # Get packed per-layer embeddings: (num_tokens, total_ple_dim)

802 - per_layer_embeds = self.embed_tokens_per_layer(per_layer_inputs_tokens)

803 -

804 - # Apply embed_scale (sqrt of per-layer hidden dim)

805 - per_layer_embeds = per_layer_embeds * self.embed_scale_per_layer

806 -

807 - # Reshape to (num_tokens, num_layers, hidden_size_per_layer_input)

808 - per_layer_embeds = per_layer_embeds.reshape(

809 - *input_ids.shape,

810 - self.config.num_hidden_layers,

811 - self.hidden_size_per_layer_input,

812 - )

813 - return per_layer_embeds

1058 + return self.self_decoder.get_per_layer_inputs(input_ids)

1059

1060 def project_per_layer_inputs(

1061 self,

1062 inputs_embeds: torch.Tensor,

1063 per_layer_inputs: torch.Tensor | None,

819 - ) -> torch.Tensor:

1064 + ) -> torch.Tensor | None:

1065 """Project inputs_embeds and combine with per_layer_inputs.

1066

1067 Steps:

@@ -826,29 +1071,94 @@ class Gemma4Model(nn.Module):

1071 4. Normalize with per_layer_projection_norm

1072 5. Combine: (projection + per_layer_inputs) * 1/sqrt(2)

1073 """

829 - if self.per_layer_model_projection is None:

830 - return None

1074 + return self.self_decoder.project_per_layer_inputs(

1075 + inputs_embeds, per_layer_inputs

1076 + )

1077

832 - # Project from hidden_size to total_ple_dim

833 - # Scaled projection: output = linear(input, weight) * scale

834 - per_layer_projection = self.per_layer_model_projection(inputs_embeds)

835 - per_layer_projection = per_layer_projection * self.per_layer_projection_scale

1078 + def fast_prefill_forward(

1079 + self,

1080 + input_ids: torch.Tensor | None,

1081 + positions: torch.Tensor,

1082 + inputs_embeds: torch.Tensor | None = None,

1083 + per_layer_inputs: torch.Tensor | None = None,

1084 + **kwargs,

1085 + ) -> torch.Tensor:

1086 + logits_indices_padded, num_logits_indices = None, None

1087 + attn_metadata = get_forward_context().attn_metadata

1088

837 - # Reshape to (num_tokens, num_layers, hidden_size_per_layer_input)

838 - per_layer_projection = per_layer_projection.reshape(

839 - *inputs_embeds.shape[:-1],

840 - self.config.num_hidden_layers,

841 - self.hidden_size_per_layer_input,

1089 + if attn_metadata is not None:

1090 + assert isinstance(attn_metadata, dict)

1091 + layer_attn_metadata = attn_metadata[

1092 + self.layers[-1].self_attn.attn.layer_name

1093 + ]

1094 + if isinstance(layer_attn_metadata, KVSharingFastPrefillMetadata):

1095 + logits_indices_padded = layer_attn_metadata.logits_indices_padded

1096 + num_logits_indices = layer_attn_metadata.num_logits_indices

1097 +

1098 + batch_size = positions.size(0)

1099 + self.positions[:batch_size].copy_(positions)

1100 + self_decoder_hidden_states, per_layer_inputs = self.self_decoder(

1101 + input_ids=input_ids,

1102 + positions=self.positions[:batch_size],

1103 + inputs_embeds=inputs_embeds,

1104 + per_layer_inputs=per_layer_inputs,

1105 + **kwargs,

1106 )

1107

844 - # Normalize

845 - per_layer_projection = self.per_layer_projection_norm(per_layer_projection)

1108 + if logits_indices_padded is None:

1109 + logits_indices_padded = torch.arange(

1110 + batch_size,

1111 + dtype=positions.dtype,

1112 + device=positions.device,

1113 + )

1114

847 - if per_layer_inputs is None:

848 - return per_layer_projection

1115 + # NOTE: Keep .clone() until fix in

1116 + # https://github.com/vllm-project/vllm/pull/22282

1117 + hidden_states = self_decoder_hidden_states.clone()

1118

850 - # Combine: (projection + per_layer_inputs) * scale

851 - return (per_layer_projection + per_layer_inputs) * self.per_layer_input_scale

1119 + num_padded = logits_indices_padded.size(0)

1120 + self.positions[:num_padded].copy_(positions[logits_indices_padded])

1121 + self.hidden_states[:num_padded].copy_(

1122 + self_decoder_hidden_states[logits_indices_padded]

1123 + )

1124 + if self.per_layer_inputs is not None and per_layer_inputs is not None:

1125 + self.per_layer_inputs[:num_padded].copy_(

1126 + per_layer_inputs[logits_indices_padded]

1127 + )

1128 +

1129 + # Update batch_descriptor so the cross-decoder's piecewise

1130 + # CUDAGraphWrapper dispatches to the correct (reduced) batch size.

1131 + forward_context = get_forward_context()

1132 + orig_batch_desc = forward_context.batch_descriptor

1133 + if orig_batch_desc is not None:

1134 + forward_context.batch_descriptor = replace(

1135 + orig_batch_desc, num_tokens=num_padded

1136 + )

1137 +

1138 + cross_per_layer = (

1139 + self.per_layer_inputs[:num_padded]

1140 + if self.per_layer_inputs is not None

1141 + else None

1142 + )

1143 + cross_hidden_states = self.cross_decoder(

1144 + self.positions[:num_padded],

1145 + self.hidden_states[:num_padded],

1146 + cross_per_layer,

1147 + **kwargs,

1148 + )

1149 +

1150 + # Restore the original batch_descriptor

1151 + forward_context.batch_descriptor = orig_batch_desc

1152 +

1153 + if num_logits_indices is not None:

1154 + assert num_logits_indices > 0

1155 + hidden_states[logits_indices_padded[:num_logits_indices]] = (

1156 + cross_hidden_states[:num_logits_indices]

1157 + )

1158 + else:

1159 + hidden_states = cross_hidden_states

1160 +

1161 + return hidden_states

1162

1163 def forward(

1164 self,

@@ -859,6 +1169,18 @@ class Gemma4Model(nn.Module):

1169 per_layer_inputs: torch.Tensor | None = None,

1170 **kwargs,

1171 ) -> torch.Tensor | IntermediateTensors:

1172 + if self.fast_prefill_enabled:

1173 + hidden_states = self.fast_prefill_forward(

1174 + input_ids,

1175 + positions,

1176 + inputs_embeds,

1177 + per_layer_inputs,

1178 + **kwargs,

1179 + )

1180 + hidden_states = self.norm(hidden_states)

1181 + return hidden_states

1182 +

1183 + # Normal (non-fast-prefill) path with PP support

1184 if get_pp_group().is_first_rank:

1185 if inputs_embeds is not None:

1186 hidden_states = inputs_embeds

[Gemma4] Enable Fast Prefill Optimization (#38879)