tract_transformers/
lib.rs

1pub mod ops;
2mod rewriter;
3use std::collections::HashSet;
4
5use rewriter::*;
6use tract_nnef::internal::*;
7
8register_simple_model_transform!("detect_apply_rope", ApplyRopeTransform);
9register_simple_model_transform!("detect_diag_gather", DetectDiagGatherTransform);
10register_simple_model_transform!("detect_scaled_masked_softmax", ScaledMaskedSoftmaxTransform);
11register_simple_model_transform!("detect_kv_cache", KeyValueCacheTransform);
12register_simple_model_transform!(
13    "detect_sdpa_kv_cache_broadcast",
14    SdpaFuseKvCacheBroadcastTransform
15);
16register_simple_model_transform!("unfold_kv_cache", UnfoldKeyValueCacheTransform);
17register_simple_model_transform!(
18    "fuse_inplace_kv_sdpa",
19    ops::inplace_kv_cache::InPlaceKvSdpaTransform
20);
21register_simple_model_transform!("transformers_detect_all", TransformersTransform);
22
23pub fn register(registry: &mut Registry) {
24    ops::apply_rope::register(registry);
25    ops::scaled_masked_softmax::register(registry);
26    ops::sdpa::register(registry);
27    ops::dyn_kv_cache::register(registry);
28    ops::window_kv_cache::register(registry);
29    ops::kv_quant::register(registry);
30}
31
32pub trait WithTractTransformers {
33    fn enable_tract_transformers(&mut self);
34    fn with_tract_transformers(self) -> Self;
35}
36
37impl WithTractTransformers for tract_nnef::framework::Nnef {
38    fn enable_tract_transformers(&mut self) {
39        self.registries.push(tract_transformers_registry());
40    }
41
42    fn with_tract_transformers(mut self) -> Self {
43        self.enable_tract_transformers();
44        self
45    }
46}
47
48pub fn tract_transformers_registry() -> Registry {
49    let mut reg = Registry::new("tract_transformers")
50        .with_doc("Extension `tract_transformers` extends NNEF with operators")
51        .with_doc("for transformer networks.")
52        .with_doc("")
53        .with_doc("Add `extension tract_transformers` to `graph.nnef`");
54
55    register(&mut reg);
56    reg
57}
58
59pub fn figure_out_causal_llm_b_s_p(
60    model: &TypedModel,
61) -> TractResult<(Option<Symbol>, Option<Symbol>, Option<Symbol>)> {
62    // expectations:
63    // - one input is for tokens, so integer dt (i64 ?) and typically of shape S or 1,S, or B,S
64    // - other inputs are kv cache, some kind of float. shape features both S and P, and B if B is present in tokens
65    let token_input = model
66        .inputs
67        .iter()
68        .position(|i| model.outlet_fact(*i).unwrap().datum_type.is_integer())
69        .context("No token input found")?;
70    let tokens_symbols = model.input_fact(token_input)?.shape.volume().symbols();
71    let kv_symbols = if let Some(kv_input) =
72        model.inputs.iter().position(|i| model.outlet_fact(*i).unwrap().datum_type.is_float())
73    {
74        model.input_fact(kv_input)?.shape.volume().symbols()
75    } else {
76        // Look for KVCache Op
77        let dummy_session_state = TurnState::default();
78        let mut symbols = HashSet::new();
79        for node in &model.nodes {
80            if let Some((_, fact)) =
81                node.op.state(&dummy_session_state, 0)?.and_then(|state| state.init_tensor_fact())
82            {
83                symbols = fact.shape.volume().symbols();
84                break;
85            }
86        }
87        symbols
88    };
89
90    let b = tokens_symbols.intersection(&kv_symbols).cloned().collect::<HashSet<_>>();
91    let s = tokens_symbols.difference(&b).cloned().collect::<HashSet<_>>();
92    let p = kv_symbols.difference(&b).cloned().collect::<HashSet<_>>();
93    Ok((b.into_iter().next(), s.into_iter().next(), p.into_iter().next()))
94}
95
96pub fn memory_arena_hints_for_causal_llm(model: &TypedModel) -> TractResult<SymbolValues> {
97    let (b, s, p) = figure_out_causal_llm_b_s_p(model)?;
98    let mut values = SymbolValues::default()
99        .with(&s.context("Could not determine sequence_len (S)")?, 1024)
100        .with(&p.context("Could not determine past_sequence_len (P)")?, 0);
101    if let Some(b) = b {
102        values = values.with(&b, 1);
103    }
104    Ok(values)
105}
tract_transformers/lib.rs

tract_transformers/
lib.rs
