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

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