vllm.v1.hisparse.binding

Worker-side wiring of HiSparse caches to attention layers.

Functions:

• allocate_hisparse_kv_caches –

  Allocate the host pool separately from the shared device backing.

• bind_hisparse_kv_caches –

  Bind existing cache storage and block tables; return the bound handles.

• init_hisparse_kv_cache –

  Allocate and bind HiSparse caches within the caller's allocation context.

• resolve_hisparse_specs –

  Resolve a common kernel block size for HiSparse MLA specs, and add the

allocate_hisparse_kv_caches(kv_cache_config, device, layout, kernel_block_sizes, host_pool)

Allocate the host pool separately from the shared device backing.

Source code in vllm/v1/hisparse/binding.py

def allocate_hisparse_kv_caches(
    kv_cache_config: KVCacheConfig,
    device: torch.device,
    layout: KVCacheLayout,
    kernel_block_sizes: list[int],
    host_pool: HiSparseHostPool,
) -> dict[str, torch.Tensor]:
    """Allocate the host pool separately from the shared device backing."""
    device_config = copy(kv_cache_config)
    device_config.kv_cache_tensors = [
        tensor
        for tensor in kv_cache_config.kv_cache_tensors
        if not tensor.host_resident
    ]
    kv_caches = allocate_kv_cache(device_config, device, layout, kernel_block_sizes)
    host_tensors = [
        tensor for tensor in kv_cache_config.kv_cache_tensors if tensor.host_resident
    ]
    (host_size,) = {tensor.size for tensor in host_tensors}
    backing = host_pool.allocate(host_size)
    (host_group_id,) = kv_cache_config.host_group_ids
    host_spec = kv_cache_config.kv_cache_groups[host_group_id].kv_cache_spec
    for tensor in host_tensors:
        spec = (
            host_spec.kv_cache_specs[tensor.layers[0]]
            if isinstance(host_spec, UniformTypeKVCacheSpecs)
            else host_spec
        )
        kernel_block_size = kernel_block_sizes[host_group_id]
        if isinstance(spec, MLAAttentionSpec) and spec.storage_block_size is not None:
            kernel_block_size = spec.storage_block_size
        views = create_kv_cache_views(
            backing,
            spec,
            kv_cache_config.num_blocks_of(tensor),
            layout,
            tensor,
            kernel_block_size=kernel_block_size,
        )
        kv_caches.update(zip(tensor.layers, views))
    return kv_caches

bind_hisparse_kv_caches(*, forward_context, kv_cache_config, kv_caches, block_tables, host_pool)

Bind existing cache storage and block tables; return the bound handles.

Source code in vllm/v1/hisparse/binding.py

def bind_hisparse_kv_caches(
    *,
    forward_context: dict[str, Any],
    kv_cache_config: KVCacheConfig,
    kv_caches: dict[str, torch.Tensor],
    block_tables: "BlockTables",
    host_pool: HiSparseHostPool,
) -> list[HiSparseCacheHandle]:
    """Bind existing cache storage and block tables; return the bound handles."""
    assert host_pool.registered is not None
    tensor_configs = {
        name: (
            tensor_config,
            tensor_config.offset + layer_index * tensor_config.layer_stride,
        )
        for tensor_config in kv_cache_config.kv_cache_tensors
        for layer_index, name in enumerate(tensor_config.layers)
    }
    resident_source_index = 0
    for group_id, group in enumerate(kv_cache_config.kv_cache_groups):
        if not isinstance(group.kv_cache_spec, HiSparseResidentSpec):
            continue
        for cache_name in group.layer_names:
            assert cache_name.endswith(HISPARSE_RESIDENT_SUFFIX)
            layer_name = cache_name[: -len(HISPARSE_RESIDENT_SUFFIX)]
            tensor_config, byte_offset = tensor_configs[cache_name]
            assert not tensor_config.host_resident
            cache_handle = _get_hisparse_cache(forward_context, layer_name)
            cache_handle.bind_cache(
                kv_caches[cache_name],
                byte_offset=byte_offset,
                block_stride=tensor_config.block_stride,
                num_blocks=kv_cache_config.num_blocks,
                block_size=group.kv_cache_spec.block_size,
                block_table=block_tables.input_block_tables[group_id],
                slot_mapping=block_tables.slot_mappings[group_id],
            )
            assert cache_handle.view is not None
            kv_caches[cache_name] = cache_handle.view.cache
            cache_handle.runtime.resident_source_index = resident_source_index
        resident_source_index += 1

    hot_backing: torch.Tensor | None = None
    cache_handles: list[HiSparseCacheHandle] = []
    for group_id, group in enumerate(kv_cache_config.kv_cache_groups):
        if not isinstance(group.kv_cache_spec, HiSparseHotSpec):
            continue
        for cache_name in group.layer_names:
            assert cache_name.endswith(HISPARSE_HOT_SUFFIX)
            raw_tensor = kv_caches[cache_name]
            if hot_backing is None:
                hot_backing = raw_tensor
            elif hot_backing.untyped_storage().data_ptr() != (
                raw_tensor.untyped_storage().data_ptr()
            ):
                raise RuntimeError("HiSparse hot tensors must share one GPU backing.")
            layer_name = cache_name[: -len(HISPARSE_HOT_SUFFIX)]
            cache_handle = _get_hisparse_cache(forward_context, layer_name)
            tensor_config, byte_offset = tensor_configs[cache_name]
            assert not tensor_config.host_resident
            cache_handle.runtime.bind_hot_cache(
                raw_tensor,
                byte_offset=byte_offset,
                block_stride=tensor_config.block_stride,
                num_blocks=kv_cache_config.num_blocks,
                block_size=group.kv_cache_spec.block_size,
                block_table=block_tables.input_block_tables[group_id],
            )
            resident = cache_handle.view
            hot = cache_handle.runtime.hot
            assert resident is not None
            if (
                resident.cache.untyped_storage().data_ptr()
                != hot.cache.untyped_storage().data_ptr()
                or resident.cache.data_ptr() != hot.cache.data_ptr()
                or resident.cache.stride() != hot.cache.stride()
            ):
                raise RuntimeError("HiSparse resident and hot layouts must match.")
            source_cache = kv_caches[layer_name]
            assert source_cache.untyped_storage().data_ptr() == (
                host_pool.registered.untyped_storage().data_ptr()
            )
            cache_handle.draft_layer = forward_context[layer_name].is_draft_layer
            cache_handle.runtime.shared_host_region = host_pool.shared_region
            cache_handle.runtime.bind_source_cache(
                source_cache,
                registered_host_pool=host_pool.registered,
            )
            cache_handles.append(cache_handle)
            # The hot slab is raw storage owned by the runtime, not a layer
            # cache: nothing downstream binds or registers it.
            del kv_caches[cache_name]

    if hot_backing is None or not cache_handles:
        raise RuntimeError("HiSparse found no hot-cache handles.")
    (source_group_id,) = kv_cache_config.host_group_ids
    for cache_handle in cache_handles:
        cache_handle.source_block_table = block_tables.input_block_tables[
            source_group_id
        ]
        cache_handle.mirror_slot_mapping = block_tables.slot_mappings[source_group_id]
    return cache_handles

