vllm.v1.worker.utils

Classes:

• AttentionGroup –
• EncoderTimingStats –

  Per-request timing statistics for encoder forward pass.

• KVBlockZeroer –

  Manages efficient zeroing of KV cache blocks via a Triton kernel.

Functions:

• add_kv_sharing_layers_to_kv_cache_groups –

  Sets up KV cache sharing by reusing the allocated KV caches in kv_caches

• allocate_kv_cache –

  Allocate the KV cache and view it as [B, H, N, C] per layer.

• bind_kv_cache –

  Bind the allocated KV cache to both ModelRunner and forward context so

• bind_kv_cache_to_layers –

  Bind layer caches and share ReplaySSM trackers in model-layer order.

• clear_layer_kv_caches –

  Detach the KV/state cache tensors installed by bind_kv_cache().

• customize_attention_spec –

  The backend's spec. Blocks larger than every kernel block size a backend

• get_uniform_decode_token_count –

  Per-request token count of a uniform decode batch, or None.

• is_uniform_query_len –

  Whether every request in the batch has the same query length.

• map_kv_caches_to_kernel_blocks –

  The views to bind to layers. The manager views in kv_caches stay for

• prepare_kernel_block_sizes –

  Generate kernel_block_sizes that matches each block_size.

• request_memory –

  Calculate the amount of memory required by vLLM, then validate

• sanity_check_mm_encoder_outputs –

  Perform sanity checks for the result of

• select_common_block_size –

  Select a block size that is supported by all backends and is a factor of

AttentionGroup dataclass

Methods:

• map_kv_cache_to_kernel_blocks –

  Views a packed cache in manager blocks as kernel blocks.

• map_kv_cache_to_manager_blocks –

  Views a packed cache in kernel blocks as manager blocks.

Source code in vllm/v1/worker/utils.py

@dataclass(eq=False)
class AttentionGroup:
    backend: type[AttentionBackend]
    layer_names: list[str]
    kv_cache_spec: KVCacheSpec
    kv_cache_group_id: int
    # Set when the backend needs blocks smaller than the KV cache manager's: its
    # block tables and KV cache views then address kernel blocks, kernel block j
    # of block b being b * kernel_block_stride + j. The stride is set when a
    # packed KV cache is bound (other layers' pages sit between its blocks);
    # otherwise it stays 0 and the kernel blocks per block are used.
    kernel_block_size: int | None = None
    kernel_block_stride: int = 0
    # Persistent per ubatch: CUDA graphs capture the mapped block table.
    kernel_block_table: torch.Tensor | None = None
    kernel_block_offsets: torch.Tensor | None = None
    # When ubatching is enabled we will have a metadata builder for each ubatch
    # so that if they use internal persistent buffers for cudagraphs, and they
    # won't have to worry about conflicting with the other ubatches.
    metadata_builders: list[AttentionMetadataBuilder] = field(
        default_factory=lambda: []
    )

    def create_metadata_builders(
        self,
        vllm_config,
        device,
        kernel_block_size: int | None = None,
        num_metadata_builders: int = 1,
    ):
        spec = self.kv_cache_spec
        if kernel_block_size is not None and isinstance(spec, AttentionSpec):
            # A backend in a packed KV cache group may need smaller blocks than
            # the group is split into.
            kernel_block_size = select_common_block_size(
                kernel_block_size, [self.backend], [spec]
            )
            if kernel_block_size != spec.block_size:
                self.kernel_block_size = kernel_block_size
        if kernel_block_size is None:
            kv_cache_spec_builder = self.kv_cache_spec
        else:
            kv_cache_spec_builder = self.kv_cache_spec.copy_with_new_block_size(
                kernel_block_size
            )
        builder_cls = self.backend.get_builder_cls()
        builder_kwargs = {}
        if builder_cls.requires_block_table_width or self.kernel_block_size is not None:
            max_num_blocks = self.kv_cache_spec.max_num_blocks_per_req(
                vllm_config, vllm_config.model_config.max_model_len
            )
            width = get_block_table_width(max_num_blocks, self.kv_cache_spec.block_size)
            if kernel_block_size is not None:
                width *= self.kv_cache_spec.block_size // kernel_block_size
            if builder_cls.requires_block_table_width:
                builder_kwargs["block_table_width"] = width
        if self.kernel_block_size is not None:
            self.kernel_block_table = torch.zeros(
                num_metadata_builders,
                vllm_config.scheduler_config.max_num_seqs,
                width,
                dtype=torch.int32,
                device=device,
            )
            self.kernel_block_offsets = torch.arange(
                spec.block_size // self.kernel_block_size,
                dtype=torch.int32,
                device=device,
            )
        self.metadata_builders = [
            builder_cls(
                kv_cache_spec_builder,
                self.layer_names,
                vllm_config,
                device,
                **builder_kwargs,
            )
            for _ in range(num_metadata_builders)
        ]

    def build_metadata(
        self,
        common_attn_metadata: CommonAttentionMetadata,
        ubatch_idx: int = 0,
        common_prefix_len: int = 0,
        **kwargs,
    ) -> Any:
        return self.get_metadata_builder(ubatch_idx).build(
            common_prefix_len,
            self._with_kernel_block_table(common_attn_metadata, ubatch_idx),
            **kwargs,
        )

    def build_metadata_for_cudagraph_capture(
        self, common_attn_metadata: CommonAttentionMetadata, ubatch_idx: int = 0
    ) -> Any:
        return self.get_metadata_builder(ubatch_idx).build_for_cudagraph_capture(
            self._with_kernel_block_table(common_attn_metadata, ubatch_idx)
        )

    def build_metadata_for_drafting(
        self, common_attn_metadata: CommonAttentionMetadata, draft_index: int
    ) -> Any:
        return self.get_metadata_builder().build_for_drafting(
            self._with_kernel_block_table(common_attn_metadata), draft_index
        )

    def _with_kernel_block_table(
        self, common_attn_metadata: CommonAttentionMetadata, ubatch_idx: int = 0
    ) -> CommonAttentionMetadata:
        if self.kernel_block_size is None:
            return common_attn_metadata
        return replace(
            common_attn_metadata,
            block_table_tensor=self.map_to_kernel_block_table(
                common_attn_metadata.block_table_tensor, ubatch_idx
            ),
        )

    def map_to_kernel_block_table(
        self, block_table: torch.Tensor, ubatch_idx: int = 0
    ) -> torch.Tensor:
        if self.kernel_block_size is None:
            return block_table
        assert self.kernel_block_table is not None
        blocks_per_kv_block = self.kv_cache_spec.block_size // self.kernel_block_size
        rows, cols = block_table.shape
        out = self.kernel_block_table[ubatch_idx, :rows, : cols * blocks_per_kv_block]
        torch.add(
            self.kernel_block_offsets,
            block_table.unsqueeze(-1),
            alpha=self.kernel_block_stride or blocks_per_kv_block,
            out=out.unflatten(1, (cols, -1)),
        )
        return out

    def map_kv_cache_to_kernel_blocks(self, kv_cache: torch.Tensor) -> torch.Tensor:
        """Views a packed cache in manager blocks as kernel blocks."""
        assert self.kernel_block_size is not None
        spec = self.kv_cache_spec
        kernel_rows = spec.get_num_kernel_states(self.kernel_block_size)
        if kv_cache.shape[-2] == kernel_rows:
            return kv_cache
        # Token and pooled caches are written through manager-block slot
        # mappings, so only compressed caches (written from their block table)
        # can be mapped.
        if (
            spec.tokens_per_state <= 1
            or kv_cache.shape[1] != 1
            or kv_cache.stride(0) % (kernel_rows * kv_cache.stride(-2))
        ):
            raise ValueError(
                f"{self.backend.get_name()} needs {self.kernel_block_size}-token "
                f"kernel blocks, which packed {spec.block_size}-token KV cache "
                f"blocks cannot be split into. Reduce --block-size to "
                f"{self.kernel_block_size} or set VLLM_KV_CACHE_LAYOUT to a "
                "layer-compact layout (e.g. LBNHC)."
            )
        # Kernel block j of block b is b * kernel_block_stride + j. The view
        # spans the other layers' pages between blocks (never addressing them),
        # which ``view`` cannot express.
        page_stride = kernel_rows * kv_cache.stride(-2)
        self.kernel_block_stride = kv_cache.stride(0) // page_stride
        num_pages = (kv_cache.shape[0] - 1) * self.kernel_block_stride + (
            kv_cache.shape[-2] // kernel_rows
        )
        return kv_cache.as_strided(
            (num_pages, *kv_cache.shape[1:-2], kernel_rows, kv_cache.shape[-1]),
            (page_stride, *kv_cache.stride()[1:]),
        )

    def map_kv_cache_to_manager_blocks(
        self, kv_cache: torch.Tensor, num_blocks: int
    ) -> torch.Tensor:
        """Views a packed cache in kernel blocks as manager blocks."""
        assert self.kernel_block_stride
        spec = self.kv_cache_spec
        rows = spec.get_num_kernel_states(spec.block_size)
        return kv_cache.as_strided(
            (num_blocks, *kv_cache.shape[1:-2], rows, kv_cache.shape[-1]),
            (kv_cache.stride(0) * self.kernel_block_stride, *kv_cache.stride()[1:]),
        )

    def get_metadata_builder(self, ubatch_id: int = 0) -> AttentionMetadataBuilder:
        assert len(self.metadata_builders) > ubatch_id
        return self.metadata_builders[ubatch_id]

    @property
    def supports_draft_decode_metadata_update(self) -> bool:
        return self.get_metadata_builder().supports_draft_decode_metadata_update

    def update_draft_decode_metadata(
        self,
        attn_metadata: Mapping[str, Any],
    ) -> None:
        metadata = attn_metadata[self.layer_names[0]]
        self.get_metadata_builder().update_draft_decode_metadata(metadata)

