rudnn/attention/kernel_ir/launch/
args.rs

1use ruda_kernel::dsl::prelude::*;
2use ruda_kernel::library::tensor::r#virtual::VirtualTensorOperations;
3use ruda_kernel::library::tensor::r#virtual::VirtualTensorOperationsExpand;
4use ruda_kernel::dsl as kernel_dsl;
5
6use crate::attention::kernel_ir::definition::{AttentionBlueprint, AttentionProblem};
7
8/// Create the input runtime arguments for a attention kernel that works on concrete inputs and
9/// output (not fused).
10pub trait ConcreteInputsFactory: LaunchArg {
11    fn create<R: Runtime>(
12        query: TensorBinding<R>,
13        key: TensorBinding<R>,
14        value: TensorBinding<R>,
15        mask: Option<TensorBinding<R>>,
16        selection: &AttentionBlueprint,
17        problem: &AttentionProblem,
18    ) -> Self::RuntimeArg<R>;
19}
20
21/// Create the output runtime argument for a attention kernel that works on concrete inputs and
22/// output (not fused).
23pub trait ConcreteOutputFactory: LaunchArg {
24    fn create<R: Runtime>(
25        out: TensorBinding<R>,
26        selection: &AttentionBlueprint,
27        problem: &AttentionProblem,
28    ) -> Self::RuntimeArg<R>;
29}
30
31pub trait FloatLine: 'static {
32    type T: Float;
33    type N: Size;
34}
35
36impl<T: Float, N: Size> FloatLine for (T, N) {
37    type T = T;
38    type N = N;
39}
40
41pub trait NumericLine: 'static {
42    type T: Numeric;
43    type N: Size;
44}
45
46impl<T: Numeric, N: Size> NumericLine for (T, N) {
47    type T = T;
48    type N = N;
49}
50
51#[ruda]
52/// Arguments for the attention algorithm.
53pub trait AttentionArgs: Send + Sync + 'static + Clone {
54    /// Type used for the input.
55    type Input<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine>: LaunchArg + RudaType;
56    /// Type used for the output.
57    type Output<O: FloatLine>: LaunchArg + RudaType;
58    /// Inner state that is used to create tensor inputs and
59    /// [tensor outputs](TensorOutput) .
60    type State<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>: RudaType;
61
62    /// Init the state.
63    fn init_state<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
64        input: &Self::Input<Q, K, V, M>,
65        output: &mut Self::Output<O>,
66    ) -> Self::State<Q, K, V, M, O>;
67
68    /// Whether the mask argument is present. Returns `Option` to allow matching at
69    /// comptime
70    fn has_mask<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
71        state: &Self::State<Q, K, V, M, O>,
72    ) -> ComptimeOption<()>;
73
74    /// Read the vector of the query tensor using the state at the given coordinate.
75    fn read_query<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
76        state: &Self::State<Q, K, V, M, O>,
77        coordinate: usize,
78    ) -> Vector<Q::T, Q::N>;
79    /// Read the vector of the key tensor using the state at the given coordinate.
80    fn read_key<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
81        state: &Self::State<Q, K, V, M, O>,
82        coordinate: usize,
83    ) -> Vector<K::T, K::N>;
84    /// Read the vector of the value tensor using the state at the given coordinate.
85    fn read_value<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
86        state: &Self::State<Q, K, V, M, O>,
87        coordinate: usize,
88    ) -> Vector<V::T, V::N>;
89    /// Read the vector of the mask tensor using the state at the given coordinate.
90    fn read_mask<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
91        state: &Self::State<Q, K, V, M, O>,
92        coordinate: usize,
93    ) -> Vector<M::T, M::N>;
94
95    /// Read the vector of the query tensor using the state at the given coordinate.
96    fn read_window_query<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
97        state: &Self::State<Q, K, V, M, O>,
98        start: usize,
99        end: usize,
100    ) -> Slice<Vector<Q::T, Q::N>>;
101    /// Read the vector of the key tensor using the state at the given coordinate.
102    fn read_window_key<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
103        state: &Self::State<Q, K, V, M, O>,
104        start: usize,
105        end: usize,
106    ) -> Slice<Vector<K::T, K::N>>;
107    /// Read the vector of the value tensor using the state at the given coordinate.
108    fn read_window_value<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
109        state: &Self::State<Q, K, V, M, O>,
110        start: usize,
111        end: usize,
112    ) -> Slice<Vector<V::T, V::N>>;
113    /// Read the vector of the mask tensor using the state at the given coordinate.
114    fn read_window_mask<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
115        state: &Self::State<Q, K, V, M, O>,
116        start: usize,
117        end: usize,
118    ) -> Slice<Vector<M::T, M::N>>;
119
120    /// Reinterpret query as tensor map
121    fn as_tensor_map_query<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
122        state: &Self::State<Q, K, V, M, O>,
123    ) -> ComptimeOption<TensorMap<Q::T, Tiled>>;
124    /// Reinterpret key as tensor map
125    fn as_tensor_map_key<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
126        state: &Self::State<Q, K, V, M, O>,
127    ) -> ComptimeOption<TensorMap<K::T, Tiled>>;
128    /// Reinterpret value as tensor map
129    fn as_tensor_map_value<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
130        state: &Self::State<Q, K, V, M, O>,
131    ) -> ComptimeOption<TensorMap<V::T, Tiled>>;
132    /// Reinterpret mask as tensor map
133    fn as_tensor_map_mask<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
134        state: &Self::State<Q, K, V, M, O>,
135    ) -> ComptimeOption<TensorMap<M::T, Tiled>>;
136
137    /// Write the vector to the output at the given coordinate using the state.
138    fn write_out<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
139        state: &mut Self::State<Q, K, V, M, O>,
140        coordinate: usize,
141        val: Vector<O::T, O::N>,
142    );
143
144    /// Get the rank of the query tensor using the state.
145    fn rank_query<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
146        state: &Self::State<Q, K, V, M, O>,
147    ) -> usize;
148    /// Get the rank of the key tensor using the state.
149    fn rank_key<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
150        state: &Self::State<Q, K, V, M, O>,
151    ) -> usize;
152    /// Get the rank of the value tensor using the state.
153    fn rank_value<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
154        state: &Self::State<Q, K, V, M, O>,
155    ) -> usize;
156    /// Get the rank of the mask tensor using the state.
157    fn rank_mask<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
158        state: &Self::State<Q, K, V, M, O>,
159    ) -> usize;
160    /// Get the rank of the out tensor using the state.
161    fn rank_out<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
162        state: &Self::State<Q, K, V, M, O>,
163    ) -> usize;
164
165    /// Get the length of the query tensor using the state.
166    fn len_query<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
167        state: &Self::State<Q, K, V, M, O>,
168    ) -> usize;
169    /// Get the length of the key tensor using the state.
170    fn len_key<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
171        state: &Self::State<Q, K, V, M, O>,
172    ) -> usize;
173    /// Get the length of the value tensor using the state.
174    fn len_value<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
175        state: &Self::State<Q, K, V, M, O>,
176    ) -> usize;
177    /// Get the length of the mask tensor using the state.
178    fn len_mask<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
179        state: &Self::State<Q, K, V, M, O>,
180    ) -> usize;
181    /// Get the length of the out tensor using the state.
182    fn len_out<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
183        state: &Self::State<Q, K, V, M, O>,
184    ) -> usize;
185
186    /// Get the buffer length of the query tensor using the state.
187    fn buffer_len_query<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
188        state: &Self::State<Q, K, V, M, O>,
189    ) -> usize;
190    /// Get the buffer length of the key tensor using the state.
191    fn buffer_len_key<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
192        state: &Self::State<Q, K, V, M, O>,
193    ) -> usize;
194    /// Get the buffer length of the value tensor using the state.
195    fn buffer_len_value<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
196        state: &Self::State<Q, K, V, M, O>,
197    ) -> usize;
198    /// Get the buffer length of the mask tensor using the state.
199    fn buffer_len_mask<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
200        state: &Self::State<Q, K, V, M, O>,
201    ) -> usize;
202    /// Get the buffer length of the out tensor using the state.
203    fn buffer_len_out<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
204        state: &Self::State<Q, K, V, M, O>,
205    ) -> usize;
206
207    /// Get the shape of the query tensor using the state.
208    fn shape_query<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
209        state: &Self::State<Q, K, V, M, O>,
210        axis: usize,
211    ) -> usize;
212    /// Get the shape of the key tensor using the state.
213    fn shape_key<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
214        state: &Self::State<Q, K, V, M, O>,
215        axis: usize,
216    ) -> usize;
217    /// Get the shape of the value tensor using the state.
218    fn shape_value<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
219        state: &Self::State<Q, K, V, M, O>,
220        axis: usize,
221    ) -> usize;
222    /// Get the shape of the mask tensor using the state.
223    fn shape_mask<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
224        state: &Self::State<Q, K, V, M, O>,
225        axis: usize,
226    ) -> usize;
227    /// Get the shape of the out tensor using the state.
228    fn shape_out<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
229        state: &Self::State<Q, K, V, M, O>,
230        axis: usize,
231    ) -> usize;
232
233    /// Get the stride of the query tensor using the state.
234    fn stride_query<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
235        state: &Self::State<Q, K, V, M, O>,
236        axis: usize,
237    ) -> usize;
238    /// Get the stride of the key tensor using the state.
239    fn stride_key<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
240        state: &Self::State<Q, K, V, M, O>,
241        axis: usize,
242    ) -> usize;
243    /// Get the stride of the value tensor using the state.
244    fn stride_value<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
245        state: &Self::State<Q, K, V, M, O>,
246        axis: usize,
247    ) -> usize;
248    /// Get the stride of the mask tensor using the state.
249    fn stride_mask<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
250        state: &Self::State<Q, K, V, M, O>,
251        axis: usize,
252    ) -> usize;
253    /// Get the stride of the out tensor using the state.
254    fn stride_out<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
255        state: &Self::State<Q, K, V, M, O>,
256        axis: usize,
257    ) -> usize;
258
259    /// Get the vector size of the query tensor using the state.
260    fn vector_size_query<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
261        state: &Self::State<Q, K, V, M, O>,
262    ) -> comptime_type!(VectorSize);
263    /// Get the vector size of the key tensor using the state.
264    fn vector_size_key<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
265        state: &Self::State<Q, K, V, M, O>,
266    ) -> comptime_type!(VectorSize);
267    /// Get the vector size of the value tensor using the state.
268    fn vector_size_value<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
269        state: &Self::State<Q, K, V, M, O>,
270    ) -> comptime_type!(VectorSize);
271    /// Get the vector size of the mask tensor using the state.
272    fn vector_size_mask<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
273        state: &Self::State<Q, K, V, M, O>,
274    ) -> comptime_type!(VectorSize);
275    /// Get the vector size of the out tensor using the state.
276    fn vector_size_out<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
277        state: &Self::State<Q, K, V, M, O>,
278    ) -> comptime_type!(VectorSize);
279}
280
281/// Tensor input representation.
282///
283/// You can use the tensor input as if it was a pointer to the actually tensor.
284pub struct TensorQuery<
285    Q: FloatLine,
286    K: FloatLine,
287    V: FloatLine,
288    M: NumericLine,
289    O: FloatLine,
290    GA: AttentionArgs,
291> {
292    state: *const GA::State<Q, K, V, M, O>,
293}
294
295pub struct TensorKey<
296    Q: FloatLine,
297    K: FloatLine,
298    V: FloatLine,
299    M: NumericLine,
300    O: FloatLine,
301    GA: AttentionArgs,
302> {
303    state: *const GA::State<Q, K, V, M, O>,
304}
305
306pub struct TensorValue<
307    Q: FloatLine,
308    K: FloatLine,
309    V: FloatLine,
310    M: NumericLine,
311    O: FloatLine,
312    GA: AttentionArgs,
313> {
314    state: *const GA::State<Q, K, V, M, O>,
315}
316
317pub struct TensorMask<
318    Q: FloatLine,
319    K: FloatLine,
320    V: FloatLine,
321    M: NumericLine,
322    O: FloatLine,
323    GA: AttentionArgs,
324> {
325    state: *const GA::State<Q, K, V, M, O>,
326}
327
328impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
329    VirtualTensorOperations<Q::T, Q::N> for TensorQuery<Q, K, V, M, O, MA>
330{
331}
332
333impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
334    VirtualTensorOperations<K::T, K::N> for TensorKey<Q, K, V, M, O, MA>
335{
336}
337
338impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
339    VirtualTensorOperations<V::T, V::N> for TensorValue<Q, K, V, M, O, MA>
340{
341}
342
343impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
344    VirtualTensorOperations<M::T, M::N> for TensorMask<Q, K, V, M, O, MA>
345{
346}
347
348impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
349    VirtualTensorOperations<O::T, O::N> for TensorOutput<Q, K, V, M, O, MA>
350{
351}
352
353impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
354    VirtualTensorOperationsExpand<O::T, O::N> for TensorOutputExpand<Q, K, V, M, O, MA>
355{
356    fn __expand_read_method(
357        &self,
358        _scope: &mut Scope,
359        _index: NativeExpand<usize>,
360    ) -> NativeExpand<Vector<O::T, O::N>> {
361        panic!("Can't read output tensor");
362    }
363
364    fn __expand_read_window_method(
365        &self,
366        _context: &mut Scope,
367        _start: NativeExpand<usize>,
368        _end: NativeExpand<usize>,
369    ) -> SliceExpand<Vector<O::T, O::N>, ReadOnly> {
370        panic!("Can't read output tensor");
371    }
372
373    fn __expand_write_method(
374        &self,
375        scope: &mut Scope,
376        index: NativeExpand<usize>,
377        val: NativeExpand<Vector<O::T, O::N>>,
378    ) {
379        TensorOutputExpand::__expand_write_method(self.clone(), scope, index, val)
380    }
381
382    fn __expand_shape_method(
383        &self,
384        scope: &mut Scope,
385        axis: NativeExpand<usize>,
386    ) -> NativeExpand<usize> {
387        TensorOutputExpand::__expand_shape_method(self.clone(), scope, axis)
388    }
389
390    fn __expand_stride_method(
391        &self,
392        scope: &mut Scope,
393        axis: NativeExpand<usize>,
394    ) -> NativeExpand<usize> {
395        TensorOutputExpand::__expand_stride_method(self.clone(), scope, axis)
396    }
397
398    fn __expand_rank_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
399        TensorOutputExpand::__expand_rank_method(self.clone(), scope)
400    }
401
402    fn __expand_len_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
403        TensorOutputExpand::__expand_len_method(self.clone(), scope)
404    }
405
406    fn __expand_buffer_len_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
407        TensorOutputExpand::__expand_buffer_len_method(self.clone(), scope)
408    }
409
410    fn __expand_as_tensor_map_method(
411        &self,
412        scope: &mut Scope,
413    ) -> ComptimeOptionExpand<TensorMap<O::T, Tiled>> {
414        ComptimeOption::__expand_new_None(scope)
415    }
416}
417
418impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
419    Vectorized for TensorOutput<Q, K, V, M, O, MA>
420{
421}
422impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
423    VectorizedExpand for TensorOutputExpand<Q, K, V, M, O, MA>
424{
425    fn vector_size(&self) -> VectorSize {
426        let mut scope = Scope::root(false);
427        TensorOutputExpand::__expand_vector_size_method(self.clone(), &mut scope)
428    }
429}
430
431impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
432    VirtualTensorOperationsExpand<Q::T, Q::N> for TensorQueryExpand<Q, K, V, M, O, MA>
433{
434    fn __expand_read_method(
435        &self,
436        scope: &mut Scope,
437        index: NativeExpand<usize>,
438    ) -> NativeExpand<Vector<Q::T, Q::N>> {
439        TensorQueryExpand::__expand_read_method(self.clone(), scope, index)
440    }
441    fn __expand_read_window_method(
442        &self,
443        context: &mut Scope,
444        start: NativeExpand<usize>,
445        end: NativeExpand<usize>,
446    ) -> SliceExpand<Vector<Q::T, Q::N>, ReadOnly> {
447        TensorQueryExpand::__expand_read_window_method(self.clone(), context, start, end)
448    }
449
450    fn __expand_write_method(
451        &self,
452        _scope: &mut Scope,
453        _index: NativeExpand<usize>,
454        _val: NativeExpand<Vector<Q::T, Q::N>>,
455    ) {
456        panic!("Can't write to input tensor");
457    }
458
459    fn __expand_shape_method(
460        &self,
461        scope: &mut Scope,
462        axis: NativeExpand<usize>,
463    ) -> NativeExpand<usize> {
464        TensorQueryExpand::__expand_shape_method(self.clone(), scope, axis)
465    }
466
467    fn __expand_stride_method(
468        &self,
469        scope: &mut Scope,
470        axis: NativeExpand<usize>,
471    ) -> NativeExpand<usize> {
472        TensorQueryExpand::__expand_stride_method(self.clone(), scope, axis)
473    }
474
475    fn __expand_rank_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
476        TensorQueryExpand::__expand_rank_method(self.clone(), scope)
477    }
478
479    fn __expand_len_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
480        TensorQueryExpand::__expand_len_method(self.clone(), scope)
481    }
482
483    fn __expand_buffer_len_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
484        TensorQueryExpand::__expand_buffer_len_method(self.clone(), scope)
485    }
486
487    fn __expand_as_tensor_map_method(
488        &self,
489        scope: &mut Scope,
490    ) -> ComptimeOptionExpand<TensorMap<Q::T, Tiled>> {
491        TensorQueryExpand::__expand_as_tensor_map_method(self.clone(), scope)
492    }
493}
494
495impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
496    Vectorized for TensorQuery<Q, K, V, M, O, MA>
497{
498}
499impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
500    VectorizedExpand for TensorQueryExpand<Q, K, V, M, O, MA>
501{
502    fn vector_size(&self) -> VectorSize {
503        let mut scope = Scope::root(false);
504        TensorQueryExpand::__expand_vector_size_method(self.clone(), &mut scope)
505    }
506}
507
508impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
509    VirtualTensorOperationsExpand<K::T, K::N> for TensorKeyExpand<Q, K, V, M, O, MA>
510{
511    fn __expand_read_method(
512        &self,
513        scope: &mut Scope,
514        index: NativeExpand<usize>,
515    ) -> NativeExpand<Vector<K::T, K::N>> {
516        TensorKeyExpand::__expand_read_method(self.clone(), scope, index)
517    }
518    fn __expand_read_window_method(
519        &self,
520        context: &mut Scope,
521        start: NativeExpand<usize>,
522        end: NativeExpand<usize>,
523    ) -> SliceExpand<Vector<K::T, K::N>, ReadOnly> {
524        TensorKeyExpand::__expand_read_window_method(self.clone(), context, start, end)
525    }
526
527    fn __expand_write_method(
528        &self,
529        _scope: &mut Scope,
530        _index: NativeExpand<usize>,
531        _val: NativeExpand<Vector<K::T, K::N>>,
532    ) {
533        panic!("Can't write to input tensor");
534    }
535
536    fn __expand_shape_method(
537        &self,
538        scope: &mut Scope,
539        axis: NativeExpand<usize>,
540    ) -> NativeExpand<usize> {
541        TensorKeyExpand::__expand_shape_method(self.clone(), scope, axis)
542    }
543
544    fn __expand_stride_method(
545        &self,
546        scope: &mut Scope,
547        axis: NativeExpand<usize>,
548    ) -> NativeExpand<usize> {
549        TensorKeyExpand::__expand_stride_method(self.clone(), scope, axis)
550    }
551
552    fn __expand_rank_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
553        TensorKeyExpand::__expand_rank_method(self.clone(), scope)
554    }
555
556    fn __expand_len_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
557        TensorKeyExpand::__expand_len_method(self.clone(), scope)
558    }
559
560    fn __expand_buffer_len_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
561        TensorKeyExpand::__expand_buffer_len_method(self.clone(), scope)
562    }
563
564    fn __expand_as_tensor_map_method(
565        &self,
566        scope: &mut Scope,
567    ) -> ComptimeOptionExpand<TensorMap<K::T, Tiled>> {
568        TensorKeyExpand::__expand_as_tensor_map_method(self.clone(), scope)
569    }
570}
571
572impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
573    Vectorized for TensorKey<Q, K, V, M, O, MA>
574{
575}
576impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
577    VectorizedExpand for TensorKeyExpand<Q, K, V, M, O, MA>
578{
579    fn vector_size(&self) -> VectorSize {
580        let mut scope = Scope::root(false);
581        TensorKeyExpand::__expand_vector_size_method(self.clone(), &mut scope)
582    }
583}
584
585impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
586    VirtualTensorOperationsExpand<V::T, V::N> for TensorValueExpand<Q, K, V, M, O, MA>
587{
588    fn __expand_read_method(
589        &self,
590        scope: &mut Scope,
591        index: NativeExpand<usize>,
592    ) -> NativeExpand<Vector<V::T, V::N>> {
593        TensorValueExpand::__expand_read_method(self.clone(), scope, index)
594    }
595    fn __expand_read_window_method(
596        &self,
597        context: &mut Scope,
598        start: NativeExpand<usize>,
599        end: NativeExpand<usize>,
600    ) -> SliceExpand<Vector<V::T, V::N>, ReadOnly> {
601        TensorValueExpand::__expand_read_window_method(self.clone(), context, start, end)
602    }
603
604    fn __expand_write_method(
605        &self,
606        _scope: &mut Scope,
607        _index: NativeExpand<usize>,
608        _val: NativeExpand<Vector<V::T, V::N>>,
609    ) {
610        panic!("Can't write to input tensor");
611    }
612
613    fn __expand_shape_method(
614        &self,
615        scope: &mut Scope,
616        axis: NativeExpand<usize>,
617    ) -> NativeExpand<usize> {
618        TensorValueExpand::__expand_shape_method(self.clone(), scope, axis)
619    }
620
621    fn __expand_stride_method(
622        &self,
623        scope: &mut Scope,
624        axis: NativeExpand<usize>,
625    ) -> NativeExpand<usize> {
626        TensorValueExpand::__expand_stride_method(self.clone(), scope, axis)
627    }
628
629    fn __expand_rank_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
630        TensorValueExpand::__expand_rank_method(self.clone(), scope)
631    }
632
633    fn __expand_len_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
634        TensorValueExpand::__expand_len_method(self.clone(), scope)
635    }
636
637    fn __expand_buffer_len_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
638        TensorValueExpand::__expand_buffer_len_method(self.clone(), scope)
639    }
640
641    fn __expand_as_tensor_map_method(
642        &self,
643        scope: &mut Scope,
644    ) -> ComptimeOptionExpand<TensorMap<V::T, Tiled>> {
645        TensorValueExpand::__expand_as_tensor_map_method(self.clone(), scope)
646    }
647}
648
649impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
650    Vectorized for TensorValue<Q, K, V, M, O, MA>
651{
652}
653impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
654    VectorizedExpand for TensorValueExpand<Q, K, V, M, O, MA>
655{
656    fn vector_size(&self) -> VectorSize {
657        let mut scope = Scope::root(false);
658        TensorValueExpand::__expand_vector_size_method(self.clone(), &mut scope)
659    }
660}
661
662impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
663    VirtualTensorOperationsExpand<M::T, M::N> for TensorMaskExpand<Q, K, V, M, O, MA>
664{
665    fn __expand_read_method(
666        &self,
667        scope: &mut Scope,
668        index: NativeExpand<usize>,
669    ) -> NativeExpand<Vector<M::T, M::N>> {
670        TensorMaskExpand::__expand_read_method(self.clone(), scope, index)
671    }
672    fn __expand_read_window_method(
673        &self,
674        context: &mut Scope,
675        start: NativeExpand<usize>,
676        end: NativeExpand<usize>,
677    ) -> SliceExpand<Vector<M::T, M::N>, ReadOnly> {
678        TensorMaskExpand::__expand_read_window_method(self.clone(), context, start, end)
679    }
680
681    fn __expand_write_method(
682        &self,
683        _scope: &mut Scope,
684        _index: NativeExpand<usize>,
685        _val: NativeExpand<Vector<M::T, M::N>>,
686    ) {
687        panic!("Can't write to input tensor");
688    }
689
690    fn __expand_shape_method(
691        &self,
692        scope: &mut Scope,
693        axis: NativeExpand<usize>,
694    ) -> NativeExpand<usize> {
695        TensorMaskExpand::__expand_shape_method(self.clone(), scope, axis)
696    }
697
698    fn __expand_stride_method(
699        &self,
700        scope: &mut Scope,
701        axis: NativeExpand<usize>,
702    ) -> NativeExpand<usize> {
703        TensorMaskExpand::__expand_stride_method(self.clone(), scope, axis)
704    }
705
706    fn __expand_rank_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
707        TensorMaskExpand::__expand_rank_method(self.clone(), scope)
708    }
709
710    fn __expand_len_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
711        TensorMaskExpand::__expand_len_method(self.clone(), scope)
712    }
713
714    fn __expand_buffer_len_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
715        TensorMaskExpand::__expand_buffer_len_method(self.clone(), scope)
716    }
717
718    fn __expand_as_tensor_map_method(
719        &self,
720        scope: &mut Scope,
721    ) -> ComptimeOptionExpand<TensorMap<M::T, Tiled>> {
722        TensorMaskExpand::__expand_as_tensor_map_method(self.clone(), scope)
723    }
724}
725
726impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
727    Vectorized for TensorMask<Q, K, V, M, O, MA>
728{
729}
730impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
731    VectorizedExpand for TensorMaskExpand<Q, K, V, M, O, MA>
732{
733    fn vector_size(&self) -> VectorSize {
734        let mut scope = Scope::root(false);
735        TensorMaskExpand::__expand_vector_size_method(self.clone(), &mut scope)
736    }
737}
738
739/// Tensor output representation.
740///
741/// You can use the tensor output as if it was a pointer to the actually tensor.
742///
743/// # Warning
744/// # Warning
745///
746/// There is no mutability guarantee.
747pub struct TensorOutput<
748    Q: FloatLine,
749    K: FloatLine,
750    V: FloatLine,
751    M: NumericLine,
752    O: FloatLine,
753    GA: AttentionArgs,
754> {
755    state: *mut GA::State<Q, K, V, M, O>,
756}
757
758/// Expand type for tensor input.
759pub struct TensorQueryExpand<
760    Q: FloatLine,
761    K: FloatLine,
762    V: FloatLine,
763    M: NumericLine,
764    O: FloatLine,
765    GA: AttentionArgs,
766> {
767    state: <GA::State<Q, K, V, M, O> as RudaType>::ExpandType,
768}
769
770pub struct TensorKeyExpand<
771    Q: FloatLine,
772    K: FloatLine,
773    V: FloatLine,
774    M: NumericLine,
775    O: FloatLine,
776    GA: AttentionArgs,
777> {
778    state: <GA::State<Q, K, V, M, O> as RudaType>::ExpandType,
779}
780
781pub struct TensorValueExpand<
782    Q: FloatLine,
783    K: FloatLine,
784    V: FloatLine,
785    M: NumericLine,
786    O: FloatLine,
787    GA: AttentionArgs,
788> {
789    state: <GA::State<Q, K, V, M, O> as RudaType>::ExpandType,
790}
791
792pub struct TensorMaskExpand<
793    Q: FloatLine,
794    K: FloatLine,
795    V: FloatLine,
796    M: NumericLine,
797    O: FloatLine,
798    GA: AttentionArgs,
799> {
800    state: <GA::State<Q, K, V, M, O> as RudaType>::ExpandType,
801}
802
803/// Expand type for [tensor output](TensorOutput).
804pub struct TensorOutputExpand<
805    Q: FloatLine,
806    K: FloatLine,
807    V: FloatLine,
808    M: NumericLine,
809    O: FloatLine,
810    GA: AttentionArgs,
811> {
812    state: <GA::State<Q, K, V, M, O> as RudaType>::ExpandType,
813}
814
815#[ruda]
816impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
817    TensorQuery<Q, K, V, M, O, MA>
818{
819    /// Create a tensor input from the state and the ident.
820    pub fn new(state: &MA::State<Q, K, V, M, O>) -> TensorQuery<Q, K, V, M, O, MA> {
821        TensorQuery::<Q, K, V, M, O, MA> { state }
822    }
823
824    //// Read the tensor at the given coordinate.
825    pub fn read_window(&self, start: usize, end: usize) -> Slice<Vector<Q::T, Q::N>> {
826        unsafe { MA::read_window_query(&(*self.state), start, end) }
827    }
828
829    /// Read the tensor at the given coordinate.
830    pub fn read(&self, coordinate: usize) -> Vector<Q::T, Q::N> {
831        unsafe { MA::read_query(&(*self.state), coordinate) }
832    }
833
834    /// Get the shape of the tensor at the given axis.
835    pub fn shape(&self, axis: usize) -> usize {
836        unsafe { MA::shape_query(&(*self.state), axis) }
837    }
838
839    /// Get the stride of the tensor at the given axis.
840    pub fn stride(&self, axis: usize) -> usize {
841        unsafe { MA::stride_query(&(*self.state), axis) }
842    }
843
844    /// Get the rank of the tensor.
845    pub fn rank(&self) -> usize {
846        unsafe { MA::rank_query(&(*self.state)) }
847    }
848
849    /// Get the length of the tensor.
850    #[allow(clippy::len_without_is_empty)]
851    pub fn len(&self) -> usize {
852        unsafe { MA::len_query(&(*self.state)) }
853    }
854
855    /// Get the buffer length of the tensor.
856    pub fn buffer_len(&self) -> usize {
857        unsafe { MA::buffer_len_query(&(*self.state)) }
858    }
859
860    /// Get the buffer length of the tensor.
861    pub fn as_tensor_map(&self) -> ComptimeOption<TensorMap<Q::T, Tiled>> {
862        unsafe { MA::as_tensor_map_query(&(*self.state)) }
863    }
864
865    /// Get the vector size of the tensor.
866    pub fn vector_size(&self) -> comptime_type!(VectorSize) {
867        unsafe { MA::vector_size_query(&(*self.state)) }
868    }
869}
870
871#[ruda]
872impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
873    TensorKey<Q, K, V, M, O, MA>
874{
875    /// Create a tensor input from the state and the ident.
876    pub fn new(state: &MA::State<Q, K, V, M, O>) -> TensorKey<Q, K, V, M, O, MA> {
877        TensorKey::<Q, K, V, M, O, MA> { state }
878    }
879
880    //// Read the tensor at the given coordinate.
881    pub fn read_window(&self, start: usize, end: usize) -> Slice<Vector<K::T, K::N>> {
882        unsafe { MA::read_window_key(&(*self.state), start, end) }
883    }
884
885    /// Read the tensor at the given coordinate.
886    pub fn read(&self, coordinate: usize) -> Vector<K::T, K::N> {
887        unsafe { MA::read_key(&(*self.state), coordinate) }
888    }
889
890    /// Get the shape of the tensor at the given axis.
891    pub fn shape(&self, axis: usize) -> usize {
892        unsafe { MA::shape_key(&(*self.state), axis) }
893    }
894
895    /// Get the stride of the tensor at the given axis.
896    pub fn stride(&self, axis: usize) -> usize {
897        unsafe { MA::stride_key(&(*self.state), axis) }
898    }
899
900    /// Get the rank of the tensor.
901    pub fn rank(&self) -> usize {
902        unsafe { MA::rank_key(&(*self.state)) }
903    }
904
905    /// Get the length of the tensor.
906    #[allow(clippy::len_without_is_empty)]
907    pub fn len(&self) -> usize {
908        unsafe { MA::len_key(&(*self.state)) }
909    }
910
911    /// Get the buffer length of the tensor.
912    pub fn buffer_len(&self) -> usize {
913        unsafe { MA::buffer_len_key(&(*self.state)) }
914    }
915
916    /// Get the buffer length of the tensor.
917    pub fn as_tensor_map(&self) -> ComptimeOption<TensorMap<K::T, Tiled>> {
918        unsafe { MA::as_tensor_map_key(&(*self.state)) }
919    }
920
921    /// Get the vector size of the tensor.
922    pub fn vector_size(&self) -> comptime_type!(VectorSize) {
923        unsafe { MA::vector_size_key(&(*self.state)) }
924    }
925}
926
927#[ruda]
928impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
929    TensorValue<Q, K, V, M, O, MA>
930{
931    /// Create a tensor input from the state and the ident.
932    pub fn new(state: &MA::State<Q, K, V, M, O>) -> TensorValue<Q, K, V, M, O, MA> {
933        TensorValue::<Q, K, V, M, O, MA> { state }
934    }
935
936    //// Read the tensor at the given coordinate.
937    pub fn read_window(&self, start: usize, end: usize) -> Slice<Vector<V::T, V::N>> {
938        unsafe { MA::read_window_value(&(*self.state), start, end) }
939    }
940
941    /// Read the tensor at the given coordinate.
942    pub fn read(&self, coordinate: usize) -> Vector<V::T, V::N> {
943        unsafe { MA::read_value(&(*self.state), coordinate) }
944    }
945
946    /// Get the shape of the tensor at the given axis.
947    pub fn shape(&self, axis: usize) -> usize {
948        unsafe { MA::shape_value(&(*self.state), axis) }
949    }
950
951    /// Get the stride of the tensor at the given axis.
952    pub fn stride(&self, axis: usize) -> usize {
953        unsafe { MA::stride_value(&(*self.state), axis) }
954    }
955
956    /// Get the rank of the tensor.
957    pub fn rank(&self) -> usize {
958        unsafe { MA::rank_value(&(*self.state)) }
959    }
960
961    /// Get the length of the tensor.
962    #[allow(clippy::len_without_is_empty)]
963    pub fn len(&self) -> usize {
964        unsafe { MA::len_value(&(*self.state)) }
965    }
966
967    /// Get the buffer length of the tensor.
968    pub fn buffer_len(&self) -> usize {
969        unsafe { MA::buffer_len_value(&(*self.state)) }
970    }
971
972    /// Get the buffer length of the tensor.
973    pub fn as_tensor_map(&self) -> ComptimeOption<TensorMap<V::T, Tiled>> {
974        unsafe { MA::as_tensor_map_value(&(*self.state)) }
975    }
976
977    /// Get the vector size of the tensor.
978    pub fn vector_size(&self) -> comptime_type!(VectorSize) {
979        unsafe { MA::vector_size_value(&(*self.state)) }
980    }
981}
982
983#[ruda]
984impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, MA: AttentionArgs>
985    TensorMask<Q, K, V, M, O, MA>
986{
987    /// Create a tensor input from the state and the ident.
988    pub fn new(state: &MA::State<Q, K, V, M, O>) -> TensorMask<Q, K, V, M, O, MA> {
989        TensorMask::<Q, K, V, M, O, MA> { state }
990    }
991
992    //// Read the tensor at the given coordinate.
993    pub fn read_window(&self, start: usize, end: usize) -> Slice<Vector<M::T, M::N>> {
994        unsafe { MA::read_window_mask(&(*self.state), start, end) }
995    }
996
997    /// Read the tensor at the given coordinate.
998    pub fn read(&self, coordinate: usize) -> Vector<M::T, M::N> {
999        unsafe { MA::read_mask(&(*self.state), coordinate) }
1000    }
1001
1002    /// Get the shape of the tensor at the given axis.
1003    pub fn shape(&self, axis: usize) -> usize {
1004        unsafe { MA::shape_mask(&(*self.state), axis) }
1005    }
1006
1007    /// Get the stride of the tensor at the given axis.
1008    pub fn stride(&self, axis: usize) -> usize {
1009        unsafe { MA::stride_mask(&(*self.state), axis) }
1010    }
1011
1012    /// Get the rank of the tensor.
1013    pub fn rank(&self) -> usize {
1014        unsafe { MA::rank_mask(&(*self.state)) }
1015    }
1016
1017    /// Get the length of the tensor.
1018    #[allow(clippy::len_without_is_empty)]
1019    pub fn len(&self) -> usize {
1020        unsafe { MA::len_mask(&(*self.state)) }
1021    }
1022
1023    /// Get the buffer length of the tensor.
1024    pub fn buffer_len(&self) -> usize {
1025        unsafe { MA::buffer_len_mask(&(*self.state)) }
1026    }
1027
1028    /// Get the buffer length of the tensor.
1029    pub fn as_tensor_map(&self) -> ComptimeOption<TensorMap<M::T, Tiled>> {
1030        unsafe { MA::as_tensor_map_mask(&(*self.state)) }
1031    }
1032
1033    /// Get the vector size of the tensor.
1034    pub fn vector_size(&self) -> comptime_type!(VectorSize) {
1035        unsafe { MA::vector_size_mask(&(*self.state)) }
1036    }
1037}
1038
1039#[ruda]
1040impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1041    TensorOutput<Q, K, V, M, O, GA>
1042{
1043    /// Create a [tensor output](TensorOutput) from the state.
1044    pub fn new(state: &mut GA::State<Q, K, V, M, O>) -> TensorOutput<Q, K, V, M, O, GA> {
1045        TensorOutput::<Q, K, V, M, O, GA> { state }
1046    }
1047
1048    /// Write the val to tensor at the given coordinate.
1049    pub fn write(&self, coordinate: usize, val: Vector<O::T, O::N>) {
1050        unsafe { GA::write_out(&mut (*self.state), coordinate, val) }
1051    }
1052
1053    /// Get the shape of the tensor at the given axis.
1054    pub fn shape(&self, axis: usize) -> usize {
1055        unsafe { GA::shape_out(&(*self.state), axis) }
1056    }
1057
1058    /// Get the stride of the tensor at the given axis.
1059    pub fn stride(&self, dim: usize) -> usize {
1060        unsafe { GA::stride_out(&(*self.state), dim) }
1061    }
1062
1063    /// Get the rank of the tensor.
1064    pub fn rank(&self) -> usize {
1065        unsafe { GA::rank_out(&(*self.state)) }
1066    }
1067
1068    /// Get the length of the tensor.
1069    #[allow(clippy::len_without_is_empty)]
1070    pub fn len(&self) -> usize {
1071        unsafe { GA::len_out(&(*self.state)) }
1072    }
1073
1074    /// Get the buffer length of the tensor.
1075    pub fn buffer_len(&self) -> usize {
1076        unsafe { GA::buffer_len_out(&(*self.state)) }
1077    }
1078
1079    /// Get the vector size of the tensor.
1080    pub fn vector_size(&self) -> comptime_type!(VectorSize) {
1081        unsafe { GA::vector_size_out(&(*self.state)) }
1082    }
1083}
1084
1085#[derive(Clone)]
1086/// Type implementing [AttentionArgs] where all inputs and the output are materialized tensors.
1087///
1088/// Other types might implement [AttentionArgs] for fused matrix multiplication kernels.
1089pub struct TensorArgs;
1090
1091#[derive(RudaLaunch, RudaType)]
1092/// Input representation for [TensorArgs] implementing [AttentionArgs].
1093pub struct TensorInputs<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine> {
1094    pub query: Tensor<Vector<Q::T, Q::N>>,
1095    pub key: Tensor<Vector<K::T, K::N>>,
1096    pub value: Tensor<Vector<V::T, V::N>>,
1097    pub mask: ComptimeOption<Tensor<Vector<M::T, M::N>>>,
1098}
1099
1100impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine> ConcreteInputsFactory
1101    for TensorInputs<Q, K, V, M>
1102{
1103    fn create<R: Runtime>(
1104        query: TensorBinding<R>,
1105        key: TensorBinding<R>,
1106        value: TensorBinding<R>,
1107        mask: Option<TensorBinding<R>>,
1108        _selection: &AttentionBlueprint,
1109        _problem: &AttentionProblem,
1110    ) -> Self::RuntimeArg<R> {
1111        TensorInputsLaunch::new(
1112            query.into_tensor_arg(),
1113            key.into_tensor_arg(),
1114            value.into_tensor_arg(),
1115            match mask {
1116                Some(mask) => ComptimeOptionArgs::Some(mask.into_tensor_arg()),
1117                None => ComptimeOptionArgs::None,
1118            },
1119        )
1120    }
1121}
1122
1123impl<EG: Numeric, EGS: Size> ConcreteOutputFactory for Tensor<Vector<EG, EGS>> {
1124    fn create<'a, R: Runtime>(
1125        out: TensorBinding<R>,
1126        _selection: &AttentionBlueprint,
1127        _problem: &AttentionProblem,
1128    ) -> Self::RuntimeArg<R> {
1129        out.into_tensor_arg()
1130    }
1131}
1132
1133#[derive(RudaType)]
1134pub struct AttentionState<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine> {
1135    pub query: *const Tensor<Vector<Q::T, Q::N>>,
1136    pub key: *const Tensor<Vector<K::T, K::N>>,
1137    pub value: *const Tensor<Vector<V::T, V::N>>,
1138    pub mask: ComptimeOption<*const Tensor<Vector<M::T, M::N>>>,
1139    pub output: *mut Tensor<Vector<O::T, O::N>>,
1140}
1141
1142#[ruda]
1143impl AttentionArgs for TensorArgs {
1144    type Input<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine> = TensorInputs<Q, K, V, M>;
1145    type Output<O: FloatLine> = Tensor<Vector<O::T, O::N>>;
1146    type State<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine> =
1147        AttentionState<Q, K, V, M, O>;
1148
1149    fn init_state<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1150        input: &Self::Input<Q, K, V, M>,
1151        output: &mut Self::Output<O>,
1152    ) -> Self::State<Q, K, V, M, O> {
1153        let mask = input.mask.as_ref().map(|mask| {
1154            let ptr: *const Tensor<Vector<M::T, M::N>> = mask;
1155            ptr
1156        });
1157
1158        AttentionState::<Q, K, V, M, O> {
1159            query: &input.query,
1160            key: &input.key,
1161            value: &input.value,
1162            mask,
1163            output,
1164        }
1165    }
1166
1167    fn has_mask<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1168        state: &Self::State<Q, K, V, M, O>,
1169    ) -> ComptimeOption<()> {
1170        state.mask.as_ref().map(|_| ())
1171    }
1172
1173    fn read_query<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1174        state: &Self::State<Q, K, V, M, O>,
1175        coordinate: usize,
1176    ) -> Vector<Q::T, Q::N> {
1177        unsafe { (*state.query)[coordinate] }
1178    }
1179
1180    fn read_key<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1181        state: &Self::State<Q, K, V, M, O>,
1182        coordinate: usize,
1183    ) -> Vector<K::T, K::N> {
1184        unsafe { (*state.key)[coordinate] }
1185    }
1186
1187    fn read_value<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1188        state: &Self::State<Q, K, V, M, O>,
1189        coordinate: usize,
1190    ) -> Vector<V::T, V::N> {
1191        unsafe { (*state.value)[coordinate] }
1192    }
1193
1194    fn read_mask<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1195        state: &Self::State<Q, K, V, M, O>,
1196        coordinate: usize,
1197    ) -> Vector<M::T, M::N> {
1198        unsafe { (*state.mask.unwrap())[coordinate] }
1199    }
1200
1201    fn read_window_query<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1202        state: &Self::State<Q, K, V, M, O>,
1203        start: usize,
1204        end: usize,
1205    ) -> Slice<Vector<Q::T, Q::N>> {
1206        unsafe { (*state.query).slice(start, end) }
1207    }
1208
1209    fn read_window_key<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1210        state: &Self::State<Q, K, V, M, O>,
1211        start: usize,
1212        end: usize,
1213    ) -> Slice<Vector<K::T, K::N>> {
1214        unsafe { (*state.key).slice(start, end) }
1215    }
1216
1217    fn read_window_value<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1218        state: &Self::State<Q, K, V, M, O>,
1219        start: usize,
1220        end: usize,
1221    ) -> Slice<Vector<V::T, V::N>> {
1222        unsafe { (*state.value).slice(start, end) }
1223    }
1224
1225    fn read_window_mask<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1226        state: &Self::State<Q, K, V, M, O>,
1227        start: usize,
1228        end: usize,
1229    ) -> Slice<Vector<M::T, M::N>> {
1230        unsafe { (*state.mask.unwrap()).slice(start, end) }
1231    }
1232
1233    fn as_tensor_map_query<
1234        Q: FloatLine,
1235        K: FloatLine,
1236        V: FloatLine,
1237        M: NumericLine,
1238        O: FloatLine,
1239    >(
1240        _state: &Self::State<Q, K, V, M, O>,
1241    ) -> ComptimeOption<TensorMap<Q::T, Tiled>> {
1242        ComptimeOption::new_None()
1243    }
1244
1245    fn as_tensor_map_key<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1246        _state: &Self::State<Q, K, V, M, O>,
1247    ) -> ComptimeOption<TensorMap<K::T, Tiled>> {
1248        ComptimeOption::new_None()
1249    }
1250
1251    fn as_tensor_map_value<
1252        Q: FloatLine,
1253        K: FloatLine,
1254        V: FloatLine,
1255        M: NumericLine,
1256        O: FloatLine,
1257    >(
1258        _state: &Self::State<Q, K, V, M, O>,
1259    ) -> ComptimeOption<TensorMap<V::T, Tiled>> {
1260        ComptimeOption::new_None()
1261    }
1262
1263    fn as_tensor_map_mask<
1264        Q: FloatLine,
1265        K: FloatLine,
1266        V: FloatLine,
1267        M: NumericLine,
1268        O: FloatLine,
1269    >(
1270        _state: &Self::State<Q, K, V, M, O>,
1271    ) -> ComptimeOption<TensorMap<M::T, Tiled>> {
1272        ComptimeOption::new_None()
1273    }
1274
1275    fn shape_query<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1276        state: &Self::State<Q, K, V, M, O>,
1277        dim: usize,
1278    ) -> usize {
1279        unsafe { (*state.query).shape(dim) }
1280    }
1281
1282    fn shape_key<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1283        state: &Self::State<Q, K, V, M, O>,
1284        dim: usize,
1285    ) -> usize {
1286        unsafe { (*state.key).shape(dim) }
1287    }
1288
1289    fn shape_value<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1290        state: &Self::State<Q, K, V, M, O>,
1291        dim: usize,
1292    ) -> usize {
1293        unsafe { (*state.value).shape(dim) }
1294    }
1295
1296    fn shape_mask<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1297        state: &Self::State<Q, K, V, M, O>,
1298        dim: usize,
1299    ) -> usize {
1300        unsafe { (*state.mask.unwrap()).shape(dim) }
1301    }
1302
1303    fn shape_out<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1304        state: &Self::State<Q, K, V, M, O>,
1305        dim: usize,
1306    ) -> usize {
1307        unsafe { (*state.output).shape(dim) }
1308    }
1309
1310    fn stride_query<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1311        state: &Self::State<Q, K, V, M, O>,
1312        dim: usize,
1313    ) -> usize {
1314        unsafe { (*state.query).stride(dim) }
1315    }
1316
1317    fn stride_key<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1318        state: &Self::State<Q, K, V, M, O>,
1319        dim: usize,
1320    ) -> usize {
1321        unsafe { (*state.key).stride(dim) }
1322    }
1323
1324    fn stride_value<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1325        state: &Self::State<Q, K, V, M, O>,
1326        dim: usize,
1327    ) -> usize {
1328        unsafe { (*state.value).stride(dim) }
1329    }
1330
1331    fn stride_mask<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1332        state: &Self::State<Q, K, V, M, O>,
1333        dim: usize,
1334    ) -> usize {
1335        unsafe { (*state.mask.unwrap()).stride(dim) }
1336    }
1337
1338    fn stride_out<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1339        state: &Self::State<Q, K, V, M, O>,
1340        dim: usize,
1341    ) -> usize {
1342        unsafe { (*state.output).stride(dim) }
1343    }
1344
1345    fn write_out<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1346        state: &mut Self::State<Q, K, V, M, O>,
1347        coordinate: usize,
1348        val: Vector<O::T, O::N>,
1349    ) {
1350        unsafe { (*state.output)[coordinate] = val }
1351    }
1352
1353    fn rank_query<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1354        state: &Self::State<Q, K, V, M, O>,
1355    ) -> usize {
1356        unsafe { (*state.query).rank() }
1357    }
1358
1359    fn rank_key<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1360        state: &Self::State<Q, K, V, M, O>,
1361    ) -> usize {
1362        unsafe { (*state.key).rank() }
1363    }
1364
1365    fn rank_value<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1366        state: &Self::State<Q, K, V, M, O>,
1367    ) -> usize {
1368        unsafe { (*state.value).rank() }
1369    }
1370
1371    fn rank_mask<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1372        state: &Self::State<Q, K, V, M, O>,
1373    ) -> usize {
1374        unsafe { (*state.mask.unwrap()).rank() }
1375    }
1376
1377    fn rank_out<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1378        state: &Self::State<Q, K, V, M, O>,
1379    ) -> usize {
1380        unsafe { (*state.output).rank() }
1381    }
1382
1383    fn len_query<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1384        state: &Self::State<Q, K, V, M, O>,
1385    ) -> usize {
1386        unsafe { (*state.query).len() }
1387    }
1388
1389    fn len_key<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1390        state: &Self::State<Q, K, V, M, O>,
1391    ) -> usize {
1392        unsafe { (*state.key).len() }
1393    }
1394
1395    fn len_value<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1396        state: &Self::State<Q, K, V, M, O>,
1397    ) -> usize {
1398        unsafe { (*state.value).len() }
1399    }
1400
1401    fn len_mask<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1402        state: &Self::State<Q, K, V, M, O>,
1403    ) -> usize {
1404        unsafe { (*state.mask.unwrap()).len() }
1405    }
1406
1407    fn len_out<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1408        state: &Self::State<Q, K, V, M, O>,
1409    ) -> usize {
1410        unsafe { (*state.output).len() }
1411    }
1412
1413    fn buffer_len_query<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1414        state: &Self::State<Q, K, V, M, O>,
1415    ) -> usize {
1416        unsafe { (*state.query).buffer_len() }
1417    }
1418
1419    fn buffer_len_key<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1420        state: &Self::State<Q, K, V, M, O>,
1421    ) -> usize {
1422        unsafe { (*state.key).buffer_len() }
1423    }
1424
1425    fn buffer_len_value<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1426        state: &Self::State<Q, K, V, M, O>,
1427    ) -> usize {
1428        unsafe { (*state.value).buffer_len() }
1429    }
1430
1431    fn buffer_len_mask<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1432        state: &Self::State<Q, K, V, M, O>,
1433    ) -> usize {
1434        unsafe { (*state.mask.unwrap()).buffer_len() }
1435    }
1436
1437    fn buffer_len_out<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1438        state: &Self::State<Q, K, V, M, O>,
1439    ) -> usize {
1440        unsafe { (*state.output).buffer_len() }
1441    }
1442
1443    fn vector_size_query<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1444        state: &Self::State<Q, K, V, M, O>,
1445    ) -> comptime_type!(usize) {
1446        unsafe { (*state.query).vector_size() }
1447    }
1448
1449    fn vector_size_key<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1450        state: &Self::State<Q, K, V, M, O>,
1451    ) -> comptime_type!(usize) {
1452        unsafe { (*state.key).vector_size() }
1453    }
1454
1455    fn vector_size_value<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1456        state: &Self::State<Q, K, V, M, O>,
1457    ) -> comptime_type!(usize) {
1458        unsafe { (*state.value).vector_size() }
1459    }
1460
1461    fn vector_size_mask<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1462        state: &Self::State<Q, K, V, M, O>,
1463    ) -> comptime_type!(usize) {
1464        unsafe { (*state.mask.unwrap()).vector_size() }
1465    }
1466
1467    fn vector_size_out<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine>(
1468        state: &Self::State<Q, K, V, M, O>,
1469    ) -> comptime_type!(usize) {
1470        unsafe { (*state.output).vector_size() }
1471    }
1472}
1473
1474mod __query {
1475    use super::*;
1476
1477    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1478        RudaType for TensorQuery<Q, K, V, M, O, GA>
1479    {
1480        type ExpandType = TensorQueryExpand<Q, K, V, M, O, GA>;
1481    }
1482
1483    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1484        Clone for TensorQueryExpand<Q, K, V, M, O, GA>
1485    {
1486        fn clone(&self) -> Self {
1487            Self {
1488                state: self.state.clone(),
1489            }
1490        }
1491    }
1492
1493    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1494        IntoMut for TensorQueryExpand<Q, K, V, M, O, GA>
1495    {
1496        fn into_mut(mut self, scope: &mut Scope) -> Self {
1497            self.state = self.state.into_mut(scope);
1498            self
1499        }
1500    }
1501    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1502        RudaDebug for TensorQueryExpand<Q, K, V, M, O, GA>
1503    {
1504        fn set_debug_name(&self, scope: &mut Scope, name: &'static str) {
1505            self.state.set_debug_name(scope, name);
1506        }
1507    }
1508    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1509        Clone for TensorQuery<Q, K, V, M, O, GA>
1510    {
1511        fn clone(&self) -> Self {
1512            *self
1513        }
1514    }
1515    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1516        Copy for TensorQuery<Q, K, V, M, O, GA>
1517    {
1518    }
1519}
1520
1521mod __key {
1522    use super::*;
1523
1524    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1525        RudaType for TensorKey<Q, K, V, M, O, GA>
1526    {
1527        type ExpandType = TensorKeyExpand<Q, K, V, M, O, GA>;
1528    }
1529
1530    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1531        Clone for TensorKeyExpand<Q, K, V, M, O, GA>
1532    {
1533        fn clone(&self) -> Self {
1534            Self {
1535                state: self.state.clone(),
1536            }
1537        }
1538    }
1539
1540    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1541        IntoMut for TensorKeyExpand<Q, K, V, M, O, GA>
1542    {
1543        fn into_mut(mut self, scope: &mut Scope) -> Self {
1544            self.state = self.state.into_mut(scope);
1545            self
1546        }
1547    }
1548    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1549        RudaDebug for TensorKeyExpand<Q, K, V, M, O, GA>
1550    {
1551        fn set_debug_name(&self, scope: &mut Scope, name: &'static str) {
1552            self.state.set_debug_name(scope, name);
1553        }
1554    }
1555    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1556        Clone for TensorKey<Q, K, V, M, O, GA>
1557    {
1558        fn clone(&self) -> Self {
1559            *self
1560        }
1561    }
1562    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1563        Copy for TensorKey<Q, K, V, M, O, GA>
1564    {
1565    }
1566}
1567
1568mod __value {
1569    use super::*;
1570
1571    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1572        RudaType for TensorValue<Q, K, V, M, O, GA>
1573    {
1574        type ExpandType = TensorValueExpand<Q, K, V, M, O, GA>;
1575    }
1576
1577    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1578        Clone for TensorValueExpand<Q, K, V, M, O, GA>
1579    {
1580        fn clone(&self) -> Self {
1581            Self {
1582                state: self.state.clone(),
1583            }
1584        }
1585    }
1586
1587    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1588        IntoMut for TensorValueExpand<Q, K, V, M, O, GA>
1589    {
1590        fn into_mut(mut self, scope: &mut Scope) -> Self {
1591            self.state = self.state.into_mut(scope);
1592            self
1593        }
1594    }
1595    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1596        RudaDebug for TensorValueExpand<Q, K, V, M, O, GA>
1597    {
1598        fn set_debug_name(&self, scope: &mut Scope, name: &'static str) {
1599            self.state.set_debug_name(scope, name);
1600        }
1601    }
1602    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1603        Clone for TensorValue<Q, K, V, M, O, GA>
1604    {
1605        fn clone(&self) -> Self {
1606            *self
1607        }
1608    }
1609    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1610        Copy for TensorValue<Q, K, V, M, O, GA>
1611    {
1612    }
1613}
1614
1615mod __mask {
1616    use super::*;
1617
1618    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1619        RudaType for TensorMask<Q, K, V, M, O, GA>
1620    {
1621        type ExpandType = TensorMaskExpand<Q, K, V, M, O, GA>;
1622    }
1623
1624    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1625        Clone for TensorMaskExpand<Q, K, V, M, O, GA>
1626    {
1627        fn clone(&self) -> Self {
1628            Self {
1629                state: self.state.clone(),
1630            }
1631        }
1632    }
1633
1634    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1635        IntoMut for TensorMaskExpand<Q, K, V, M, O, GA>
1636    {
1637        fn into_mut(mut self, scope: &mut Scope) -> Self {
1638            self.state = self.state.into_mut(scope);
1639            self
1640        }
1641    }
1642    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1643        RudaDebug for TensorMaskExpand<Q, K, V, M, O, GA>
1644    {
1645        fn set_debug_name(&self, scope: &mut Scope, name: &'static str) {
1646            self.state.set_debug_name(scope, name);
1647        }
1648    }
1649    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1650        Clone for TensorMask<Q, K, V, M, O, GA>
1651    {
1652        fn clone(&self) -> Self {
1653            *self
1654        }
1655    }
1656    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1657        Copy for TensorMask<Q, K, V, M, O, GA>
1658    {
1659    }
1660}
1661
1662mod __output {
1663    use super::*;
1664
1665    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1666        RudaType for TensorOutput<Q, K, V, M, O, GA>
1667    {
1668        type ExpandType = TensorOutputExpand<Q, K, V, M, O, GA>;
1669    }
1670
1671    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1672        Clone for TensorOutput<Q, K, V, M, O, GA>
1673    {
1674        fn clone(&self) -> Self {
1675            *self
1676        }
1677    }
1678
1679    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1680        Clone for TensorOutputExpand<Q, K, V, M, O, GA>
1681    {
1682        fn clone(&self) -> Self {
1683            Self {
1684                state: self.state.clone(),
1685            }
1686        }
1687    }
1688
1689    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1690        IntoMut for TensorOutputExpand<Q, K, V, M, O, GA>
1691    {
1692        fn into_mut(mut self, scope: &mut Scope) -> Self {
1693            self.state = self.state.into_mut(scope);
1694            self
1695        }
1696    }
1697
1698    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1699        RudaDebug for TensorOutputExpand<Q, K, V, M, O, GA>
1700    {
1701        fn set_debug_name(&self, scope: &mut Scope, name: &'static str) {
1702            self.state.set_debug_name(scope, name);
1703        }
1704    }
1705
1706    impl<Q: FloatLine, K: FloatLine, V: FloatLine, M: NumericLine, O: FloatLine, GA: AttentionArgs>
1707        Copy for TensorOutput<Q, K, V, M, O, GA>
1708    {
1709    }
1710}
rudnn/attention/kernel_ir/launch/args.rs

rudnn/attention/kernel_ir/launch/
args.rs
