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

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

[V1] Enable prefill optimization for Gemma3n (#22628)

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

Yong Hoon Shin committed Aug 28, 2025 at 14:54 UTC cb293f6a790d555d6d7ced872118ff029bd828e8

9 files changed +583 -228

tests/v1/e2e/test_kv_sharing_fast_prefill.py

-57

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 import random

5 -from typing import Optional, Union

6 import pytest

7 import torch

@@ -10,12 +9,6 @@ import torch

9 from vllm import LLM, SamplingParams

10 from vllm.config import CompilationConfig, CompilationLevel

11 from vllm.distributed import cleanup_dist_env_and_memory

13 -from vllm.forward_context import get_forward_context

14 -from vllm.model_executor.models.gemma3n_mm import (

15 - Gemma3nForConditionalGeneration)

16 -from vllm.model_executor.models.registry import ModelRegistry

17 -from vllm.model_executor.models.utils import extract_layer_index

18 -from vllm.sequence import IntermediateTensors

13 from ...utils import fork_new_process_for_each_test

@@ -23,54 +16,6 @@ from ...utils import fork_new_process_for_each_test

16 SEED = 42

26 -class TestGemma3nForConditionalGeneration(Gemma3nForConditionalGeneration):

27 -

28 - def forward(

29 - self,

30 - input_ids: torch.Tensor,

31 - positions: torch.Tensor,

32 - intermediate_tensors: Optional[IntermediateTensors] = None,

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

34 - **kwargs,

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

36 - hidden_states = super().forward(input_ids, positions,

37 - intermediate_tensors, inputs_embeds,

38 - **kwargs)

39 - attn_metadata = get_forward_context().attn_metadata

40 - # attn_metadata is None during dummy runs

41 - if (attn_metadata is not None

42 - and self.language_model.cache_config.kv_sharing_fast_prefill):

43 - assert isinstance(attn_metadata, dict) # true in V1

44 - # Gemma3n-E2B has 30 layers, with last 20 layers being

45 - # cross-decoder layers. Check attention metadata is correct

46 - for layer_name, metadata in attn_metadata.items():

47 - layer_idx = extract_layer_index(layer_name)

48 - if layer_idx >= 20:

49 - assert hasattr(metadata, 'logits_indices_padded')

50 - assert hasattr(metadata, 'num_logits_indices')

51 - else:

52 - assert not hasattr(metadata, 'logits_indices_padded')

53 - assert not hasattr(metadata, 'num_logits_indices')

54 -

55 - # Last layer will be a KV sharing layer

56 - layer_attn_metadata = attn_metadata[

57 - self.language_model.model.layers[-1].self_attn.attn.layer_name]

58 - logits_indices_padded = (layer_attn_metadata.logits_indices_padded)

59 - assert logits_indices_padded is not None

60 - num_logits_indices = layer_attn_metadata.num_logits_indices

61 - assert num_logits_indices > 0

62 - # Reset hidden states to random values and

63 - # only set logits at logits_indices to valid values

64 - # Because logits_indices are the only positions that are used

65 - # for output token sampling, this still produces same outputs

66 - logits_hs = hidden_states[logits_indices_padded]

67 - hidden_states = torch.randn_like(hidden_states)

68 - gen_indices = logits_indices_padded[:num_logits_indices]

69 - hidden_states[gen_indices] = logits_hs[:num_logits_indices]

70 -

71 - return hidden_states

72 -

73 -

19 @pytest.fixture

20 def test_prompts():

21 """

@@ -124,8 +69,6 @@ def test_kv_sharing_fast_prefill(

69 enforce_eager: bool,

70 test_prompts: list[str],

71 ):

127 - ModelRegistry.register_model("Gemma3nForConditionalGeneration",

128 - TestGemma3nForConditionalGeneration)

72 sampling_params = SamplingParams(temperature=0.0, max_tokens=100)

73 compilation_config = CompilationConfig(

74 # This allows vLLM compilation backend to handle allocating and

vllm/config/cache.py

+7 -5

@@ -145,12 +145,19 @@ class CacheConfig:

145

146 self._verify_cache_dtype()

147 self._verify_prefix_caching()

148 + self._verify_kv_sharing_fast_prefill()

149

150 def metrics_info(self):

151 # convert cache_config to dict(key: str, value: str) for prometheus

152 # metrics info

153 return {key: str(value) for key, value in self.__dict__.items()}

154

155 + def _verify_kv_sharing_fast_prefill(self) -> None:

156 + if self.kv_sharing_fast_prefill and not envs.VLLM_USE_V1:

157 + raise NotImplementedError(

158 + "Fast prefill optimization for KV sharing is not supported "

159 + "in V0 currently.")

160 +

161 @model_validator(mode='after')

162 def _verify_args(self) -> Self:

163 if self.cpu_offload_gb < 0:

@@ -162,11 +169,6 @@ class CacheConfig:

169 "GPU memory utilization must be less than 1.0. Got "

170 f"{self.gpu_memory_utilization}.")

171

165 - if self.kv_sharing_fast_prefill:

166 - logger.warning_once(

167 - "--kv-sharing-fast-prefill is currently work in progress "

168 - "and not functional yet (i.e. no prefill savings)")

169 -

172 return self

173

174 def _verify_cache_dtype(self) -> None:

vllm/model_executor/models/gemma3n.py

+354 -65

@@ -23,9 +23,11 @@ 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

27 from vllm.compilation.decorators import support_torch_compile

28 from vllm.config import CacheConfig, VllmConfig

29 from vllm.distributed import get_tensor_model_parallel_world_size

30 +from vllm.forward_context import get_forward_context

31 from vllm.logger import init_logger

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

33 GeluAndMul,

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

47 default_weight_loader, maybe_remap_kv_scale_name)

48 from vllm.model_executor.sampling_metadata import SamplingMetadata

49 from vllm.sequence import IntermediateTensors

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

52 from .interfaces import SupportsQuant

53 from .utils import (AutoWeightsLoader, extract_layer_index,

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

536 return corrected_predictions

537

538

536 -@support_torch_compile

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)

711 class Gemma3nTextModel(nn.Module, SupportsQuant):

712

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

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

717 quant_config = vllm_config.quant_config

718 self.config = config

719 self.quant_config = quant_config

546 -

720 self.embed_tokens = VocabParallelEmbedding(

721 config.vocab_size,

722 config.hidden_size,

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

786 lambda prefix: Gemma3nDecoderLayer(

787 config, cache_config, quant_config, prefix=prefix),

788 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 +

820 self.norm = RMSNorm(

821 config.hidden_size,

822 eps=config.rms_norm_eps,

823 )

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

824

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

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

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 + )

847

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

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

849 + return self.self_decoder.get_input_embeddings(input_ids)

850

637 - def forward(

851 + def fast_prefill_forward(

852 self,

639 - input_ids: Optional[torch.Tensor],

853 + input_ids: torch.Tensor,

854 positions: torch.Tensor,

641 - per_layer_inputs: torch.Tensor,

642 - intermediate_tensors: Optional[IntermediateTensors] = None,

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

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

857 **kwargs,

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)

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

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,

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 + )

891 +

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,

913 )

657 - per_layer_projection = self.per_layer_projection_norm(

658 - per_layer_projection)

914

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

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])

920 else:

665 - per_layer_inputs = per_layer_projection

921 + hidden_states = cross_decoder_hidden_states

922

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)

923 + return hidden_states

924

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 - )

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

947

948 + def altup_unembed(

949 + self,

950 + hidden_states: torch.Tensor,

951 + ) -> torch.Tensor:

952 # Altup unembed.

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

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

954 dim=-1,

955 keepdim=True)**0.5

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

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

696 - hidden_states[i])

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

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

958 + hidden_states[..., i])

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

960 dim=-1,

961 keepdim=True)**0.5

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

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

963 new_magnitude, self.eps)

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

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

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

967

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)

994 return self.norm(hidden_states)

995

996 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.get_per_layer_input_embeddings(

623 + per_layer_inputs = self.language_model.model.self_decoder.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,

vllm/v1/attention/backends/utils.py

+126 -13

@@ -4,11 +4,13 @@ import abc

4 import enum

5 import functools

6 from abc import abstractmethod

7 -from dataclasses import dataclass, make_dataclass

8 -from typing import TYPE_CHECKING, Any, ClassVar, Generic, Optional, TypeVar

7 +from dataclasses import dataclass, fields, make_dataclass

8 +from typing import (TYPE_CHECKING, Any, ClassVar, Generic, Optional, Protocol,

9 + TypeVar)

11 import numpy as np

12 import torch

13 +from typing_extensions import runtime_checkable

15 from vllm.config import VllmConfig, get_layers_from_vllm_config

16 from vllm.utils import cdiv

@@ -19,7 +21,8 @@ if TYPE_CHECKING:

21 from vllm.v1.worker.gpu_input_batch import InputBatch

23 import vllm.envs as envs

22 -from vllm.attention.backends.abstract import AttentionBackend

24 +from vllm.attention.backends.abstract import (AttentionBackend,

25 + AttentionMetadata)

26 from vllm.attention.layer import Attention

27 from vllm.distributed.kv_transfer.kv_connector.utils import (

28 get_kv_connector_cache_layout)

@@ -65,6 +68,10 @@ class CommonAttentionMetadata:

69 causal: bool = True

71 + # Needed by FastPrefillAttentionBuilder

72 + logits_indices_padded: Optional[torch.Tensor] = None

73 + num_logits_indices: Optional[int] = None

74 +

76 @dataclass

77 class UbatchSlice:

@@ -542,6 +549,69 @@ def make_local_attention_virtual_batches(

549 )

550

551

552 +def make_kv_sharing_fast_prefill_common_attn_metadata(

553 + common_attn_metadata: CommonAttentionMetadata,

554 +) -> CommonAttentionMetadata:

555 + if common_attn_metadata.max_query_len == 1:

556 + # All requests are decode (assume 1 token for now)

557 + # Skip computing fast prefill path

558 + return common_attn_metadata

559 +

560 + assert common_attn_metadata.logits_indices_padded is not None

561 + assert common_attn_metadata.num_logits_indices is not None

562 +

563 + logits_indices_padded = common_attn_metadata.logits_indices_padded

564 + num_logits_indices = common_attn_metadata.num_logits_indices

565 + # Get rid of CUDAGraph padding, if any

566 + logits_indices = logits_indices_padded[:num_logits_indices]

567 + num_reqs = common_attn_metadata.num_reqs

568 + query_start_loc = common_attn_metadata.query_start_loc

569 + seq_lens = common_attn_metadata.seq_lens

570 + # Example inputs

571 + # num_reqs: 3

572 + # generation_indices: [14, 18, 19, 27]

573 + # query_start_loc: [0, 15, 20, 28]

574 + # seq_lens: [41, 31, 40]

575 +

576 + # Find how many decode indices belong to each request

577 + # request_ids: [0, 1, 1, 2]

578 + request_ids = torch.bucketize(logits_indices,

579 + query_start_loc[1:],

580 + right=True)

581 +

582 + # Figure out how many tokens are in each request

583 + # num_decode_tokens: [1, 2, 1]

584 + num_decode_tokens = torch.bincount(request_ids, minlength=num_reqs)

585 +

586 + # Calculate new query_start_loc with tokens in generation_indices

587 + # decode_query_start_loc: [0, 1, 3, 4]

588 + decode_query_start_loc = torch.empty(num_reqs + 1,

589 + device=query_start_loc.device,

590 + dtype=query_start_loc.dtype)

591 +

592 + decode_query_start_loc[0] = 0

593 + decode_query_start_loc[1:] = torch.cumsum(num_decode_tokens, dim=0)

594 + decode_max_query_len = int(num_decode_tokens.max().item())

595 + total_num_decode_tokens = int(num_decode_tokens.sum().item())

596 +

597 + common_attn_metadata = CommonAttentionMetadata(

598 + query_start_loc=decode_query_start_loc,

599 + query_start_loc_cpu=decode_query_start_loc.to("cpu",

600 + non_blocking=True),

601 + seq_lens=seq_lens,

602 + seq_lens_cpu=seq_lens.to("cpu", non_blocking=True),

603 + num_computed_tokens_cpu=common_attn_metadata.num_computed_tokens_cpu,

604 + num_reqs=num_reqs,

605 + num_actual_tokens=total_num_decode_tokens,

606 + max_query_len=decode_max_query_len,

607 + max_seq_len=common_attn_metadata.max_seq_len,

608 + block_table_tensor=common_attn_metadata.block_table_tensor,

609 + slot_mapping=common_attn_metadata.slot_mapping,

610 + causal=True,

611 + )

612 + return common_attn_metadata

613 +

614 +

615 def subclass_attention_backend(

616 name_prefix: str, attention_backend_cls: type[AttentionBackend],

617 builder_cls: type[AttentionMetadataBuilder[M]]

@@ -679,13 +749,56 @@ def subclass_attention_metadata(

749 return Wrapped

750

751

682 -def make_kv_sharing_fast_prefill_attention_metadata(

683 - metadata_cls: Any, ) -> Any:

684 - """

685 - Return a new subclass of `metadata_cls` for fast prefill

686 - """

687 - return subclass_attention_metadata(

688 - name_prefix="KVSharingFastPrefill",

689 - metadata_cls=metadata_cls,

690 - fields=KV_SHARING_FAST_PREFILL_METADATA_FIELDS,

691 - )

752 +@runtime_checkable

753 +class KVSharingFastPrefillMetadata(Protocol):

754 + logits_indices_padded: torch.Tensor

755 + num_logits_indices: int

756 +

757 +

758 +def create_fast_prefill_custom_backend(

759 + prefix: str,

760 + underlying_attn_backend: AttentionBackend,

761 +) -> type[AttentionBackend]:

762 +

763 + underlying_builder = underlying_attn_backend.get_builder_cls()

764 +

765 + class FastPrefillAttentionBuilder(underlying_builder): # type: ignore

766 +

767 + def build(self,

768 + common_prefix_len: int,

769 + common_attn_metadata: CommonAttentionMetadata,

770 + fast_build: bool = False) -> AttentionMetadata:

771 + new_common_attn_metadata =\

772 + make_kv_sharing_fast_prefill_common_attn_metadata(common_attn_metadata)

773 + metadata = super().build(common_prefix_len,

774 + new_common_attn_metadata, fast_build)

775 +

776 + class KVSharingFastPrefillAttentionMetadata(

777 + metadata.__class__, # type: ignore

778 + KVSharingFastPrefillMetadata):

779 +

780 + def __init__(self, metadata, common_attn_metadata):

781 + # Shallow copy all fields in metadata cls

782 + for field in fields(metadata.__class__):

783 + setattr(self, field.name,

784 + getattr(metadata, field.name))

785 +

786 + # Set additional fields that will be used in model code

787 + assert (common_attn_metadata.logits_indices_padded

788 + is not None

789 + and common_attn_metadata.num_logits_indices

790 + is not None)

791 + self.logits_indices_padded = \

792 + common_attn_metadata.logits_indices_padded

793 + self.num_logits_indices = \

794 + common_attn_metadata.num_logits_indices

795 +

796 + return KVSharingFastPrefillAttentionMetadata(

797 + metadata, common_attn_metadata)

798 +

799 + attn_backend = subclass_attention_backend(

800 + name_prefix=prefix,

801 + attention_backend_cls=underlying_attn_backend,

802 + builder_cls=FastPrefillAttentionBuilder)

803 +

804 + return attn_backend

vllm/v1/engine/async_llm.py

@@ -335,6 +335,13 @@ class AsyncLLM(EngineClient):

335 returning the RequestOutput back to the caller.

336 """

337

338 + if (self.vllm_config.cache_config.kv_sharing_fast_prefill

339 + and sampling_params.prompt_logprobs):

340 + raise ValueError(

341 + "--kv-sharing-fast-prefill produces incorrect logprobs for "

342 + "prompt tokens, please disable it when the requests need "

343 + "prompt logprobs")

344 +

345 try:

346 # We start the output_handler on the first call to generate() so

347 # we can call __init__ before the event loop, which enables us

vllm/v1/worker/gpu_model_runner.py

+53 -43

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -import dataclasses

4 import gc

5 import itertools

6 import time

@@ -58,7 +57,7 @@ from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,

57 supports_dynamo)

58 from vllm.v1.attention.backends.utils import (

59 AttentionCGSupport, AttentionMetadataBuilder, CommonAttentionMetadata,

61 - make_kv_sharing_fast_prefill_attention_metadata,

60 + create_fast_prefill_custom_backend,

61 reorder_batch_to_split_decodes_and_prefills)

62 from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher

63 from vllm.v1.kv_cache_interface import (AttentionSpec,

@@ -84,9 +83,10 @@ from vllm.v1.worker.kv_connector_model_runner_mixin import (

83 KVConnectorModelRunnerMixin, KVConnectorOutput)

84 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin

87 -from .utils import (AttentionGroup, MultiModalBudget, bind_kv_cache,

88 - gather_mm_placeholders, initialize_kv_cache_for_kv_sharing,

89 - sanity_check_mm_encoder_outputs, scatter_mm_placeholders)

86 +from .utils import (AttentionGroup, MultiModalBudget,

87 + add_kv_sharing_layers_to_kv_cache_groups, bind_kv_cache,

88 + gather_mm_placeholders, sanity_check_mm_encoder_outputs,

89 + scatter_mm_placeholders)

91 if TYPE_CHECKING:

92 import xgrammar as xgr

@@ -860,6 +860,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

860 max_seq_len=max_seq_len,

861 block_table_tensor=blk_table_tensor,

862 slot_mapping=slot_mapping,

863 + logits_indices_padded=logits_indices_padded,

864 + num_logits_indices=logits_indices.size(0),

865 causal=True,

866 )

867

@@ -884,28 +886,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

886 common_attn_metadata=common_attn_metadata,

887 ))

