vllm.v1.worker.gpu.attn_utils ¶

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,
    page_aligned_layers: set[str],
) -> 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
        if not page_aligned_layers.isdisjoint(kv_tensor.shared_by):
            # These KV-first views already address blocks by page, just like
            # blocks-first views, so they must not be restrided onto the
            # unaligned KV-first storage layout.
            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]
                )

def _kv_first_layers_sharing_pool_with_mamba(
    attn_groups: Sequence[AttentionGroup],
    kernel_block_sizes: list[int],
    cache_dtype: str,
    kv_cache_config: "KVCacheConfig | None",
) -> set[str]:
    """KV-first attention layers that share an allocation with Mamba state.

    Mamba addresses its state by page, so these layers have to lay their blocks
    out page by page as well; otherwise block ids from the two groups resolve
    to overlapping bytes.
    """
    if kv_cache_config is None:
        return set()

    kv_first_layers: set[str] = set()
    mamba_layers: set[str] = set()
    for group in attn_groups:
        kv_cache_spec = group.kv_cache_spec
        if isinstance(kv_cache_spec, MambaSpec):
            mamba_layers.update(group.layer_names)
            continue
        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,
        )
        if block_dim != 0:
            kv_first_layers.update(group.layer_names)

    page_aligned: set[str] = set()
    for kv_tensor in kv_cache_config.kv_cache_tensors:
        if mamba_layers.isdisjoint(kv_tensor.shared_by):
            continue
        page_aligned.update(kv_first_layers.intersection(kv_tensor.shared_by))
    return page_aligned

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

def get_query_lens_mismatch_unsupported_backend(
    attn_groups: list[list[AttentionGroup]],
) -> 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 not group.backend.supports_device_cpu_query_lens_mismatch():
                return group.backend.__name__
    return None