metrale_model_layers/layers/ops/
prefill_attn_fa128.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-27: `AttnFa128Kernels`: the bit-identical 128-row twins of the BF16 prefill
4//! flash-attention kernels (`kernels/gb10/common/attn_prefill_fa128.cu`):
5//! `attn_prefill_fa128_paged` for `attn_prefill_paged_64` and `attn_prefill_fa128` for
6//! `attn_prefill_64`. Each row gets the original's arithmetic in the original's order; the
7//! twins read each K/V tile once per 128 query rows, keep P in registers and let every warp
8//! run both QK^T and P V. On GB10 they are 1.5-2.3x faster (microtest
9//! `attn_prefill_fa128_microtest`: 0 differing values against both originals).
10//!
11//! Owner: model-layers ops.
12//! Invariants: a launcher returns Ok(false), launching nothing, whenever the twin does not
13//! apply (handle 0, head_dim != 256, a sliding window, no causal mask, or a head count that
14//! is not a multiple of the KV heads); the caller then runs the original.
15
16use anyhow::Result;
17use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
18use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
19
20/// 2026-09-27: Query rows per CTA, and the 96 KiB of dynamic shared memory (Q 128 rows, one
21/// K and one V tile of 32 rows, 512 bytes per row).
22const ROWS: u32 = 128;
23const SMEM: u32 = (128 + 32 + 32) * 512;
24
25/// 2026-09-27: The two twins' handles; 0 when the target lacks the module or
26/// `METRALE_NO_ATTN_FA128` is present (the originals then run, for A/B measurement).
27#[derive(Clone, Copy, Debug)]
28pub struct AttnFa128Kernels {
29    paged: KernelHandle,
30    contiguous: KernelHandle,
31}
32
33impl AttnFa128Kernels {
34    #[track_caller]
35    pub fn resolve(gpu: &dyn GpuBackend) -> Self {
36        const MODULE: &str = "attn_prefill_fa128";
37        if std::env::var_os("METRALE_NO_ATTN_FA128").is_some() {
38            return Self {
39                paged: KernelHandle(0),
40                contiguous: KernelHandle(0),
41            };
42        }
43        Self {
44            paged: crate::layers::try_target_kernel(gpu, MODULE, "attn_prefill_fa128_paged"),
45            contiguous: crate::layers::try_target_kernel(gpu, MODULE, "attn_prefill_fa128"),
46        }
47    }
48
49    fn applies(
50        kernel: KernelHandle,
51        head_dim: u32,
52        num_q_heads: u32,
53        num_kv_heads: u32,
54        causal: bool,
55        sliding_window: u32,
56    ) -> bool {
57        kernel.0 != 0
58            && head_dim == 256
59            && causal
60            && sliding_window == 0
61            && num_kv_heads > 0
62            && num_q_heads.is_multiple_of(num_kv_heads)
63    }
64
65    /// 2026-09-27: In place of `ops::prefill_attention_paged_64` (causal, BF16 paged cache).
66    #[allow(clippy::too_many_arguments)]
67    pub fn paged(
68        &self,
69        gpu: &dyn GpuBackend,
70        q: DevicePtr,
71        k_cache: DevicePtr,
72        v_cache: DevicePtr,
73        output: DevicePtr,
74        block_table: DevicePtr,
75        q_len: u32,
76        kv_len: u32,
77        q_offset: u32,
78        num_q_heads: u32,
79        num_kv_heads: u32,
80        head_dim: u32,
81        cache_block_size: u32,
82        sliding_window: u32,
83        inv_sqrt_d: f32,
84        stream: u64,
85    ) -> Result<bool> {
86        if q_len == 0
87            || !Self::applies(
88                self.paged,
89                head_dim,
90                num_q_heads,
91                num_kv_heads,
92                true,
93                sliding_window,
94            )
95        {
96            return Ok(false);
97        }
98        KernelLaunch::new(gpu, self.paged)
99            .grid([num_q_heads, div_ceil(q_len, ROWS), 1])
100            .block([256, 1, 1])
101            .shared_mem(SMEM)
102            .arg_ptr(q)
103            .arg_ptr(k_cache)
104            .arg_ptr(v_cache)
105            .arg_ptr(output)
106            .arg_ptr(block_table)
107            .arg_u32(q_len)
108            .arg_u32(kv_len)
109            .arg_u32(q_offset)
110            .arg_u32(num_q_heads)
111            .arg_u32(num_kv_heads)
112            .arg_u32(cache_block_size)
113            .arg_f32(inv_sqrt_d)
114            .launch(stream)?;
115        Ok(true)
116    }
117
118    /// 2026-09-27: In place of `ops::prefill_attention_64` (contiguous BF16 Q/K/V).
119    #[allow(clippy::too_many_arguments)]
120    pub fn contiguous(
121        &self,
122        gpu: &dyn GpuBackend,
123        q: DevicePtr,
124        k: DevicePtr,
125        v: DevicePtr,
126        output: DevicePtr,
127        seq_len: u32,
128        batch: u32,
129        num_q_heads: u32,
130        num_kv_heads: u32,
131        head_dim: u32,
132        inv_sqrt_d: f32,
133        causal: bool,
134        sliding_window: u32,
135        stream: u64,
136    ) -> Result<bool> {
137        if seq_len == 0
138            || batch == 0
139            || !Self::applies(
140                self.contiguous,
141                head_dim,
142                num_q_heads,
143                num_kv_heads,
144                causal,
145                sliding_window,
146            )
147        {
148            return Ok(false);
149        }
150        KernelLaunch::new(gpu, self.contiguous)
151            .grid([num_q_heads, div_ceil(seq_len, ROWS), batch])
152            .block([256, 1, 1])
153            .shared_mem(SMEM)
154            .arg_ptr(q)
155            .arg_ptr(k)
156            .arg_ptr(v)
157            .arg_ptr(output)
158            .arg_u32(seq_len)
159            .arg_u32(num_q_heads)
160            .arg_u32(num_kv_heads)
161            .arg_f32(inv_sqrt_d)
162            .launch(stream)?;
163        Ok(true)
164    }
165}
metrale_model_layers/layers/ops/prefill_attn_fa128.rs

metrale_model_layers/layers/ops/
prefill_attn_fa128.rs