888

887 - fast_prefill_metadata = attn_metadata_i

888 - if (self.cache_config.kv_sharing_fast_prefill

889 - and self.kv_sharing_fast_prefill_eligible_layers):

890 - # Dynamically create a a dataclass type that inherits

891 - # from attention metadata type but includes additional

892 - # fields logits_indices_padded and num_logits_indices

893 - # which are required for prefill truncation

894 - fast_prefill_metadata_type = (

895 - make_kv_sharing_fast_prefill_attention_metadata(

896 - metadata_cls=type(attn_metadata_i), ))

897 - fast_prefill_metadata = fast_prefill_metadata_type(

898 - **dataclasses.asdict(attn_metadata_i),

899 - logits_indices_padded=logits_indices_padded,

900 - num_logits_indices=logits_indices.size(0),

901 - )

902 -

889 for layer_name in attn_group.layer_names:

904 - if (self.cache_config.kv_sharing_fast_prefill

905 - and layer_name

906 - in self.kv_sharing_fast_prefill_eligible_layers):

907 - attn_metadata[layer_name] = fast_prefill_metadata

908 - continue

890 attn_metadata[layer_name] = attn_metadata_i

891

892 # Hot-Swap lora model

@@ -1484,6 +1465,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1465 return self.kv_connector_no_forward(scheduler_output,

1466 self.vllm_config)

