vllm.v1.worker.gpu.attn_utils

Classes:

• AttentionCGSupportInfo –
• FastPrefillBatchMetadata –

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

• FastPrefillHelper –

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

Functions:

• bind_kv_caches –

  Bind per-layer KV cache views to the attention layers.

• get_attn_cg_support –

  Return the weakest CUDA graph support among the checked layers.

• get_kv_sharing_fast_prefill_eligible_layers –

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

• get_query_lens_mismatch_unsupported_backend –

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

• get_varlen_cudagraph_unsupported_backend –

  Name the first backend whose FULL cudagraphs cannot replay decode batches

• init_attn_backend –

  maps_slot_mapping: the caller maps every slot mapping it builds to

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

@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

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

@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

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",
        num_active_loras: int = 0,
    ) -> FastPrefillBatchMetadata | None:
        if (
            not has_prefill
            or batch_desc.cg_mode == CUDAGraphMode.FULL
            or batch_desc.num_ubatches > 1
            # LoRA metadata describes the full batch, not the logits-only batch.
            or num_active_loras > 0
        ):
            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=num_active_loras,
        )
        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()),
        )

bind_kv_caches(kv_caches, forward_context, kv_cache_config, vllm_config, attn_groups=())

Bind per-layer KV cache views to the attention layers.

Source code in vllm/v1/worker/gpu/attn_utils.py

def bind_kv_caches(
    kv_caches: dict[str, torch.Tensor],
    forward_context: dict[str, Any],
    kv_cache_config: KVCacheConfig,
    vllm_config: VllmConfig,
    attn_groups: Iterable[AttentionGroup] = (),
) -> dict[str, Any]:
    """Bind per-layer KV cache views to the attention layers."""
    for layer_name, target in get_shared_kv_cache_layers(vllm_config).items():
        kv_caches[layer_name] = kv_caches[target]
    bindable_caches = map_kv_caches_to_kernel_blocks(
        {name: cache for name, cache in kv_caches.items() if name in forward_context},
        attn_groups,
    )
    bind_kv_cache_to_layers(bindable_caches, forward_context)
    return kv_caches

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

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

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

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

get_varlen_cudagraph_unsupported_backend(attn_groups, vllm_config, max_query_len, checked_layer_names=None)

Name the first backend whose FULL cudagraphs cannot replay decode batches of 1 to max_query_len tokens per request, with its bound, if any.

ALWAYS builders replay any batch; the others need a bound of at least max_query_len. See AttentionMetadataBuilder.get_varlen_cudagraph_max_query_len().

Source code in vllm/v1/worker/gpu/attn_utils.py

def get_varlen_cudagraph_unsupported_backend(
    attn_groups: list[list[AttentionGroup]],
    vllm_config: VllmConfig,
    max_query_len: int,
    checked_layer_names: set[str] | None = None,
) -> tuple[str, int | None] | None:
    """Name the first backend whose FULL cudagraphs cannot replay decode batches
    of 1 to max_query_len tokens per request, with its bound, if any.

    ALWAYS builders replay any batch; the others need a bound of at least
    max_query_len. See AttentionMetadataBuilder.get_varlen_cudagraph_max_query_len().
    """
    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)
            spec = group.kv_cache_spec
            support = builder.get_cudagraph_support(vllm_config, spec)
            if support == AttentionCGSupport.ALWAYS:
                continue
            bound = (
                # NEVER means never, whatever an inherited override says.
                None
                if support == AttentionCGSupport.NEVER
                else builder.get_varlen_cudagraph_max_query_len(vllm_config, spec)
            )
            if bound is None or bound < max_query_len:
                return group.backend.__name__, bound
    return None

init_attn_backend(kv_cache_config, vllm_config, device, active_layer_names=None, draft_layer_names=None, maps_slot_mapping=False, slot_mapped_layers=())

maps_slot_mapping: the caller maps every slot mapping it builds to kernel blocks (AttentionGroup.map_to_kernel_slot_mapping); or only for the slot_mapped_layers, which another runner (a speculator) writes.

Source code in vllm/v1/worker/gpu/attn_utils.py

