metrale_model_layers/layers/ops/
prefill_attn_main_b.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the paged prefill flash-attention kernels over
4//! NVFP4, HDIM=512 BF16 and TurboQuant KV caches, and the 64-row BF16, FP8 and
5//! NVFP4 variants. The BF16-K / packed-V launchers live in
6//! `prefill_attn_main_b_turbov.rs`.
7//!
8//! Owner: model-layers ops.
9//! Invariants: none beyond the types.
10
11#![allow(unused_imports)]
12
13use anyhow::Result;
14use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
15use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
16
17use crate::layers::moe;
18use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
19
20use super::*;
21#[path = "prefill_attn_main_b_turbov.rs"]
22mod bf16k_turbov;
23pub use bf16k_turbov::{
24    prefill_attention_paged_bf16k_turbo2v_64, prefill_attention_paged_bf16k_turbo3v_64,
25    prefill_attention_paged_bf16k_turbo4v_64,
26};
27
28/// 2026-09-25: Causal paged prefill over an NVFP4 KV cache, 32 query rows per
29/// block.
30#[allow(clippy::too_many_arguments)]
31pub fn prefill_attention_paged_nvfp4(
32    gpu: &dyn GpuBackend,
33    kernel: KernelHandle,
34    q: DevicePtr,
35    k_cache: DevicePtr,
36    v_cache: DevicePtr,
37    output: DevicePtr,
38    block_table: DevicePtr,
39    q_len: u32,
40    kv_len: u32,
41    q_offset: u32,
42    num_q_heads: u32,
43    num_kv_heads: u32,
44    head_dim: u32,
45    cache_block_size: u32,
46    sliding_window: u32,
47    inv_sqrt_d: f32,
48    block_stride_bytes: u64,
49    data_section_bytes: u64,
50    stream: u64,
51) -> Result<()> {
52    let br = 32u32;
53    KernelLaunch::new(gpu, kernel)
54        .grid([num_q_heads, div_ceil(q_len, br), 1])
55        .block([128, 1, 1])
56        .arg_ptr(q)
57        .arg_ptr(k_cache)
58        .arg_ptr(v_cache)
59        .arg_ptr(output)
60        .arg_ptr(block_table)
61        .arg_u32(q_len)
62        .arg_u32(kv_len)
63        .arg_u32(q_offset)
64        .arg_u32(num_q_heads)
65        .arg_u32(num_kv_heads)
66        .arg_u32(head_dim)
67        .arg_u32(cache_block_size)
68        .arg_u32(sliding_window)
69        .arg_u32(1u32)
70        .arg_f32(inv_sqrt_d)
71        .arg_u64(block_stride_bytes)
72        .arg_u64(data_section_bytes)
73        .launch(stream)
74}
75
76/// 2026-09-25: Causal paged prefill for `head_dim == 512` over a BF16 KV cache:
77/// 32 query rows and 8 warps per block, with the 101,120 bytes of dynamic shared
78/// memory `attn_prefill_paged_512.cu` requires.
79#[allow(clippy::too_many_arguments)]
80pub fn prefill_attention_paged_512(
81    gpu: &dyn GpuBackend,
82    kernel: KernelHandle,
83    q: DevicePtr,
84    k_cache: DevicePtr,
85    v_cache: DevicePtr,
86    output: DevicePtr,
87    block_table: DevicePtr,
88    q_len: u32,
89    kv_len: u32,
90    q_offset: u32,
91    num_q_heads: u32,
92    num_kv_heads: u32,
93    head_dim: u32,
94    cache_block_size: u32,
95    sliding_window: u32,
96    inv_sqrt_d: f32,
97    stream: u64,
98) -> Result<()> {
99    let br = 32u32;
100    KernelLaunch::new(gpu, kernel)
101        .grid([num_q_heads, div_ceil(q_len, br), 1])
102        .block([256, 1, 1])
103        .shared_mem(101_120)
104        .arg_ptr(q)
105        .arg_ptr(k_cache)
106        .arg_ptr(v_cache)
107        .arg_ptr(output)
108        .arg_ptr(block_table)
109        .arg_u32(q_len)
110        .arg_u32(kv_len)
111        .arg_u32(q_offset)
112        .arg_u32(num_q_heads)
113        .arg_u32(num_kv_heads)
114        .arg_u32(head_dim)
115        .arg_u32(cache_block_size)
116        .arg_u32(sliding_window)
117        .arg_u32(1u32)
118        .arg_f32(inv_sqrt_d)
119        .launch(stream)
120}
121
122/// 2026-09-25: [`prefill_attention_paged`] with 64 query rows per block (32 on
123/// the SCALE and HIP targets).
124#[allow(clippy::too_many_arguments)]
125pub fn prefill_attention_paged_64(
126    gpu: &dyn GpuBackend,
127    kernel: KernelHandle,
128    q: DevicePtr,
129    k_cache: DevicePtr,
130    v_cache: DevicePtr,
131    output: DevicePtr,
132    block_table: DevicePtr,
133    q_len: u32,
134    kv_len: u32,
135    q_offset: u32,
136    num_q_heads: u32,
137    num_kv_heads: u32,
138    head_dim: u32,
139    cache_block_size: u32,
140    sliding_window: u32,
141    inv_sqrt_d: f32,
142    stream: u64,
143) -> Result<()> {
144    // 2026-09-25: The grid's rows per block must equal the kernel's `BR64`, or rows
145    // 32..63 of every 64-row band are never written. `BR64` is 32 in the gb10
146    // `prefill_paged_compute.cuh` under `__SCALE__` and always in the strix-hip
147    // one; `metrale_scale` is set for exactly `strix` and `strix-hip`.
148    let br = if cfg!(metrale_scale) { 32u32 } else { 64u32 };
149    KernelLaunch::new(gpu, kernel)
150        .grid([num_q_heads, div_ceil(q_len, br), 1])
151        .block([256, 1, 1])
152        .arg_ptr(q)
153        .arg_ptr(k_cache)
154        .arg_ptr(v_cache)
155        .arg_ptr(output)
156        .arg_ptr(block_table)
157        .arg_u32(q_len)
158        .arg_u32(kv_len)
159        .arg_u32(q_offset)
160        .arg_u32(num_q_heads)
161        .arg_u32(num_kv_heads)
162        .arg_u32(head_dim)
163        .arg_u32(cache_block_size)
164        .arg_u32(sliding_window)
165        .arg_u32(1u32)
166        .arg_f32(inv_sqrt_d)
167        .launch(stream)
168}
169
170/// 2026-09-25: [`prefill_attention_paged_fp8`] with 64 query rows per block (32
171/// on the SCALE and HIP targets).
172#[allow(clippy::too_many_arguments)]
173pub fn prefill_attention_paged_fp8_64(
174    gpu: &dyn GpuBackend,
175    kernel: KernelHandle,
176    q: DevicePtr,
177    k_cache: DevicePtr,
178    v_cache: DevicePtr,
179    output: DevicePtr,
180    block_table: DevicePtr,
181    q_len: u32,
182    kv_len: u32,
183    q_offset: u32,
184    num_q_heads: u32,
185    num_kv_heads: u32,
186    head_dim: u32,
187    cache_block_size: u32,
188    sliding_window: u32,
189    inv_sqrt_d: f32,
190    k_scale: f32,
191    v_scale: f32,
192    cache_stride: u64,
193    stream: u64,
194) -> Result<()> {
195    // 2026-09-25: The grid's rows per block must equal the kernel's `BR64`, or rows
196    // 32..63 of every 64-row band are never written. `BR64` is 32 in the gb10
197    // `prefill_paged_compute.cuh` under `__SCALE__` and always in the strix-hip
198    // one; `metrale_scale` is set for exactly `strix` and `strix-hip`.
199    let br = if cfg!(metrale_scale) { 32u32 } else { 64u32 };
200    KernelLaunch::new(gpu, kernel)
201        .grid([num_q_heads, div_ceil(q_len, br), 1])
202        .block([256, 1, 1])
203        .arg_ptr(q)
204        .arg_ptr(k_cache)
205        .arg_ptr(v_cache)
206        .arg_ptr(output)
207        .arg_ptr(block_table)
208        .arg_u32(q_len)
209        .arg_u32(kv_len)
210        .arg_u32(q_offset)
211        .arg_u32(num_q_heads)
212        .arg_u32(num_kv_heads)
213        .arg_u32(head_dim)
214        .arg_u32(cache_block_size)
215        .arg_u32(sliding_window)
216        .arg_u32(1u32)
217        .arg_f32(inv_sqrt_d)
218        .arg_f32(k_scale)
219        .arg_f32(v_scale)
220        .arg_u64(cache_stride)
221        .launch(stream)
222}
223
224/// 2026-09-25: Causal paged prefill over a TurboQuant KV cache (same format for K
225/// and V), 64 query rows per block (32 on the SCALE and HIP targets). The
226/// turbo8, turbo4 and turbo3 `_64` entries share this argument list; the caller
227/// picks one through `kernel` and passes that cache's block stride and
228/// data-section size.
229#[allow(clippy::too_many_arguments)]
230pub fn prefill_attention_paged_turbo_64(
231    gpu: &dyn GpuBackend,
232    kernel: KernelHandle,
233    q: DevicePtr,
234    k_cache: DevicePtr,
235    v_cache: DevicePtr,
236    output: DevicePtr,
237    block_table: DevicePtr,
238    q_len: u32,
239    kv_len: u32,
240    q_offset: u32,
241    num_q_heads: u32,
242    num_kv_heads: u32,
243    head_dim: u32,
244    cache_block_size: u32,
245    sliding_window: u32,
246    inv_sqrt_d: f32,
247    block_stride_bytes: u64,
248    data_section_bytes: u64,
249    stream: u64,
250) -> Result<()> {
251    // 2026-09-25: The grid's rows per block must equal the kernel's `BR64`, or rows
252    // 32..63 of every 64-row band are never written. `BR64` is 32 in the gb10
253    // `prefill_paged_compute.cuh` under `__SCALE__` and always in the strix-hip
254    // one; `metrale_scale` is set for exactly `strix` and `strix-hip`.
255    let br = if cfg!(metrale_scale) { 32u32 } else { 64u32 };
256    KernelLaunch::new(gpu, kernel)
257        .grid([num_q_heads, div_ceil(q_len, br), 1])
258        .block([256, 1, 1])
259        .arg_ptr(q)
260        .arg_ptr(k_cache)
261        .arg_ptr(v_cache)
262        .arg_ptr(output)
263        .arg_ptr(block_table)
264        .arg_u32(q_len)
265        .arg_u32(kv_len)
266        .arg_u32(q_offset)
267        .arg_u32(num_q_heads)
268        .arg_u32(num_kv_heads)
269        .arg_u32(head_dim)
270        .arg_u32(cache_block_size)
271        .arg_u32(sliding_window)
272        .arg_u32(1u32)
273        .arg_f32(inv_sqrt_d)
274        .arg_u64(block_stride_bytes)
275        .arg_u64(data_section_bytes)
276        .launch(stream)
277}
278
279/// 2026-09-25: The turbo2 counterpart of [`prefill_attention_paged_turbo_64`].
280pub fn prefill_attention_paged_turbo2_64(
281    gpu: &dyn GpuBackend,
282    kernel: KernelHandle,
283    q: DevicePtr,
284    k_cache: DevicePtr,
285    v_cache: DevicePtr,
286    output: DevicePtr,
287    block_table: DevicePtr,
288    q_len: u32,
289    kv_len: u32,
290    q_offset: u32,
291    num_q_heads: u32,
292    num_kv_heads: u32,
293    head_dim: u32,
294    cache_block_size: u32,
295    sliding_window: u32,
296    inv_sqrt_d: f32,
297    block_stride_bytes: u64,
298    data_section_bytes: u64,
299    stream: u64,
300) -> Result<()> {
301    // 2026-09-26: The handle is the 32-row entry `attn_prefill_paged_turbo2`
302    // (`qwen3_attention/init_prefill_kernels.rs`), so the grid uses 32 rows and 128
303    // threads like the other 32-row paged launchers.
304    let br = 32u32;
305    KernelLaunch::new(gpu, kernel)
306        .grid([num_q_heads, div_ceil(q_len, br), 1])
307        .block([128, 1, 1])
308        .arg_ptr(q)
309        .arg_ptr(k_cache)
310        .arg_ptr(v_cache)
311        .arg_ptr(output)
312        .arg_ptr(block_table)
313        .arg_u32(q_len)
314        .arg_u32(kv_len)
315        .arg_u32(q_offset)
316        .arg_u32(num_q_heads)
317        .arg_u32(num_kv_heads)
318        .arg_u32(head_dim)
319        .arg_u32(cache_block_size)
320        .arg_u32(sliding_window)
321        .arg_u32(1u32)
322        .arg_f32(inv_sqrt_d)
323        .arg_u64(block_stride_bytes)
324        .arg_u64(data_section_bytes)
325        .launch(stream)
326}
327
328/// 2026-09-25: [`prefill_attention_paged_nvfp4`] with 64 query rows per block
329/// (32 on the SCALE and HIP targets).
330pub fn prefill_attention_paged_nvfp4_64(
331    gpu: &dyn GpuBackend,
332    kernel: KernelHandle,
333    q: DevicePtr,
334    k_cache: DevicePtr,
335    v_cache: DevicePtr,
336    output: DevicePtr,
337    block_table: DevicePtr,
338    q_len: u32,
339    kv_len: u32,
340    q_offset: u32,
341    num_q_heads: u32,
342    num_kv_heads: u32,
343    head_dim: u32,
344    cache_block_size: u32,
345    sliding_window: u32,
346    inv_sqrt_d: f32,
347    block_stride_bytes: u64,
348    data_section_bytes: u64,
349    stream: u64,
350) -> Result<()> {
351    // 2026-09-25: The grid's rows per block must equal the kernel's `BR64`, or rows
352    // 32..63 of every 64-row band are never written. `BR64` is 32 in the gb10
353    // `prefill_paged_compute.cuh` under `__SCALE__` and always in the strix-hip
354    // one; `metrale_scale` is set for exactly `strix` and `strix-hip`.
355    let br = if cfg!(metrale_scale) { 32u32 } else { 64u32 };
356    KernelLaunch::new(gpu, kernel)
357        .grid([num_q_heads, div_ceil(q_len, br), 1])
358        .block([256, 1, 1])
359        .arg_ptr(q)
360        .arg_ptr(k_cache)
361        .arg_ptr(v_cache)
362        .arg_ptr(output)
363        .arg_ptr(block_table)
364        .arg_u32(q_len)
365        .arg_u32(kv_len)
366        .arg_u32(q_offset)
367        .arg_u32(num_q_heads)
368        .arg_u32(num_kv_heads)
369        .arg_u32(head_dim)
370        .arg_u32(cache_block_size)
371        .arg_u32(sliding_window)
372        .arg_u32(1u32)
373        .arg_f32(inv_sqrt_d)
374        .arg_u64(block_stride_bytes)
375        .arg_u64(data_section_bytes)
376        .launch(stream)
377}
metrale_model_layers/layers/ops/prefill_attn_main_b.rs

metrale_model_layers/layers/ops/
prefill_attn_main_b.rs
