77d906995c · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -26,6 +26,7 @@ from vllm.attention import Attention

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

29 +from vllm.forward_context import get_forward_context

30 from vllm.logger import init_logger

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

32 GeluAndMul,

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

45 from vllm.model_executor.model_loader.weight_utils import (

46 default_weight_loader, maybe_remap_kv_scale_name)

47 from vllm.sequence import IntermediateTensors

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

49

50 from .interfaces import SupportsQuant

51 from .utils import (AutoWeightsLoader, extract_layer_index,

@@ -51,6 +53,8 @@ from .utils import (AutoWeightsLoader, extract_layer_index,

53

54 logger = init_logger(__name__)

55

56 +EPS = torch.tensor(torch.finfo().min)

57 +

58

59 class Gemma3nAltUp(nn.Module):

60 """Alternating updates (Altup)

@@ -532,16 +536,29 @@ class Gemma3nDecoderLayer(nn.Module):

536 return corrected_predictions

537

538

535 -@support_torch_compile

536 -class Gemma3nTextModel(nn.Module, SupportsQuant):

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

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

547 + def __init__(

548 + self,

549 + *,

550 + vllm_config: VllmConfig,

551 + prefix: str = "",

552 + decoder_layers: list[Gemma3nDecoderLayer],

553 + layer_idx_start: int,

554 + ):

555 super().__init__()

556 + self.decoder_layers = decoder_layers

557 + self.layer_idx_start = layer_idx_start

558 +

559 config = vllm_config.model_config.hf_config

541 - cache_config = vllm_config.cache_config

542 - quant_config = vllm_config.quant_config

560 self.config = config

544 - self.quant_config = quant_config

561 + quant_config = vllm_config.quant_config

562

563 self.embed_tokens = VocabParallelEmbedding(

564 config.vocab_size,

@@ -594,32 +611,6 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):

611 prefix=f"{prefix}.altup_projections.{idx-1}",

612 ) for idx in range(1, self.config.altup_num_inputs)

613 ])

597 - self.altup_unembed_projections = nn.ModuleList([

598 - ColumnParallelLinear(

599 - config.hidden_size,

600 - config.hidden_size,

601 - bias=False,

602 - gather_output=True,

603 - return_bias=False,

604 - quant_config=quant_config,

605 - prefix=f"{prefix}.altup_unembed_projections.{idx-1}",

606 - ) for idx in range(1, self.config.altup_num_inputs)

607 - ])

608 -

609 - # Transformer blocks.

610 - self.start_layer, self.end_layer, self.layers = make_layers(

611 - config.num_hidden_layers,

612 - lambda prefix: Gemma3nDecoderLayer(

613 - config, cache_config, quant_config, prefix=prefix),

614 - prefix=f"{prefix}.layers")

615 - self.norm = RMSNorm(

616 - config.hidden_size,

617 - eps=config.rms_norm_eps,

618 - )

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

620 -

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

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

614

615 def get_per_layer_input_embeddings(

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

@@ -633,20 +624,11 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):

624 return self.embed_tokens_per_layer(

625 per_layer_inputs_tokens) * self.embed_scale_per_layer

626

636 - def forward(

627 + def get_per_layer_inputs(

628 self,

638 - input_ids: Optional[torch.Tensor],

639 - positions: torch.Tensor,

640 - per_layer_inputs: torch.Tensor,

641 - intermediate_tensors: Optional[IntermediateTensors] = None,

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

643 - **kwargs,

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

645 - if inputs_embeds is not None:

646 - hidden_states_0 = inputs_embeds

647 - else:

648 - hidden_states_0 = self.get_input_embeddings(input_ids)

649 -

629 + hidden_states_0: torch.Tensor,

630 + per_layer_inputs: Optional[torch.Tensor],

631 + ) -> torch.Tensor:

632 per_layer_projection = self.per_layer_model_projection(hidden_states_0)

633 per_layer_projection = per_layer_projection.reshape(

634 *hidden_states_0.shape[:-1],

@@ -655,14 +637,18 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):

637 )

638 per_layer_projection = self.per_layer_projection_norm(

639 per_layer_projection)

658 -

640 if per_layer_inputs is not None:

641 # Profiling run does not compute per_layer_inputs

642 per_layer_inputs = per_layer_projection + per_layer_inputs

643 per_layer_inputs *= self.per_layer_input_scale

644 else:

645 per_layer_inputs = per_layer_projection

646 + return per_layer_inputs

647

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

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

650 +

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

652 # Altup embed.

653 hidden_states = [hidden_states_0] * self.config.altup_num_inputs

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

@@ -673,11 +659,77 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):

659 dim=-1,

660 keepdim=True)**0.5

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

676 - new_magnitude, self.eps)

677 - hidden_states = torch.stack(hidden_states, dim=0)

662 + new_magnitude, EPS)

663 + hidden_states = torch.stack(hidden_states, dim=-1)

664 + return hidden_states

665 +

666 + def forward(

667 + self,

668 + input_ids: torch.Tensor,

669 + positions: torch.Tensor,

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

671 + per_layer_inputs: Optional[torch.Tensor] = None,

672 + **kwargs,

673 + ) -> tuple[torch.Tensor, torch.Tensor]:

674 + if inputs_embeds is not None:

675 + hidden_states_0 = inputs_embeds

676 + else:

677 + hidden_states_0 = self.get_input_embeddings(input_ids)

678 +

679 + adjusted_per_layer_inputs = self.get_per_layer_inputs(

680 + hidden_states_0, per_layer_inputs)

681 + hidden_states = self.altup_embed(hidden_states_0)

682 +

683 + # [altnum_inputs, num_tokens, hidden_size]

684 + hidden_states = hidden_states.permute(2, 0, 1)

685 +

686 + for idx, layer in enumerate(self.decoder_layers):

687 + layer_idx = idx + self.layer_idx_start

688 + # [altup_num_inputs, num_tokens, hidden_size]

689 + hidden_states = layer(

690 + positions=positions,

691 + hidden_states=hidden_states,

692 + per_layer_input=adjusted_per_layer_inputs[:, layer_idx, :],

693 + **kwargs,

694 + )

695 +

696 + # [num_tokens, hidden_size, altnum_inputs]

697 + hidden_states = hidden_states.permute(1, 2, 0)

698 +

699 + return hidden_states, adjusted_per_layer_inputs

700 +

701 +

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

703 +@support_torch_compile(enable_if=lambda vllm_config: vllm_config.cache_config.

704 + kv_sharing_fast_prefill)

705 +class Gemma3nCrossDecoder(nn.Module):

706 + """

707 + Cross-decoder layers

708 + """

709 +

710 + def __init__(

711 + self,

712 + *,

713 + vllm_config: VllmConfig,

714 + prefix: str = "",

715 + decoder_layers: list[Gemma3nDecoderLayer],

716 + layer_idx_start: int,

717 + ):

718 + super().__init__()

719 + self.decoder_layers = decoder_layers

720 + self.layer_idx_start = layer_idx_start

721

679 - # Transformer blocks.

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

722 + def forward(

723 + self,

724 + positions: torch.Tensor,

725 + hidden_states: torch.Tensor,

726 + per_layer_inputs: torch.Tensor,

727 + **kwargs,

728 + ) -> torch.Tensor:

729 + # [altnum_inputs, num_tokens, hidden_size]

730 + hidden_states = hidden_states.permute(2, 0, 1)

731 + for idx, layer in enumerate(self.decoder_layers):

732 + layer_idx = idx + self.layer_idx_start

733 # [altup_num_inputs, num_tokens, hidden_size]

734 hidden_states = layer(

735 positions=positions,

@@ -685,22 +737,249 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):

737 per_layer_input=per_layer_inputs[:, layer_idx, :],

738 **kwargs,

739 )

740 + # [num_tokens, hidden_size, altnum_inputs]

741 + hidden_states = hidden_states.permute(1, 2, 0)

742 + return hidden_states

743 +

744 +

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

746 +@support_torch_compile(enable_if=lambda vllm_config: not vllm_config.

747 + cache_config.kv_sharing_fast_prefill)

748 +class Gemma3nTextModel(nn.Module, SupportsQuant):

749 +

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

751 + super().__init__()

752 + config = vllm_config.model_config.hf_config

753 + cache_config = vllm_config.cache_config

754 + quant_config = vllm_config.quant_config

755 + self.config = config

756 + self.quant_config = quant_config

757 +

758 + self.altup_unembed_projections = nn.ModuleList([

759 + ColumnParallelLinear(

760 + config.hidden_size,

761 + config.hidden_size,

762 + bias=False,

763 + gather_output=True,

764 + return_bias=False,

765 + quant_config=quant_config,

766 + prefix=f"{prefix}.altup_unembed_projections.{idx-1}",

767 + ) for idx in range(1, self.config.altup_num_inputs)

768 + ])

769 +

770 + # Allocate config.num_kv_shared_layers layers for self-decoder

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

772 + config.num_hidden_layers,

773 + lambda prefix: Gemma3nDecoderLayer(

774 + config, cache_config, quant_config, prefix=prefix),

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

776 +

777 + first_kv_shared_layer_idx = (config.num_hidden_layers -

778 + config.num_kv_shared_layers)

779 +

780 + # NOTE(sarckk): importing this top level seems to cause issues

781 + # during running of tests.

782 + from vllm.compilation.backends import set_model_tag

783 +

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

785 + with set_model_tag("self_decoder"):

786 + self.self_decoder = Gemma3nSelfDecoder(

787 + vllm_config=vllm_config,

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

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

790 + layer_idx_start=0,

791 + )

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

793 + with set_model_tag("cross_decoder"):

794 + self.cross_decoder = Gemma3nCrossDecoder(

795 + vllm_config=vllm_config,

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

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

798 + layer_idx_start=first_kv_shared_layer_idx,

799 + )

800 +

801 + self.norm = RMSNorm(

802 + config.hidden_size,

803 + eps=config.rms_norm_eps,

804 + )

805 +

806 + self.fast_prefill_enabled = cache_config.kv_sharing_fast_prefill

807 +

808 + if self.fast_prefill_enabled:

809 + # Allocate static buffers for CUDAGraph

810 + # TODO(sarckk): Extract this functionality to interface

811 + max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens

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

813 + self.positions = torch.zeros(max_num_tokens,

814 + dtype=torch.int64,

815 + device=device)

816 + self.hidden_states = torch.zeros(

817 + (max_num_tokens, config.hidden_size,

818 + self.config.altup_num_inputs),

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

820 + device=device,

821 + )

822 + self.per_layer_inputs = torch.zeros(

823 + (max_num_tokens, self.config.num_hidden_layers,

824 + self.config.hidden_size_per_layer_input),

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

826 + device=device,

827 + )

828 +

829 + @property

830 + def embed_tokens(self):

831 + return self.self_decoder.embed_tokens

832 +

833 + def get_per_layer_input_embeddings(

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

835 + return self.self_decoder.get_per_layer_input_embeddings(input_ids)

836 +

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

838 + return self.self_decoder.get_input_embeddings(input_ids)

839 +

840 + def fast_prefill_forward(

841 + self,

842 + input_ids: torch.Tensor,

843 + positions: torch.Tensor,

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

845 + per_layer_inputs: Optional[torch.Tensor] = None,

846 + **kwargs,

847 + ) -> torch.Tensor:

848 + logits_indices_padded, num_logits_indices = None, None

849 + attn_metadata = get_forward_context().attn_metadata

850 +

851 + # attn_metadata is None during dummy runs

852 + if (self.fast_prefill_enabled and attn_metadata is not None):

853 + assert isinstance(attn_metadata, dict)

854 + # Last layer is a KV sharing layer

855 + layer_attn_metadata = attn_metadata[

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

857 + if (isinstance(layer_attn_metadata, KVSharingFastPrefillMetadata)):

858 + logits_indices_padded = (

859 + layer_attn_metadata.logits_indices_padded)

860 + num_logits_indices = layer_attn_metadata.num_logits_indices

861 +

862 + # Copy inputs for cudagraph

863 + batch_size = positions.size(0)

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

865 + self_decoder_hidden_states, per_layer_inputs_adjusted = \

866 + self.self_decoder(

867 + input_ids=input_ids,

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

869 + inputs_embeds=inputs_embeds,

870 + per_layer_inputs=per_layer_inputs,

871 + **kwargs,

872 + )

873 +

874 + if logits_indices_padded is None:

875 + logits_indices_padded = torch.arange(

876 + positions.size(0),

877 + dtype=positions.dtype,

878 + device=positions.device,

879 + )

880 +

881 + # NOTE(sarckk): There is currently a bug caused by

882 + # vLLM converting output of last piecewise CUDA graph

883 + # to weakref, causing memory to be prematurely freed

884 + # when there are multiple compilation units

885 + # Keep .clone() until fix in

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

887 + hidden_states = self_decoder_hidden_states.clone()

888 +

889 + # Copy inputs for cudagraph

890 + num_padded_logits_indices = logits_indices_padded.size(0)

891 + self.positions[:num_padded_logits_indices].copy_(

892 + positions[logits_indices_padded])

893 + self.hidden_states[:num_padded_logits_indices].copy_(

894 + self_decoder_hidden_states[logits_indices_padded])

895 + self.per_layer_inputs[:num_padded_logits_indices].copy_(

896 + per_layer_inputs_adjusted[logits_indices_padded])

897 + cross_decoder_hidden_states = self.cross_decoder(

898 + positions=self.positions[:num_padded_logits_indices],

899 + hidden_states=self.hidden_states[:num_padded_logits_indices],

900 + per_layer_inputs=self.per_layer_inputs[:num_padded_logits_indices],

901 + **kwargs,

902 + )

903 +

904 + if num_logits_indices is not None:

905 + assert num_logits_indices > 0

906 + # Merge cross-decoder and self-decoder hidden states

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

908 + cross_decoder_hidden_states[:num_logits_indices])

909 + else:

910 + hidden_states = cross_decoder_hidden_states

911 +

912 + return hidden_states

913 +

914 + def normal_forward(

915 + self,

916 + input_ids: torch.Tensor,

917 + positions: torch.Tensor,

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

919 + per_layer_inputs: Optional[torch.Tensor] = None,

920 + **kwargs,

921 + ) -> torch.Tensor:

922 + hidden_states, per_layer_inputs = self.self_decoder(

923 + input_ids=input_ids,

924 + positions=positions,

925 + inputs_embeds=inputs_embeds,

926 + per_layer_inputs=per_layer_inputs,

927 + **kwargs,

928 + )

929 + hidden_states = self.cross_decoder(

930 + positions=positions,

931 + hidden_states=hidden_states,

932 + per_layer_inputs=per_layer_inputs,

933 + **kwargs,

934 + )

935 + return hidden_states

936

937 + def altup_unembed(

938 + self,

939 + hidden_states: torch.Tensor,

940 + ) -> torch.Tensor:

941 # Altup unembed.

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

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

943 dim=-1,

944 keepdim=True)**0.5

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

694 - hidden_states[i] = self.altup_unembed_projections[i - 1](

695 - hidden_states[i])

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

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

947 + hidden_states[..., i])

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

949 dim=-1,

950 keepdim=True)**0.5

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

700 - new_magnitude, self.eps)

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

702 - hidden_states = torch.mean(hidden_states, dim=0)

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

952 + new_magnitude, EPS)

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

954 + hidden_states = torch.mean(hidden_states, dim=-1)

955 + return hidden_states

956

957 + def forward(

958 + self,

959 + input_ids: Optional[torch.Tensor],

960 + positions: torch.Tensor,

961 + per_layer_inputs: Optional[torch.Tensor] = None,

962 + intermediate_tensors: Optional[IntermediateTensors] = None,

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

964 + **kwargs,

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

966 + if self.fast_prefill_enabled:

967 + hidden_states = self.fast_prefill_forward(

968 + input_ids,

969 + positions,

970 + inputs_embeds,

971 + per_layer_inputs,

972 + **kwargs,

973 + )

974 + else:

975 + hidden_states = self.normal_forward(

976 + input_ids,

977 + positions,

978 + inputs_embeds,

979 + per_layer_inputs,

980 + **kwargs,

981 + )

982 + hidden_states = self.altup_unembed(hidden_states)

983 return self.norm(hidden_states)

984

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

@@ -716,6 +995,13 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):

995 params_dict = dict(self.named_parameters())

996 loaded_params: set[str] = set()

997 for name, loaded_weight in weights:

998 + # decoder layer weights, altup_unembed_projections and rmsnorm

999 + # are initialized in text model, others are in self decoder

1000 + if (not name.startswith('layers')

1001 + and not name.startswith('altup_unembed_projections')

1002 + and not name.startswith('norm')):

1003 + name = f"self_decoder.{name}"

1004 +

1005 if (self.quant_config is not None and

1006 (scale_name := self.quant_config.get_cache_scale(name))):

1007 # Loading kv cache scales for compressed-tensors quantization

[KV sharing] Re-land Gemma3n model changes from #22628 (#24357)