rlx_ir/ops/
attention.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Attention builders: SDPA with custom or kernel-synthesized
6//! masks (plan #53).
7
8use crate::op::MaskKind;
9use crate::{Graph, NodeId, Op, Shape};
10
11/// Build an [`Op::Attention`] with optional score scale and logit softcap.
12/// `v_head_dim = None` ⇒ V/output width equals `head_dim` (the common case).
13pub fn attention_kind_op(
14    num_heads: usize,
15    head_dim: usize,
16    v_head_dim: Option<usize>,
17    mask_kind: MaskKind,
18    score_scale: Option<f32>,
19    attn_logit_softcap: Option<f32>,
20) -> Op {
21    Op::Attention {
22        num_heads,
23        head_dim,
24        v_head_dim,
25        mask_kind,
26        score_scale,
27        attn_logit_softcap,
28    }
29}
30
31impl Graph {
32    /// Scaled dot-product attention with a custom (caller-supplied) mask.
33    /// Equivalent to `attention_kind(.., MaskKind::Custom, ..)`.
34    pub fn attention(
35        &mut self,
36        q: NodeId,
37        k: NodeId,
38        v: NodeId,
39        mask: NodeId,
40        num_heads: usize,
41        head_dim: usize,
42        shape: Shape,
43    ) -> NodeId {
44        self.attention_opts(q, k, v, mask, num_heads, head_dim, shape, None, None)
45    }
46
47    /// Like [`Self::attention`] with optional score scale and logit softcap.
48    pub fn attention_opts(
49        &mut self,
50        q: NodeId,
51        k: NodeId,
52        v: NodeId,
53        mask: NodeId,
54        num_heads: usize,
55        head_dim: usize,
56        shape: Shape,
57        score_scale: Option<f32>,
58        attn_logit_softcap: Option<f32>,
59    ) -> NodeId {
60        self.push(
61            attention_kind_op(
62                num_heads,
63                head_dim,
64                None,
65                MaskKind::Custom,
66                score_scale,
67                attn_logit_softcap,
68            ),
69            vec![q, k, v, mask],
70            shape,
71            None,
72        )
73    }
74
75    /// Scaled dot-product attention with a kernel-synthesized mask
76    /// (`None` / `Causal` / `SlidingWindow`). Inputs are Q, K, V only —
77    /// no mask tensor is allocated or read in the inner loop. Use
78    /// `MaskKind::None` for a single un-padded sequence.
79    pub fn attention_kind(
80        &mut self,
81        q: NodeId,
82        k: NodeId,
83        v: NodeId,
84        num_heads: usize,
85        head_dim: usize,
86        mask_kind: MaskKind,
87        shape: Shape,
88    ) -> NodeId {
89        self.attention_kind_opts(q, k, v, num_heads, head_dim, mask_kind, shape, None, None)
90    }
91
92    /// Like [`Self::attention_kind`] with optional score scale and logit softcap.
93    pub fn attention_kind_opts(
94        &mut self,
95        q: NodeId,
96        k: NodeId,
97        v: NodeId,
98        num_heads: usize,
99        head_dim: usize,
100        mask_kind: MaskKind,
101        shape: Shape,
102        score_scale: Option<f32>,
103        attn_logit_softcap: Option<f32>,
104    ) -> NodeId {
105        debug_assert!(
106            !matches!(mask_kind, MaskKind::Custom | MaskKind::Bias),
107            "attention_kind() requires a non-tensor MaskKind; use attention() for Custom or attention_bias() for Bias"
108        );
109        self.push(
110            attention_kind_op(
111                num_heads,
112                head_dim,
113                None,
114                mask_kind,
115                score_scale,
116                attn_logit_softcap,
117            ),
118            vec![q, k, v],
119            shape,
120            None,
121        )
122    }
123
124    /// Scaled dot-product attention with an additive bias tensor of shape
125    /// `[batch, num_heads, query_len, key_len]` added to the
126    /// `QK^T · scale` scores before softmax. Lets boxRPB / per-query
127    /// position biases reuse the fast `Op::Attention` kernel path.
128    pub fn attention_bias(
129        &mut self,
130        q: NodeId,
131        k: NodeId,
132        v: NodeId,
133        bias: NodeId,
134        num_heads: usize,
135        head_dim: usize,
136        shape: Shape,
137    ) -> NodeId {
138        self.attention_bias_opts(q, k, v, bias, num_heads, head_dim, shape, None, None)
139    }
140
141    /// Like [`Self::attention_bias`] with optional score scale and logit softcap
142    /// (e.g. Gemma's `score_scale = 1.0` + final-attn-logit softcap, which the
143    /// bias-less variant would silently drop).
144    #[allow(clippy::too_many_arguments)]
145    pub fn attention_bias_opts(
146        &mut self,
147        q: NodeId,
148        k: NodeId,
149        v: NodeId,
150        bias: NodeId,
151        num_heads: usize,
152        head_dim: usize,
153        shape: Shape,
154        score_scale: Option<f32>,
155        attn_logit_softcap: Option<f32>,
156    ) -> NodeId {
157        self.push(
158            attention_kind_op(
159                num_heads,
160                head_dim,
161                None,
162                MaskKind::Bias,
163                score_scale,
164                attn_logit_softcap,
165            ),
166            vec![q, k, v, bias],
167            shape,
168            None,
169        )
170    }
171
172    /// Asymmetric SDPA where V/output per-head width (`v_head_dim`) differs from
173    /// the Q/K score width (`head_dim`). Q/K/V only (kernel-synthesized mask).
174    /// Used by MLA to skip zero-padding V up to `head_dim`. `shape` is the
175    /// output `[.., num_heads * v_head_dim]`.
176    #[allow(clippy::too_many_arguments)]
177    pub fn attention_kind_vdim(
178        &mut self,
179        q: NodeId,
180        k: NodeId,
181        v: NodeId,
182        num_heads: usize,
183        head_dim: usize,
184        v_head_dim: usize,
185        mask_kind: MaskKind,
186        shape: Shape,
187    ) -> NodeId {
188        debug_assert!(
189            !matches!(mask_kind, MaskKind::Custom | MaskKind::Bias),
190            "attention_kind_vdim() requires a non-tensor MaskKind"
191        );
192        self.push(
193            attention_kind_op(num_heads, head_dim, Some(v_head_dim), mask_kind, None, None),
194            vec![q, k, v],
195            shape,
196            None,
197        )
198    }
199}
rlx_ir/ops/attention.rs

rlx_ir/ops/
attention.rs
