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

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

Revert gemma3n fast prefill changes (#23897)

Signed-off-by: Yong Hoon Shin <yhshin@meta.com>

Yong Hoon Shin committed Aug 29, 2025 at 12:16 UTC 8c3e199998cc5b1225328f2de01a7443fbb4f3cd

3 files changed +67 -355

tests/v1/e2e/test_kv_sharing_fast_prefill.py

@@ -64,6 +64,7 @@ def cleanup(llm: LLM, compilation_config: CompilationConfig):

65 @fork_new_process_for_each_test

66 @pytest.mark.parametrize("enforce_eager", [True])

67 +@pytest.mark.skip(reason="Disable until Gemma3n supports fast prefill")

68 def test_kv_sharing_fast_prefill(

69 monkeypatch: pytest.MonkeyPatch,

70 enforce_eager: bool,

vllm/model_executor/models/gemma3n.py

+65 -354

@@ -23,11 +23,9 @@ from torch import nn

23 from transformers.models.gemma3n.configuration_gemma3n import Gemma3nTextConfig

25 from vllm.attention import Attention

26 -from vllm.compilation.backends import set_model_tag

26 from vllm.compilation.decorators import support_torch_compile

27 from vllm.config import CacheConfig, VllmConfig

28 from vllm.distributed import get_tensor_model_parallel_world_size

30 -from vllm.forward_context import get_forward_context

29 from vllm.logger import init_logger

30 from vllm.model_executor.layers.activation import (_ACTIVATION_REGISTRY,

31 GeluAndMul,

@@ -47,7 +45,6 @@ from vllm.model_executor.model_loader.weight_utils import (

45 default_weight_loader, maybe_remap_kv_scale_name)

46 from vllm.model_executor.sampling_metadata import SamplingMetadata

47 from vllm.sequence import IntermediateTensors

50 -from vllm.v1.attention.backends.utils import KVSharingFastPrefillMetadata

49 from .interfaces import SupportsQuant

50 from .utils import (AutoWeightsLoader, extract_layer_index,

@@ -536,178 +533,7 @@ class Gemma3nDecoderLayer(nn.Module):

533 return corrected_predictions

534

535

539 -# This enables torch.compile if --kv-sharing-fast-prefill passed

540 -@support_torch_compile(enable_if=lambda vllm_config: vllm_config.cache_config.

541 - kv_sharing_fast_prefill)

542 -class Gemma3nSelfDecoder(nn.Module):

543 - """

544 - Includes altup embedding and self decoder layers

545 - """

546 -

547 - def __init__(

548 - self,

549 - *,

550 - vllm_config: VllmConfig,

551 - prefix: str = "",

552 - decoder_layers: list[Gemma3nDecoderLayer],

553 - layer_idx_start: int,

554 - per_layer_model_projection: ColumnParallelLinear,

555 - embed_scale_per_layer: torch.Tensor,

556 - embed_tokens_per_layer: VocabParallelEmbedding,

557 - per_layer_projection_norm: RMSNorm,

558 - per_layer_input_scale: torch.Tensor,

559 - altup_projections: nn.ModuleList,

560 - eps: torch.Tensor,

561 - embed_tokens: VocabParallelEmbedding,

562 - embed_scale: torch.Tensor,

563 - ):

564 - super().__init__()

565 - self.decoder_layers = decoder_layers

566 - self.layer_idx_start = layer_idx_start

567 - self.per_layer_model_projection = per_layer_model_projection

568 - self.config = vllm_config.model_config.hf_config

569 - self.embed_scale_per_layer = embed_scale_per_layer

570 - self.embed_tokens_per_layer = embed_tokens_per_layer

571 - self.per_layer_projection_norm = per_layer_projection_norm

572 - self.per_layer_input_scale = per_layer_input_scale

573 - self.altup_projections = altup_projections

574 - self.eps = eps

575 - self.embed_tokens = embed_tokens

576 - self.embed_scale = embed_scale

577 -

578 - def get_per_layer_input_embeddings(

579 - self, input_ids: torch.Tensor) -> torch.Tensor:

580 - # Deal with the fact that vocab_size_per_layer_input < vocab_size

581 - # which causes us to have some out of vocab tokens by setting

582 - # those token ids to 0. This matches the HF implementation.

583 - per_layer_inputs_mask = torch.logical_and(

584 - input_ids >= 0, input_ids < self.config.vocab_size_per_layer_input)

585 - per_layer_inputs_tokens = torch.where(per_layer_inputs_mask, input_ids,

586 - torch.zeros_like(input_ids))

587 - return self.embed_tokens_per_layer(

588 - per_layer_inputs_tokens) * self.embed_scale_per_layer

589 -

590 - def get_per_layer_inputs(

591 - self,

592 - hidden_states_0: torch.Tensor,

593 - per_layer_inputs: Optional[torch.Tensor],

594 - ) -> torch.Tensor:

595 - per_layer_projection = self.per_layer_model_projection(hidden_states_0)

596 - per_layer_projection = per_layer_projection.reshape(

597 - *hidden_states_0.shape[:-1],

598 - self.config.num_hidden_layers,

599 - self.config.hidden_size_per_layer_input,

600 - )

601 - per_layer_projection = self.per_layer_projection_norm(

602 - per_layer_projection)

603 - if per_layer_inputs is not None:

604 - # Profiling run does not compute per_layer_inputs

605 - per_layer_inputs = per_layer_projection + per_layer_inputs

606 - per_layer_inputs *= self.per_layer_input_scale

607 - else:

608 - per_layer_inputs = per_layer_projection

609 - return per_layer_inputs

610 -

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

612 - return self.embed_tokens(input_ids) * self.embed_scale

613 -

614 - def altup_embed(self, hidden_states_0: torch.Tensor) -> torch.Tensor:

615 - # Altup embed.

616 - hidden_states = [hidden_states_0] * self.config.altup_num_inputs

617 - target_magnitude = torch.mean(hidden_states_0**2, dim=-1,

618 - keepdim=True)**0.5

619 - for i in range(1, self.config.altup_num_inputs):

620 - hidden_states[i] = self.altup_projections[i - 1](hidden_states[i])

621 - new_magnitude = torch.mean(hidden_states[i]**2,

622 - dim=-1,

623 - keepdim=True)**0.5

624 - hidden_states[i] *= target_magnitude / torch.maximum(

625 - new_magnitude, self.eps)

626 - hidden_states = torch.stack(hidden_states, dim=-1)

627 - return hidden_states

628 -

629 - def forward(

630 - self,

631 - input_ids: torch.Tensor,

632 - positions: torch.Tensor,

633 - inputs_embeds: Optional[torch.Tensor] = None,

634 - per_layer_inputs: Optional[torch.Tensor] = None,

635 - **kwargs,

636 - ) -> tuple[torch.Tensor, torch.Tensor]:

637 - if inputs_embeds is not None:

638 - hidden_states_0 = inputs_embeds

639 - else:

640 - hidden_states_0 = self.get_input_embeddings(input_ids)

641 -

642 - adjusted_per_layer_inputs = self.get_per_layer_inputs(

643 - hidden_states_0, per_layer_inputs)

644 - hidden_states = self.altup_embed(hidden_states_0)

645 -

646 - # [altnum_inputs, num_tokens, hidden_size]

647 - hidden_states = hidden_states.permute(2, 0, 1)

648 -

649 - for idx, layer in enumerate(self.decoder_layers):

650 - layer_idx = idx + self.layer_idx_start

651 - # [altup_num_inputs, num_tokens, hidden_size]

652 - hidden_states = layer(

653 - positions=positions,

654 - hidden_states=hidden_states,

655 - per_layer_input=adjusted_per_layer_inputs[:, layer_idx, :],

656 - **kwargs,

657 - )

658 -

659 - # [num_tokens, hidden_size, altnum_inputs]

660 - hidden_states = hidden_states.permute(1, 2, 0)

661 -

662 - return hidden_states, adjusted_per_layer_inputs

663 -

664 -

665 -# This enables torch.compile if --kv-sharing-fast-prefill passed

666 -@support_torch_compile(enable_if=lambda vllm_config: vllm_config.cache_config.

667 - kv_sharing_fast_prefill)

668 -class Gemma3nCrossDecoder(nn.Module):

669 - """

670 - Cross-decoder layers

671 - """

672 -

673 - def __init__(

674 - self,

675 - *,

676 - vllm_config: VllmConfig,

677 - prefix: str = "",

678 - decoder_layers: list[Gemma3nDecoderLayer],

679 - layer_idx_start: int,

680 - ):

681 - super().__init__()

682 - self.decoder_layers = decoder_layers

683 - self.layer_idx_start = layer_idx_start

684 -

685 - def forward(

686 - self,

687 - positions: torch.Tensor,

688 - hidden_states: torch.Tensor,

689 - per_layer_inputs: torch.Tensor,

690 - **kwargs,

691 - ) -> torch.Tensor:

692 - # [altnum_inputs, num_tokens, hidden_size]

693 - hidden_states = hidden_states.permute(2, 0, 1)

694 - for idx, layer in enumerate(self.decoder_layers):

695 - layer_idx = idx + self.layer_idx_start

696 - # [altup_num_inputs, num_tokens, hidden_size]

697 - hidden_states = layer(

698 - positions=positions,

699 - hidden_states=hidden_states,

700 - per_layer_input=per_layer_inputs[:, layer_idx, :],

701 - **kwargs,

702 - )

703 - # [num_tokens, hidden_size, altnum_inputs]

704 - hidden_states = hidden_states.permute(1, 2, 0)

705 - return hidden_states

706 -

707 -

708 -# This disables torch.compile if --kv-sharing-fast-prefill passed

709 -@support_torch_compile(enable_if=lambda vllm_config: not vllm_config.

710 - cache_config.kv_sharing_fast_prefill)

536 +@support_torch_compile

537 class Gemma3nTextModel(nn.Module, SupportsQuant):

538

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

@@ -717,6 +543,7 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):

543 quant_config = vllm_config.quant_config

544 self.config = config

545 self.quant_config = quant_config

546 +

547 self.embed_tokens = VocabParallelEmbedding(

548 config.vocab_size,

549 config.hidden_size,

@@ -786,211 +613,95 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):

613 lambda prefix: Gemma3nDecoderLayer(

614 config, cache_config, quant_config, prefix=prefix),

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

789 -

790 - self.eps = torch.tensor(torch.finfo().min)

791 -

792 - first_kv_shared_layer_idx = (config.num_hidden_layers -

793 - config.num_kv_shared_layers)

794 - # Layer idx 0-19 are self-decoder layers in You Only Cache Once (YOCO)

795 - with set_model_tag("self_decoder"):

796 - self.self_decoder = Gemma3nSelfDecoder(

797 - vllm_config=vllm_config,

798 - prefix=f"{prefix}.self_decoder",

799 - decoder_layers=self.layers[:first_kv_shared_layer_idx],

800 - layer_idx_start=0,

801 - per_layer_model_projection=self.per_layer_model_projection,

802 - embed_scale_per_layer=self.embed_scale_per_layer,

803 - embed_tokens_per_layer=self.embed_tokens_per_layer,

804 - per_layer_projection_norm=self.per_layer_projection_norm,

805 - per_layer_input_scale=self.per_layer_input_scale,

806 - altup_projections=self.altup_projections,

807 - eps=self.eps,

808 - embed_tokens=self.embed_tokens,

809 - embed_scale=self.embed_scale,

810 - )

811 - # Layer idx 20-30 are cross-decoder layers in YOCO

812 - with set_model_tag("cross_decoder"):

813 - self.cross_decoder = Gemma3nCrossDecoder(

814 - vllm_config=vllm_config,

815 - prefix=f"{prefix}.cross_decoder",

816 - decoder_layers=self.layers[first_kv_shared_layer_idx:],

817 - layer_idx_start=first_kv_shared_layer_idx,

818 - )

819 -

616 self.norm = RMSNorm(

617 config.hidden_size,

618 eps=config.rms_norm_eps,

619 )

824 -

825 - self.fast_prefill_enabled = cache_config.kv_sharing_fast_prefill

826 -

827 - if self.fast_prefill_enabled:

828 - # Allocate static buffers for CUDAGraph

829 - # TODO(sarckk): Extract this functionality to interface

830 - max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens

831 - device = next(self.parameters()).device

832 - self.positions = torch.zeros(max_num_tokens,

833 - dtype=torch.int64,

834 - device=device)

835 - self.hidden_states = torch.zeros(

836 - (max_num_tokens, config.hidden_size,

837 - self.config.altup_num_inputs),

838 - dtype=self.embed_tokens.weight.dtype,

839 - device=device,

840 - )

841 - self.per_layer_inputs = torch.zeros(

842 - (max_num_tokens, self.config.num_hidden_layers,

843 - self.config.hidden_size_per_layer_input),

844 - dtype=self.embed_tokens.weight.dtype,

845 - device=device,

846 - )

620 + self.eps = torch.tensor(torch.finfo().min)

621

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

849 - return self.self_decoder.get_input_embeddings(input_ids)

623 + return self.embed_tokens(input_ids) * self.embed_scale

624 +

625 + def get_per_layer_input_embeddings(

626 + self, input_ids: torch.Tensor) -> torch.Tensor:

627 + # Deal with the fact that vocab_size_per_layer_input < vocab_size

628 + # which causes us to have some out of vocab tokens by setting

629 + # those token ids to 0. This matches the HF implementation.

630 + per_layer_inputs_mask = torch.logical_and(

631 + input_ids >= 0, input_ids < self.config.vocab_size_per_layer_input)

632 + per_layer_inputs_tokens = torch.where(per_layer_inputs_mask, input_ids,

633 + torch.zeros_like(input_ids))

634 + return self.embed_tokens_per_layer(

635 + per_layer_inputs_tokens) * self.embed_scale_per_layer

636

851 - def fast_prefill_forward(

637 + def forward(

638 self,

853 - input_ids: torch.Tensor,

639 + input_ids: Optional[torch.Tensor],

640 positions: torch.Tensor,

641 + per_layer_inputs: torch.Tensor,

642 + intermediate_tensors: Optional[IntermediateTensors] = None,

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

856 - per_layer_inputs: Optional[torch.Tensor] = None,

644 **kwargs,

858 - ) -> torch.Tensor:

859 - logits_indices_padded, num_logits_indices = None, None

860 - attn_metadata = get_forward_context().attn_metadata

861 -

862 - # attn_metadata is None during dummy runs

863 - if (self.fast_prefill_enabled and attn_metadata is not None):

864 - assert isinstance(attn_metadata, dict)

865 - # Last layer is a KV sharing layer

866 - layer_attn_metadata = attn_metadata[

867 - self.layers[-1].self_attn.attn.layer_name]

868 - if (isinstance(layer_attn_metadata, KVSharingFastPrefillMetadata)):

869 - logits_indices_padded = (

870 - layer_attn_metadata.logits_indices_padded)

871 - num_logits_indices = layer_attn_metadata.num_logits_indices

872 -

873 - # Copy inputs for cudagraph

874 - batch_size = positions.size(0)

875 - self.positions[:batch_size].copy_(positions)

876 - self_decoder_hidden_states, per_layer_inputs_adjusted = \

877 - self.self_decoder(

878 - input_ids=input_ids,

879 - positions=self.positions[:batch_size],

880 - inputs_embeds=inputs_embeds,

881 - per_layer_inputs=per_layer_inputs,

882 - **kwargs,

883 - )

884 -

885 - if logits_indices_padded is None:

886 - logits_indices_padded = torch.arange(

887 - positions.size(0),

888 - dtype=positions.dtype,

889 - device=positions.device,

890 - )

645 + ) -> Union[torch.Tensor, IntermediateTensors]:

646 + if inputs_embeds is not None:

647 + hidden_states_0 = inputs_embeds

648 + else:

649 + hidden_states_0 = self.get_input_embeddings(input_ids)

650

892 - # NOTE(sarckk): There is currently a bug caused by

893 - # vLLM converting output of last piecewise CUDA graph

894 - # to weakref, causing memory to be prematurely freed

895 - # when there are multiple compilation units

896 - # Keep .clone() until fix in

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

898 - hidden_states = self_decoder_hidden_states.clone()

899 -

900 - # Copy inputs for cudagraph

901 - num_padded_logits_indices = logits_indices_padded.size(0)

902 - self.positions[:num_padded_logits_indices].copy_(

903 - positions[logits_indices_padded])

904 - self.hidden_states[:num_padded_logits_indices].copy_(

905 - self_decoder_hidden_states[logits_indices_padded])

906 - self.per_layer_inputs[:num_padded_logits_indices].copy_(

907 - per_layer_inputs_adjusted[logits_indices_padded])

908 - cross_decoder_hidden_states = self.cross_decoder(

909 - positions=self.positions[:num_padded_logits_indices],

910 - hidden_states=self.hidden_states[:num_padded_logits_indices],

911 - per_layer_inputs=self.per_layer_inputs[:num_padded_logits_indices],

912 - **kwargs,

651 + per_layer_projection = self.per_layer_model_projection(hidden_states_0)

652 + per_layer_projection = per_layer_projection.reshape(

653 + *hidden_states_0.shape[:-1],

654 + self.config.num_hidden_layers,

655 + self.config.hidden_size_per_layer_input,

656 )

657 + per_layer_projection = self.per_layer_projection_norm(

658 + per_layer_projection)

659

915 - if num_logits_indices is not None:

916 - assert num_logits_indices > 0

917 - # Merge cross-decoder and self-decoder hidden states

918 - hidden_states[logits_indices_padded[:num_logits_indices]] = (

919 - cross_decoder_hidden_states[:num_logits_indices])

660 + if per_layer_inputs is not None:

661 + # Profiling run does not compute per_layer_inputs

662 + per_layer_inputs = per_layer_projection + per_layer_inputs

663 + per_layer_inputs *= self.per_layer_input_scale

664 else:

921 - hidden_states = cross_decoder_hidden_states

665 + per_layer_inputs = per_layer_projection

666

923 - return hidden_states

667 + # Altup embed.

668 + hidden_states = [hidden_states_0] * self.config.altup_num_inputs

669 + target_magnitude = torch.mean(hidden_states_0**2, dim=-1,

670 + keepdim=True)**0.5

671 + for i in range(1, self.config.altup_num_inputs):

672 + hidden_states[i] = self.altup_projections[i - 1](hidden_states[i])

673 + new_magnitude = torch.mean(hidden_states[i]**2,

674 + dim=-1,

675 + keepdim=True)**0.5

676 + hidden_states[i] *= target_magnitude / torch.maximum(

677 + new_magnitude, self.eps)

678 + hidden_states = torch.stack(hidden_states, dim=0)

679

925 - def normal_forward(

926 - self,

927 - input_ids: torch.Tensor,

928 - positions: torch.Tensor,

929 - inputs_embeds: Optional[torch.Tensor] = None,

930 - per_layer_inputs: Optional[torch.Tensor] = None,

931 - **kwargs,

932 - ) -> torch.Tensor:

933 - hidden_states, per_layer_inputs = self.self_decoder(

934 - input_ids=input_ids,

935 - positions=positions,

936 - inputs_embeds=inputs_embeds,

937 - per_layer_inputs=per_layer_inputs,

938 - **kwargs,

939 - )

940 - hidden_states = self.cross_decoder(

941 - positions=positions,

942 - hidden_states=hidden_states,

943 - per_layer_inputs=per_layer_inputs,

944 - **kwargs,

945 - )

946 - return hidden_states

680 + # Transformer blocks.

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

682 + # [altup_num_inputs, num_tokens, hidden_size]

683 + hidden_states = layer(

684 + positions=positions,

685 + hidden_states=hidden_states,

686 + per_layer_input=per_layer_inputs[:, layer_idx, :],

687 + **kwargs,

688 + )

689

948 - def altup_unembed(

949 - self,

950 - hidden_states: torch.Tensor,

951 - ) -> torch.Tensor:

690 # Altup unembed.

953 - target_magnitude = torch.mean(hidden_states[..., 0]**2,

691 + target_magnitude = torch.mean(hidden_states[0]**2,

692 dim=-1,

693 keepdim=True)**0.5

694 for i in range(1, self.config.altup_num_inputs):

957 - hidden_states[..., i] = self.altup_unembed_projections[i - 1](

958 - hidden_states[..., i])

959 - new_magnitude = torch.mean(hidden_states[..., i]**2,

695 + hidden_states[i] = self.altup_unembed_projections[i - 1](

696 + hidden_states[i])

697 + new_magnitude = torch.mean(hidden_states[i]**2,

698 dim=-1,

699 keepdim=True)**0.5

962 - hidden_states[..., i] *= target_magnitude / torch.maximum(

700 + hidden_states[i] *= target_magnitude / torch.maximum(

701 new_magnitude, self.eps)

964 - # [num_tokens,hidden_size, altup_num_inputs] -> [num_tokens,hidden_size]

965 - hidden_states = torch.mean(hidden_states, dim=-1)

966 - return hidden_states

702 + # [altup_num_inputs,num_tokens,hidden_size] -> [num_tokens,hidden_size]

703 + hidden_states = torch.mean(hidden_states, dim=0)

704

968 - def forward(

969 - self,

970 - input_ids: Optional[torch.Tensor],

971 - positions: torch.Tensor,

972 - per_layer_inputs: Optional[torch.Tensor] = None,

973 - intermediate_tensors: Optional[IntermediateTensors] = None,

974 - inputs_embeds: Optional[torch.Tensor] = None,

975 - **kwargs,

976 - ) -> Union[torch.Tensor, IntermediateTensors]:

977 - if self.fast_prefill_enabled:

978 - hidden_states = self.fast_prefill_forward(

979 - input_ids,

980 - positions,

981 - inputs_embeds,

982 - per_layer_inputs,

983 - **kwargs,

984 - )

985 - else:

986 - hidden_states = self.normal_forward(

987 - input_ids,

988 - positions,

989 - inputs_embeds,

990 - per_layer_inputs,

991 - **kwargs,

992 - )

993 - hidden_states = self.altup_unembed(hidden_states)

705 return self.norm(hidden_states)

706

707 def load_weights(self, weights: Iterable[tuple[str,

vllm/model_executor/models/gemma3n_mm.py

+1 -1

@@ -620,7 +620,7 @@ class Gemma3nForConditionalGeneration(nn.Module, SupportsMultiModal):

620 # NOTE (NickLucche) Each pass needs tokens to compute PLE so we cache

621 # them here, as the model forward has only access to the input_embeds.

622 if input_ids is not None:

623 - per_layer_inputs = self.language_model.model.self_decoder.get_per_layer_input_embeddings(

623 + per_layer_inputs = self.language_model.model.get_per_layer_input_embeddings(

624 input_ids)

625 per_layer_inputs = per_layer_inputs.reshape(

626 -1, self.config.text_config.num_hidden_layers,