burn_cubecl/kernel/attention/
base.rs

1#[cfg(feature = "autotune")]
2use crate::kernel::attention::attention_autotune;
3use crate::{CubeBackend, ops::numeric::empty_device_dtype, tensor::CubeTensor};
4use burn_backend::cubecl::dtype_to_storage_type;
5use burn_backend::{
6    DType, Shape,
7    ops::{AttentionModuleOptions, attention::attention_fallback},
8};
9use cubek::attention::forward::{
10    definition::{
11        AccumulatorPrecision, AttentionGlobalTypes, AttentionOptions, AttentionSetupError,
12    },
13    launch,
14    routines::blackbox_accelerated::BlackboxAcceleratedStrategy,
15};
16
17#[derive(Debug)]
18/// Strategy used to select which attention implementation to run.
19pub enum AttentionStrategy {
20    /// Flash Attention using accelerated inner matmuls.
21    FlashBlackboxAccelerated(BlackboxAcceleratedStrategy),
22
23    /// Flash Attention using unit inner matmuls.
24    FlashUnit,
25
26    /// Fallback implementation using multiple separate kernels.
27    Fallback,
28
29    /// Automatically benchmark and select the best strategy at runtime.
30    #[cfg(feature = "autotune")]
31    Autotune,
32}
33
34impl Default for AttentionStrategy {
35    fn default() -> Self {
36        // if autotune is enabled, default to autotune
37        #[cfg(feature = "autotune")]
38        return AttentionStrategy::Autotune;
39
40        // if autotune is disabled, default to fallback to make sure it runs
41        #[cfg(not(feature = "autotune"))]
42        AttentionStrategy::Fallback
43    }
44}
45
46#[allow(clippy::too_many_arguments)]
47/// Launch an attention kernel with given strategy
48pub fn attention(
49    query: CubeTensor,
50    key: CubeTensor,
51    value: CubeTensor,
52    mask: Option<CubeTensor>,
53    attn_bias: Option<CubeTensor>,
54    options: AttentionModuleOptions,
55    strategy: AttentionStrategy,
56) -> Result<CubeTensor, AttentionSetupError> {
57    // Resolve the flash launch strategy; the non-flash arms answer directly.
58    let flash = match strategy {
59        AttentionStrategy::FlashBlackboxAccelerated(strategy) => {
60            launch::Strategy::BlackboxAccelerated(launch::BlueprintStrategy::Inferred(strategy))
61        }
62        AttentionStrategy::FlashUnit => {
63            launch::Strategy::Unit(launch::BlueprintStrategy::Inferred(()))
64        }
65        AttentionStrategy::Fallback => {
66            return Ok(attention_fallback::<CubeBackend>(
67                query, key, value, mask, attn_bias, options,
68            ));
69        }
70        #[cfg(feature = "autotune")]
71        AttentionStrategy::Autotune => {
72            return Ok(attention_autotune(
73                query, key, value, mask, attn_bias, options,
74            ));
75        }
76    };
77
78    // The flash routines carry hard launch constraints — the accelerated
79    // stage's `seq_q` must divide the problem's, and the tile head dim must
80    // divide the problem's — while attention autotune caches its selection per
81    // *anchored* key: a winner benchmarked on one raw `seq_q` (or `head_dim`)
82    // must still run on the bucket's other raw shapes. So a flash strategy that
83    // can't launch degrades to the fallback here, the same way the matmul
84    // dispatch degrades a constrained routine to the unit kernel; the fallback
85    // (separate kernels, materialized scores) has no such constraint and always
86    // runs. Only an `InvalidConfig` degrades — availability and other errors
87    // surface unchanged. `options` is `Copy`; the tensors are cheap handle
88    // clones, taken only so the originals survive for the fallback.
89    match flash_attention(
90        query.clone(),
91        key.clone(),
92        value.clone(),
93        mask.clone(),
94        attn_bias.clone(),
95        options,
96        flash,
97    ) {
98        Err(AttentionSetupError::InvalidConfig(_)) => Ok(attention_fallback::<CubeBackend>(
99            query, key, value, mask, attn_bias, options,
100        )),
101        other => other,
102    }
103}
104
105#[allow(clippy::too_many_arguments)]
106/// Launch a flash attention kernel
107pub fn flash_attention(
108    query: CubeTensor,
109    key: CubeTensor,
110    value: CubeTensor,
111    mask: Option<CubeTensor>,
112    _attn_bias: Option<CubeTensor>,
113    options: AttentionModuleOptions,
114    strategy: launch::Strategy,
115) -> Result<CubeTensor, AttentionSetupError> {
116    let client = query.client.clone();
117    let out = init_attention_output(&query, &value);
118
119    let dtypes = AttentionGlobalTypes {
120        query: dtype_to_storage_type(query.dtype),
121        key: dtype_to_storage_type(key.dtype),
122        value: dtype_to_storage_type(value.dtype),
123        mask: dtype_to_storage_type(mask.as_ref().map(|m| m.dtype).unwrap_or(DType::U8)),
124        out: dtype_to_storage_type(out.dtype),
125    };
126
127    launch::launch_ref(
128        strategy,
129        &client,
130        query.binding(),
131        key.binding(),
132        value.binding(),
133        mask.map(|mask| mask.binding()),
134        out.clone().binding(),
135        &dtypes,
136        AttentionOptions {
137            causal: options.is_causal,
138            accumulator_precision: AccumulatorPrecision::Strict(cubecl::ir::ElemType::Float(
139                cubecl::ir::FloatKind::F32,
140            )),
141        },
142    )?;
143
144    Ok(out)
145}
146
147pub(crate) fn init_attention_output(query: &CubeTensor, value: &CubeTensor) -> CubeTensor {
148    let num_batches = query.meta.shape[0];
149    let num_heads = query.meta.shape[1];
150    let seq_q = query.meta.shape[2];
151    let val_dim = value.meta.shape[3];
152    let out_shape = Shape::new([num_batches, num_heads, seq_q, val_dim]);
153
154    empty_device_dtype(
155        query.client.clone(),
156        query.device.clone(),
157        out_shape,
158        query.dtype,
159    )
160}
burn_cubecl/kernel/attention/base.rs

burn_cubecl/kernel/attention/
base.rs