def init_attn_backend(
    kv_cache_config: KVCacheConfig,
    vllm_config: VllmConfig,
    device: torch.device,
    active_layer_names: set[str] | None = None,
    draft_layer_names: set[str] | None = None,
    maps_slot_mapping: bool = False,
    slot_mapped_layers: Collection[str] = (),
) -> tuple[list[list[AttentionGroup]], AttentionCGSupportInfo, list[int]]:
    """`maps_slot_mapping`: the caller maps every slot mapping it builds to
    kernel blocks (AttentionGroup.map_to_kernel_slot_mapping); or only for the
    `slot_mapped_layers`, which another runner (a speculator) writes."""
    # Phase 1: discover attention groups for each kv cache group.
    attn_groups: list[list[AttentionGroup]] = []

    # Add KV-sharing layers to their target's kv cache group so they are
    # discovered alongside the target layer in Phase 1 below.
    shared_kv_cache_layers = get_shared_kv_cache_layers(vllm_config)
    add_kv_sharing_layers_to_kv_cache_groups(
        shared_kv_cache_layers, kv_cache_config.kv_cache_groups
    )
    fast_prefill_eligible_layers = get_kv_sharing_fast_prefill_eligible_layers(
        vllm_config, draft_layer_names
    )

    # Phase 1: discover attention groups for each kv cache group.
    for kv_cache_group_id, kv_cache_group_spec in enumerate(
        kv_cache_config.kv_cache_groups
    ):
        layer_names = kv_cache_group_spec.layer_names
        if not kv_cache_group_spec.kv_cache_spec.has_layer_views:
            attn_groups.append([])
            continue
        if active_layer_names is not None:
            layer_names = list(active_layer_names.intersection(layer_names))

        layer_type = cast(type[Any], AttentionLayerBase)
        attn_layers = get_layers_from_vllm_config(vllm_config, layer_type, layer_names)

        group_map: dict[
            tuple[tuple[str, str], KVCacheSpec, int, bool], AttentionGroup
        ] = {}
        group_order: list[tuple[tuple[str, str], KVCacheSpec, int, bool]] = []

        for layer_name in layer_names:
            attn_backend = attn_layers[layer_name].get_attn_backend()
            if (
                envs.VLLM_BATCH_INVARIANT
                and not attn_backend.supports_batch_invariance()
            ):
                raise RuntimeError(
                    "VLLM batch_invariant mode is not supported for "
                    f"{attn_backend.get_name()}."
                )
            if layer_name in fast_prefill_eligible_layers:
                attn_backend = create_fast_prefill_custom_backend(
                    "FastPrefill", attn_backend
                )

            layer_kv_cache_spec: KVCacheSpec = kv_cache_group_spec.kv_cache_spec
            if isinstance(layer_kv_cache_spec, UniformTypeKVCacheSpecs):
                # KV-sharing layers have no spec of their own; use the target's.
                spec_layer_name = shared_kv_cache_layers.get(layer_name, layer_name)
                layer_kv_cache_spec = layer_kv_cache_spec.kv_cache_specs[
                    spec_layer_name
                ]

            # Split on per-rank num_heads_q so layers with different Q-head
            # counts (e.g. a spec-decode draft head and its target) get separate
            # metadata builders, and on whether a speculator maps the layer's
            # slot mappings.
            num_heads_q = getattr(attn_layers[layer_name], "num_heads", 0)
            maps = maps_slot_mapping or layer_name in slot_mapped_layers
            key = (
                attn_backend.full_cls_name(),
                layer_kv_cache_spec,
                num_heads_q,
                maps,
            )
            if key not in group_map:
                group_map[key] = AttentionGroup(
                    attn_backend,
                    [layer_name],
                    layer_kv_cache_spec,
                    kv_cache_group_id,
                    maps_slot_mapping=maps,
                )
                group_order.append(key)
            else:
                group_map[key].layer_names.append(layer_name)

        attn_groups.append([group_map[key] for key in group_order])

    # Phase 2: pick a kernel block size per kv cache group that is supported
    # by all backends within that group.
    kernel_block_sizes = prepare_kernel_block_sizes(kv_cache_config, attn_groups)

    if current_platform.is_xpu():
        for kv_cache_group_id, selected_kernel_block_size in enumerate(
            kernel_block_sizes
        ):
            if selected_kernel_block_size % 64 != 0:
                for attn_group in attn_groups[kv_cache_group_id]:
                    logger.warning_once(
                        "On XPU platform, current selected kernel block size is "
                        "not fully optimized: kernel_block_size=%s, backend=%s, "
                        "kv_cache_spec=%s",
                        selected_kernel_block_size,
                        attn_group.backend,
                        repr(attn_group.kv_cache_spec),
                    )
    # Phase 3: create metadata builders and determine cudagraph support.
    attn_backend_workspace: torch.Tensor | None = None
    for kv_cache_group_id, groups in enumerate(attn_groups):
        kernel_block_size: int | None = None
        if kv_cache_group_id < len(kernel_block_sizes):
            kernel_block_size = kernel_block_sizes[kv_cache_group_id]
        for group in groups:
            group.create_metadata_builders(
                vllm_config=vllm_config,
                device=device,
                kernel_block_size=kernel_block_size,
                # Microbatches build attention metadata concurrently, and some
                # builders keep the prepared metadata on themselves (MLA stores
                # it on the prefill backend), so each ubatch needs its own.
                num_metadata_builders=get_num_ubatches(vllm_config.parallel_config),
            )
            # The microbatches' builders share the workspace: they all issue
            # attention on the one compute stream the threads hand off, so the
            # buffer is written serially, as it already is across steps.
            for builder in group.metadata_builders:
                if attn_backend_workspace is None:
                    if hasattr(builder, "_get_workspace_buffer"):
                        attn_backend_workspace = builder._get_workspace_buffer()
                elif hasattr(builder, "set_workspace_buffer"):
                    builder.set_workspace_buffer(attn_backend_workspace)
    attn_cg_support_info = get_attn_cg_support(attn_groups, vllm_config)
    return attn_groups, attn_cg_support_info, kernel_block_sizes