Skip to content

vllm.v1.worker.gpu.attn_utils

Classes:

Functions:

AttentionCGSupportInfo dataclass

Methods:

  • narrow –

    Return an info tightened by support if it is more restrictive.

Source code in vllm/v1/worker/gpu/attn_utils.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
@dataclass(frozen=True)
class AttentionCGSupportInfo:
    min_cg_support: AttentionCGSupport = AttentionCGSupport.ALWAYS
    min_cg_attn_backend: str | None = None

    def narrow(
        self, support: AttentionCGSupport, backend: str | None
    ) -> "AttentionCGSupportInfo":
        """Return an info tightened by ``support`` if it is more restrictive.

        Lets attention groups built outside ``init_attn_backend`` (e.g.
        encoder-only layers) contribute to the runner's cudagraph decision.
        """
        if support.value < self.min_cg_support.value:
            return AttentionCGSupportInfo(support, backend)
        return self

narrow(support, backend)

Return an info tightened by support if it is more restrictive.

Lets attention groups built outside init_attn_backend (e.g. encoder-only layers) contribute to the runner's cudagraph decision.

Source code in vllm/v1/worker/gpu/attn_utils.py 
49
50
51
52
53
54
55
56
57
58
59
def narrow(
    self, support: AttentionCGSupport, backend: str | None
) -> "AttentionCGSupportInfo":
    """Return an info tightened by ``support`` if it is more restrictive.

    Lets attention groups built outside ``init_attn_backend`` (e.g.
    encoder-only layers) contribute to the runner's cudagraph decision.
    """
    if support.value < self.min_cg_support.value:
        return AttentionCGSupportInfo(support, backend)
    return self

FastPrefillBatchMetadata dataclass

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

Source code in vllm/v1/worker/gpu/attn_utils.py 
62
63
64
65
66
67
68
@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, and stages the logits indices it runs on.

Source code in vllm/v1/worker/gpu/attn_utils.py 
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
class FastPrefillHelper:
    """Decides per step whether to arm the KV-sharing fast prefill path, and
    stages the logits indices it runs on.
    """

    def __init__(self, cudagraph_manager: "CudaGraphManager", max_num_tokens: int):
        self.max_num_tokens = max_num_tokens
        self.cudagraph_manager = cudagraph_manager
        self.logits_indices_buf = torch.zeros(
            max_num_tokens, dtype=torch.int32, device=cudagraph_manager.device
        )

    def prepare(
        self,
        logits_indices: torch.Tensor,
        num_reqs: int,
        cu_num_logits_np: np.ndarray,
        has_prefill: bool,
        batch_desc: "BatchExecutionDescriptor",
    ) -> FastPrefillBatchMetadata | None:
        if (
            not has_prefill
            or batch_desc.cg_mode == CUDAGraphMode.FULL
            or batch_desc.num_ubatches > 1
        ):
            return None
        buf = self.logits_indices_buf
        num_logits = logits_indices.shape[0]
        assert num_logits > 0
        buf[:num_logits].copy_(logits_indices)
        # There might be leftover indices in buf[num_logits:] from previous
        # iterations. Broadcast the scalar GPU-side to keep them valid.
        buf[num_logits:] = logits_indices[-1]
        # Pad so the model's KV-sharing layers run their reduced (logits-only)
        # batch at a captured piecewise cudagraph size. Pre-capture, or with
        # cudagraphs off, dispatch returns the unpadded count.
        desc = self.cudagraph_manager.dispatch(
            num_reqs=num_reqs,
            num_tokens=num_logits,
            uniform_token_count=None,
            num_active_loras=0,
        )
        num_logits_padded = min(desc.num_tokens, self.max_num_tokens)
        return FastPrefillBatchMetadata(
            logits_indices_padded=buf[:num_logits_padded],
            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()),
        )

compute_mm_prefix_ranges(req_ids, mm_features, sliding_window=None)

Compute PrefixLM bidirectional ranges for multimodal tokens.

Ranges exceeding sliding_window are skipped to prevent early tokens from attending across the entire image span.

