tract_transformers/ops/
dyn_kv_cache.rs

1use std::str::FromStr;
2
3use tract_nnef::internal::*;
4use tract_nnef::prelude::tract_itertools::Itertools;
5use tract_nnef::ser::{datum_type, tdims};
6use tract_nnef::tract_core::ops::array::TypedConcat;
7use tract_nnef::tract_core::ops::source::TypedSource;
8
9pub fn register(registry: &mut Registry) {
10    registry.register_dumper(ser_dyn_kv_cache);
11    registry.register_primitive(
12        "tract_transformers_dyn_kv_cache",
13        &[
14            TypeName::Scalar.tensor().named("input"),
15            TypeName::String.named("name"),
16            TypeName::Integer.named("axis"),
17            TypeName::String.named("datum_type"),
18            TypeName::Integer.array().named("past_sequence_shape"),
19            TypeName::Integer.array().named("input_sequence_shape"),
20        ],
21        &[("output", TypeName::Scalar.tensor())],
22        de_dyn_kv_cache,
23    );
24}
25
26fn ser_dyn_kv_cache(
27    ast: &mut IntoAst,
28    node: &TypedNode,
29    op: &DynKeyValueCache,
30) -> TractResult<Option<Arc<RValue>>> {
31    let input = ast.mapping[&node.inputs[0]].clone();
32    Ok(Some(invocation(
33        "tract_transformers_dyn_kv_cache",
34        &[input],
35        &[
36            ("name", string(&op.name)),
37            ("axis", numeric(op.axis)),
38            ("datum_type", datum_type(op.past_sequence_fact.datum_type)),
39            ("past_sequence_shape", tdims(op.past_sequence_fact.shape.dims())),
40            ("input_sequence_shape", tdims(op.input_sequence_fact.shape.dims())),
41        ],
42    )))
43}
44
45fn de_dyn_kv_cache(
46    builder: &mut ModelBuilder,
47    invocation: &ResolvedInvocation,
48) -> TractResult<Value> {
49    let input = invocation.named_arg_as(builder, "input")?;
50    let name: String = invocation.named_arg_as(builder, "name")?;
51    let axis: usize = invocation.named_arg_as(builder, "axis")?;
52    let dt = DatumType::from_str(&invocation.named_arg_as::<String>(builder, "datum_type")?)?;
53    let past_sequence_shape: TVec<TDim> = builder
54        .allowing_new_symbols(|builder| invocation.named_arg_as(builder, "past_sequence_shape"))?;
55    let input_sequence_shape: TVec<TDim> = builder
56        .allowing_new_symbols(|builder| invocation.named_arg_as(builder, "input_sequence_shape"))?;
57    builder.wire(
58        DynKeyValueCache {
59            name,
60            axis,
61            past_sequence_fact: dt.fact(&*past_sequence_shape),
62            input_sequence_fact: dt.fact(&*input_sequence_shape),
63        },
64        &[input],
65    )
66}
67
68#[derive(Debug, Clone)]
69pub struct DynKeyValueCacheState {
70    name: String,
71    axis: usize,
72    past_sequence_fact: TypedFact,
73    kv_cache: Option<TValue>,
74}
75
76impl DynKeyValueCacheState {
77    pub fn resolve_symbols(
78        state: &mut TurnState,
79        fact: TypedFact,
80        concrete_shape: Option<&[usize]>,
81    ) -> TractResult<()> {
82        let unresolved = fact
83            .shape
84            .iter()
85            .enumerate()
86            .filter_map(|(ax, symb)| match symb {
87                TDim::Sym(s) if state.resolved_symbols.get(s).is_none() => Some((ax, s)),
88                _ => None,
89            })
90            .collect_vec();
91
92        if unresolved.is_empty() {
93            return Ok(());
94        }
95
96        ensure!(unresolved.len() == 1);
97        let (ax, sym) = unresolved[0];
98        if let Some(shape) = concrete_shape {
99            ensure!(ax < shape.len());
100            state.resolved_symbols.set(sym, shape[ax] as i64);
101        } else {
102            state.resolved_symbols.set(sym, 0);
103        }
104
105        if state.scenario.is_none() {
106            state.scenario = sym.scope().unwrap().guess_scenario(&state.resolved_symbols)?;
107        }
108        Ok(())
109    }
110
111    pub fn truncate(&mut self, len: usize) -> TractResult<()> {
112        if let Some(t) = self.kv_cache.as_mut() {
113            *t = t.slice(self.axis, 0, len)?.into_tvalue();
114        } else {
115            bail!("Can not truncate a zero-len kv-cache value");
116        }
117        Ok(())
118    }
119}
120
121impl OpState for DynKeyValueCacheState {
122    fn load_from(
123        &mut self,
124        state: &mut TurnState,
125        states: &mut dyn Iterator<Item = tract_nnef::prelude::TValue>,
126    ) -> TractResult<()> {
127        // KV Cache fact is always at index 0
128        let kv_cache_init = states.next().context("Not enough state initializers")?;
129        Self::resolve_symbols(state, self.past_sequence_fact.clone(), Some(kv_cache_init.shape()))?;
130        self.kv_cache = Some(kv_cache_init.clone());
131
132        Ok(())
133    }
134
135    fn save_to(&self, states: &mut Vec<TValue>) -> TractResult<()> {
136        if let Some(kv_cache) = &self.kv_cache {
137            states.push(kv_cache.clone());
138            Ok(())
139        } else {
140            bail!("KV cache {} was never initialized", self.name)
141        }
142    }
143
144    fn init_tensor_fact(&self) -> Option<(String, TypedFact)> {
145        Some((self.name.clone(), self.past_sequence_fact.clone()))
146    }
147
148    fn has_init_tensor_fact(&self) -> bool {
149        true
150    }
151
152    fn resolve_symbols(&mut self, state: &mut TurnState) -> TractResult<()> {
153        let shape = self.kv_cache.as_ref().map(|kv_cache| kv_cache.shape());
154        Self::resolve_symbols(state, self.past_sequence_fact.clone(), shape)
155    }
156
157    fn eval(
158        &mut self,
159        _ctx: &EvalContext,
160        _op: &dyn Op,
161        inputs: TVec<TValue>,
162    ) -> TractResult<TVec<TValue>> {
163        let input = args_1!(inputs);
164        // build output
165        let output = if let Some(curr) = self.kv_cache.take() {
166            TypedConcat { axis: self.axis }
167                .eval(&EvalContext::out_of_plan(), tvec![curr, input])?
168                .remove(0)
169        } else {
170            input
171        };
172        self.kv_cache = Some(output.clone());
173
174        Ok(tvec!(output))
175    }
176
177    fn reset_lanes(&mut self, _lanes: &[LaneId]) -> TractResult<()> {
178        bail!("DynKeyValueCache is not lane-aware: the cache has no lane axis")
179    }
180}
181
182#[derive(Clone, Debug, PartialEq, Eq)]
183pub struct DynKeyValueCache {
184    pub name: String,
185    pub axis: usize,
186    pub past_sequence_fact: TypedFact,
187    pub input_sequence_fact: TypedFact,
188}
189
190impl Op for DynKeyValueCache {
191    fn name(&self) -> StaticName {
192        "DynamicKeyValueCache".to_string().into()
193    }
194
195    op_as_typed_op!();
196}
197
198impl EvalOp for DynKeyValueCache {
199    not_out_of_plan!();
200
201    fn state(&self, _ctx: &EvalContext) -> TractResult<Option<Box<dyn OpState>>> {
202        Ok(Some(Box::new(DynKeyValueCacheState {
203            name: self.name.clone(),
204            axis: self.axis,
205            past_sequence_fact: self.past_sequence_fact.clone(),
206            kv_cache: None,
207        })))
208    }
209}
210
211impl TypedOp for DynKeyValueCache {
212    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
213        ensure!(inputs.len() == 1);
214        let input = inputs[0];
215        let mut fact = input.without_value();
216
217        fact.shape.set(
218            self.axis,
219            self.past_sequence_fact.shape.dims()[self.axis].clone()
220                + self.input_sequence_fact.shape.dims()[self.axis].clone(),
221        );
222        Ok(tvec!(fact))
223    }
224
225    fn cost(&self, _inputs: &[&TypedFact]) -> TractResult<TVec<(Cost, TDim)>> {
226        let token_volume = self
227            .past_sequence_fact
228            .shape
229            .iter()
230            .enumerate()
231            .filter(|(axis, _d)| *axis != self.axis)
232            .map(|(_axis, d)| d)
233            .product::<TDim>();
234        Ok(tvec!((Cost::Custom(false, "KVCacheValuesPerToken".to_string()), token_volume)))
235    }
236
237    as_op!();
238}
239
240/// Reverse of `replace_kv_cache`: replaces a DynKeyValueCache node with Source + Concat,
241/// restoring KV cache state as explicit model inputs and outputs.
242pub fn unfold_kv_cache(target: &mut TypedModel, kv_node_id: usize) -> TractResult<()> {
243    let node = target.node(kv_node_id);
244    let op = node.op_as::<DynKeyValueCache>().context("Not a DynKeyValueCache node")?;
245    let name = op.name.clone();
246    let axis = op.axis;
247    let past_fact = op.past_sequence_fact.clone();
248    let input_fact = op.input_sequence_fact.clone();
249    let existing_input = node.inputs[0];
250
251    // Add a new Source node for the past KV cache
252    let source_outlet = target.add_source(&name, past_fact)?;
253
254    // Compute output fact for the Concat
255    let mut output_fact = input_fact.clone();
256    output_fact.shape.set(
257        axis,
258        target.outlet_fact(source_outlet)?.shape.dims()[axis].clone()
259            + input_fact.shape.dims()[axis].clone(),
260    );
261
262    // Replace DynKeyValueCache op with TypedConcat
263    let kv_node = target.node_mut(kv_node_id);
264    kv_node.name = format!("{name}_concat");
265    kv_node.op = Box::new(TypedConcat { axis });
266    kv_node.outputs[0].fact = output_fact;
267
268    // Rewire: Concat takes [source, existing_input] as inputs
269    // Currently the node has [existing_input] at slot 0
270    // We need [source_outlet, existing_input] at slots [0, 1]
271    kv_node.inputs = vec![source_outlet, existing_input];
272
273    // Update successor info on the source node
274    target.nodes[source_outlet.node].outputs[source_outlet.slot]
275        .successors
276        .push(InletId::new(kv_node_id, 0));
277
278    // Update the existing input's successor slot from 0 to 1
279    target.nodes[existing_input.node].outputs[existing_input.slot].successors.iter_mut().for_each(
280        |succ| {
281            if succ.node == kv_node_id && succ.slot == 0 {
282                succ.slot = 1;
283            }
284        },
285    );
286
287    // Add the Concat output to model outputs and label it so runtimes preserve the name
288    let concat_outlet = OutletId::new(kv_node_id, 0);
289    target.outputs.push(concat_outlet);
290    target.set_outlet_label(concat_outlet, format!("{name}_concat"))?;
291
292    Ok(())
293}
294
295/// Search pattern => Input -> Concat -> Output
296/// Return type is for using rule-ensure macro
297pub fn replace_kv_cache(target: &mut TypedModel, source_node_id: usize) -> TractResult<Option<()>> {
298    assert!(target.node(source_node_id).op_is::<TypedSource>());
299    let (concat_node_id, non_source_input_id, axis, input_facts) = {
300        rule_if_some!(concat_node = target.next_node(target.node(source_node_id)));
301
302        // Check KV Cache Pattern
303        rule_if!(
304            concat_node.op_is::<TypedConcat>()
305                && concat_node.inputs.len() == 2
306                && concat_node.outputs.len() == 1
307                && target.outputs.contains(&concat_node.id.into())
308        );
309
310        let concat_in_facts = target.node_input_facts(concat_node.id)?;
311
312        // Check on shapes
313        let concat_in_shapes = [concat_in_facts[0].shape.dims(), concat_in_facts[1].shape.dims()];
314        let rank = concat_in_shapes[0].len();
315        let axes = (0..rank)
316            .filter(|ax| concat_in_shapes[0][*ax] != concat_in_shapes[1][*ax])
317            .collect_vec();
318        ensure!(axes.len() == 1);
319
320        let axis = axes[0];
321        rule_if!(
322            matches!(concat_in_shapes[0][axis], TDim::Sym(_))
323                && matches!(concat_in_shapes[1][axis], TDim::Sym(_))
324        );
325        let mut facts = [concat_in_facts[0].clone(), concat_in_facts[1].clone()];
326        if concat_node.inputs[0].node == source_node_id {
327            (concat_node.id, concat_node.inputs[1].node, axis, facts)
328        } else if concat_node.inputs[1].node == source_node_id {
329            facts.swap(0, 1);
330            (concat_node.id, concat_node.inputs[0].node, axis, facts)
331        } else {
332            return Ok(None);
333        }
334    };
335
336    {
337        // Replace Concat by KVCache
338        let name = target.node_names().collect_vec()[source_node_id].to_string();
339        let concat_node = target.node_mut(concat_node_id);
340        concat_node.op = Box::new(DynKeyValueCache {
341            name: name.clone(),
342            axis,
343            past_sequence_fact: input_facts[0].clone(),
344            input_sequence_fact: input_facts[1].clone(),
345        });
346        concat_node.name = name;
347        concat_node.inputs.retain(|input| input != &source_node_id.into());
348    }
349
350    {
351        // Replace Source by Dummy Op for it to be cleaned later
352        let dummy_op = target.create_dummy();
353        let source_node = target.node_mut(source_node_id);
354        source_node.outputs[0].successors.clear();
355        source_node.op = dummy_op;
356    }
357    {
358        // Non-source input is usually the second input of Concat. Rewire it to the only input of the new KVCache Op
359        let non_source_input = target.node_mut(non_source_input_id);
360        non_source_input.outputs.iter_mut().for_each(|output| {
361            output.successors.iter_mut().for_each(|succ| {
362                if succ.node == concat_node_id {
363                    succ.slot = 0
364                }
365            })
366        });
367    }
368
369    // Clean model I/Os
370    target.outputs.retain(|output| output.node != concat_node_id);
371    target.inputs.retain(|input| input.node != source_node_id);
372    target.outlet_labels.remove(&concat_node_id.into());
373    Ok(None)
374}
375
376#[cfg(test)]
377mod tests {
378    use super::*;
379    use tract_num_traits::AsPrimitive;
380    use tract_num_traits::Zero;
381
382    fn run_test_case<F: Datum + Zero + Copy>(
383        input_shapes: &[Vec<usize>],
384        axis: usize,
385    ) -> TractResult<()>
386    where
387        usize: AsPrimitive<F>,
388    {
389        let first_shape = &input_shapes[0];
390        ensure!(input_shapes.iter().all(|shape| (shape.len() == first_shape.len())
391            && (shape[..axis] == first_shape[..axis])
392            && (if axis != (shape.len() - 1) {
393                shape[(axis + 1)..] == first_shape[(axis + 1)..]
394            } else {
395                true
396            })));
397
398        let op_name = "test".to_string();
399        let dummy_model = TypedModel::default();
400
401        let make_shape =
402            |sym: &str| {
403                input_shapes[0]
404                    .iter()
405                    .enumerate()
406                    .map(|(i, &dim)| {
407                        if i == axis {
408                            TDim::Sym(dummy_model.sym(sym))
409                        } else {
410                            TDim::Val(dim as _)
411                        }
412                    })
413                    .collect::<TVec<TDim>>()
414            };
415
416        let past_shape = make_shape("P");
417        let input_shape = make_shape("S");
418
419        let op = DynKeyValueCache {
420            name: op_name.clone(),
421            past_sequence_fact: TypedFact::dt_shape(F::datum_type(), past_shape),
422            input_sequence_fact: TypedFact::dt_shape(F::datum_type(), input_shape),
423            axis,
424        };
425
426        let mut turn = TurnState::default();
427        let mut state = op.state(&EvalContext::out_of_plan())?.unwrap();
428
429        let mut inputs = tvec![];
430
431        // Init state with first shape
432        let shape = &input_shapes[0];
433        let len = shape.iter().product::<usize>();
434        let input = Tensor::from_shape(shape, &(0..len).map(|f| f.as_()).collect::<Vec<F>>())?;
435        inputs.push(input.clone().into_tvalue());
436
437        let mut state_initializers = vec![input.into()].into_iter();
438
439        state.load_from(&mut turn, &mut state_initializers)?;
440
441        for shape in input_shapes {
442            let len = shape.iter().product::<usize>();
443            let input = Tensor::from_shape(shape, &(0..len).map(|f| f.as_()).collect::<Vec<F>>())?;
444            inputs.push(input.clone().into_tvalue());
445            state.eval(&EvalContext::out_of_plan(), &op, tvec!(input.clone().into()))?[0]
446                .clone()
447                .into_tensor();
448        }
449
450        let mut curr_states = vec![];
451        state.save_to(&mut curr_states)?;
452        let output = curr_states.remove(0);
453
454        let reference = &TypedConcat { axis }.eval(&EvalContext::out_of_plan(), inputs)?[0];
455        output.close_enough(&reference.clone().into_tensor(), Approximation::Close)?;
456        Ok(())
457    }
458
459    #[test]
460    fn test_dyn_kv_cache() -> TractResult<()> {
461        run_test_case::<f32>(&[vec![2, 2]], 0)?;
462        run_test_case::<f32>(&[vec![2, 2], vec![4, 2]], 0)?;
463        run_test_case::<f32>(&[vec![2, 2], vec![2, 1], vec![2, 3]], 1)?;
464        Ok(())
465    }
466
467    // Guards against `has_init_tensor_fact` (the allocation-free predicate used
468    // on the per-run symbol-resolution hot path) drifting out of sync with
469    // `init_tensor_fact`. If they disagree, `resolve_symbols` would silently stop
470    // running for this op.
471    #[test]
472    fn has_init_tensor_fact_matches_init_tensor_fact() -> TractResult<()> {
473        let model = TypedModel::default();
474        let past: TVec<TDim> = tvec![1.to_dim(), model.sym("P").into(), 64.to_dim()];
475        let input: TVec<TDim> = tvec![1.to_dim(), model.sym("S").into(), 64.to_dim()];
476        let op = DynKeyValueCache {
477            name: "kv_cache_0".to_string(),
478            axis: 1,
479            past_sequence_fact: f32::fact(&past),
480            input_sequence_fact: f32::fact(&input),
481        };
482        let _turn = TurnState::default();
483        let state = op.state(&EvalContext::out_of_plan())?.unwrap();
484        assert!(state.has_init_tensor_fact());
485        assert_eq!(state.has_init_tensor_fact(), state.init_tensor_fact().is_some());
486        Ok(())
487    }
488
489    #[test]
490    fn test_unfold_kv_cache() -> TractResult<()> {
491        // Build a model with DynKeyValueCache
492        let mut model = TypedModel::default();
493        let s = model.sym("S");
494        let p = model.sym("P");
495
496        let input_shape: TVec<TDim> = tvec![1.to_dim(), s.into(), 64.to_dim()];
497        let past_shape: TVec<TDim> = tvec![1.to_dim(), p.into(), 64.to_dim()];
498
499        let input = model.add_source("input", f32::fact(&input_shape))?;
500        let op = DynKeyValueCache {
501            name: "kv_cache_0".to_string(),
502            axis: 1,
503            past_sequence_fact: f32::fact(&past_shape),
504            input_sequence_fact: f32::fact(&input_shape),
505        };
506        let out = model.wire_node("kv_cache", op, &[input])?;
507        model.select_output_outlets(&out)?;
508
509        // Model should have 1 input (input), 1 output (kv_cache)
510        assert_eq!(model.inputs.len(), 1);
511        assert_eq!(model.outputs.len(), 1);
512        assert!(model.node(1).op_is::<DynKeyValueCache>());
513
514        // Unfold
515        unfold_kv_cache(&mut model, 1)?;
516
517        // After unfold: 2 inputs (input + kv_cache_0 source), 2 outputs (original + concat)
518        assert_eq!(model.inputs.len(), 2);
519        assert_eq!(model.outputs.len(), 2);
520
521        // The KV cache node should now be a Concat
522        assert!(model.node(1).op_is::<TypedConcat>());
523        let concat = model.node(1).op_as::<TypedConcat>().unwrap();
524        assert_eq!(concat.axis, 1);
525
526        // The new source node should exist
527        let source_node_id = model.inputs[1].node;
528        assert!(model.node(source_node_id).op_is::<TypedSource>());
529        assert_eq!(model.node(source_node_id).name, "kv_cache_0");
530
531        // Concat should have 2 inputs: [source, input]
532        assert_eq!(model.node(1).inputs.len(), 2);
533        assert_eq!(model.node(1).inputs[0].node, source_node_id);
534        assert_eq!(model.node(1).inputs[1].node, 0); // original input
535
536        Ok(())
537    }
538
539    #[test]
540    fn test_fold_unfold_round_trip() -> TractResult<()> {
541        use crate::rewriter::KeyValueCacheTransform;
542        use tract_nnef::tract_core::transform::ModelTransform;
543
544        // Build a model with Source + Concat (the pre-fold pattern)
545        let mut model = TypedModel::default();
546        let s = model.sym("S");
547        let p = model.sym("P");
548
549        let input_shape: TVec<TDim> = tvec![1.to_dim(), s.into(), 64.to_dim()];
550        let past_shape: TVec<TDim> = tvec![1.to_dim(), p.into(), 64.to_dim()];
551
552        let past = model.add_source("kv_past", f32::fact(&past_shape))?;
553        let input = model.add_source("input", f32::fact(&input_shape))?;
554        let concat = model.wire_node("concat", TypedConcat { axis: 1 }, &[past, input])?;
555        model.select_output_outlets(&concat)?;
556
557        let orig_input_count = model.inputs.len();
558        let orig_output_count = model.outputs.len();
559
560        // Fold: Source + Concat -> DynKeyValueCache
561        KeyValueCacheTransform.transform(&mut model)?;
562        assert_eq!(model.inputs.len(), orig_input_count - 1); // past source removed
563        assert_eq!(model.outputs.len(), orig_output_count - 1); // concat output removed
564
565        // Find the DynKeyValueCache node
566        let kv_node_id = model.nodes().iter().find(|n| n.op_is::<DynKeyValueCache>()).unwrap().id;
567
568        // Unfold: DynKeyValueCache -> Source + Concat
569        unfold_kv_cache(&mut model, kv_node_id)?;
570
571        // Should be back to original structure
572        assert_eq!(model.inputs.len(), orig_input_count);
573        assert_eq!(model.outputs.len(), orig_output_count);
574
575        // Verify it's a Concat again
576        let concat_node = model.nodes().iter().find(|n| n.op_is::<TypedConcat>()).unwrap();
577        assert_eq!(concat_node.op_as::<TypedConcat>().unwrap().axis, 1);
578        assert_eq!(concat_node.inputs.len(), 2);
579
580        Ok(())
581    }
582
583    #[test]
584    fn test_dyn_kv_cache_nnef_round_trip() -> TractResult<()> {
585        use crate::WithTractTransformers;
586
587        let mut model = TypedModel::default();
588        let s = model.sym("S");
589        let p = model.sym("P");
590
591        let input_shape: TVec<TDim> = tvec![1.to_dim(), s.into(), 64.to_dim()];
592        let past_shape: TVec<TDim> = tvec![1.to_dim(), p.into(), 64.to_dim()];
593
594        let input = model.add_source("input", f32::fact(&input_shape))?;
595        let op = DynKeyValueCache {
596            name: "kv_cache_0".to_string(),
597            axis: 1,
598            past_sequence_fact: f32::fact(&past_shape),
599            input_sequence_fact: f32::fact(&input_shape),
600        };
601        let out = model.wire_node("kv_cache", op, &[input])?;
602        model.select_output_outlets(&out)?;
603
604        let nnef = tract_nnef::nnef().with_tract_transformers();
605        let mut buffer = vec![];
606        nnef.write_to_tar(&model, &mut buffer)?;
607        let reloaded = nnef.model_for_read(&mut &*buffer)?;
608
609        assert_eq!(reloaded.nodes().len(), model.nodes().len());
610        let reloaded_kv = reloaded.node(1);
611        let reloaded_op = reloaded_kv.op_as::<DynKeyValueCache>().unwrap();
612        assert_eq!(reloaded_op.name, "kv_cache_0");
613        assert_eq!(reloaded_op.axis, 1);
614        assert_eq!(reloaded_op.past_sequence_fact.datum_type, DatumType::F32);
615        assert_eq!(reloaded_op.past_sequence_fact.shape.rank(), 3);
616        assert_eq!(reloaded_op.input_sequence_fact.datum_type, DatumType::F32);
617        assert_eq!(reloaded_op.input_sequence_fact.shape.rank(), 3);
618        Ok(())
619    }
620}
tract_transformers/ops/dyn_kv_cache.rs

tract_transformers/ops/
dyn_kv_cache.rs
