1use crate::{DType, QuantCodecError};
2
3#[cfg(feature = "serde")]
4use serde::{Deserialize, Serialize};
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
7#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
8pub enum KvRole {
9 Key,
10 Value,
11}
12
13#[derive(Debug, Clone, PartialEq, Eq, Hash)]
14#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
15pub enum KvLayout {
16 LayersHeadsTokensDim,
17 LayersTokensHeadsDim,
18 RuntimeSpecific(String),
19}
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
22#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
23pub struct LayerId(pub u32);
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
26#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
27pub struct HeadId(pub u32);
28
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
30#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
31pub struct TokenSpan {
32 pub start: u64,
33 pub end: u64,
34}
35
36impl TokenSpan {
37 pub fn new(start: u64, end: u64) -> Result<Self, QuantCodecError> {
38 if start >= end {
39 return Err(QuantCodecError::InvalidTokenSpan { start, end });
40 }
41 Ok(Self { start, end })
42 }
43
44 pub fn len(self) -> u64 {
45 self.end - self.start
46 }
47
48 pub fn is_empty(self) -> bool {
49 self.start >= self.end
50 }
51}
52
53#[derive(Debug, Clone, PartialEq, Eq, Hash)]
54#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
55pub enum KvAttentionKind {
56 Mha,
57 Mqa,
58 Gqa,
59 Unsupported(String),
60}
61
62#[derive(Debug, Clone, PartialEq, Eq, Hash)]
63#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
64pub struct KvCacheShapeV2 {
65 pub batch: u32,
66 pub layers: u32,
67 pub num_q_heads: u32,
68 pub num_kv_heads: u32,
69 pub seq_len: u64,
70 pub head_dim: u32,
71 pub layout: KvLayout,
72 pub dtype: DType,
73 pub attention_kind: KvAttentionKind,
74}
75
76impl KvCacheShapeV2 {
77 #[allow(clippy::too_many_arguments)]
78 pub fn new(
79 batch: u32,
80 layers: u32,
81 num_q_heads: u32,
82 num_kv_heads: u32,
83 seq_len: u64,
84 head_dim: u32,
85 layout: KvLayout,
86 dtype: DType,
87 attention_kind: KvAttentionKind,
88 ) -> Result<Self, QuantCodecError> {
89 let shape = Self {
90 batch,
91 layers,
92 num_q_heads,
93 num_kv_heads,
94 seq_len,
95 head_dim,
96 layout,
97 dtype,
98 attention_kind,
99 };
100 shape.validate()?;
101 Ok(shape)
102 }
103
104 pub fn mha(
105 batch: u32,
106 layers: u32,
107 heads: u32,
108 seq_len: u64,
109 head_dim: u32,
110 layout: KvLayout,
111 dtype: DType,
112 ) -> Result<Self, QuantCodecError> {
113 Self::new(
114 batch,
115 layers,
116 heads,
117 heads,
118 seq_len,
119 head_dim,
120 layout,
121 dtype,
122 KvAttentionKind::Mha,
123 )
124 }
125
126 pub fn mqa(
127 batch: u32,
128 layers: u32,
129 num_q_heads: u32,
130 seq_len: u64,
131 head_dim: u32,
132 layout: KvLayout,
133 dtype: DType,
134 ) -> Result<Self, QuantCodecError> {
135 Self::new(
136 batch,
137 layers,
138 num_q_heads,
139 1,
140 seq_len,
141 head_dim,
142 layout,
143 dtype,
144 KvAttentionKind::Mqa,
145 )
146 }
147
148 #[allow(clippy::too_many_arguments)]
149 pub fn gqa(
150 batch: u32,
151 layers: u32,
152 num_q_heads: u32,
153 num_kv_heads: u32,
154 seq_len: u64,
155 head_dim: u32,
156 layout: KvLayout,
157 dtype: DType,
158 ) -> Result<Self, QuantCodecError> {
159 Self::new(
160 batch,
161 layers,
162 num_q_heads,
163 num_kv_heads,
164 seq_len,
165 head_dim,
166 layout,
167 dtype,
168 KvAttentionKind::Gqa,
169 )
170 }
171
172 pub fn validate(&self) -> Result<(), QuantCodecError> {
173 if self.batch == 0 {
174 return Err(invalid_shape("batch must be greater than zero"));
175 }
176 if self.layers == 0 {
177 return Err(invalid_shape("layers must be greater than zero"));
178 }
179 if self.num_q_heads == 0 {
180 return Err(invalid_shape("num_q_heads must be greater than zero"));
181 }
182 if self.num_kv_heads == 0 {
183 return Err(invalid_shape("num_kv_heads must be greater than zero"));
184 }
185 if self.seq_len == 0 {
186 return Err(invalid_shape("seq_len must be greater than zero"));
187 }
188 if self.head_dim == 0 {
189 return Err(invalid_shape("head_dim must be greater than zero"));
190 }
191 if matches!(&self.layout, KvLayout::RuntimeSpecific(s) if s.is_empty()) {
192 return Err(invalid_shape(
193 "runtime-specific layout label cannot be empty",
194 ));
195 }
196 match &self.attention_kind {
197 KvAttentionKind::Mha => {
198 if self.num_q_heads != self.num_kv_heads {
199 return Err(invalid_shape(
200 "MHA requires num_q_heads equal to num_kv_heads",
201 ));
202 }
203 }
204 KvAttentionKind::Mqa => {
205 if self.num_kv_heads != 1 || self.num_q_heads <= 1 {
206 return Err(invalid_shape(
207 "MQA requires num_kv_heads == 1 and num_q_heads > 1",
208 ));
209 }
210 }
211 KvAttentionKind::Gqa => {
212 if self.num_q_heads <= self.num_kv_heads || self.num_kv_heads <= 1 {
213 return Err(invalid_shape(
214 "GQA requires num_q_heads > num_kv_heads and num_kv_heads > 1",
215 ));
216 }
217 if self.num_q_heads % self.num_kv_heads != 0 {
218 return Err(invalid_shape(
219 "GQA requires num_q_heads divisible by num_kv_heads",
220 ));
221 }
222 }
223 KvAttentionKind::Unsupported(label) => {
224 if label.trim().is_empty() {
225 return Err(invalid_shape("unsupported attention label cannot be empty"));
226 }
227 return Err(invalid_shape(format!(
228 "unsupported attention kind {label} is not adapter-owned"
229 )));
230 }
231 }
232 Ok(())
233 }
234}
235
236#[derive(Debug, Clone, PartialEq, Eq, Hash)]
237#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
238pub struct KvTensorShape {
239 pub layers: u32,
240 pub key_heads: u32,
241 pub value_heads: u32,
242 pub seq_len: u64,
243 pub head_dim: u32,
244 pub layout: KvLayout,
245 pub dtype: DType,
246}
247
248impl KvTensorShape {
249 pub fn gqa(
250 layers: u32,
251 key_heads: u32,
252 value_heads: u32,
253 seq_len: u64,
254 head_dim: u32,
255 layout: KvLayout,
256 dtype: DType,
257 ) -> Result<Self, QuantCodecError> {
258 let shape = Self {
259 layers,
260 key_heads,
261 value_heads,
262 seq_len,
263 head_dim,
264 layout,
265 dtype,
266 };
267 shape.validate()?;
268 Ok(shape)
269 }
270
271 pub fn validate(&self) -> Result<(), QuantCodecError> {
272 if self.layers == 0 {
273 return Err(invalid_shape("layers must be greater than zero"));
274 }
275 if self.key_heads == 0 {
276 return Err(invalid_shape("key_heads must be greater than zero"));
277 }
278 if self.value_heads == 0 {
279 return Err(invalid_shape("value_heads must be greater than zero"));
280 }
281 if self.seq_len == 0 {
282 return Err(invalid_shape("seq_len must be greater than zero"));
283 }
284 if self.head_dim == 0 {
285 return Err(invalid_shape("head_dim must be greater than zero"));
286 }
287 if matches!(&self.layout, KvLayout::RuntimeSpecific(s) if s.is_empty()) {
288 return Err(invalid_shape(
289 "runtime-specific layout label cannot be empty",
290 ));
291 }
292 Ok(())
293 }
294
295 pub fn heads_for(&self, role: KvRole) -> u32 {
296 match role {
297 KvRole::Key => self.key_heads,
298 KvRole::Value => self.value_heads,
299 }
300 }
301
302 pub fn layer_element_count(&self, role: KvRole) -> Result<usize, QuantCodecError> {
303 checked_usize_product(
304 &[
305 self.heads_for(role) as u64,
306 self.seq_len,
307 self.head_dim as u64,
308 ],
309 "layer element count",
310 )
311 }
312
313 pub fn total_element_count(&self, role: KvRole) -> Result<usize, QuantCodecError> {
314 checked_usize_product(
315 &[
316 self.layers as u64,
317 self.heads_for(role) as u64,
318 self.seq_len,
319 self.head_dim as u64,
320 ],
321 "total element count",
322 )
323 }
324
325 pub fn validate_span(&self, span: TokenSpan) -> Result<(), QuantCodecError> {
326 if span.is_empty() || span.end > self.seq_len {
327 return Err(QuantCodecError::InvalidTokenSpan {
328 start: span.start,
329 end: span.end,
330 });
331 }
332 Ok(())
333 }
334}
335
336#[derive(Debug, Clone, PartialEq, Eq, Hash)]
337#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
338pub struct KvSliceRequest {
339 pub layer: LayerId,
340 pub role: KvRole,
341 pub token_span: TokenSpan,
342 pub head: Option<HeadId>,
343}
344
345impl KvSliceRequest {
346 pub fn layer_span(layer: LayerId, token_span: TokenSpan) -> Self {
347 Self {
348 layer,
349 role: KvRole::Key,
350 token_span,
351 head: None,
352 }
353 }
354
355 pub fn for_role(mut self, role: KvRole) -> Self {
356 self.role = role;
357 self
358 }
359
360 pub fn for_head(mut self, head: HeadId) -> Self {
361 self.head = Some(head);
362 self
363 }
364
365 pub fn validate_for_shape(&self, shape: &KvTensorShape) -> Result<(), QuantCodecError> {
366 shape.validate()?;
367 if self.layer.0 >= shape.layers {
368 return Err(QuantCodecError::ShapeMismatch {
369 reason: format!(
370 "requested layer {} but shape has {} layers",
371 self.layer.0, shape.layers
372 ),
373 });
374 }
375 shape.validate_span(self.token_span)?;
376 if let Some(head) = self.head {
377 let heads = shape.heads_for(self.role);
378 if head.0 >= heads {
379 return Err(QuantCodecError::ShapeMismatch {
380 reason: format!("requested head {} but role has {} heads", head.0, heads),
381 });
382 }
383 }
384 Ok(())
385 }
386}
387
388fn checked_usize_product(values: &[u64], context: &'static str) -> Result<usize, QuantCodecError> {
389 let mut product = 1u64;
390 for value in values {
391 product = product
392 .checked_mul(*value)
393 .ok_or(QuantCodecError::IntegerOverflow { context })?;
394 }
395 usize::try_from(product).map_err(|_| QuantCodecError::IntegerOverflow { context })
396}
397
398fn invalid_shape(reason: impl Into<String>) -> QuantCodecError {
399 QuantCodecError::InvalidShape {
400 reason: reason.into(),
401 }
402}