tract_gpu/ops/
dyn_kv_cache.rs

1use crate::fact::DeviceTypedFactExt;
2use crate::tensor::{DeviceTensor, DeviceTensorExt, IntoDevice};
3use derive_new::new;
4use tract_core::internal::*;
5use tract_core::ops::OpStateFreeze;
6use tract_transformers::ops::dyn_kv_cache::{DynKeyValueCache, DynKeyValueCacheState};
7
8#[derive(Debug, Clone, new)]
9pub struct GpuDynKVCacheState {
10    node_id: usize,
11    name: String,
12    axis: usize,
13    past_sequence_fact: TypedFact,
14    kv_cache: Option<TValue>,
15}
16
17impl OpState for GpuDynKVCacheState {
18    fn load_from(
19        &mut self,
20        state: &mut TurnState,
21        states: &mut dyn Iterator<Item = TValue>,
22    ) -> TractResult<()> {
23        let kv_cache = states.next().context("Not enough state initializers")?;
24        DynKeyValueCacheState::resolve_symbols(
25            state,
26            self.past_sequence_fact.clone(),
27            Some(kv_cache.shape()),
28        )?;
29        self.kv_cache = Some(kv_cache.into_tensor().into_device()?.into_tensor().into_tvalue());
30        Ok(())
31    }
32
33    fn save_to(&self, states: &mut Vec<TValue>) -> TractResult<()> {
34        if let Some(kv_cache) = &self.kv_cache {
35            states.push(kv_cache.to_device_tensor()?.to_host()?.into_tensor().into_tvalue());
36            Ok(())
37        } else {
38            bail!("KV cache {} was never initialized", self.name)
39        }
40    }
41
42    fn init_tensor_fact(&self) -> Option<(String, TypedFact)> {
43        Some((self.name.clone(), self.past_sequence_fact.clone()))
44    }
45
46    fn has_init_tensor_fact(&self) -> bool {
47        true
48    }
49
50    fn resolve_symbols(&mut self, state: &mut TurnState) -> TractResult<()> {
51        let shape = self
52            .kv_cache
53            .as_ref()
54            .map(|kv_cache| kv_cache.to_device_tensor().expect("Expected GPU Tensor").shape());
55        DynKeyValueCacheState::resolve_symbols(state, self.past_sequence_fact.clone(), shape)
56    }
57
58    fn eval(
59        &mut self,
60        session: &mut TurnState,
61        op: &dyn Op,
62        inputs: TVec<TValue>,
63    ) -> TractResult<TVec<TValue>> {
64        ensure!(inputs.len() == 1);
65        let mut op_inputs = TVec::new();
66
67        if let Some(kv_cache) = self.kv_cache.take() {
68            op_inputs.push(kv_cache);
69        }
70
71        op_inputs.push(inputs.into_iter().next().unwrap());
72
73        let gpu_op =
74            op.downcast_ref::<GpuDynKVCache>().ok_or_else(|| format_err!("Wrong Op type"))?;
75        let axis = gpu_op.axis;
76
77        let inputs =
78            op_inputs.iter().map(|it| it.to_device_tensor()).collect::<TractResult<TVec<_>>>()?;
79        let mut output_shape = inputs[0].shape().to_vec();
80        output_shape[axis] = inputs.iter().map(|it| it.shape()[axis]).sum();
81        let output = crate::session_handler::make_tensor_for_node(
82            session,
83            self.node_id,
84            inputs[0].datum_type(),
85            &output_shape,
86        )?;
87
88        // Concat inputs into output
89        let ctx = crate::device::get_context()?;
90        let mut cursor = 0usize;
91        for input in &inputs {
92            let slice_len = input.shape()[axis];
93            if slice_len == 0 {
94                continue;
95            }
96            let dst_offset =
97                cursor * output.strides()[axis] as usize * output.datum_type().size_of();
98            ctx.copy_nd(
99                input,
100                0,
101                input.strides(),
102                &output,
103                dst_offset,
104                input.shape(),
105                output.strides(),
106            )?;
107            cursor += slice_len;
108        }
109
110        let res = output.into_tensor().into_tvalue();
111        self.kv_cache = Some(res.clone());
112        Ok(tvec!(res))
113    }
114}
115
116impl GpuDynKVCacheState {
117    pub fn truncate(&mut self, len: usize) -> TractResult<()> {
118        if let Some(v) = &mut self.kv_cache {
119            let mut t: Tensor = v.to_device_tensor()?.to_host()?.into_tensor();
120            t = t.slice(self.axis, 0, len)?;
121            *v = t.into_device()?.into_tensor().into_tvalue();
122        }
123        Ok(())
124    }
125}
126
127#[derive(Debug, Clone)]
128pub struct FrozenGpuDynKVCacheState {
129    node_id: usize,
130    name: String,
131    axis: usize,
132    past_sequence_fact: TypedFact,
133    kv_cache: Option<DeviceTensor>,
134}
135
136impl OpStateFreeze for GpuDynKVCacheState {
137    fn freeze(&self) -> Box<dyn FrozenOpState + 'static> {
138        Box::new(FrozenGpuDynKVCacheState {
139            node_id: self.node_id,
140            name: self.name.clone(),
141            axis: self.axis,
142            past_sequence_fact: self.past_sequence_fact.clone(),
143            kv_cache: self.kv_cache.clone().map(|t| t.to_device_tensor().cloned().unwrap()),
144        })
145    }
146
147    fn freeze_into(self: Box<Self>) -> Box<dyn FrozenOpState> {
148        Box::new(FrozenGpuDynKVCacheState {
149            node_id: self.node_id,
150            name: self.name,
151            axis: self.axis,
152            past_sequence_fact: self.past_sequence_fact,
153            kv_cache: self.kv_cache.map(|t| t.to_device_tensor().cloned().unwrap()),
154        })
155    }
156}
157
158impl FrozenOpState for FrozenGpuDynKVCacheState {
159    fn unfreeze(&self) -> Box<dyn OpState> {
160        Box::new(GpuDynKVCacheState {
161            node_id: self.node_id,
162            name: self.name.clone(),
163            axis: self.axis,
164            past_sequence_fact: self.past_sequence_fact.clone(),
165            kv_cache: self.kv_cache.clone().map(|t| t.into_tensor().into_tvalue()),
166        })
167    }
168}
169
170#[derive(Clone)]
171pub struct GpuDynKVCache {
172    pub name: String,
173    pub past_sequence_fact: TypedFact,
174    pub input_sequence_fact: TypedFact,
175    pub axis: usize,
176}
177
178impl GpuDynKVCache {
179    pub fn from_tract_transformers(op: &DynKeyValueCache) -> Self {
180        Self {
181            name: op.name.clone(),
182            axis: op.axis,
183            past_sequence_fact: op.past_sequence_fact.clone(),
184            input_sequence_fact: op.input_sequence_fact.clone(),
185        }
186    }
187}
188
189impl std::fmt::Debug for GpuDynKVCache {
190    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
191        write!(f, "GpuDynKVCache({}, axis={})", self.name, self.axis)
192    }
193}
194
195impl PartialEq for GpuDynKVCache {
196    fn eq(&self, other: &Self) -> bool {
197        self.name == other.name
198            && self.axis == other.axis
199            && self.past_sequence_fact == other.past_sequence_fact
200            && self.input_sequence_fact == other.input_sequence_fact
201    }
202}
203
204impl Eq for GpuDynKVCache {}
205
206impl std::hash::Hash for GpuDynKVCache {
207    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
208        self.name.hash(state);
209        self.axis.hash(state);
210    }
211}
212
213impl Op for GpuDynKVCache {
214    fn name(&self) -> StaticName {
215        "GpuDynKVCache".into()
216    }
217
218    fn info(&self) -> TractResult<Vec<String>> {
219        Ok(vec![format!("axis: {}", self.axis)])
220    }
221
222    op_as_typed_op!();
223}
224
225impl EvalOp for GpuDynKVCache {
226    fn is_stateless(&self) -> bool {
227        false
228    }
229
230    fn state(&self, _session: &TurnState, node_id: usize) -> TractResult<Option<Box<dyn OpState>>> {
231        Ok(Some(Box::new(GpuDynKVCacheState::new(
232            node_id,
233            self.name.clone(),
234            self.axis,
235            self.past_sequence_fact.clone(),
236            None,
237        ))))
238    }
239}
240
241impl TypedOp for GpuDynKVCache {
242    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
243        ensure!(inputs.len() == 1);
244        let mut facts = crate::utils::facts_to_device_facts(inputs, |facts| {
245            let mut fact = facts[0].without_value();
246            fact.shape.set(
247                self.axis,
248                self.past_sequence_fact.shape.dims()[self.axis].clone()
249                    + self.input_sequence_fact.shape.dims()[self.axis].clone(),
250            );
251            Ok(tvec!(fact))
252        })
253        .with_context(|| format!("Error while computing facts for {:?}", self.name()))?;
254        facts[0].as_device_fact_mut().unwrap().state_owned = true;
255        Ok(facts)
256    }
257
258    as_op!();
259}
tract_gpu/ops/dyn_kv_cache.rs

tract_gpu/ops/
dyn_kv_cache.rs
