metrale_model_layers/layers/ops/
prefill_attn_b.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the NVFP4 KV cache write, the BF16 absmax used by
4//! FP8 KV calibration, and the paged decode attention kernels for NVFP4, FP8
5//! and BF16 caches, with their split-K partial and reduce passes.
6//!
7//! Owner: model-layers ops.
8//! Invariants: none beyond the types.
9
10#![allow(unused_imports)]
11
12use anyhow::Result;
13use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
14use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
15
16use crate::layers::moe;
17use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
18
19use super::*;
20
21/// 2026-09-25: Write each token's K and V into the paged NVFP4 cache at its
22/// `slot_mapping` slot: E2M1 values with one E4M3 scale per 16. One block per
23/// token.
24pub fn reshape_and_cache_nvfp4(
25    gpu: &dyn GpuBackend,
26    kernel: KernelHandle,
27    key: DevicePtr,
28    value: DevicePtr,
29    k_cache: DevicePtr,
30    v_cache: DevicePtr,
31    slot_mapping: DevicePtr,
32    num_tokens: u32,
33    num_kv_heads: u32,
34    head_dim: u32,
35    block_size: u32,
36    key_stride: u32,
37    value_stride: u32,
38    block_stride_bytes: u64,
39    data_section_bytes: u64,
40    stream: u64,
41) -> Result<()> {
42    KernelLaunch::new(gpu, kernel)
43        .grid([num_tokens, 1, 1])
44        .block([256, 1, 1])
45        .arg_ptr(key)
46        .arg_ptr(value)
47        .arg_ptr(k_cache)
48        .arg_ptr(v_cache)
49        .arg_ptr(slot_mapping)
50        .arg_u32(num_kv_heads)
51        .arg_u32(head_dim)
52        .arg_u32(block_size)
53        .arg_u32(key_stride)
54        .arg_u32(value_stride)
55        .arg_u64(block_stride_bytes)
56        .arg_u64(data_section_bytes)
57        .launch(stream)
58}
59
60/// 2026-09-25: Max absolute value of `n_elems` BF16 values, folded into the f32
61/// at `out_max` with an atomic max. The caller zeroes `out_max` first; the FP8
62/// KV calibration (`fp8_calibration.rs`) does so before each K/V pair.
63pub fn bf16_absmax(
64    gpu: &dyn GpuBackend,
65    kernel: KernelHandle,
66    data: DevicePtr,
67    out_max: DevicePtr,
68    n_elems: u32,
69    stream: u64,
70) -> Result<()> {
71    // 2026-09-25: At most 256 blocks; the kernel's grid-stride loop covers the rest.
72    let grid_x = (n_elems as u64).div_ceil(256 * 2).min(256) as u32;
73    KernelLaunch::new(gpu, kernel)
74        .grid([grid_x, 1, 1])
75        .block([256, 1, 1])
76        .arg_ptr(data)
77        .arg_ptr(out_max)
78        .arg_u32(n_elems)
79        .launch(stream)
80}
81
82/// 2026-09-25: Paged decode attention over an NVFP4 KV cache, one block per
83/// `(q_head, seq)`.
84pub fn paged_decode_attn_nvfp4(
85    gpu: &dyn GpuBackend,
86    kernel: KernelHandle,
87    q: DevicePtr,
88    k_cache: DevicePtr,
89    v_cache: DevicePtr,
90    output: DevicePtr,
91    block_tables: DevicePtr,
92    seq_lens: DevicePtr,
93    max_blocks_per_seq: u32,
94    num_seqs: u32,
95    num_q_heads: u32,
96    num_kv_heads: u32,
97    head_dim: u32,
98    block_size: u32,
99    inv_sqrt_d: f32,
100    q_stride: u32,
101    block_stride_bytes: u64,
102    data_section_bytes: u64,
103    stream: u64,
104) -> Result<()> {
105    KernelLaunch::new(gpu, kernel)
106        .grid([num_q_heads, num_seqs, 1])
107        .block([256, 1, 1])
108        .arg_ptr(q)
109        .arg_ptr(k_cache)
110        .arg_ptr(v_cache)
111        .arg_ptr(output)
112        .arg_ptr(block_tables)
113        .arg_ptr(seq_lens)
114        .arg_u32(max_blocks_per_seq)
115        .arg_u32(num_q_heads)
116        .arg_u32(num_kv_heads)
117        .arg_u32(head_dim)
118        .arg_u32(block_size)
119        .arg_f32(inv_sqrt_d)
120        .arg_u32(q_stride)
121        .arg_u64(block_stride_bytes)
122        .arg_u64(data_section_bytes)
123        .launch(stream)
124}
125
126/// 2026-09-25: Split-K paged decode attention over an NVFP4 KV cache: each
127/// `(q_head, seq)` sequence is cut into `num_splits` parts, and each block
128/// writes its partial output and softmax state to `workspace`.
129#[allow(clippy::too_many_arguments)]
130pub fn paged_decode_attn_splitk_nvfp4(
131    gpu: &dyn GpuBackend,
132    kernel: KernelHandle,
133    q: DevicePtr,
134    k_cache: DevicePtr,
135    v_cache: DevicePtr,
136    workspace: DevicePtr,
137    block_tables: DevicePtr,
138    seq_lens: DevicePtr,
139    max_blocks_per_seq: u32,
140    num_q_heads: u32,
141    num_kv_heads: u32,
142    head_dim: u32,
143    block_size: u32,
144    inv_sqrt_d: f32,
145    num_splits: u32,
146    q_stride: u32,
147    block_stride_bytes: u64,
148    data_section_bytes: u64,
149    num_seqs: u32,
150    stream: u64,
151) -> Result<()> {
152    KernelLaunch::new(gpu, kernel)
153        .grid([num_q_heads, num_splits, num_seqs])
154        .block([256, 1, 1])
155        .arg_ptr(q)
156        .arg_ptr(k_cache)
157        .arg_ptr(v_cache)
158        .arg_ptr(workspace)
159        .arg_ptr(block_tables)
160        .arg_ptr(seq_lens)
161        .arg_u32(max_blocks_per_seq)
162        .arg_u32(num_q_heads)
163        .arg_u32(num_kv_heads)
164        .arg_u32(head_dim)
165        .arg_u32(block_size)
166        .arg_f32(inv_sqrt_d)
167        .arg_u32(num_splits)
168        .arg_u32(q_stride)
169        .arg_u64(block_stride_bytes)
170        .arg_u64(data_section_bytes)
171        .launch(stream)
172}
173
174/// 2026-09-25: Combine the split-K partials in `workspace` into the BF16 output.
175#[allow(clippy::too_many_arguments)]
176pub fn paged_decode_attn_reduce_nvfp4(
177    gpu: &dyn GpuBackend,
178    kernel: KernelHandle,
179    workspace: DevicePtr,
180    output: DevicePtr,
181    seq_lens: DevicePtr,
182    num_q_heads: u32,
183    head_dim: u32,
184    num_splits: u32,
185    num_seqs: u32,
186    stream: u64,
187) -> Result<()> {
188    KernelLaunch::new(gpu, kernel)
189        .grid([num_q_heads, num_seqs, 1])
190        .block([32, 1, 1])
191        .arg_ptr(workspace)
192        .arg_ptr(output)
193        .arg_ptr(seq_lens)
194        .arg_u32(num_q_heads)
195        .arg_u32(head_dim)
196        .arg_u32(num_splits)
197        .launch(stream)
198}
199
200/// 2026-09-25: The FP8-KV counterpart of [`paged_decode_attn_splitk_nvfp4`].
201#[allow(clippy::too_many_arguments)]
202pub fn paged_decode_attn_splitk_fp8(
203    gpu: &dyn GpuBackend,
204    kernel: KernelHandle,
205    q: DevicePtr,
206    k_cache: DevicePtr,
207    v_cache: DevicePtr,
208    workspace: DevicePtr,
209    block_tables: DevicePtr,
210    seq_lens: DevicePtr,
211    max_blocks_per_seq: u32,
212    num_q_heads: u32,
213    num_kv_heads: u32,
214    head_dim: u32,
215    block_size: u32,
216    inv_sqrt_d: f32,
217    num_splits: u32,
218    k_scale: f32,
219    v_scale: f32,
220    q_stride: u32,
221    cache_stride: u64,
222    num_seqs: u32,
223    sliding_window: u32,
224    stream: u64,
225) -> Result<()> {
226    KernelLaunch::new(gpu, kernel)
227        .grid([num_q_heads, num_splits, num_seqs])
228        .block([256, 1, 1])
229        .arg_ptr(q)
230        .arg_ptr(k_cache)
231        .arg_ptr(v_cache)
232        .arg_ptr(workspace)
233        .arg_ptr(block_tables)
234        .arg_ptr(seq_lens)
235        .arg_u32(max_blocks_per_seq)
236        .arg_u32(num_q_heads)
237        .arg_u32(num_kv_heads)
238        .arg_u32(head_dim)
239        .arg_u32(block_size)
240        .arg_f32(inv_sqrt_d)
241        .arg_u32(num_splits)
242        .arg_f32(k_scale)
243        .arg_f32(v_scale)
244        .arg_u32(q_stride)
245        .arg_u64(cache_stride)
246        .arg_u32(sliding_window)
247        .launch(stream)
248}
249
250/// 2026-09-25: The BF16-KV counterpart of [`paged_decode_attn_splitk_nvfp4`]
251/// (`kernels/hopper/common/paged_decode_bf16_splitk_hopper.cu`). It takes no
252/// `cache_stride`: the cache is `[blocks, block_size, kv_heads, head_dim]` and
253/// the kernel computes the page stride, as [`paged_decode_attn_bf16`] does.
254#[allow(clippy::too_many_arguments)]
255pub fn paged_decode_attn_splitk_bf16(
256    gpu: &dyn GpuBackend,
257    kernel: KernelHandle,
258    q: DevicePtr,
259    k_cache: DevicePtr,
260    v_cache: DevicePtr,
261    workspace: DevicePtr,
262    block_tables: DevicePtr,
263    seq_lens: DevicePtr,
264    max_blocks_per_seq: u32,
265    num_q_heads: u32,
266    num_kv_heads: u32,
267    head_dim: u32,
268    block_size: u32,
269    inv_sqrt_d: f32,
270    num_splits: u32,
271    q_stride: u32,
272    num_seqs: u32,
273    sliding_window: u32,
274    stream: u64,
275) -> Result<()> {
276    KernelLaunch::new(gpu, kernel)
277        .grid([num_q_heads, num_splits, num_seqs])
278        .block([256, 1, 1])
279        .arg_ptr(q)
280        .arg_ptr(k_cache)
281        .arg_ptr(v_cache)
282        .arg_ptr(workspace)
283        .arg_ptr(block_tables)
284        .arg_ptr(seq_lens)
285        .arg_u32(max_blocks_per_seq)
286        .arg_u32(num_q_heads)
287        .arg_u32(num_kv_heads)
288        .arg_u32(head_dim)
289        .arg_u32(block_size)
290        .arg_f32(inv_sqrt_d)
291        .arg_u32(num_splits)
292        .arg_u32(q_stride)
293        .arg_u32(sliding_window)
294        .launch(stream)
295}
296
297/// 2026-09-25: Combine split-K partials into the BF16 output. The workspace is F32
298/// `[head_dim values, m, l]` per split whatever the cache type, so the BF16
299/// split-K's reduce handle is launched through this function too
300/// (`qwen3_attention/decode/splitk_dispatch.rs`).
301#[allow(clippy::too_many_arguments)]
302pub fn paged_decode_attn_reduce_fp8(
303    gpu: &dyn GpuBackend,
304    kernel: KernelHandle,
305    workspace: DevicePtr,
306    output: DevicePtr,
307    seq_lens: DevicePtr,
308    num_q_heads: u32,
309    head_dim: u32,
310    num_splits: u32,
311    num_seqs: u32,
312    stream: u64,
313) -> Result<()> {
314    KernelLaunch::new(gpu, kernel)
315        .grid([num_q_heads, num_seqs, 1])
316        .block([32, 1, 1])
317        .arg_ptr(workspace)
318        .arg_ptr(output)
319        .arg_ptr(seq_lens)
320        .arg_u32(num_q_heads)
321        .arg_u32(head_dim)
322        .arg_u32(num_splits)
323        .launch(stream)
324}
metrale_model_layers/layers/ops/prefill_attn_b.rs

metrale_model_layers/layers/ops/
prefill_attn_b.rs