map_kv_cache_to_kernel_blocks(kv_cache)

Views a packed cache in manager blocks as kernel blocks.

Source code in vllm/v1/worker/utils.py

def map_kv_cache_to_kernel_blocks(self, kv_cache: torch.Tensor) -> torch.Tensor:
    """Views a packed cache in manager blocks as kernel blocks."""
    assert self.kernel_block_size is not None
    spec = self.kv_cache_spec
    kernel_rows = spec.get_num_kernel_states(self.kernel_block_size)
    if kv_cache.shape[-2] == kernel_rows:
        return kv_cache
    # Token and pooled caches are written through manager-block slot
    # mappings, so only compressed caches (written from their block table)
    # can be mapped.
    if (
        spec.tokens_per_state <= 1
        or kv_cache.shape[1] != 1
        or kv_cache.stride(0) % (kernel_rows * kv_cache.stride(-2))
    ):
        raise ValueError(
            f"{self.backend.get_name()} needs {self.kernel_block_size}-token "
            f"kernel blocks, which packed {spec.block_size}-token KV cache "
            f"blocks cannot be split into. Reduce --block-size to "
            f"{self.kernel_block_size} or set VLLM_KV_CACHE_LAYOUT to a "
            "layer-compact layout (e.g. LBNHC)."
        )
    # Kernel block j of block b is b * kernel_block_stride + j. The view
    # spans the other layers' pages between blocks (never addressing them),
    # which ``view`` cannot express.
    page_stride = kernel_rows * kv_cache.stride(-2)
    self.kernel_block_stride = kv_cache.stride(0) // page_stride
    num_pages = (kv_cache.shape[0] - 1) * self.kernel_block_stride + (
        kv_cache.shape[-2] // kernel_rows
    )
    return kv_cache.as_strided(
        (num_pages, *kv_cache.shape[1:-2], kernel_rows, kv_cache.shape[-1]),
        (page_stride, *kv_cache.stride()[1:]),
    )

map_kv_cache_to_manager_blocks(kv_cache, num_blocks)

Views a packed cache in kernel blocks as manager blocks.

Source code in vllm/v1/worker/utils.py

def map_kv_cache_to_manager_blocks(
    self, kv_cache: torch.Tensor, num_blocks: int
) -> torch.Tensor:
    """Views a packed cache in kernel blocks as manager blocks."""
    assert self.kernel_block_stride
    spec = self.kv_cache_spec
    rows = spec.get_num_kernel_states(spec.block_size)
    return kv_cache.as_strided(
        (num_blocks, *kv_cache.shape[1:-2], rows, kv_cache.shape[-1]),
        (kv_cache.stride(0) * self.kernel_block_stride, *kv_cache.stride()[1:]),
    )

EncoderTimingStats dataclass

Per-request timing statistics for encoder forward pass.

Attributes:

• encoder_forward_secs (float) –

  Time spent in vision encoder forward pass (seconds).

• num_encoder_calls (int) –

  Number of times encoder was called for this request.

Source code in vllm/v1/worker/utils.py

@dataclass
class EncoderTimingStats:
    """Per-request timing statistics for encoder forward pass."""

    encoder_forward_secs: float = 0.0
    """Time spent in vision encoder forward pass (seconds)."""

    num_encoder_calls: int = 0
    """Number of times encoder was called for this request."""

    def to_dict(self) -> dict[str, float | int]:
        return {
            "encoder_forward_secs": self.encoder_forward_secs,
            "num_encoder_calls": self.num_encoder_calls,
        }

encoder_forward_secs = 0.0 class-attribute instance-attribute

Time spent in vision encoder forward pass (seconds).

num_encoder_calls = 0 class-attribute instance-attribute

