Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / infer / kv_cache.py: 0%
162 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""KV cache container for generation."""
16from dataclasses import dataclass
17from typing import Any, Iterable, List, Optional, Tuple
19import torch
21PastKeyValues = List[Tuple[torch.Tensor, torch.Tensor]]
24def _validate_pair_shapes(key: torch.Tensor, value: torch.Tensor) -> None:
25 """Validate one key/value cache tensor pair."""
26 if not isinstance(key, torch.Tensor) or not isinstance(value, torch.Tensor):
27 raise ValueError("key and value must be tensors")
28 if key.ndim != 4 or value.ndim != 4:
29 raise ValueError("key and value must have shape (batch, heads, seq, dim)")
30 if key.shape != value.shape:
31 raise ValueError("key and value batch/heads/seq/dim dimensions must match")
34def detach_and_validate_past_key_values(past_key_values: Iterable) -> PastKeyValues:
35 """Return detached tuple KV tensors after validating their shapes."""
36 values = []
37 for item in past_key_values:
38 if not isinstance(item, (tuple, list)) or len(item) != 2:
39 raise ValueError("each cache entry must be a (key, value) pair")
40 key, value = item
41 _validate_pair_shapes(key, value)
42 values.append((key.detach(), value.detach()))
43 return values
46@dataclass(frozen=True)
47class SequenceShardInfo:
48 """Sequence range held by one context-parallel rank."""
50 rank: int
51 world_size: int
52 start: int
53 end: int
54 global_seq_len: int
56 @property
57 def local_seq_len(self) -> int:
58 """Return the sequence length stored by this rank."""
59 return self.end - self.start
62def get_sequence_shard_info(
63 global_seq_len: int,
64 rank: int,
65 world_size: int,
66) -> SequenceShardInfo:
67 """Return the contiguous sequence shard range for a CP rank."""
68 if global_seq_len < 0:
69 raise ValueError("global_seq_len must be >= 0")
70 if world_size <= 0:
71 raise ValueError("world_size must be > 0")
72 if rank < 0 or rank >= world_size:
73 raise ValueError("rank must be in [0, world_size)")
74 base = global_seq_len // world_size
75 remainder = global_seq_len % world_size
76 start = rank * base + min(rank, remainder)
77 end = start + base + (1 if rank < remainder else 0)
78 return SequenceShardInfo(
79 rank=rank,
80 world_size=world_size,
81 start=start,
82 end=end,
83 global_seq_len=global_seq_len,
84 )
87def shard_past_key_values(
88 past_key_values: Iterable,
89 rank: int,
90 world_size: int,
91 global_seq_len: Optional[int] = None,
92) -> Tuple[PastKeyValues, SequenceShardInfo]:
93 """Shard full past key values on the sequence dimension for CP cache."""
94 values = detach_and_validate_past_key_values(past_key_values)
95 if not values:
96 seq_len = 0 if global_seq_len is None else global_seq_len
97 else:
98 seq_len = values[0][0].shape[-2]
99 if global_seq_len is None:
100 global_seq_len = seq_len
101 if seq_len != global_seq_len:
102 raise ValueError("global_seq_len must match full cache sequence length")
103 shard_info = get_sequence_shard_info(global_seq_len, rank, world_size)
104 sharded = [
105 (
106 key.narrow(-2, shard_info.start, shard_info.local_seq_len).contiguous(),
107 value.narrow(-2, shard_info.start, shard_info.local_seq_len).contiguous(),
108 )
109 for key, value in values
110 ]
111 return sharded, shard_info
114class KVCache:
115 """Stores per-layer key/value tensors."""
117 def __init__(self):
118 self.past_key_values: Optional[Any] = None
120 @property
121 def is_empty(self) -> bool:
122 """Check whether no usable KV cache is stored."""
123 return self.past_key_values is None or (
124 isinstance(self.past_key_values, list) and len(self.past_key_values) == 0
125 )
127 def clear(self) -> None:
128 """Drop all cached tensors."""
129 self.past_key_values = None
131 def update(self, past_key_values: Optional[Iterable]) -> None:
132 """Replace the cache with detached past key values."""
133 if past_key_values is None:
134 return
135 if self._is_opaque_cache(past_key_values):
136 self.past_key_values = past_key_values
137 return
138 values = self._detach_and_validate(past_key_values)
139 self.past_key_values = None if not values else values
141 def merge(self, past_key_values: Optional[Iterable]) -> None:
142 """Append incremental key/value tensors on the sequence dimension."""
143 if past_key_values is None:
144 return
145 if self._is_opaque_cache(past_key_values):
146 self.past_key_values = past_key_values
147 return
148 new_values = self._detach_and_validate(past_key_values)
149 if not new_values:
150 return
151 if self.past_key_values is None:
152 self.past_key_values = new_values
153 return
154 if len(self.past_key_values) != len(new_values):
155 raise ValueError("past_key_values layer count mismatch")
156 merged = []
157 for (old_k, old_v), (new_k, new_v) in zip(self.past_key_values, new_values):
158 self._validate_pair_shapes(old_k, old_v)
159 self._validate_pair_shapes(new_k, new_v)
160 if old_k.shape[:-2] != new_k.shape[:-2] or old_k.shape[-1] != new_k.shape[-1]:
161 raise ValueError("key cache shape mismatch")
162 if old_v.shape[:-2] != new_v.shape[:-2] or old_v.shape[-1] != new_v.shape[-1]:
163 raise ValueError("value cache shape mismatch")
164 merged.append((
165 torch.cat([old_k, new_k], dim=-2),
166 torch.cat([old_v, new_v], dim=-2),
167 ))
168 self.past_key_values = merged
170 @classmethod
171 def _detach_and_validate(cls, past_key_values: Iterable) -> PastKeyValues:
172 return detach_and_validate_past_key_values(past_key_values)
174 @staticmethod
175 def _validate_pair_shapes(key: torch.Tensor, value: torch.Tensor) -> None:
176 _validate_pair_shapes(key, value)
178 @staticmethod
179 def _is_opaque_cache(past_key_values) -> bool:
180 return hasattr(past_key_values, "get_seq_length") and not isinstance(
181 past_key_values, (list, tuple),
182 )
185class ContextParallelKVCache(KVCache):
186 """Stores a local sequence shard of generation KV cache."""
188 def __init__(self, rank: int, world_size: int):
189 super().__init__()
190 if world_size <= 0:
191 raise ValueError("world_size must be > 0")
192 if rank < 0 or rank >= world_size:
193 raise ValueError("rank must be in [0, world_size)")
194 self.rank = rank
195 self.world_size = world_size
196 self.shard_info = get_sequence_shard_info(0, rank, world_size)
198 def update_full(self, past_key_values: Optional[Iterable]) -> None:
199 """Shard full prefill K/V cache and store only this rank's sequence slice."""
200 if past_key_values is None:
201 return
202 sharded, shard_info = shard_past_key_values(
203 past_key_values,
204 rank=self.rank,
205 world_size=self.world_size,
206 )
207 self.past_key_values = sharded
208 self.shard_info = shard_info
210 def update_local(
211 self,
212 past_key_values: Optional[Iterable],
213 shard_info: SequenceShardInfo,
214 ) -> None:
215 """Store K/V tensors that are already local to this CP rank."""
216 if past_key_values is None:
217 return
218 self._validate_shard_info(shard_info)
219 values = self._detach_and_validate(past_key_values)
220 self._validate_local_seq_len(values, shard_info.local_seq_len)
221 self.past_key_values = values
222 self.shard_info = shard_info
224 def merge_local(
225 self,
226 past_key_values: Optional[Iterable],
227 global_seq_len: Optional[int] = None,
228 ) -> None:
229 """Append local incremental K/V tensors and advance global sequence metadata."""
230 if past_key_values is None:
231 return
232 new_values = self._detach_and_validate(past_key_values)
233 if self.past_key_values is None:
234 if global_seq_len is None and self.world_size > 1:
235 raise ValueError("global_seq_len is required for initial CP local cache")
236 inferred_global = (
237 self.shard_info.global_seq_len + new_values[0][0].shape[-2]
238 if global_seq_len is None and new_values
239 else global_seq_len
240 )
241 shard_info = get_sequence_shard_info(
242 0 if inferred_global is None else inferred_global,
243 self.rank,
244 self.world_size,
245 )
246 self._validate_local_seq_len(new_values, shard_info.local_seq_len)
247 self.past_key_values = new_values
248 self.shard_info = shard_info
249 return
250 old_local_seq_len = self.shard_info.local_seq_len
251 next_global_seq_len = (
252 self.shard_info.global_seq_len + new_values[0][0].shape[-2]
253 if global_seq_len is None and new_values
254 else global_seq_len
255 )
256 if next_global_seq_len is None:
257 raise ValueError("global_seq_len is required for empty incremental cache")
258 shard_info = get_sequence_shard_info(
259 next_global_seq_len,
260 self.rank,
261 self.world_size,
262 )
263 expected_growth = shard_info.local_seq_len - old_local_seq_len
264 actual_growth = new_values[0][0].shape[-2] if new_values else 0
265 if expected_growth != actual_growth:
266 raise ValueError("local cache growth does not match CP shard metadata")
267 super().merge(new_values)
268 self.shard_info = shard_info
270 def clear(self) -> None:
271 """Drop all cached tensors and reset CP sequence metadata."""
272 super().clear()
273 self.shard_info = get_sequence_shard_info(0, self.rank, self.world_size)
275 def _validate_shard_info(self, shard_info: SequenceShardInfo) -> None:
276 if shard_info.rank != self.rank or shard_info.world_size != self.world_size:
277 raise ValueError("shard_info does not match this CP cache")
279 @staticmethod
280 def _validate_local_seq_len(values: PastKeyValues, local_seq_len: int) -> None:
281 for key, value in values:
282 if key.shape[-2] != local_seq_len or value.shape[-2] != local_seq_len:
283 raise ValueError("local cache sequence length does not match shard_info")