tract_gpu/ops/
dyn_kv_cache.rs

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

tract_gpu/ops/
dyn_kv_cache.rs
