| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Normalized configuration consumed by native offloading backends.""" |
| 4 | |
| 5 | from collections.abc import Mapping |
| 6 | from dataclasses import dataclass |
| 7 | from typing import Any |
| 8 | |
| 9 | |
| 10 | @dataclass(frozen=True) |
| 11 | class OffloadingGroupConfig: |
| 12 | # Total token span covered by one block across all workers |
| 13 | # (accounts for context parallelism). |
| 14 | tokens_per_block: int |
| 15 | # Layer names belonging to this group. |
| 16 | layer_names: tuple[str, ...] |
| 17 | # Original KVCacheConfig group index. |
| 18 | group_id: int |
| 19 | |
| 20 | |
| 21 | @dataclass(frozen=True) |
| 22 | class OffloadingModelConfig: |
| 23 | # Model identifier (e.g. HuggingFace model path). |
| 24 | name: str |
| 25 | # KV cache data type (e.g. "float16"). |
| 26 | dtype: str |
| 27 | |
| 28 | |
| 29 | @dataclass(frozen=True) |
| 30 | class OffloadingCacheConfig: |
| 31 | # Tokens per block hash. |
| 32 | tokens_per_hash: int |
| 33 | # Blocks coalesced into one offload chunk. |
| 34 | blocks_per_chunk: int |
| 35 | |
| 36 | |
| 37 | @dataclass(frozen=True) |
| 38 | class OffloadingParallelConfig: |
| 39 | # Worker index in [0, world_size). 0 on the scheduler side. |
| 40 | rank: int |
| 41 | # Total number of workers. |
| 42 | world_size: int |
| 43 | # Tensor parallel size. |
| 44 | tp_size: int |
| 45 | # Pipeline parallel size. |
| 46 | pp_size: int |
| 47 | # Prefill context parallel size. |
| 48 | pcp_size: int |
| 49 | # Decode context parallel size. |
| 50 | dcp_size: int |
| 51 | # Data parallel replica index of this engine. |
| 52 | data_parallel_index: int |
| 53 | # Number of data parallel replicas. |
| 54 | data_parallel_size: int |
| 55 | # Local rank of the data parallel group, set only in SPMD mode. |
| 56 | data_parallel_rank_local: int | None |
| 57 | # True when the bytes that will be persisted for a block are portable |
| 58 | # across parallelism configurations: for the direct layout, concatenating |
| 59 | # the block's data across all workers in rank order yields the same bytes |
| 60 | # under any topology; for the canonical layout, the canonical page itself |
| 61 | # is topology-free. |
| 62 | is_parallelism_agnostic: bool |
| 63 | |
| 64 | |
| 65 | @dataclass(frozen=True) |
| 66 | class OffloadingConfig: |
| 67 | groups: tuple[OffloadingGroupConfig, ...] |
| 68 | # KV bytes stored by one worker per block. |
| 69 | worker_kv_bytes_per_block: int |
| 70 | # Whether the scheduler emits KV cache events. When true, |
| 71 | # the offloading backend should emit events as well. |
| 72 | enable_kv_cache_events: bool |
| 73 | # Offloading-specific configuration from kv_connector_extra_config. |
| 74 | extra_config: Mapping[str, Any] |
| 75 | # Unique identifier for this engine, distinct per DP rank. |
| 76 | engine_id: str |
| 77 | model: OffloadingModelConfig |
| 78 | cache: OffloadingCacheConfig |
| 79 | parallel: OffloadingParallelConfig |
| 80 | # True when the offloaded bytes of every worker are expected to be |
| 81 | # byte-identical per block (pure-MLA model, single-node TP-only |
| 82 | # parallelism), enabling a single-copy host layout in backends that |
| 83 | # support it. Aggregate layout decision; per-layer replication metadata |
| 84 | # is planned for CanonicalKVCacheRef (#48408). |
| 85 | replicated_layout: bool = False |
| 86 | # True when the canonical per-layer host byte layout was requested via |
| 87 | # kv_connector_extra_config; certified per-layer at worker registration. |
| 88 | canonical_layout: bool = False |
| 89 | # Resolved KVCacheLayout name of the worker KV cache. |
| 90 | kv_cache_layout: str | None = None |