metrale_model_layers/layers/ops/
prefill_attn_fp8k.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 an
4//! FP8 K cache and a Lloyd-Max packed V cache (`attn_prefill_paged_fp8k_turbo{2,3,4}v`
5//! in `kernels/gb10/common/`). K is dequantized with the per-tensor `k_scale`; V
6//! holds 2, 3 or 4 bits per value plus one E4M3 scale per 16 values, laid out as
7//! a data section then a scale section per block (`v_block_stride_bytes`,
8//! `v_data_section_bytes`).
9//!
10//! Owner: model-layers ops.
11//! Invariants: none beyond the types.
12
13#![allow(unused_imports)]
14
15use anyhow::Result;
16use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
17use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
18
19/// 2026-09-25: FP8 K with 3-bit V.
20#[allow(clippy::too_many_arguments)]
21pub fn prefill_attention_paged_fp8k_turbo3v_64(
22    gpu: &dyn GpuBackend,
23    kernel: KernelHandle,
24    q: DevicePtr,
25    k_cache: DevicePtr,
26    v_cache: DevicePtr,
27    output: DevicePtr,
28    block_table: DevicePtr,
29    q_len: u32,
30    kv_len: u32,
31    q_offset: u32,
32    num_q_heads: u32,
33    num_kv_heads: u32,
34    head_dim: u32,
35    cache_block_size: u32,
36    sliding_window: u32,
37    inv_sqrt_d: f32,
38    k_scale: f32,
39    v_block_stride_bytes: u64,
40    v_data_section_bytes: u64,
41    stream: u64,
42) -> Result<()> {
43    let br = 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_f32(k_scale)
63        .arg_u64(v_block_stride_bytes)
64        .arg_u64(v_data_section_bytes)
65        .launch(stream)
66}
67
68/// 2026-09-25: FP8 K with 4-bit V.
69#[allow(clippy::too_many_arguments)]
70pub fn prefill_attention_paged_fp8k_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    k_scale: f32,
88    v_block_stride_bytes: u64,
89    v_data_section_bytes: u64,
90    stream: u64,
91) -> Result<()> {
92    let br = 64u32;
93    KernelLaunch::new(gpu, kernel)
94        .grid([num_q_heads, div_ceil(q_len, br), 1])
95        .block([256, 1, 1])
96        .arg_ptr(q)
97        .arg_ptr(k_cache)
98        .arg_ptr(v_cache)
99        .arg_ptr(output)
100        .arg_ptr(block_table)
101        .arg_u32(q_len)
102        .arg_u32(kv_len)
103        .arg_u32(q_offset)
104        .arg_u32(num_q_heads)
105        .arg_u32(num_kv_heads)
106        .arg_u32(head_dim)
107        .arg_u32(cache_block_size)
108        .arg_u32(sliding_window)
109        .arg_u32(1u32)
110        .arg_f32(inv_sqrt_d)
111        .arg_f32(k_scale)
112        .arg_u64(v_block_stride_bytes)
113        .arg_u64(v_data_section_bytes)
114        .launch(stream)
115}
116
117/// 2026-09-25: FP8 K with 2-bit V.
118#[allow(clippy::too_many_arguments)]
119pub fn prefill_attention_paged_fp8k_turbo2v_64(
120    gpu: &dyn GpuBackend,
121    kernel: KernelHandle,
122    q: DevicePtr,
123    k_cache: DevicePtr,
124    v_cache: DevicePtr,
125    output: DevicePtr,
126    block_table: DevicePtr,
127    q_len: u32,
128    kv_len: u32,
129    q_offset: u32,
130    num_q_heads: u32,
131    num_kv_heads: u32,
132    head_dim: u32,
133    cache_block_size: u32,
134    sliding_window: u32,
135    inv_sqrt_d: f32,
136    k_scale: f32,
137    v_block_stride_bytes: u64,
138    v_data_section_bytes: u64,
139    stream: u64,
140) -> Result<()> {
141    let br = 64u32;
142    KernelLaunch::new(gpu, kernel)
143        .grid([num_q_heads, div_ceil(q_len, br), 1])
144        .block([256, 1, 1])
145        .arg_ptr(q)
146        .arg_ptr(k_cache)
147        .arg_ptr(v_cache)
148        .arg_ptr(output)
149        .arg_ptr(block_table)
150        .arg_u32(q_len)
151        .arg_u32(kv_len)
152        .arg_u32(q_offset)
153        .arg_u32(num_q_heads)
154        .arg_u32(num_kv_heads)
155        .arg_u32(head_dim)
156        .arg_u32(cache_block_size)
157        .arg_u32(sliding_window)
158        .arg_u32(1u32)
159        .arg_f32(inv_sqrt_d)
160        .arg_f32(k_scale)
161        .arg_u64(v_block_stride_bytes)
162        .arg_u64(v_data_section_bytes)
163        .launch(stream)
164}
metrale_model_layers/layers/ops/prefill_attn_fp8k.rs

metrale_model_layers/layers/ops/
prefill_attn_fp8k.rs
