Merge branch 'main' into rhypot

Author: vaijira

Cargo.toml

@@ -11,7 +11,7 @@ edition = "2024"
 license = "MIT OR Apache-2.0"
 readme = "README.md"
 rust-version = "1.88"
-version = "0.9.0-pre.1"
+version = "0.9.0-pre.2"
 
 [workspace.dependencies]
 bitflags = { version = "2.9.1", features = ["serde"] }

crates/cubecl-attention/Cargo.toml

@@ -19,12 +19,12 @@ attention_tests = []
 
 [dependencies]
 bytemuck = { workspace = true }
-cubecl-common = { path = "../cubecl-common", version = "0.9.0-pre.1", default-features = false }
-cubecl-core = { path = "../cubecl-core", version = "0.9.0-pre.1", default-features = false }
-cubecl-runtime = { path = "../cubecl-runtime", version = "0.9.0-pre.1", default-features = false }
-cubecl-std = { path = "../cubecl-std", version = "0.9.0-pre.1", default-features = false }
-cubecl-matmul = { path = "../cubecl-matmul", version = "0.9.0-pre.1", default-features = false }
-cubecl-random = { path = "../cubecl-random", version = "0.9.0-pre.1", default-features = false }
+cubecl-common = { path = "../cubecl-common", version = "0.9.0-pre.2", default-features = false }
+cubecl-core = { path = "../cubecl-core", version = "0.9.0-pre.2", default-features = false }
+cubecl-runtime = { path = "../cubecl-runtime", version = "0.9.0-pre.2", default-features = false }
+cubecl-std = { path = "../cubecl-std", version = "0.9.0-pre.2", default-features = false }
+cubecl-matmul = { path = "../cubecl-matmul", version = "0.9.0-pre.2", default-features = false }
+cubecl-random = { path = "../cubecl-random", version = "0.9.0-pre.2", default-features = false }
 half = { workspace = true, features = ["bytemuck"] }
 pretty_assertions = { workspace = true, optional = true }
 serde = { workspace = true }

crates/cubecl-attention/src/components/batch/hypercube/base.rs

@@ -26,8 +26,11 @@ impl HypercubeConfig {
         selection: &AttentionSelection,
     ) -> CubeCountPlan {
         CubeCountPlan {
-            inner: (problem.seq_q as u32)
-                .div_ceil(selection.tiling_scheme.elements_in_stage_seq_q()),
+            inner: (problem.seq_q as u32).div_ceil(
+                selection.tiling_scheme.tile_size.seq_q
+                    * selection.tiling_scheme.partition_size.seq_q
+                    * selection.tiling_scheme.stage_size.seq_q,
+            ),
             outer: (problem.batch * problem.num_heads) as u32,
         }
     }

crates/cubecl-attention/src/components/batch/simple/attention.rs

@@ -10,6 +10,7 @@ use crate::components::{
         BatchAttention, BatchAttentionConfig, CubeCountInput, simple::config::SimpleBatchConfig,
     },
     global::{GlobalAttention, GlobalAttentionConfig as _},
