metrale_model_layers/layers/ops/
prefill_attn_turbok.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the 64-row paged prefill attention kernels with a
4//! quantized K cache and a differently quantized V cache
5//! (`attn_prefill_paged_turbo{4k_turbo3v,4k_turbo8v,3k_turbo8v}` in
6//! `kernels/gb10/common/`). Turbo3/Turbo4 are 3/4-bit Lloyd-Max codes with an
7//! E4M3 scale per 16 values; Turbo8 is FP8 E4M3 with a BF16 scale per 16. The
8//! two caches have separate byte layouts, so each launcher takes a block stride
9//! and data-section size for K and for V. All launch with
10//! `causal_mask_enabled = 1`.
11//!
12//! Owner: model-layers ops.
13//! Invariants: none beyond the types.
14
15#![allow(unused_imports)]
16
17use anyhow::Result;
18use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
19use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
20
21/// 2026-09-25: 4-bit K with 3-bit V.
22#[allow(clippy::too_many_arguments)]
23pub fn prefill_attention_paged_turbo4k_turbo3v_64(
24    gpu: &dyn GpuBackend,
25    kernel: KernelHandle,
26    q: DevicePtr,
27    k_cache: DevicePtr,
28    v_cache: DevicePtr,
29    output: DevicePtr,
30    block_table: DevicePtr,
31    q_len: u32,
32    kv_len: u32,
33    q_offset: u32,
34    num_q_heads: u32,
35    num_kv_heads: u32,
36    head_dim: u32,
37    cache_block_size: u32,
38    sliding_window: u32,
39    inv_sqrt_d: f32,
40    k_block_stride_bytes: u64,
41    k_data_section_bytes: u64,
42    v_block_stride_bytes: u64,
43    v_data_section_bytes: u64,
44    stream: u64,
45) -> Result<()> {
46    let br = 64u32;
47    KernelLaunch::new(gpu, kernel)
48        .grid([num_q_heads, div_ceil(q_len, br), 1])
49        .block([256, 1, 1])
50        .arg_ptr(q)
51        .arg_ptr(k_cache)
52        .arg_ptr(v_cache)
53        .arg_ptr(output)
54        .arg_ptr(block_table)
55        .arg_u32(q_len)
56        .arg_u32(kv_len)
57        .arg_u32(q_offset)
58        .arg_u32(num_q_heads)
59        .arg_u32(num_kv_heads)
60        .arg_u32(head_dim)
61        .arg_u32(cache_block_size)
62        .arg_u32(sliding_window)
63        .arg_u32(1u32)
64        .arg_f32(inv_sqrt_d)
65        .arg_u64(k_block_stride_bytes)
66        .arg_u64(k_data_section_bytes)
67        .arg_u64(v_block_stride_bytes)
68        .arg_u64(v_data_section_bytes)
69        .launch(stream)
70}
71
72/// 2026-09-25: 4-bit K with FP8 V.
73#[allow(clippy::too_many_arguments)]
74pub fn prefill_attention_paged_turbo4k_turbo8v_64(
75    gpu: &dyn GpuBackend,
76    kernel: KernelHandle,
77    q: DevicePtr,
78    k_cache: DevicePtr,
79    v_cache: DevicePtr,
80    output: DevicePtr,
81    block_table: DevicePtr,
82    q_len: u32,
83    kv_len: u32,
84    q_offset: u32,
85    num_q_heads: u32,
86    num_kv_heads: u32,
87    head_dim: u32,
88    cache_block_size: u32,
89    sliding_window: u32,
90    inv_sqrt_d: f32,
91    k_block_stride_bytes: u64,
92    k_data_section_bytes: u64,
93    v_block_stride_bytes: u64,
94    v_data_section_bytes: u64,
95    stream: u64,
96) -> Result<()> {
97    let br = 64u32;
98    KernelLaunch::new(gpu, kernel)
99        .grid([num_q_heads, div_ceil(q_len, br), 1])
100        .block([256, 1, 1])
101        .arg_ptr(q)
102        .arg_ptr(k_cache)
103        .arg_ptr(v_cache)
104        .arg_ptr(output)
105        .arg_ptr(block_table)
106        .arg_u32(q_len)
107        .arg_u32(kv_len)
108        .arg_u32(q_offset)
109        .arg_u32(num_q_heads)
110        .arg_u32(num_kv_heads)
111        .arg_u32(head_dim)
112        .arg_u32(cache_block_size)
113        .arg_u32(sliding_window)
114        .arg_u32(1u32)
115        .arg_f32(inv_sqrt_d)
116        .arg_u64(k_block_stride_bytes)
117        .arg_u64(k_data_section_bytes)
118        .arg_u64(v_block_stride_bytes)
119        .arg_u64(v_data_section_bytes)
120        .launch(stream)
121}
122
123/// 2026-09-25: 3-bit K with FP8 V.
124#[allow(clippy::too_many_arguments)]
125pub fn prefill_attention_paged_turbo3k_turbo8v_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    k_block_stride_bytes: u64,
143    k_data_section_bytes: u64,
144    v_block_stride_bytes: u64,
145    v_data_section_bytes: u64,
146    stream: u64,
147) -> Result<()> {
148    let br = 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        .arg_u64(k_block_stride_bytes)
168        .arg_u64(k_data_section_bytes)
169        .arg_u64(v_block_stride_bytes)
170        .arg_u64(v_data_section_bytes)
171        .launch(stream)
172}
metrale_model_layers/layers/ops/prefill_attn_turbok.rs

metrale_model_layers/layers/ops/
prefill_attn_turbok.rs