Source code in vllm/v1/worker/gpu/attn_utils.py 
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
def compute_mm_prefix_ranges(
    req_ids: list[str],
    mm_features: dict[str, list[MultiModalFeatureSpec]],
    sliding_window: int | None = None,
) -> dict[int, list[tuple[int, int]]]:
    """Compute PrefixLM bidirectional ranges for multimodal tokens.

    Ranges exceeding sliding_window are skipped to prevent early tokens
    from attending across the entire image span.
    """
    req_doc_ranges: dict[int, list[tuple[int, int]]] = {}
    for req_idx, req_id in enumerate(req_ids):
        image_doc_ranges = []
        for mm_feature in mm_features.get(req_id, ()):
            if mm_feature.modality not in ("image", "video"):
                continue
            for r in mm_feature.mm_position.extract_embeds_range():
                if sliding_window is not None and (r[1] - r[0] + 1) > sliding_window:
                    continue
                image_doc_ranges.append(r)
        req_doc_ranges[req_idx] = image_doc_ranges
    return req_doc_ranges

get_attn_cg_support(attn_groups, vllm_config, checked_layer_names=None)

Return the weakest CUDA graph support among the checked layers.

Source code in vllm/v1/worker/gpu/attn_utils.py 
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
def get_attn_cg_support(
    attn_groups: list[list[AttentionGroup]],
    vllm_config: VllmConfig,
    checked_layer_names: set[str] | None = None,
) -> AttentionCGSupportInfo:
    """Return the weakest CUDA graph support among the checked layers."""
    min_cg_support = AttentionCGSupport.ALWAYS
    min_cg_attn_backend = None
    for groups in attn_groups:
        for group in groups:
            if checked_layer_names is not None and checked_layer_names.isdisjoint(
                group.layer_names
            ):
                continue
            builder = group.get_metadata_builder(0)
            cg_support = builder.get_cudagraph_support(
                vllm_config,
                group.kv_cache_spec,
            )
            if cg_support.value < min_cg_support.value:
                min_cg_support = cg_support
                min_cg_attn_backend = group.backend.__name__
    return AttentionCGSupportInfo(
        min_cg_support=min_cg_support,
        min_cg_attn_backend=min_cg_attn_backend,
    )

get_kv_sharing_fast_prefill_eligible_layers(vllm_config, draft_layer_names=None)

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.

Speculator draft layers register after the target model's layers (and may themselves share KV), so they are excluded from the walk.

Source code in vllm/v1/worker/gpu/attn_utils.py 
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
def get_kv_sharing_fast_prefill_eligible_layers(
    vllm_config: VllmConfig, draft_layer_names: set[str] | None = None
) -> 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.

    Speculator draft layers register after the target model's layers (and may
    themselves share KV), so they are excluded from the walk.
    """
    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 draft_layer_names is not None and layer_name in draft_layer_names:
            continue
        if layer_name not in shared_kv_cache_layers:
            break
        eligible_layers.add(layer_name)
    return eligible_layers

get_query_lens_mismatch_unsupported_backend(attn_groups, checked_layer_names=None)

Name the first backend needing the CPU query lengths to be exact, if any.

The attention selector already excludes these when adaptive verification is enabled, but models that hard-wire their backend never consult it. See AttentionBackend.supports_device_cpu_query_lens_mismatch().

Source code in vllm/v1/worker/gpu/attn_utils.py 
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
def get_query_lens_mismatch_unsupported_backend(
    attn_groups: list[list[AttentionGroup]],
    checked_layer_names: set[str] | None = None,
) -> str | None:
    """Name the first backend needing the CPU query lengths to be exact, if any.

    The attention selector already excludes these when adaptive verification is
    enabled, but models that hard-wire their backend never consult it. See
    AttentionBackend.supports_device_cpu_query_lens_mismatch().
    """
    for groups in attn_groups:
        for group in groups:
            if checked_layer_names is not None and checked_layer_names.isdisjoint(
                group.layer_names
            ):
                continue
            if not group.backend.supports_device_cpu_query_lens_mismatch():
                return group.backend.__name__
    return None