init_hisparse_kv_cache(kv_cache_config, device, kernel_block_sizes, vllm_config, forward_context, block_tables)

Allocate and bind HiSparse caches within the caller's allocation context.

Source code in vllm/v1/hisparse/binding.py

def init_hisparse_kv_cache(
    kv_cache_config: KVCacheConfig,
    device: torch.device,
    kernel_block_sizes: list[int],
    vllm_config: VllmConfig,
    forward_context: dict[str, Any],
    block_tables: "BlockTables",
) -> dict[str, torch.Tensor]:
    """Allocate and bind HiSparse caches within the caller's allocation context."""
    host_pool = HiSparseHostPool(vllm_config, kv_cache_config)
    try:
        kv_caches = allocate_hisparse_kv_caches(
            kv_cache_config,
            device,
            vllm_config.cache_config.get_resolved_kv_cache_layout(),
            kernel_block_sizes,
            host_pool,
        )
        cache_handles = bind_hisparse_kv_caches(
            forward_context=forward_context,
            kv_cache_config=kv_cache_config,
            kv_caches=kv_caches,
            block_tables=block_tables,
            host_pool=host_pool,
        )
        initialize_hisparse_runtime_buffers(
            cache_handles,
            max_num_reqs=vllm_config.scheduler_config.max_num_seqs,
            max_num_batched_tokens=vllm_config.scheduler_config.max_num_batched_tokens,
        )
        return kv_caches
    except Exception:
        if host_pool.shared_region is not None:
            host_pool.shared_region.cleanup()
        raise

resolve_hisparse_specs(vllm_config, kv_cache_spec, attn_layers)

Resolve a common kernel block size for HiSparse MLA specs, and add the resident/hot specs HiSparse derives from them.

Source code in vllm/v1/hisparse/binding.py

def resolve_hisparse_specs(
    vllm_config: VllmConfig,
    kv_cache_spec: dict[str, KVCacheSpec],
    attn_layers: Mapping[str, "AttentionLayerBase"],
) -> dict[str, KVCacheSpec]:
    """Resolve a common kernel block size for HiSparse MLA specs, and add the
    resident/hot specs HiSparse derives from them."""
    mla_specs = {
        name: spec
        for name, spec in kv_cache_spec.items()
        if isinstance(spec, MLAAttentionSpec)
    }
    if not mla_specs:
        return kv_cache_spec
    block_sizes = {spec.block_size for spec in mla_specs.values()}
    if len(block_sizes) != 1:
        raise ValueError("HiSparse requires one scheduler block size.")
    backends = [attn_layers[name].get_attn_backend() for name in mla_specs]
    try:
        block_size = select_common_block_size(block_sizes.pop(), backends)
    except ValueError as error:
        raise ValueError(
            "HiSparse requires a GPU block size supported by every sparse "
            f"attention and indexer backend: {error}"
        ) from error
    kv_cache_spec = kv_cache_spec | {
        name: spec.copy_with_new_block_size(block_size)
        for name, spec in mla_specs.items()
    }
    groups = get_hisparse_kv_cache_groups(vllm_config, kv_cache_spec)
    assert groups is not None
    return kv_cache_spec | {
        layer_name: group.kv_cache_spec
        for group in groups
        if isinstance(group.kv_cache_spec, (HiSparseResidentSpec, HiSparseHotSpec))
        for layer_name in group.layer_names
    }