Skip to content

vllm.v1.worker.cpu.attn_utils

Classes:

Functions:

FastPrefillBatchMetadata dataclass

Per-step inputs for the KV-sharing fast prefill path.

Source code in vllm/v1/worker/cpu/attn_utils.py 
43
44
45
46
47
48
49
@dataclass(frozen=True)
class FastPrefillBatchMetadata:
    """Per-step inputs for the KV-sharing fast prefill path."""

    logits_indices_padded: torch.Tensor
    num_logits_indices: int
    max_logits_per_req: int

FastPrefillHelper

Decides per step whether to arm the KV-sharing fast prefill path.

Source code in vllm/v1/worker/cpu/attn_utils.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
class FastPrefillHelper:
    """Decides per step whether to arm the KV-sharing fast prefill path."""

    def prepare(
        self,
        logits_indices: torch.Tensor,
        cu_num_logits_np: np.ndarray,
        has_prefill: bool,
    ) -> FastPrefillBatchMetadata | None:
        if not has_prefill:
            return None

        num_logits = logits_indices.shape[0]
        assert num_logits > 0
        return FastPrefillBatchMetadata(
            logits_indices_padded=logits_indices,
            num_logits_indices=num_logits,
            # Largest per-request logits count, known on the host.
            max_logits_per_req=int(np.diff(cu_num_logits_np).max()),
        )

get_kv_sharing_fast_prefill_eligible_layers(vllm_config)

Trailing run of KV-sharing layers, eligible for fast prefill.

In You Only Cache Once (https://arxiv.org/abs/2405.05254) or other similar KV sharing setups, only the layers that generate KV caches are involved in the prefill phase, enabling prefill to early exit. Layers are registered in execution order, so the eligible layers are the contiguous suffix of KV-sharing layers.

Source code in vllm/v1/worker/cpu/attn_utils.py 
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
def get_kv_sharing_fast_prefill_eligible_layers(
    vllm_config: VllmConfig,
) -> set[str]:
    """Trailing run of KV-sharing layers, eligible for fast prefill.

    In You Only Cache Once (https://arxiv.org/abs/2405.05254) or other similar
    KV sharing setups, only the layers that generate KV caches are involved in
    the prefill phase, enabling prefill to early exit. Layers are registered in
    execution order, so the eligible layers are the contiguous suffix of
    KV-sharing layers.
    """
    if not vllm_config.cache_config.kv_sharing_fast_prefill:
        return set()
    shared_kv_cache_layers = get_shared_kv_cache_layers(vllm_config)
    if not shared_kv_cache_layers:
        return set()
    eligible_layers: set[str] = set()
    attn_layers = get_layers_from_vllm_config(vllm_config, Attention)
    for layer_name in reversed(attn_layers):
        if layer_name not in shared_kv_cache_layers:
            break
        eligible_layers.add(layer_name)
    return eligible_layers