+    stage::StageAttentionConfig as _,
 };
 
 pub struct SimpleBatchAttention<AP: AttentionPrecision, GA: GlobalAttention<AP>> {
@@ -35,7 +36,7 @@ impl<GA: GlobalAttention<AP>, AP: AttentionPrecision> BatchAttention<AP>
         let q_index = CUBE_POS_X;
         let batch_index = CUBE_POS_Y;
 
-        let stage_q_offset = q_index * global_config.tiling_scheme().elements_in_stage_seq_q();
+        let stage_q_offset = q_index * global_config.stage_config().elements_in_stage_seq_q();
 
         // Assume [batch, num_heads, seq_*, head_dim] layout
         let seq_q = query.shape(2);

crates/cubecl-attention/src/components/batch/simple/config.rs

@@ -10,7 +10,6 @@ use crate::components::{
 pub struct SimpleBatchConfig<G: GlobalAttentionConfig> {
     global_config: G,
     hypercube_config: HypercubeConfig,
-    seq_kv: u32,
 }
 
 impl<G: GlobalAttentionConfig> BatchAttentionConfig for SimpleBatchConfig<G> {
@@ -30,11 +29,10 @@ impl<G: GlobalAttentionConfig> BatchAttentionConfig for SimpleBatchConfig<G> {
 }
 
 impl<G: GlobalAttentionConfig> SimpleBatchConfig<G> {
-    pub fn new(global_config: G, hypercube_config: HypercubeConfig, seq_kv: u32) -> Self {
+    pub fn new(global_config: G, hypercube_config: HypercubeConfig) -> Self {
         Self {
             global_config,
             hypercube_config,
-            seq_kv,
         }
     }
 

crates/cubecl-attention/src/components/batch/simple/setup.rs

@@ -36,7 +36,6 @@ impl<GA: GlobalAttentionFamily> BatchAttentionFamily for SimpleBatchAttentionFam
             selection
                 .hypercube_selection
                 .to_hypercube_config(problem, client.properties().hardware.max_cube_count.clone()),
-            problem.seq_kv as u32,
         )
         .validate(problem)
     }

crates/cubecl-attention/src/components/global/base.rs

@@ -2,13 +2,12 @@ use cubecl_core as cubecl;
 use cubecl_core::prelude::*;
 
 use crate::components::{AttentionElems, global::simple::AttentionWriter};
-use cubecl_matmul::components::{global::memory::GlobalMemoryConfig, stage::StageMemoryConfig};
 use cubecl_std::{CubeOption, tensor::r#virtual::VirtualTensor};
 
 use crate::components::{
-    AttentionIdent, AttentionLineSizes, AttentionPrecision, AttentionProblem, AttentionSelection,
-    AttentionSetupError, AttentionTilingScheme, AvailableLineSizes, attention_types::*,
-    global::simple::QueryReader, stage::StageAttentionConfig,
+    AttentionLineSizes, AttentionPrecision, AttentionProblem, AttentionSelection,
+    AttentionSetupError, AvailableLineSizes, attention_types::*, global::simple::QueryReader,
+    stage::StageAttentionConfig,
 };
 use std::{fmt::Debug, hash::Hash};
 
@@ -107,14 +106,5 @@ pub trait GlobalAttentionConfig:
     type StageConfig: StageAttentionConfig;
 
     fn stage_config(&self) -> Self::StageConfig;
-    fn key_stage_memory_config(&self) -> StageMemoryConfig;
-    fn value_stage_memory_config(&self) -> StageMemoryConfig;
-
     fn cube_dim(&self) -> CubeDim;
-    fn plane_dim(&self) -> u32;
-    fn global_memory_config(&self, ident: AttentionIdent) -> GlobalMemoryConfig;
-
-    fn tiling_scheme(&self) -> AttentionTilingScheme;
-
-    fn causal_mask(&self) -> bool;
 }

crates/cubecl-attention/src/components/global/layout.rs

@@ -44,7 +44,7 @@ impl Layout for AttentionGlobalLayout {
     type SourceCoordinates = Coords1d;
 
     fn to_source_pos(&self, coords: Self::Coordinates) -> u32 {
-        let line_size = comptime![self.config.line_size()];
+        let line_size = comptime![self.config.line_size];
         let (row, col) = coords;
         let idx = self.batch_offset + row * self.stride_row + col * self.stride_col;
 
@@ -62,10 +62,7 @@ impl Layout for AttentionGlobalLayout {
     fn is_in_bounds(&self, pos: Self::Coordinates) -> bool {
         let (row, col) = pos;
 
-        match comptime!((
-            self.config.check_row_bounds(),
-            self.config.check_col_bounds()
-        )) {
+        match comptime!((self.config.check_row_bounds, self.config.check_col_bounds)) {
             (true, true) => row < self.rows && col < self.columns,
             (true, false) => row < self.rows,
             (false, true) => col < self.columns,

crates/cubecl-attention/src/components/global/simple/attention.rs

@@ -7,14 +7,15 @@ use cubecl_std::{CubeOption, CubeOptionExpand};
 use std::marker::PhantomData;
 
 use crate::components::attention_types::*;
-use crate::components::global::base::GlobalAttentionConfig;
+use crate::components::global::simple::QueryReader;
 use crate::components::global::simple::{AttentionWriter, AttentionWriterExpand, MaskReader};
 use crate::components::global::{AttentionGlobalLayout, simple::DummyKeyValueReader};
-use crate::components::stage::{AttentionPartitioner, AttentionTilingLayout, StageAttention};
-use crate::components::{AttentionIdent, global::simple::QueryReader};
+use crate::components::stage::{
+    AttentionPartitioner, AttentionTilingLayout, StageAttention, StageAttentionConfig as _,
+};
 use crate::components::{
     AttentionPrecision,
-    global::{GlobalAttention, simple::config::SimpleGlobalConfig},
+    global::{GlobalAttention, simple::config::SimpleGlobalAttentionConfig},
 };
 
 pub struct SimpleGlobalAttention<AP: AttentionPrecision, SA: StageAttention<AP>> {
@@ -32,13 +33,13 @@ impl<
     AP: AttentionPrecision,
 > GlobalAttention<AP> for SimpleGlobalAttention<AP, SA>
 {
-    type KeyReader = DummyKeyValueReader<KG<AP>, KS<AP>, Self::Config>;
-    type ValueReader = DummyKeyValueReader<VG<AP>, VS<AP>, Self::Config>;
+    type KeyReader = DummyKeyValueReader<KG<AP>, KS<AP>>;
+    type ValueReader = DummyKeyValueReader<VG<AP>, VS<AP>>;
     type MaskReader = MaskReader<AP>;
 
     type Writer = <SA::Partitioner as AttentionPartitioner>::Writer<OS<AP>, OG<AP>>;
 
-    type Config = SimpleGlobalConfig<SA::Config>;
+    type Config = SimpleGlobalAttentionConfig<SA::Config>;
 
     fn execute(
         query_reader: QueryReader<AP>,
@@ -51,32 +52,32 @@ impl<
         #[comptime] config: Self::Config,
     ) {
         // Init staging shared memories
-        let mut key_stage = key_reader.init_stage(config.key_stage_memory_config());
-        let mut value_stage = value_reader.init_stage(config.value_stage_memory_config());
+        let mut key_stage = key_reader.init_stage();
+        let mut value_stage = value_reader.init_stage();
 
         // Load queries which stay alive in registers for all the kernel
-        let mut query_registers = SA::init_query(config.stage_config());
-        SA::read_query(&query_reader, &mut query_registers, config.stage_config());
+        let mut query_registers = SA::init_query(config.stage_config);
+        SA::read_query(&query_reader, &mut query_registers, config.stage_config);
 
         // Init registers that will change inside global loop
-        let mut key_value_registers = SA::init_key_value(config.stage_config());
+        let mut key_value_registers = SA::init_key_value(config.stage_config);
         let mut mask_registers =
-            SA::init_mask(CubeOption::new_Some((seq_q, seq_kv)), config.stage_config());
-        let mut softmax_registers = SA::init_softmax(config.stage_config());
-        let mut accumulator_registers = SA::init_accumulator(config.stage_config());
+            SA::init_mask(CubeOption::new_Some((seq_q, seq_kv)), config.stage_config);
+        let mut softmax_registers = SA::init_softmax(config.stage_config);
+        let mut accumulator_registers = SA::init_accumulator(config.stage_config);
 
         // Init running state
-        let mut stage_state = SA::init_state(config.stage_config());
+        let mut stage_state = SA::init_state(config.stage_config);
 
         // Define number of global iterations
         let num_stage_iterations =
-            seq_kv.div_ceil(config.tiling_scheme().elements_in_partition_seq_kv());
+            seq_kv.div_ceil(config.stage_config.elements_in_partition_seq_kv());
 
         // Global loop over seq_kv
         for _ in 0..num_stage_iterations {
             // Put key and value into stage
-            key_reader.read_global(&mut key_stage, config);
-            value_reader.read_global(&mut value_stage, config);
+            key_reader.read_global(&mut key_stage);
+            value_reader.read_global(&mut value_stage);
 
             sync_cube();
 
@@ -91,7 +92,7 @@ impl<
                 &mut softmax_registers,
                 &mut accumulator_registers,
                 &mut stage_state,
-                config.stage_config(),
+                config.stage_config,
             );
 
             sync_cube();
@@ -103,19 +104,15 @@ impl<
         }
 
         // Accumulators must be rescaled using running state
-        SA::rescale(
-            &mut accumulator_registers,
-            stage_state,
-            config.stage_config(),
-        );
+        SA::rescale(&mut accumulator_registers, stage_state, config.stage_config);
 
         // Write accumulators to output
         let mut out_stage = writer.stage();
         SA::write::<Self::Writer, Self::Config>(
             &accumulator_registers,
             &mut out_stage,
             &mut writer,
-            config.stage_config(),
+            config.stage_config,
         )
     }
 
@@ -125,11 +122,7 @@ impl<
         query: VirtualTensor<QG<AP>>,
         #[comptime] config: Self::Config,
     ) -> QueryReader<AP> {
-        let layout = AttentionGlobalLayout::new(
-            &query,
-            batch_index,
-            config.global_memory_config(AttentionIdent::Query),
-        );
+        let layout = AttentionGlobalLayout::new(&query, batch_index, config.query_gmem_config);
 
         QueryReader::<AP>::new(stage_q_offset, query.view(layout))
     }
@@ -139,27 +132,21 @@ impl<
         key: VirtualTensor<KG<AP>>,
         #[comptime] config: Self::Config,
     ) -> Self::KeyReader {
-        let step = reduction_step::<Self::Config>(config);
-        let layout = AttentionGlobalLayout::new(
-            &key,
-            batch_index,
-            config.global_memory_config(AttentionIdent::Key),
-        );
-        DummyKeyValueReader::new(key.view(layout), step, AttentionIdent::Key)
+        let step = config.stage_config.elements_in_partition_seq_kv().runtime();
+        let layout =
+            AttentionGlobalLayout::new(&key, batch_index, config.key_reader_config.gmem_config);
+        DummyKeyValueReader::new(key.view(layout), step, config.key_reader_config)
     }
 
     fn init_value_reader(
         batch_index: u32,
         value: VirtualTensor<VG<AP>>,
         #[comptime] config: Self::Config,
     ) -> Self::ValueReader {
-        let step = reduction_step::<Self::Config>(config);
-        let layout = AttentionGlobalLayout::new(
-            &value,
-            batch_index,
-            config.global_memory_config(AttentionIdent::Value),
-        );
-        DummyKeyValueReader::new(value.view(layout), step, AttentionIdent::Value)
+        let step = config.stage_config.elements_in_partition_seq_kv().runtime();
+        let layout =
+            AttentionGlobalLayout::new(&value, batch_index, config.value_reader_config.gmem_config);
+        DummyKeyValueReader::new(value.view(layout), step, config.value_reader_config)
     }
 
     fn init_mask_reader(
@@ -169,24 +156,22 @@ impl<
         seq_kv_shape: u32,
         #[comptime] config: Self::Config,
     ) -> Self::MaskReader {
-        let step = reduction_step::<Self::Config>(config);
+        let step = config.stage_config.elements_in_partition_seq_kv().runtime();
         let partition_q_offset = <SA::Partitioner as AttentionPartitioner>::seq_q_index()
-            * config.tiling_scheme().elements_in_partition_seq_q();
+            * config.stage_config.elements_in_partition_seq_q();
 
         match mask {
             CubeOption::Some(mask) => {
-                let layout = AttentionGlobalLayout::new(
-                    &mask,
-                    batch_index,
-                    config.global_memory_config(AttentionIdent::Mask),
-                );
+                let layout =
+                    AttentionGlobalLayout::new(&mask, batch_index, config.mask_gmem_config);
 
                 MaskReader::new_materialized(
                     stage_q_offset,
                     partition_q_offset,
                     mask.view(layout),
                     step,
                     seq_kv_shape,
+                    config.mask_gmem_config.view_direction,
                 )
             }
             CubeOption::None => MaskReader::new_logical(stage_q_offset + partition_q_offset, step),
@@ -199,22 +184,13 @@ impl<
         out: VirtualTensor<OG<AP>, ReadWrite>,
         #[comptime] config: Self::Config,
     ) -> Self::Writer {
-        let conf = config.global_memory_config(AttentionIdent::Out);
-        let layout = AttentionGlobalLayout::new(&out, batch_index, conf);
+        let layout =
+            AttentionGlobalLayout::new(&out, batch_index, config.writer_config.gmem_config);
         let out = out.view_mut(layout);
 
-        Self::Writer::new::<SA::Config>(
+        Self::Writer::init::<SA::Config>(
             out.slice_mut_unchecked((stage_q_offset, 0), out.shape()),
-            conf,
-            config.stage_config(),
+            config.writer_config,
         )
     }
 }
-
-#[cube]
-fn reduction_step<C: GlobalAttentionConfig>(#[comptime] config: C) -> u32 {
-    config
-        .tiling_scheme()
-        .elements_in_partition_seq_kv()
-        .runtime()
-}