metrale_model_layers/layers/ops/
prefill_attn_main_b_turbov.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the paged prefill attention kernels with a BF16 K
4//! cache and a Lloyd-Max packed V cache (`attn_prefill_paged_bf16k_turbo{2,3,4}v`
5//! in `kernels/gb10/common/`, built from `prefill_paged_compute_asym.cuh`). V
6//! holds 2, 3 or 4 bits per value plus one E4M3 scale per 16 values, laid out
7//! as a data section then a scale section per block (`v_block_stride_bytes`,
8//! `v_data_section_bytes`). All three launch with `causal_mask_enabled = 1`.
9//!
10//! Owner: model-layers ops.
11//! Invariants: none beyond the types.
12
13use super::*;
14
15/// 2026-09-25: BF16 K (`[block_size, num_kv_heads, head_dim]` per block) with
16/// 3-bit V.
17#[allow(clippy::too_many_arguments)]
18pub fn prefill_attention_paged_bf16k_turbo3v_64(
19    gpu: &dyn GpuBackend,
20    kernel: KernelHandle,
21    q: DevicePtr,
22    k_cache: DevicePtr,
23    v_cache: DevicePtr,
24    output: DevicePtr,
25    block_table: DevicePtr,
26    q_len: u32,
27    kv_len: u32,
28    q_offset: u32,
29    num_q_heads: u32,
30    num_kv_heads: u32,
31    head_dim: u32,
32    cache_block_size: u32,
33    sliding_window: u32,
34    inv_sqrt_d: f32,
35    v_block_stride_bytes: u64,
36    v_data_section_bytes: u64,
37    stream: u64,
38) -> Result<()> {
39    // 2026-09-25: `prefill_paged_compute_asym.cuh` defines `BR64` as 64 on every
40    // target, so under `metrale_scale` this 32-row grid has about twice the
41    // blocks the kernel needs; a block whose `q_block * 64` reaches `q_len`
42    // returns at once.
43    let br = if cfg!(metrale_scale) { 32u32 } else { 64u32 };
44    KernelLaunch::new(gpu, kernel)
45        .grid([num_q_heads, div_ceil(q_len, br), 1])
46        .block([256, 1, 1])
47        .arg_ptr(q)
48        .arg_ptr(k_cache)
49        .arg_ptr(v_cache)
50        .arg_ptr(output)
51        .arg_ptr(block_table)
52        .arg_u32(q_len)
53        .arg_u32(kv_len)
54        .arg_u32(q_offset)
55        .arg_u32(num_q_heads)
56        .arg_u32(num_kv_heads)
57        .arg_u32(head_dim)
58        .arg_u32(cache_block_size)
59        .arg_u32(sliding_window)
60        .arg_u32(1u32)
61        .arg_f32(inv_sqrt_d)
62        .arg_u64(v_block_stride_bytes)
63        .arg_u64(v_data_section_bytes)
64        .launch(stream)
65}
66
67/// 2026-09-25: BF16 K with 4-bit V; the arguments of
68/// [`prefill_attention_paged_bf16k_turbo3v_64`].
69#[allow(clippy::too_many_arguments)]
70pub fn prefill_attention_paged_bf16k_turbo4v_64(
71    gpu: &dyn GpuBackend,
72    kernel: KernelHandle,
73    q: DevicePtr,
74    k_cache: DevicePtr,
75    v_cache: DevicePtr,
76    output: DevicePtr,
77    block_table: DevicePtr,
78    q_len: u32,
79    kv_len: u32,
80    q_offset: u32,
81    num_q_heads: u32,
82    num_kv_heads: u32,
83    head_dim: u32,
84    cache_block_size: u32,
85    sliding_window: u32,
86    inv_sqrt_d: f32,
87    v_block_stride_bytes: u64,
88    v_data_section_bytes: u64,
89    stream: u64,
90) -> Result<()> {
91    // 2026-09-25: `prefill_paged_compute_asym.cuh` defines `BR64` as 64 on every
92    // target, so under `metrale_scale` this 32-row grid has about twice the
93    // blocks the kernel needs; a block whose `q_block * 64` reaches `q_len`
94    // returns at once.
95    let br = if cfg!(metrale_scale) { 32u32 } else { 64u32 };
96    KernelLaunch::new(gpu, kernel)
97        .grid([num_q_heads, div_ceil(q_len, br), 1])
98        .block([256, 1, 1])
99        .arg_ptr(q)
100        .arg_ptr(k_cache)
101        .arg_ptr(v_cache)
102        .arg_ptr(output)
103        .arg_ptr(block_table)
104        .arg_u32(q_len)
105        .arg_u32(kv_len)
106        .arg_u32(q_offset)
107        .arg_u32(num_q_heads)
108        .arg_u32(num_kv_heads)
109        .arg_u32(head_dim)
110        .arg_u32(cache_block_size)
111        .arg_u32(sliding_window)
112        .arg_u32(1u32)
113        .arg_f32(inv_sqrt_d)
114        .arg_u64(v_block_stride_bytes)
115        .arg_u64(v_data_section_bytes)
116        .launch(stream)
117}
118
119/// 2026-09-25: BF16 K with 2-bit V; the arguments of
120/// [`prefill_attention_paged_bf16k_turbo3v_64`].
121#[allow(clippy::too_many_arguments)]
122pub fn prefill_attention_paged_bf16k_turbo2v_64(
123    gpu: &dyn GpuBackend,
124    kernel: KernelHandle,
125    q: DevicePtr,
126    k_cache: DevicePtr,
127    v_cache: DevicePtr,
128    output: DevicePtr,
129    block_table: DevicePtr,
130    q_len: u32,
131    kv_len: u32,
132    q_offset: u32,
133    num_q_heads: u32,
134    num_kv_heads: u32,
135    head_dim: u32,
136    cache_block_size: u32,
137    sliding_window: u32,
138    inv_sqrt_d: f32,
139    v_block_stride_bytes: u64,
140    v_data_section_bytes: u64,
141    stream: u64,
142) -> Result<()> {
143    // 2026-09-25: `prefill_paged_compute_asym.cuh` defines `BR64` as 64 on every
144    // target, so under `metrale_scale` this 32-row grid has about twice the
145    // blocks the kernel needs; a block whose `q_block * 64` reaches `q_len`
146    // returns at once.
147    let br = if cfg!(metrale_scale) { 32u32 } else { 64u32 };
148    KernelLaunch::new(gpu, kernel)
149        .grid([num_q_heads, div_ceil(q_len, br), 1])
150        .block([256, 1, 1])
151        .arg_ptr(q)
152        .arg_ptr(k_cache)
153        .arg_ptr(v_cache)
154        .arg_ptr(output)
155        .arg_ptr(block_table)
156        .arg_u32(q_len)
157        .arg_u32(kv_len)
158        .arg_u32(q_offset)
159        .arg_u32(num_q_heads)
160        .arg_u32(num_kv_heads)
161        .arg_u32(head_dim)
162        .arg_u32(cache_block_size)
163        .arg_u32(sliding_window)
164        .arg_u32(1u32)
165        .arg_f32(inv_sqrt_d)
166        .arg_u64(v_block_stride_bytes)
167        .arg_u64(v_data_section_bytes)
168        .launch(stream)
169}
metrale_model_layers/layers/ops/prefill_attn_main_b_turbov.rs

metrale_model_layers/layers/ops/
prefill_attn_main_b_turbov.rs
