rudnn/attention/kernel_ir/components/global/simple/reader/
query.rs

1use ruda_kernel::dsl as kernel_dsl;
2use ruda_kernel::dsl::prelude::*;
3use ruda_kernel::library::Swizzle;
4use ruda_kernel::library::tensor::View;
5use ruda_kernel::library::tensor::layout::Coords2d;
6use rublas::kernel_ir::components::global::memory::GlobalMemoryConfig;
7use ruda_kernel::tiling::tile::StridedTile;
8
9use crate::attention::kernel_ir::{
10    components::stage::AttentionPartitioner,
11    definition::attention_types::{QG, QGS},
12    definition::{AttentionPrecision, AttentionTileSize},
13};
14
15#[derive(RudaType)]
16pub struct QueryReader<AP: AttentionPrecision> {
17    query: View<Vector<QG<AP>, QGS<AP>>, Coords2d>,
18    #[ruda(comptime)]
19    gmem_config: GlobalMemoryConfig,
20}
21
22#[ruda]
23impl<AP: AttentionPrecision> QueryReader<AP> {
24    pub fn new(
25        stage_q_offset: u32,
26        query: View<Vector<QG<AP>, QGS<AP>>, Coords2d>,
27        #[comptime] gmem_config: GlobalMemoryConfig,
28    ) -> Self {
29        let query = query.slice((stage_q_offset, 0), query.shape());
30
31        QueryReader::<AP> { query, gmem_config }
32    }
33
34    pub fn get_tile<P: AttentionPartitioner>(
35        &self,
36        tile: Coords2d,
37        #[comptime] tile_size: AttentionTileSize,
38        #[comptime] partition_seq_q: u32,
39        #[comptime] partition_head_dim: u32,
40    ) -> StridedTile<QG<AP>, QGS<AP>> {
41        let (row_in_partition, col) = tile;
42
43        let row = row_in_partition + P::seq_q_index() * partition_seq_q;
44
45        let vector_size = self.gmem_config.vector_size.comptime() as u32;
46
47        let slice = self
48            .query
49            .slice(
50                (row * tile_size.seq_q, col * tile_size.head_dim),
51                (tile_size.seq_q, tile_size.head_dim).runtime(),
52            )
53            .to_linear_slice();
54
55        let start = 0;
56        let vectors_per_tile = tile_size.seq_q * tile_size.head_dim / vector_size;
57        let end = start + vectors_per_tile;
58        let vectors_per_partition_row = partition_head_dim * tile_size.head_dim / vector_size;
59
60        StridedTile::<QG<AP>, QGS<AP>>::new_strided(
61            slice,
62            start,
63            end,
64            vectors_per_partition_row,
65            Swizzle::none(),
66            self.gmem_config.matrix_layout,
67        )
68    }
69}
rudnn/attention/kernel_ir/components/global/simple/reader/query.rs

rudnn/attention/kernel_ir/components/global/simple/reader/
query.rs
