| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import hashlib |
| 5 | import json |
| 6 | |
| 7 | from vllm.distributed.kv_transfer.kv_connector.v1.offloading.canonical_mapping import ( |
| 8 | canonical_format_id, |
| 9 | ) |
| 10 | from vllm.v1.kv_offload.base import ( |
| 11 | OffloadingSpec, |
| 12 | OffloadKey, |
| 13 | get_offload_block_hash, |
| 14 | get_offload_group_idx, |
| 15 | ) |
| 16 | |
| 17 | _BASE_PATH_HASH_LEN = 12 |
| 18 | _CONFIG_FILENAME = "config.json" |
| 19 | |
| 20 | |
| 21 | class FileMapper: |
| 22 | """FileMapper maps KV blocks (given by their hash) to file names.""" |
| 23 | |
| 24 | def __init__( |
| 25 | self, |
| 26 | root_dir: str, |
| 27 | model_name: str, |
| 28 | tokens_per_hash: int, |
| 29 | blocks_per_file: int, |
| 30 | tp_size: int, |
| 31 | pp_size: int, |
| 32 | pcp_size: int, |
| 33 | dcp_size: int, |
| 34 | rank: int, |
| 35 | dtype: str, |
| 36 | kv_cache_groups: list[dict] | None = None, |
| 37 | inference_engine: str = "vllm", |
| 38 | parallel_agnostic: bool = False, |
| 39 | replicated_layout: bool = False, |
| 40 | canonical_format: str | None = None, |
| 41 | ): |
| 42 | """Initialize the file mapper. Each worker constructs its own, but |
| 43 | `config.json` is shared across workers since rank lives outside the hash. |
| 44 | When `parallel_agnostic=True`, tp/pp/pcp/dcp are forced to 1 and rank |
| 45 | to 0 so multiple parallelism layouts collapse into the same folder. |
| 46 | """ |
| 47 | if parallel_agnostic: |
| 48 | tp_size = pp_size = pcp_size = dcp_size = 1 |
| 49 | rank = 0 |
| 50 | self.rank: int = rank |
| 51 | self.fields: dict = { |
| 52 | "model_name": model_name, |
| 53 | "tokens_per_hash": tokens_per_hash, |
| 54 | "blocks_per_file": blocks_per_file, |
| 55 | "tp_size": tp_size, |
| 56 | "pp_size": pp_size, |
| 57 | "pcp_size": pcp_size, |
| 58 | "dcp_size": dcp_size, |
| 59 | "dtype": str(dtype), |
| 60 | "kv_cache_groups": kv_cache_groups or [], |
| 61 | "inference_engine": inference_engine, |
| 62 | } |
| 63 | if not parallel_agnostic: |
| 64 | self.fields["parallel_agnostic"] = False |
| 65 | # Only written when True so existing deployments' hashed fields are |
| 66 | # unchanged (False is the historical default and must not appear). |
| 67 | if replicated_layout: |
| 68 | self.fields["replicated_layout"] = True |
| 69 | # The canonical byte format is not interchangeable with the direct |
| 70 | # layout (or with other canonical format versions/families), so its |
| 71 | # identity participates in the storage namespace. |
| 72 | if canonical_format is not None: |
| 73 | self.fields["canonical_format"] = canonical_format |
| 74 | self.base_path: str = self._compute_base_path(root_dir, self.fields) |
| 75 | |
| 76 | @classmethod |
| 77 | def from_offloading_spec( |
| 78 | cls, |
| 79 | root_dir: str, |
| 80 | offloading_spec: OffloadingSpec, |
| 81 | blocks_per_file: int = 1, |
| 82 | parallel_agnostic: bool = False, |
| 83 | ) -> "FileMapper": |
| 84 | """Build a FileMapper from an OffloadingSpec.""" |
| 85 | config = offloading_spec.config |
| 86 | kv_cache_groups = [ |
| 87 | { |
| 88 | "tokens_per_block": group.tokens_per_block, |
| 89 | "layer_names": list(group.layer_names), |
| 90 | } |
| 91 | for group in config.groups |
| 92 | ] |
| 93 | parallel = config.parallel |
| 94 | canonical_format = None |
| 95 | if config.canonical_layout: |
| 96 | assert config.kv_cache_layout is not None |
| 97 | canonical_format = canonical_format_id(config.kv_cache_layout) |
| 98 | return cls( |
| 99 | root_dir=root_dir, |
| 100 | model_name=config.model.name, |
| 101 | tokens_per_hash=config.cache.tokens_per_hash, |
| 102 | blocks_per_file=blocks_per_file, |
| 103 | tp_size=parallel.tp_size, |
| 104 | pp_size=parallel.pp_size, |
| 105 | pcp_size=parallel.pcp_size, |
| 106 | dcp_size=parallel.dcp_size, |
| 107 | rank=parallel.rank, |
| 108 | dtype=config.model.dtype, |
| 109 | kv_cache_groups=kv_cache_groups, |
| 110 | parallel_agnostic=( |
| 111 | parallel_agnostic |
| 112 | and (parallel.is_parallelism_agnostic or config.replicated_layout) |
| 113 | ), |
| 114 | replicated_layout=(parallel_agnostic and config.replicated_layout), |
| 115 | canonical_format=canonical_format, |
| 116 | ) |
| 117 | |
| 118 | def get_file_name(self, key: OffloadKey) -> str: |
| 119 | """Map an OffloadKey to <base>_r<rank>/<hhh>/<hh>_g<group_idx>/<hash>.bin.""" |
| 120 | hash_hex = get_offload_block_hash(key).hex() |
| 121 | group_idx = get_offload_group_idx(key) |
| 122 | subfolder1, subfolder2 = hash_hex[:3], hash_hex[3:5] |
| 123 | return ( |
| 124 | f"{self.base_path}_r{self.rank}" |
| 125 | f"/{subfolder1}/{subfolder2}_g{group_idx}/{hash_hex}.bin" |
| 126 | ) |
| 127 | |
| 128 | def get_run_config(self) -> dict: |
| 129 | return dict(self.fields) |
| 130 | |
| 131 | def get_config_file_path(self) -> str: |
| 132 | return f"{self.base_path}/{_CONFIG_FILENAME}" |
| 133 | |
| 134 | @staticmethod |
| 135 | def _compute_base_path(root_dir: str, fields: dict) -> str: |
| 136 | """Layout: <root_dir>/<safe_model_name>_<sha256-prefix>/. |
| 137 | safe_model_name replaces '/' with '_' so HuggingFace IDs don't nest. |
| 138 | """ |
| 139 | canonical = json.dumps(fields, sort_keys=True, separators=(",", ":")) |
| 140 | digest = hashlib.sha256(canonical.encode("utf-8")).hexdigest()[ |
| 141 | :_BASE_PATH_HASH_LEN |
| 142 | ] |
| 143 | safe_model_name = fields["model_name"].replace("/", "_") |
| 144 | return f"{root_dir}/{safe_model_name}_{digest}" |