rudnn/attention/kernel_ir/launch/
base.rs

1use ruda_kernel::dsl as kernel_dsl;
2use ruda_kernel::dsl::Runtime;
3use ruda_kernel::dsl::client::ComputeClient;
4use ruda_kernel::dsl::prelude::TensorBinding;
5
6use crate::attention::kernel_ir::routines::{
7    Routine, blackbox_accelerated::BlackboxAcceleratedRoutine, unit::UnitRoutine,
8};
9use crate::attention::kernel_ir::{
10    definition::AttentionSetupError,
11    definition::{AttentionDims, AttentionGlobalTypes, AttentionOptions, AttentionProblem},
12    launch::args::{TensorArgs, TensorInputsLaunch},
13    routines::DeviceSettings,
14};
15
16use crate::attention::kernel_ir::components::batch::BatchAttentionFamily;
17
18#[derive(Debug, Clone)]
19pub enum BlueprintStrategy<R: Routine> {
20    /// Use a predefined blueprint
21    Forced(R::Blueprint),
22    /// Allows to give limited settings information, and the rest is inferred from it
23    Inferred(R::Strategy),
24}
25
26#[derive(Debug, Clone)]
27pub enum Strategy {
28    BlackboxAccelerated(BlueprintStrategy<BlackboxAcceleratedRoutine>),
29    Unit(BlueprintStrategy<UnitRoutine>),
30}
31
32#[allow(clippy::result_large_err, clippy::too_many_arguments)]
33pub fn launch_ref<R: Runtime>(
34    strategy: Strategy,
35    client: &ComputeClient<R>,
36    query: TensorBinding<R>,
37    key: TensorBinding<R>,
38    value: TensorBinding<R>,
39    mask: Option<TensorBinding<R>>,
40    out: TensorBinding<R>,
41    attention_global_types: &AttentionGlobalTypes,
42    attention_options: AttentionOptions,
43) -> Result<(), AttentionSetupError> {
44    match strategy {
45        Strategy::BlackboxAccelerated(strategy) => {
46            launch_attention::<R, BlackboxAcceleratedRoutine>(
47                client,
48                query,
49                key,
50                value,
51                mask,
52                out,
53                attention_global_types,
54                strategy,
55                attention_options,
56            )
57        }
58        Strategy::Unit(strategy) => launch_attention::<R, UnitRoutine>(
59            client,
60            query,
61            key,
62            value,
63            mask,
64            out,
65            attention_global_types,
66            strategy,
67            attention_options,
68        ),
69    }
70}
71
72#[allow(clippy::too_many_arguments)]
73pub fn launch_attention<R: Runtime, A: Routine>(
74    client: &ComputeClient<R>,
75    query: TensorBinding<R>,
76    key: TensorBinding<R>,
77    value: TensorBinding<R>,
78    mask: Option<TensorBinding<R>>,
79    out: TensorBinding<R>,
80    global_dtypes: &AttentionGlobalTypes,
81    strategy: BlueprintStrategy<A>,
82    attention_options: AttentionOptions,
83) -> Result<(), AttentionSetupError> {
84    let definition = AttentionProblem {
85        dims: AttentionDims {
86            batch: query.shape[0],
87            num_heads: query.shape[1],
88            seq_q: query.shape[2],
89            head_dim: query.shape[3],
90            seq_kv: key.shape[2],
91            val_dim: value.shape[3],
92        },
93        masked: mask.is_some(),
94        global_dtypes: global_dtypes.clone(),
95        options: attention_options,
96        address_type: [
97            query.required_address_type(global_dtypes.query.size()),
98            key.required_address_type(global_dtypes.key.size()),
99            value.required_address_type(global_dtypes.value.size()),
100            mask.as_ref()
101                .map(|mask| mask.required_address_type(global_dtypes.mask.size()))
102                .unwrap_or_default(),
103            out.required_address_type(global_dtypes.out.size()),
104        ]
105        .into_iter()
106        .max()
107        .unwrap_or_default(),
108    };
109
110    let device_settings = DeviceSettings::new(client, &definition);
111
112    let launch_info = A::prepare(&definition, &device_settings, strategy)?;
113
114    // This allows an expand_config error to be caught by the client rather than the server.
115    // Then the server can re-run expand config assuming a valid blueprint
116    <A as Routine>::BatchAttention::expand_config(
117        client.properties(),
118        launch_info.blueprint.clone(),
119        &launch_info.dtypes,
120    )?;
121
122    let result = unsafe {
123        <A as Routine>::BatchAttention::launch_unchecked::<TensorArgs, R>(
124            client,
125            launch_info.ruda_dim,
126            launch_info.ruda_count_plan.resolve(),
127            launch_info.address_type,
128            TensorInputsLaunch::new(
129                query.into_tensor_arg(),
130                key.into_tensor_arg(),
131                value.into_tensor_arg(),
132                mask.map(|it| it.into_tensor_arg()).into(),
133            ),
134            out.into_tensor_arg(),
135            ruda_kernel::tiling::ruda_count::ruda_mapping_launch(&launch_info.ruda_count_plan),
136            &launch_info.dtypes,
137            &device_settings.vector_sizes,
138            launch_info.blueprint,
139        )
140    };
141
142    match result {
143        Ok(_) => Ok(()),
144        Err(err) => Err(AttentionSetupError::Execution(err)),
145    }
146}
rudnn/attention/kernel_ir/launch/base.rs

rudnn/attention/kernel_ir/launch/
base.rs
