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 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
@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 
48
49
50
51
52
53
54
55
56
57
58
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

_align_mixed_attention_kv_cache_views(attn_groups, kv_caches, kernel_block_sizes, cache_dtype, kv_cache_config)

Align shared attention KV views when backends disagree on layout.

Encoder-decoder models can share one raw allocation between decoder self-attention (K/V-first ROCM_ATTN, block dim 1) and cross-attention (blocks-first backends, block dim 0). Keep the physical storage in the K/V-first layout expected by ROCM_ATTN, and restride the blocks-first logical views so block IDs address the same bytes.

Source code in vllm/v1/worker/gpu/attn_utils.py 
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
def _align_mixed_attention_kv_cache_views(
    attn_groups: Iterable[AttentionGroup],
    kv_caches: dict[str, Any],
    kernel_block_sizes: list[int],
    cache_dtype: str,
    kv_cache_config: KVCacheConfig,
) -> None:
    """Align shared attention KV views when backends disagree on layout.

    Encoder-decoder models can share one raw allocation between decoder
    self-attention (K/V-first ROCM_ATTN, block dim 1) and cross-attention
    (blocks-first backends, block dim 0). Keep the physical storage in the
    K/V-first layout expected by ROCM_ATTN, and restride the blocks-first
    logical views so block IDs address the same bytes.
    """
    block_dims_by_layer: dict[str, int] = {}
    for group in attn_groups:
        kv_cache_spec = group.kv_cache_spec
        if not isinstance(kv_cache_spec, AttentionSpec):
            continue
        if group.kv_cache_group_id >= len(kernel_block_sizes):
            continue
        block_dim = group.backend.get_kv_cache_block_dim(
            kernel_block_sizes[group.kv_cache_group_id],
            kv_cache_spec.num_kv_heads,
            kv_cache_spec.head_size,
            cache_dtype_str=cache_dtype,
        )
        for layer_name in group.layer_names:
            if layer_name in kv_caches:
                block_dims_by_layer[layer_name] = block_dim

    for kv_tensor in kv_cache_config.kv_cache_tensors:
        if kv_tensor.block_stride > 0:
            continue
        shared_block_dims = {
            block_dims_by_layer[layer_name]
            for layer_name in kv_tensor.shared_by
            if layer_name in block_dims_by_layer
        }
        if 0 not in shared_block_dims or 1 not in shared_block_dims:
            continue

        for layer_name in kv_tensor.shared_by:
            if block_dims_by_layer.get(layer_name) == 0:
                _restride_blocks_first_kv_cache_to_kv_first_storage(
                    kv_caches[layer_name]
                )

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 
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
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