Number of times encoder was called for this request.

KVBlockZeroer

Manages efficient zeroing of KV cache blocks via a Triton kernel.

Construct once after KV caches are allocated to precompute segment addresses, then call :meth:zero_block_ids each step to zero newly-allocated blocks.

Methods:

• __init__ –

  Precompute the absolute-address table for the Triton zeroing kernel.

• warmup –

  JIT-compile the zeroing kernel before the first real request.

• zero_block_ids –

  Zero the KV cache memory for the given block IDs.

Source code in vllm/v1/worker/utils.py

class KVBlockZeroer:
    """Manages efficient zeroing of KV cache blocks via a Triton kernel.

    Construct once after KV caches are allocated to precompute segment
    addresses, then call :meth:`zero_block_ids` each step to zero
    newly-allocated blocks.
    """

    def __init__(
        self,
        device: torch.device,
        attn_groups_iter: Iterable["AttentionGroup"],
        kernel_block_sizes: list[int],
        static_forward_context: dict[str, Any],
        num_blocks: int,
        runner_only_attn_layers: set[str] | None = None,
    ) -> None:
        """Precompute the absolute-address table for the Triton zeroing kernel.

        Each entry is the absolute byte address of a segment start on the
        GPU, so segments in different CUDA allocations work correctly.

        Per-layer views are standardized ``[B, H, N, C]`` with blocks at dim 0; dims
        physically outside B (separate head groups under LHBNC) each get their own
        segment. A segment's page spans everything inside its block, so under BHLNC it
        also covers the block's other layers -- safe, since block IDs are global pool
        indices and a newly allocated block owns its whole tile.

        Block IDs from the scheduler reference logical blocks whose size
        may differ from the kernel block size (virtual block splitting).
        Each virtual block is represented as an independent segment so its
        physical block stride and zeroed page span remain independent.

        Only AttentionSpec layers are processed; Mamba layers are skipped.
        """
        self.device = device
        self._meta: (
            tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, int] | None
        ) = None

        if runner_only_attn_layers is None:
            runner_only_attn_layers = set()
        # Overlaid layers (packed layouts) share a base address but may have
        # different page sizes; keep the widest span per address so newly
        # allocated blocks are fully zeroed for every overlaying group.
        seen_ptrs: dict[int, int] = {}
        seg_addrs: list[int] = []
        seg_block_strides: list[int] = []
        seg_page_sizes: list[int] = []

        for group in attn_groups_iter:
            spec = group.kv_cache_spec
            if not isinstance(spec, AttentionSpec):
                continue
            if group.kv_cache_group_id >= len(kernel_block_sizes):
                continue
            kernel_bs = kernel_block_sizes[group.kv_cache_group_id]
            assert spec.block_size % kernel_bs == 0
            for layer_name in group.layer_names:
                if layer_name in runner_only_attn_layers:
                    continue
                kv = static_forward_context[layer_name].kv_cache
                if not isinstance(kv, torch.Tensor):
                    continue
                if group.kernel_block_stride:
                    # Bound views address kernel blocks; zero by manager block.
                    kv = group.map_kv_cache_to_manager_blocks(kv, num_blocks)
                if kv.device.type != self.device.type:
                    continue
                dp = kv.data_ptr()

                assert kv.shape[0] % num_blocks == 0, (
                    f"{layer_name}: {kv.shape[0]} kernel blocks is not a "
                    f"multiple of {num_blocks} logical blocks"
                )
                ratio = kv.shape[0] // num_blocks

                el = kv.element_size()
                block_stride_bytes = kv.stride(0) * el
                assert block_stride_bytes % 4 == 0
                assert kv.shape[0] % ratio == 0
                outer_dims = [
                    d
                    for d in range(1, kv.ndim)
                    if kv.stride(d) * el > block_stride_bytes
                ]
                outer_strides = [kv.stride(d) * el for d in outer_dims]
                inner_dims = [d for d in range(1, kv.ndim) if d not in outer_dims]
                kernel_page_bytes = el + sum(
                    (kv.shape[d] - 1) * kv.stride(d) * el for d in inner_dims
                )
                assert kernel_page_bytes % 4 == 0
                logical_block_stride_bytes = block_stride_bytes * ratio
                for outer in iprod(*(range(kv.shape[d]) for d in outer_dims)):
                    off_bytes = sum(i * s for i, s in zip(outer, outer_strides))
                    assert (dp + off_bytes) % 4 == 0
                    for virtual_index in range(ratio):
                        addr = dp + off_bytes + virtual_index * block_stride_bytes
                        if (idx := seen_ptrs.get(addr)) is not None:
                            assert (
                                seg_block_strides[idx]
                                == logical_block_stride_bytes // 4
                            )
                            seg_page_sizes[idx] = max(
                                seg_page_sizes[idx], kernel_page_bytes // 4
                            )
                            continue
                        seen_ptrs[addr] = len(seg_addrs)
                        seg_addrs.append(addr)
                        seg_block_strides.append(logical_block_stride_bytes // 4)
                        seg_page_sizes.append(kernel_page_bytes // 4)

        if not seg_addrs:
            self._meta = None
            return

        max_page_size_el = max(seg_page_sizes)
        blk_size = min(1 << (max_page_size_el - 1).bit_length(), 1024)
        self._meta = (
            torch.tensor(seg_addrs, dtype=torch.uint64, device=self.device),
            torch.tensor(seg_block_strides, dtype=torch.int64, device=self.device),
            torch.tensor(seg_page_sizes, dtype=torch.int64, device=self.device),
            (max_page_size_el + blk_size - 1) // blk_size,
            blk_size,
            len(seg_addrs),
        )

    def zero_block_ids(self, block_ids: list[int]) -> None:
        """Zero the KV cache memory for the given block IDs."""
        if not block_ids or self._meta is None:
            return
        (
            seg_addrs,
            seg_block_strides,
            seg_page_sizes,
            max_chunks,
            blk_size,
            n_segs,
        ) = self._meta
        n_blocks = len(block_ids)
        idx = async_tensor_h2d(block_ids, device=self.device, dtype=torch.int64)
        grid = (n_blocks, n_segs, max_chunks)
        _zero_kv_blocks_kernel[grid](
            seg_addrs,
            seg_block_strides,
            seg_page_sizes,
            idx,
            BLOCK_SIZE=blk_size,
        )

    def warmup(self, num_kv_blocks: int) -> None:
        """JIT-compile the zeroing kernel before the first real request."""
        if num_kv_blocks > 0:
            self.zero_block_ids([0])

__init__(device, attn_groups_iter, kernel_block_sizes, static_forward_context, num_blocks, runner_only_attn_layers=None)

Precompute the absolute-address table for the Triton zeroing kernel.

Each entry is the absolute byte address of a segment start on the GPU, so segments in different CUDA allocations work correctly.

Per-layer views are standardized [B, H, N, C] with blocks at dim 0; dims physically outside B (separate head groups under LHBNC) each get their own segment. A segment's page spans everything inside its block, so under BHLNC it also covers the block's other layers -- safe, since block IDs are global pool indices and a newly allocated block owns its whole tile.

Block IDs from the scheduler reference logical blocks whose size may differ from the kernel block size (virtual block splitting). Each virtual block is represented as an independent segment so its physical block stride and zeroed page span remain independent.

Only AttentionSpec layers are processed; Mamba layers are skipped.

Source code in vllm/v1/worker/utils.py

def __init__(
    self,
    device: torch.device,
    attn_groups_iter: Iterable["AttentionGroup"],
    kernel_block_sizes: list[int],
    static_forward_context: dict[str, Any],
    num_blocks: int,
    runner_only_attn_layers: set[str] | None = None,
) -> None:
    """Precompute the absolute-address table for the Triton zeroing kernel.

    Each entry is the absolute byte address of a segment start on the
    GPU, so segments in different CUDA allocations work correctly.

    Per-layer views are standardized ``[B, H, N, C]`` with blocks at dim 0; dims
    physically outside B (separate head groups under LHBNC) each get their own
    segment. A segment's page spans everything inside its block, so under BHLNC it
    also covers the block's other layers -- safe, since block IDs are global pool
    indices and a newly allocated block owns its whole tile.

    Block IDs from the scheduler reference logical blocks whose size
    may differ from the kernel block size (virtual block splitting).
    Each virtual block is represented as an independent segment so its
    physical block stride and zeroed page span remain independent.

    Only AttentionSpec layers are processed; Mamba layers are skipped.
    """
    self.device = device
    self._meta: (
        tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, int] | None
    ) = None

    if runner_only_attn_layers is None:
        runner_only_attn_layers = set()
    # Overlaid layers (packed layouts) share a base address but may have
    # different page sizes; keep the widest span per address so newly
    # allocated blocks are fully zeroed for every overlaying group.
    seen_ptrs: dict[int, int] = {}
    seg_addrs: list[int] = []
    seg_block_strides: list[int] = []
    seg_page_sizes: list[int] = []

    for group in attn_groups_iter:
        spec = group.kv_cache_spec
        if not isinstance(spec, AttentionSpec):
            continue
        if group.kv_cache_group_id >= len(kernel_block_sizes):
            continue
        kernel_bs = kernel_block_sizes[group.kv_cache_group_id]
        assert spec.block_size % kernel_bs == 0
        for layer_name in group.layer_names:
            if layer_name in runner_only_attn_layers:
                continue
            kv = static_forward_context[layer_name].kv_cache
            if not isinstance(kv, torch.Tensor):
                continue
            if group.kernel_block_stride:
                # Bound views address kernel blocks; zero by manager block.
                kv = group.map_kv_cache_to_manager_blocks(kv, num_blocks)
            if kv.device.type != self.device.type:
                continue
            dp = kv.data_ptr()

            assert kv.shape[0] % num_blocks == 0, (
                f"{layer_name}: {kv.shape[0]} kernel blocks is not a "
                f"multiple of {num_blocks} logical blocks"
            )
            ratio = kv.shape[0] // num_blocks

            el = kv.element_size()
            block_stride_bytes = kv.stride(0) * el
            assert block_stride_bytes % 4 == 0
            assert kv.shape[0] % ratio == 0
            outer_dims = [
                d
                for d in range(1, kv.ndim)
                if kv.stride(d) * el > block_stride_bytes
            ]
            outer_strides = [kv.stride(d) * el for d in outer_dims]
            inner_dims = [d for d in range(1, kv.ndim) if d not in outer_dims]
            kernel_page_bytes = el + sum(
                (kv.shape[d] - 1) * kv.stride(d) * el for d in inner_dims
            )
            assert kernel_page_bytes % 4 == 0
            logical_block_stride_bytes = block_stride_bytes * ratio
            for outer in iprod(*(range(kv.shape[d]) for d in outer_dims)):
                off_bytes = sum(i * s for i, s in zip(outer, outer_strides))
                assert (dp + off_bytes) % 4 == 0
                for virtual_index in range(ratio):
                    addr = dp + off_bytes + virtual_index * block_stride_bytes
                    if (idx := seen_ptrs.get(addr)) is not None:
                        assert (
                            seg_block_strides[idx]
                            == logical_block_stride_bytes // 4
                        )
                        seg_page_sizes[idx] = max(
                            seg_page_sizes[idx], kernel_page_bytes // 4
                        )
                        continue
                    seen_ptrs[addr] = len(seg_addrs)
                    seg_addrs.append(addr)
                    seg_block_strides.append(logical_block_stride_bytes // 4)
                    seg_page_sizes.append(kernel_page_bytes // 4)

    if not seg_addrs:
        self._meta = None
        return

    max_page_size_el = max(seg_page_sizes)
    blk_size = min(1 << (max_page_size_el - 1).bit_length(), 1024)
    self._meta = (
        torch.tensor(seg_addrs, dtype=torch.uint64, device=self.device),
        torch.tensor(seg_block_strides, dtype=torch.int64, device=self.device),
        torch.tensor(seg_page_sizes, dtype=torch.int64, device=self.device),
        (max_page_size_el + blk_size - 1) // blk_size,
        blk_size,
        len(seg_addrs),
    )

warmup(num_kv_blocks)

JIT-compile the zeroing kernel before the first real request.

Source code in vllm/v1/worker/utils.py

def warmup(self, num_kv_blocks: int) -> None:
    """JIT-compile the zeroing kernel before the first real request."""
    if num_kv_blocks > 0:
        self.zero_block_ids([0])

zero_block_ids(block_ids)

Zero the KV cache memory for the given block IDs.

Source code in vllm/v1/worker/utils.py

def zero_block_ids(self, block_ids: list[int]) -> None:
    """Zero the KV cache memory for the given block IDs."""
    if not block_ids or self._meta is None:
        return
    (
        seg_addrs,
        seg_block_strides,
        seg_page_sizes,
        max_chunks,
        blk_size,
        n_segs,
    ) = self._meta
    n_blocks = len(block_ids)
    idx = async_tensor_h2d(block_ids, device=self.device, dtype=torch.int64)
    grid = (n_blocks, n_segs, max_chunks)
    _zero_kv_blocks_kernel[grid](
        seg_addrs,
        seg_block_strides,
        seg_page_sizes,
        idx,
        BLOCK_SIZE=blk_size,
    )

_block_size_is_supported(backends, block_size, specs)

Check if the block size is supported by all backends.

An exact int declaration must match exactly; a MultipleOf declaration accepts any multiple of its base.

Source code in vllm/v1/worker/utils.py

def _block_size_is_supported(
    backends: list[type[AttentionBackend]],
    block_size: int,
    specs: Sequence[KVCacheSpec | None],
) -> bool:
    """Check if the block size is supported by all backends.

    An exact ``int`` declaration must match exactly; a ``MultipleOf``
    declaration accepts any multiple of its base.
    """
    for backend, spec in zip(backends, specs, strict=True):
        is_supported = False
        for supported_size in backend.get_supported_kernel_block_sizes(spec):
            if isinstance(supported_size, int):
                if block_size == supported_size:
                    is_supported = True
            elif isinstance(supported_size, MultipleOf):
                if block_size % supported_size.base == 0:
                    is_supported = True
            else:
                raise ValueError(f"Unknown supported size: {supported_size}")
        if not is_supported:
            return False
    return True

_zero_kv_blocks_kernel(seg_addrs_ptr, seg_block_strides_ptr, seg_page_sizes_ptr, block_ids_ptr, BLOCK_SIZE)

Zero KV cache blocks across all segments in a single launch.

Each segment is a contiguous region of one block's data. Layer-compact layouts have one segment per layer buffer; dimensions physically outside the block dim (separate head groups under LHBNC) and virtual block splits each get their own segment.

Segments may have different block strides and page sizes (e.g. packed KV views or models with multiple KV cache groups like MLA + DSA indexer). Each segment's block stride determines where a logical block begins, while its page size determines how many elements are cleared.

seg_addrs_ptr holds absolute byte addresses (int64) for each segment, allowing segments to live in different CUDA allocations.

Programs are mapped directly onto a 3-D grid as (block_index, seg_index, chunk_index).

Source code in vllm/v1/worker/utils.py

@triton.jit
def _zero_kv_blocks_kernel(
    seg_addrs_ptr,
    seg_block_strides_ptr,
    seg_page_sizes_ptr,
    block_ids_ptr,
    BLOCK_SIZE: tl.constexpr,
):
    """Zero KV cache blocks across all segments in a single launch.

    Each segment is a contiguous region of one block's data.  Layer-compact
    layouts have one segment per layer buffer; dimensions physically outside
    the block dim (separate head groups under LHBNC) and virtual block splits
    each get their own segment.

    Segments may have different block strides and page sizes (e.g. packed
    KV views or models with multiple KV cache groups like MLA + DSA
    indexer). Each segment's block stride determines where a logical block
    begins, while its page size determines how many elements are cleared.

    seg_addrs_ptr holds absolute byte addresses (int64) for each segment,
    allowing segments to live in different CUDA allocations.

    Programs are mapped directly onto a 3-D grid as
    (block_index, seg_index, chunk_index).
    """
    block_index = tl.program_id(0)
    seg_index = tl.program_id(1)
    chunk_index = tl.program_id(2)
    block_stride_el = tl.load(seg_block_strides_ptr + seg_index)
    page_size_el = tl.load(seg_page_sizes_ptr + seg_index)
    chunk_offset = chunk_index.to(tl.int64) * BLOCK_SIZE
    if chunk_offset >= page_size_el:
        return
    block_id = tl.load(block_ids_ptr + block_index)
    seg_addr = tl.load(seg_addrs_ptr + seg_index)
    ptr = tl.cast(seg_addr, tl.pointer_type(tl.int32))
    block_offset = block_id.to(tl.int64) * block_stride_el.to(tl.int64)
    cols = chunk_offset + tl.arange(0, BLOCK_SIZE).to(tl.int64)
    tl.store(
        ptr + block_offset + cols,
        tl.zeros([BLOCK_SIZE], dtype=tl.int32),
        mask=cols < page_size_el,
    )

add_kv_sharing_layers_to_kv_cache_groups(shared_kv_cache_layers, kv_cache_groups, runner_only_attn_layers=None)

Sets up KV cache sharing by reusing the allocated KV caches in kv_caches for layers that do not allocate its own KV cache, based on the mapping in shared_kv_cache_layers. Adds these layers to the corresponding KV cache group, which is needed to ensure that attention metadata is assigned later.

Parameters:

• shared_kv_cache_layers

  (dict[str, str]) –

  Layer pairings for cross-layer KV sharing. If an Attention layer layer_name is in the keys of this dict, it means this layer will perform attention using the keys and values from the KV cache of shared_kv_cache_layers[layer_name].

• kv_cache_groups

  (list[KVCacheGroupSpec]) –

  The KV cache groups of the model.

• runner_only_attn_layers

  (set[str] | None, default: None ) –

  Attention layers handled by the runner only, which are excluded from the KV cache groups.

Source code in vllm/v1/worker/utils.py

def add_kv_sharing_layers_to_kv_cache_groups(
    shared_kv_cache_layers: dict[str, str],
    kv_cache_groups: list[KVCacheGroupSpec],
    runner_only_attn_layers: set[str] | None = None,
) -> None:
    """Sets up KV cache sharing by reusing the allocated KV caches in `kv_caches`
    for layers that do not allocate its own KV cache, based on the mapping in
    `shared_kv_cache_layers`. Adds these layers to the corresponding KV cache
    group, which is needed to ensure that attention metadata is assigned later.

    Args:
        shared_kv_cache_layers: Layer pairings for cross-layer KV sharing.
            If an Attention layer `layer_name` is in the keys of this dict, it
            means this layer will perform attention using the keys and values
            from the KV cache of `shared_kv_cache_layers[layer_name]`.
        kv_cache_groups: The KV cache groups of the model.
        runner_only_attn_layers: Attention layers handled by the runner only,
            which are excluded from the KV cache groups.

    """
    if not shared_kv_cache_layers:
        return

    layer_to_kv_cache_group: dict[str, KVCacheGroupSpec] = {}
    for kv_cache_group in kv_cache_groups:
        for layer_name in kv_cache_group.layer_names:
            layer_to_kv_cache_group[layer_name] = kv_cache_group

    for layer_name, target_layer_name in shared_kv_cache_layers.items():
        tgt_kv_cache_group = layer_to_kv_cache_group[target_layer_name]
        tgt_kv_cache_group.layer_names.append(layer_name)

        if runner_only_attn_layers is not None:
            runner_only_attn_layers.add(layer_name)

allocate_kv_cache(kv_cache_config, device, layout, kernel_block_sizes=None)

Allocate the KV cache and view it as [B, H, N, C] per layer.

Every KVCacheTensor places its layers in the same backing allocation: layer l of block b starts at offset + l * layer_stride + b * block_stride. Cache groups overlay each other, so tensors may address the same bytes.

Source code in vllm/v1/worker/utils.py

def allocate_kv_cache(
    kv_cache_config: KVCacheConfig,
    device: torch.device,
    layout: KVCacheLayout,
    kernel_block_sizes: list[int] | None = None,
) -> dict[str, torch.Tensor]:
    """Allocate the KV cache and view it as ``[B, H, N, C]`` per layer.

    Every KVCacheTensor places its layers in the same backing allocation: layer ``l`` of
    block ``b`` starts at ``offset + l * layer_stride + b * block_stride``. Cache
    groups overlay each other, so tensors may address the same bytes.
    """
    if not kv_cache_config.kv_cache_tensors:
        return {}

    sizes = {tensor.size for tensor in kv_cache_config.kv_cache_tensors}
    assert len(sizes) == 1, "KV cache tensors must share one backing allocation."
    raw_size = sizes.pop()
    # wvSplitKrc's process-lifetime static workspaces (csrc/rocm/skinny_gemms.cu)
    # are created lazily on the first qualifying GEMM. Force that now, before
    # the giant backing allocation below: if one landed in this segment's
    # rounding tail it would pin the whole segment at engine shutdown.
    if current_platform.is_rocm():
        warmup_rocm_skinny_gemm_workspaces(device)
        # Pad to the page granularity MoRIIO needs to register the shared
        # backing as a single RDMA memory region. Other platforms keep the
        # exact-size allocation (see #53974), so anything reading
        # storage.nbytes() has to tolerate the tail on ROCm alone.
        page_size = 4096
        buf_size = ((raw_size + page_size - 1) // page_size) * page_size
    else:
        buf_size = raw_size
    buf = torch.zeros(buf_size, dtype=torch.int8, device=device)

    kv_caches: dict[str, torch.Tensor] = {}
    for tensor in kv_cache_config.kv_cache_tensors:
        layer_name = tensor.layers[0]
        group_id, group = next(
            (group_id, group)
            for group_id, group in enumerate(kv_cache_config.kv_cache_groups)
            if layer_name in group.layer_names
        )
        spec = group.kv_cache_spec
        if isinstance(spec, UniformTypeKVCacheSpecs):
            spec = spec.kv_cache_specs[layer_name]

        if not spec.has_layer_views:
            kv_caches.update((name, buf) for name in tensor.layers)
            continue

        num_blocks = kv_cache_config.num_blocks_of(tensor)
        kernel_block_size = None
        if kernel_block_sizes is not None and group_id < len(kernel_block_sizes):
            kernel_block_size = kernel_block_sizes[group_id]
        views = create_kv_cache_views(
            buf,
            spec,
            num_blocks,
            layout,
            tensor,
            kernel_block_size=kernel_block_size,
        )
        kv_caches.update(zip(tensor.layers, views))
    return kv_caches

bind_kv_cache(kv_caches, forward_context, runner_kv_caches, num_attn_module=1, kv_cache_groups=None, layer_kv_caches=None)

Bind the allocated KV cache to both ModelRunner and forward context so that the KV cache can be used in the forward pass.

This function

1) Fills the ModelRunner's kv cache list (runner_kv_caches) with kv_caches. 2) Associates each attention layer in the forward_context with its corresponding KV cache in kv_caches.

Parameters:

• kv_caches

  (dict[str, Tensor]) –

  The allocated kv_caches with layer names as keys.

• num_attn_module

  (int, default: 1 ) –

  Number of attention modules per layer entry.

• forward_context

  (dict[str, Attention]) –

  The global forward context containing all Attention layers with layer names as keys.

• runner_kv_caches

  (list[Tensor]) –

  The kv_cache declared by ModelRunner.

• kv_cache_groups

  (Sequence[KVCacheGroupSpec] | None, default: None ) –

  The KV cache groups of the model, used to resolve layers that share a KV cache.

• layer_kv_caches

  (dict[str, Tensor] | None, default: None ) –

  The views to bind to layers when they differ from kv_caches (e.g. kernel-block views). The runner keeps kv_caches, which are addressed by scheduler block IDs.

Source code in vllm/v1/worker/utils.py

def bind_kv_cache(
    kv_caches: dict[str, torch.Tensor],
    forward_context: dict[str, Attention],
    runner_kv_caches: list[torch.Tensor],
    num_attn_module: int = 1,
    kv_cache_groups: Sequence[KVCacheGroupSpec] | None = None,
    layer_kv_caches: dict[str, torch.Tensor] | None = None,
) -> None:
    """Bind the allocated KV cache to both ModelRunner and forward context so
    that the KV cache can be used in the forward pass.

    This function:
      1) Fills the ModelRunner's kv cache list (`runner_kv_caches`) with
         kv_caches.
      2) Associates each attention layer in the `forward_context` with its
         corresponding KV cache in kv_caches.

    Args:
        kv_caches: The allocated kv_caches with layer names as keys.
        num_attn_module: Number of attention modules per layer entry.
        forward_context: The global forward context containing all Attention
            layers with layer names as keys.
        runner_kv_caches: The kv_cache declared by ModelRunner.
        kv_cache_groups: The KV cache groups of the model, used to resolve
            layers that share a KV cache.
        layer_kv_caches: The views to bind to layers when they differ from
            kv_caches (e.g. kernel-block views). The runner keeps kv_caches,
            which are addressed by scheduler block IDs.

    """
    # Bind kv_caches to ModelRunner
    assert len(runner_kv_caches) == 0

    # Convert kv_caches dict to a list of tensors in the order of layer_index.
    index2name = defaultdict(list)
    for layer_name in kv_caches:
        index2name[extract_layer_index(layer_name, num_attn_module)].append(layer_name)

    for layer_index in sorted(index2name.keys()):
        layer_names = index2name[layer_index]
        if len(layer_names) > 1:
            # One typical case is encoder-decoder model, e.g., bart.
            # The cross attention and self attention in the same decoder layer
            # has different layer_name but the same layer_index.

            # TODO - analyze where runner_kv_caches is used and the right
            # way to ensure it properly reflects multiple attention layers
            # in the same decoder block.
            current_platform.check_runner_kv_caches_multi_layer()
        for layer_name in layer_names:
            runner_kv_caches.append(kv_caches[layer_name])

    bind_kv_cache_to_layers(
        layer_kv_caches or kv_caches, forward_context, num_attn_module, kv_cache_groups
    )

bind_kv_cache_to_layers(kv_caches, forward_context, num_attn_module=1, kv_cache_groups=None)

Bind layer caches and share ReplaySSM trackers in model-layer order.

Source code in vllm/v1/worker/utils.py

def bind_kv_cache_to_layers(
    kv_caches: dict[str, torch.Tensor],
    forward_context: dict[str, Attention],
    num_attn_module: int = 1,
    kv_cache_groups: Sequence[KVCacheGroupSpec] | None = None,
) -> None:
    """Bind layer caches and share ReplaySSM trackers in model-layer order."""
    # Bind kv_caches to forward context. Each layer's bind_kv_cache unpacks
    # its raw allocation into the per-layer view(s) it needs (e.g. Mamba
    # splits conv/ssm), so the kv_caches dict can hold a single tensor per
    # layer for the KV connector to register.
    for layer_name, kv_cache in kv_caches.items():
        forward_context[layer_name].bind_kv_cache(kv_cache)

    ordered_layer_names = sorted(
        kv_caches, key=lambda name: extract_layer_index(name, num_attn_module)
    )
    share_replayssm_ring_trackers(ordered_layer_names, forward_context, kv_cache_groups)

clear_layer_kv_caches(layers)

Detach the KV/state cache tensors installed by bind_kv_cache().

The model object can outlive the runner (e.g. LLMEngine's finalizer keeps it reachable until engine deletion), so dropping the runner's references alone does not release the KV cache memory on teardown paths.

Source code in vllm/v1/worker/utils.py

def clear_layer_kv_caches(layers: Iterable[Any]) -> None:
    """Detach the KV/state cache tensors installed by bind_kv_cache().

    The model object can outlive the runner (e.g. LLMEngine's finalizer keeps
    it reachable until engine deletion), so dropping the runner's references
    alone does not release the KV cache memory on teardown paths.
    """
    for layer in layers:
        if not hasattr(layer, "kv_cache"):
            continue
        kv_cache = layer.kv_cache
        layer.kv_cache = torch.tensor([]) if isinstance(kv_cache, torch.Tensor) else []
        # Clean up quantized KV cache scale views
        # (int8_per_token_head, fp8_per_token_head)
        if hasattr(layer, "impl"):
            if hasattr(layer.impl, "_k_scale_cache"):
                layer.impl._k_scale_cache = None
            if hasattr(layer.impl, "_v_scale_cache"):
                layer.impl._v_scale_cache = None

customize_attention_spec(backend, spec)

The backend's spec. Blocks larger than every kernel block size a backend takes are split into kernel blocks, which packed blocks must keep whole.

Source code in vllm/v1/worker/utils.py

def customize_attention_spec(
    backend: type[AttentionBackend], spec: AttentionSpec
) -> AttentionSpec:
    """The backend's spec. Blocks larger than every kernel block size a backend
    takes are split into kernel blocks, which packed blocks must keep whole."""
    spec = backend.customize_spec(spec)
    sizes = backend.get_supported_kernel_block_sizes(spec)
    fixed_sizes = [s for s in sizes if isinstance(s, int)]
    if spec.block_stride_alignment is None and sizes and fixed_sizes == sizes:
        spec = replace(spec, block_stride_alignment=MultipleOf(max(fixed_sizes)))
    return spec

get_uniform_decode_token_count(num_reqs, num_tokens, max_query_len, decode_graph_eligible)

Per-request token count of a uniform decode batch, or None.

Source code in vllm/v1/worker/utils.py

def get_uniform_decode_token_count(
    num_reqs: int, num_tokens: int, max_query_len: int, decode_graph_eligible: bool
) -> int | None:
    """Per-request token count of a uniform decode batch, or None."""
    if decode_graph_eligible and is_uniform_query_len(
        num_reqs, num_tokens, max_query_len
    ):
        return max_query_len
    return None

is_uniform_query_len(num_reqs, num_tokens, max_query_len)

Whether every request in the batch has the same query length.

Shape test only; use get_uniform_decode_token_count to classify a scheduled batch, since a prompt chunk can have a decode batch's shape.

Source code in vllm/v1/worker/utils.py

def is_uniform_query_len(num_reqs: int, num_tokens: int, max_query_len: int) -> bool:
    """Whether every request in the batch has the same query length.

    Shape test only; use ``get_uniform_decode_token_count`` to classify a
    scheduled batch, since a prompt chunk can have a decode batch's shape.
    """
    return num_reqs > 0 and num_tokens == max_query_len * num_reqs

map_kv_caches_to_kernel_blocks(kv_caches, attn_groups)

The views to bind to layers. The manager views in kv_caches stay for KV connectors.

Source code in vllm/v1/worker/utils.py

def map_kv_caches_to_kernel_blocks(
    kv_caches: dict[str, torch.Tensor], attn_groups: Iterable[AttentionGroup]
) -> dict[str, torch.Tensor]:
    """The views to bind to layers. The manager views in ``kv_caches`` stay for
    KV connectors."""
    mapped = dict(kv_caches)
    for group in attn_groups:
        if group.kernel_block_size is not None:
            mapped.update(
                (name, group.map_kv_cache_to_kernel_blocks(kv_caches[name]))
                for name in group.layer_names
                if name in kv_caches
            )
    return mapped

prepare_kernel_block_sizes(kv_cache_config, attn_groups)

Generate kernel_block_sizes that matches each block_size.

For attention backends that support virtual block splitting, use the supported block sizes from the backend. For other backends (like Mamba), use the same block size (no splitting).

Parameters:

• kv_cache_config

  (KVCacheConfig) –

  The KV cache configuration.

• attn_groups

  (list[list[AttentionGroup]]) –

  Attention groups indexed by KV cache group id.

Returns:

• list[int] –

  List of kernel block sizes for each cache group.

Source code in vllm/v1/worker/utils.py

def prepare_kernel_block_sizes(
    kv_cache_config: KVCacheConfig, attn_groups: list[list[AttentionGroup]]
) -> list[int]:
    """Generate kernel_block_sizes that matches each block_size.

    For attention backends that support virtual block splitting,
    use the supported block sizes from the backend.
    For other backends (like Mamba), use the same block size (no splitting).

    Args:
        kv_cache_config: The KV cache configuration.
        attn_groups: Attention groups indexed by KV cache group id.

    Returns:
        List of kernel block sizes for each cache group.

    """
    kernel_block_sizes = []
    for kv_cache_gid, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups):
        kv_cache_spec = kv_cache_group.kv_cache_spec
        if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
            # All layers in the UniformTypeKVCacheSpecs have the same type,
            # pick an arbitrary one to dispatch.
            kv_cache_spec = next(iter(kv_cache_spec.kv_cache_specs.values()))
        if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec):
            continue
        if not kv_cache_spec.has_layer_views:
            kernel_block_sizes.append(kv_cache_spec.block_size)
        elif isinstance(kv_cache_spec, AttentionSpec):
            # This is an attention backend that supports virtual block splitting.
            kv_manager_block_size = kv_cache_group.kv_cache_spec.block_size
            group_backends = [g.backend for g in attn_groups[kv_cache_gid]]
            group_specs = [g.kv_cache_spec for g in attn_groups[kv_cache_gid]]
            # Block-outermost packing puts other layers' pages between the
            # blocks, so they are never split group-wide.
            selected_kernel_size = (
                kv_manager_block_size
                if any(
                    t.block_stride > t.layer_stride
                    for t in kv_cache_config.kv_cache_tensors
                    if not set(t.layers).isdisjoint(kv_cache_group.layer_names)
                )
                else select_common_block_size(
                    kv_manager_block_size, group_backends, group_specs
                )
            )
            kernel_block_sizes.append(selected_kernel_size)
        elif isinstance(kv_cache_spec, MambaSpec):
            # This is likely Mamba or other non-attention cache, no splitting.
            kernel_block_sizes.append(kv_cache_spec.block_size)
        else:
            raise NotImplementedError(
                f"unknown kv cache spec {kv_cache_group.kv_cache_spec}"
            )
    return kernel_block_sizes

request_memory(init_snapshot, cache_config, external_weight_memory=0)

Calculate the amount of memory required by vLLM, then validate that the current amount of free memory is sufficient for that.

Source code in vllm/v1/worker/utils.py

def request_memory(
    init_snapshot: MemorySnapshot,
    cache_config: CacheConfig,
    external_weight_memory: int = 0,
) -> int:
    """Calculate the amount of memory required by vLLM, then validate
    that the current amount of free memory is sufficient for that.
    """
    requested_memory = math.ceil(
        init_snapshot.total_memory * cache_config.gpu_memory_utilization
    )

    if external_weight_memory > 0:
        engine_memory = requested_memory - external_weight_memory
        if engine_memory <= 0:
            raise ValueError(
                f"Externally held weights on device {init_snapshot.device_} "
                f"({format_gib(external_weight_memory)}/"
                f"{format_gib(init_snapshot.total_memory)} GiB) exceed the "
                "desired GPU memory utilization "
                f"({cache_config.gpu_memory_utilization}, "
                f"{format_gib(requested_memory)} GiB). Increase GPU memory "
                "utilization or reduce GPU memory used by other processes."
            )
        if init_snapshot.free_memory < engine_memory:
            raise ValueError(
                f"Free memory on device {init_snapshot.device_} "
                f"({format_gib(init_snapshot.free_memory)}/"
                f"{format_gib(init_snapshot.total_memory)} GiB) on startup "
                "is less than the engine's budget after excluding "
                "external process's weights "
                f"({cache_config.gpu_memory_utilization}, "
                f"{format_gib(engine_memory)} GiB). Decrease GPU memory "
                "utilization or reduce GPU memory used by other processes."
            )
        logger.info_once(
            "Weights are held outside this process: of the %s GiB "
            "utilization budget, %s GiB is externally held and "
            "%s GiB remains for the engine's own allocations.",
            format_gib(requested_memory),
            format_gib(external_weight_memory),
            format_gib(engine_memory),
        )
        return engine_memory

    if init_snapshot.free_memory < requested_memory:
        raise ValueError(
            f"Free memory on device {init_snapshot.device_} "
            f"({format_gib(init_snapshot.free_memory)}/"
            f"{format_gib(init_snapshot.total_memory)} GiB) on startup "
            f"is less than desired GPU memory utilization "
            f"({cache_config.gpu_memory_utilization}, "
            f"{format_gib(requested_memory)} GiB). Decrease GPU memory "
            f"utilization or reduce GPU memory used by other processes."
        )

    return requested_memory

sanity_check_mm_encoder_outputs(mm_embeddings, expected_num_items)

Perform sanity checks for the result of vllm.model_executor.models.SupportsMultiModal.embed_multimodal.

Source code in vllm/v1/worker/utils.py

def sanity_check_mm_encoder_outputs(
    mm_embeddings: MultiModalEmbeddings,
    expected_num_items: int,
) -> None:
    """Perform sanity checks for the result of
    [`vllm.model_executor.models.SupportsMultiModal.embed_multimodal`][].
    """
    assert isinstance(mm_embeddings, (list, tuple, torch.Tensor)), (
        "Expected multimodal embeddings to be a list/tuple of 2D tensors, "
        f"or a single 3D tensor, but got {type(mm_embeddings)} "
        "instead. This is most likely due to incorrect implementation "
        "of the model's `embed_multimodal` method."
    )

    assert len(mm_embeddings) == expected_num_items, (
        "Expected number of multimodal embeddings to match number of "
        f"input items: {expected_num_items}, but got {len(mm_embeddings)=} "
        "instead. This is most likely due to incorrect implementation "
        "of the model's `embed_multimodal` method."
    )

    assert all(e.ndim == 2 for e in mm_embeddings), (
        "Expected multimodal embeddings to be a sequence of 2D tensors, "
        f"but got tensors with shapes {[e.shape for e in mm_embeddings]} "
        "instead. This is most likely due to incorrect implementation "
        "of the model's `embed_multimodal` method."
    )

select_common_block_size(kv_manager_block_size, backends, kv_cache_specs=None)

Select a block size that is supported by all backends and is a factor of kv_manager_block_size.

If kv_manager_block_size is supported by all backends, return it directly. Otherwise, return the max supported size.

Parameters:

• kv_manager_block_size

  (int) –

  Block size of KV cache.

• backends

  (list[type[AttentionBackend]]) –

  List of attention backend classes.

• kv_cache_specs

  (Sequence[KVCacheSpec | None] | None, default: None ) –

  Per-backend cache specs for spec-specific constraints.

Returns:

• int –

  The selected block size.

Raises:

• ValueError –

  If no valid block size found.

Source code in vllm/v1/worker/utils.py

def select_common_block_size(
    kv_manager_block_size: int,
    backends: list[type[AttentionBackend]],
    kv_cache_specs: Sequence[KVCacheSpec | None] | None = None,
) -> int:
    """Select a block size that is supported by all backends and is a factor of
    kv_manager_block_size.

    If kv_manager_block_size is supported by all backends, return it directly.
    Otherwise, return the max supported size.

    Args:
        kv_manager_block_size: Block size of KV cache.
        backends: List of attention backend classes.
        kv_cache_specs: Per-backend cache specs for spec-specific constraints.

    Returns:
        The selected block size.

    Raises:
        ValueError: If no valid block size found.

    """
    specs = kv_cache_specs or [None] * len(backends)
    if _block_size_is_supported(backends, kv_manager_block_size, specs):
        return kv_manager_block_size

    # MultipleOf constraints also accept the manager size if they accept a divisor.
    # Any remaining candidate must therefore be an explicit size from a backend.
    candidates = {
        size
        for backend, spec in zip(backends, specs, strict=True)
        for size in backend.get_supported_kernel_block_sizes(spec)
        if isinstance(size, int) and kv_manager_block_size % size == 0
    }

    for size in sorted(candidates, reverse=True):
        if _block_size_is_supported(backends, size, specs):
            return size
    raise ValueError(
        f"No common block size for {kv_manager_block_size} ("
        + "; ".join(
            f"{b.get_name()}: {b.get_supported_kernel_block_sizes(s)}"
            for b, s in zip(backends, specs, strict=True)
        )
        + ")."
    )