1467

1468 + if self.cache_config.kv_sharing_fast_prefill:

1469 + assert not self.input_batch.num_prompt_logprobs, (

1470 + "--kv-sharing-fast-prefill produces incorrect logprobs for "

1471 + "prompt tokens, tokens, please disable it when the requests "

1472 + "need prompt logprobs")

1473 +

1474 # Prepare the decoder inputs.

1475 (attn_metadata, logits_indices, spec_decode_metadata,

1476 num_scheduled_tokens_np, spec_decode_common_attn_metadata,

@@ -2742,6 +2729,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2729 # layer.

2730 for layer_name in layer_names:

2731 attn_backend = layers[layer_name].get_attn_backend()

2732 +

2733 + if layer_name in self.kv_sharing_fast_prefill_eligible_layers:

2734 + attn_backend = create_fast_prefill_custom_backend(

2735 + "FastPrefill",

2736 + attn_backend,

2737 + )

2738 +

2739 key = attn_backend.full_cls_name()

2740 attn_backends[key] = attn_backend

2741 attn_backend_layers[key].append(layer_name)

@@ -3074,20 +3068,40 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3068 kv_caches = self._reshape_kv_cache_tensors(kv_cache_config,

3069 kv_cache_raw_tensors)

3070

3077 - # Setup `kv_cache_config` and `kv_caches` for models

3078 - # with cross-layer KV sharing

3079 - if self.shared_kv_cache_layers:

3080 - initialize_kv_cache_for_kv_sharing(

3081 - self.shared_kv_cache_layers,

3082 - kv_cache_config.kv_cache_groups,

3083 - kv_caches,

3084 - self.attn_groups,

3085 - self.runner_only_attn_layers,

3086 - )

3071 + # Set up cross-layer KV cache sharing

3072 + for layer_name, target_layer_name in self.shared_kv_cache_layers.items(

3073 + ):

3074 + logger.debug("%s reuses KV cache of %s", layer_name,

3075 + target_layer_name)

3076 + kv_caches[layer_name] = kv_caches[target_layer_name]

3077 +

3078 + bind_kv_cache(kv_caches,

3079 + self.compilation_config.static_forward_context,

3080 + self.kv_caches)

3081 + return kv_caches

3082 +

3083 + def maybe_add_kv_sharing_layers_to_kv_cache_groups(

3084 + self, kv_cache_config: KVCacheConfig) -> None:

3085 + """

3086 + Add layers that re-use KV cache to KV cache group of its target layer.

3087 + Mapping of KV cache tensors happens in `initialize_kv_cache_tensors()`

3088 + """

3089 + if not self.shared_kv_cache_layers:

3090 + # No cross-layer KV sharing, return

3091 + return

3092 +

3093 + add_kv_sharing_layers_to_kv_cache_groups(

3094 + self.shared_kv_cache_layers,

3095 + kv_cache_config.kv_cache_groups,

3096 + self.runner_only_attn_layers,

3097 + )

3098 +

3099 + if self.cache_config.kv_sharing_fast_prefill:

3100 + # In You Only Cache Once (https://arxiv.org/abs/2405.05254) or other

3101 + # similar KV sharing setups, only the layers that generate KV caches

3102 + # are involved in the prefill phase, enabling prefill to early exit.

3103 attn_layers = get_layers_from_vllm_config(self.vllm_config,

3104 Attention)

3089 - # Iterate in reversed order and add layers that re-use KV cache

3090 - # e.g. in YOCO-like KV sharing setups (e.g. Gemma3n)

3105 for layer_name in reversed(attn_layers):

3106 if layer_name in self.shared_kv_cache_layers:

3107 self.kv_sharing_fast_prefill_eligible_layers.add(

@@ -3095,11 +3109,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3109 else:

3110 break

3111

3098 - bind_kv_cache(kv_caches,

3099 - self.compilation_config.static_forward_context,

3100 - self.kv_caches)

3101 - return kv_caches

3102 -

3112 def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:

3113 """

3114 Initialize KV cache based on `kv_cache_config`.

@@ -3111,6 +3120,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3120 self.kv_cache_config = kv_cache_config

3121 self.may_reinitialize_input_batch(kv_cache_config)

3122 self.may_add_encoder_only_layers_to_kv_cache_config()

3123 + self.maybe_add_kv_sharing_layers_to_kv_cache_groups(kv_cache_config)

3124 self.initialize_attn_backend(kv_cache_config)

3125 kv_caches = self.initialize_kv_cache_tensors(kv_cache_config)

3126

vllm/v1/worker/tpu_model_runner.py

+28 -11

@@ -55,9 +55,8 @@ from vllm.v1.worker.kv_connector_model_runner_mixin import (

55 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin

56 from vllm.v1.worker.tpu_input_batch import CachedRequestState, InputBatch

58 -from .utils import (MultiModalBudget, bind_kv_cache,

59 - initialize_kv_cache_for_kv_sharing,

60 - sanity_check_mm_encoder_outputs)

58 +from .utils import (MultiModalBudget, add_kv_sharing_layers_to_kv_cache_groups,

59 + bind_kv_cache, sanity_check_mm_encoder_outputs)

61 if TYPE_CHECKING:

62 from vllm.v1.core.sched.output import SchedulerOutput

@@ -1599,6 +1598,30 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1598 self.encoder_cache.clear()

1599 gc.collect()

1600

1601 + def maybe_setup_cross_layer_kv_sharing(

1602 + self,

1603 + kv_caches: dict[str, torch.Tensor],

1604 + kv_cache_config: KVCacheConfig,

1605 + ) -> None:

1606 + """

1607 + Add layers that re-use KV cache to KV cache group of its target layer.

1608 + Mapping of KV cache tensors happens in `initialize_kv_cache_tensors()`

1609 + """

1610 + if not self.shared_kv_cache_layers:

1611 + # No cross-layer KV sharing, return

1612 + return

1613 +

1614 + add_kv_sharing_layers_to_kv_cache_groups(

1615 + self.shared_kv_cache_layers,

1616 + kv_cache_config.kv_cache_groups,

1617 + )

1618 +

1619 + for layer_name, target_layer_name in self.shared_kv_cache_layers.items(

1620 + ):

1621 + logger.debug("%s reuses KV cache of %s", layer_name,

1622 + target_layer_name)

1623 + kv_caches[layer_name] = kv_caches[target_layer_name]

1624 +

1625 def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:

1626 """

1627 Initialize KV cache based on `kv_cache_config`.

@@ -1664,14 +1687,8 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1687 else:

1688 raise NotImplementedError

1689

1667 - # Setup `kv_cache_config` and `kv_caches` for models

1668 - # with cross-layer KV sharing

1669 - if self.shared_kv_cache_layers:

1670 - initialize_kv_cache_for_kv_sharing(

1671 - self.shared_kv_cache_layers,

1672 - kv_cache_config.kv_cache_groups,

1673 - kv_caches,

1674 - )

1690 + # Set up cross-layer KV cache sharing if needed

1691 + self.maybe_setup_cross_layer_kv_sharing(kv_caches, kv_cache_config)

1692

1693 bind_kv_cache(

1694 kv_caches,

vllm/v1/worker/utils.py

+7 -33

@@ -203,12 +203,9 @@ def gather_mm_placeholders(

203 return placeholders[is_embed]

204

205

206 -def initialize_kv_cache_for_kv_sharing(

206 +def add_kv_sharing_layers_to_kv_cache_groups(

207 shared_kv_cache_layers: dict[str, str],

208 kv_cache_groups: list[KVCacheGroupSpec],

209 - kv_caches: dict[str, torch.Tensor],

210 - # Optional for now to avoid breaking TPU

211 - attn_groups: Optional[list[list[AttentionGroup]]] = None,

209 runner_only_attn_layers: Optional[set[str]] = None,

210 ) -> None:

211 """

@@ -223,38 +220,15 @@ def initialize_kv_cache_for_kv_sharing(

220 means this layer will perform attention using the keys and values

221 from the KV cache of `shared_kv_cache_layers[layer_name]`.

222 kv_cache_groups: The KV cache groups of the model.

226 - kv_caches: The allocated kv_caches with layer names as keys.

227 - Note that layers in shared_kv_cache_layers.keys() are not

228 - originally included as it only contains layers which have its own

229 - KV cache allocation.

230 - attn_groups: Optional list of attention groups. Layers in the same KV

231 - cache group may be placed in different attention groups if they

232 - have different attention backends. Currently only provided by

233 - GPU model runner.

223 """

235 - # mapping from layer name to tuple of (kv_cache_group_idx, attn_group_idx)

236 - layer_to_attn_group_idx: dict[str, tuple[int, int]] = {}

237 - if attn_groups:

238 - for kv_cache_group_idx, kv_attn_groups in enumerate(attn_groups):

239 - for attn_group_idx, attn_group in enumerate(kv_attn_groups):

240 - for layer_name in attn_group.layer_names:

241 - layer_to_attn_group_idx[layer_name] = (kv_cache_group_idx,

242 - attn_group_idx)

243 - else:

244 - for kv_cache_group_idx, kv_cache_group in enumerate(kv_cache_groups):

245 - for layer_name in kv_cache_group.layer_names:

246 - # attn group idx default to 0 if not provided

247 - layer_to_attn_group_idx[layer_name] = (kv_cache_group_idx, 0)

224 + layer_to_kv_cache_group: dict[str, KVCacheGroupSpec] = {}

225 + for kv_cache_group in kv_cache_groups:

226 + for layer_name in kv_cache_group.layer_names:

227 + layer_to_kv_cache_group[layer_name] = kv_cache_group

228

229 for layer_name, target_layer_name in shared_kv_cache_layers.items():

250 - kv_caches[layer_name] = kv_caches[target_layer_name]

251 - kv_cache_group_idx = layer_to_attn_group_idx[target_layer_name][0]

252 - kv_cache_groups[kv_cache_group_idx].layer_names.append(layer_name)

253 -

254 - if attn_groups:

255 - attn_group_idx = layer_to_attn_group_idx[target_layer_name][1]

256 - attn_groups[kv_cache_group_idx][attn_group_idx].layer_names.append(

257 - layer_name)

230 + tgt_kv_cache_group = layer_to_kv_cache_group[target_layer_name]

231 + tgt_kv_cache_group.layer_names.append(layer_name)

232

233 if runner_only_attn_layers is not None:

234 runner_only_attn_layers.add(layer_name)