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

tract_transformers/ops/
dyn_kv_cache.rs
