metrale_model_layers/layers/ops/
prefill_attn_batched.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the batched paged-prefill attention kernels
4//! `attn_prefill_paged{,_fp8,_nvfp4}_batched` and their BR=64 siblings, built
5//! from `kernels/gb10/common/prefill_paged_compute.cuh` with `PREFILL_BATCHED`.
6//! Grid z is the stream index.
7//!
8//! The caller provides:
9//! 1. `block_table_ptrs`: a device array of per-stream block-table pointers.
10//! 2. Q and O for all streams in one buffer. With `cu_seqlens` NULL (uniform),
11//!    every stream is `q_len` long and stream `b` starts at
12//!    `b * q_len * num_q_heads * head_dim`. With `cu_seqlens` set (varlen),
13//!    streams are packed: stream `b` starts at
14//!    `cu_seqlens[b] * num_q_heads * head_dim` and is
15//!    `cu_seqlens[b+1] - cu_seqlens[b]` long; `q_len` must then be at least
16//!    the longest stream, and only sizes the grid.
17//! 3. The KV extent: `kv_lens[b]` when `kv_lens` is set, else the shared
18//!    `kv_len`. `q_offset`, `sliding_window` and the FP8/NVFP4 scales are
19//!    shared by every stream.
20//!
21//! Every launcher passes `causal_mask_enabled = 1`.
22//!
23//! Owner: model-layers ops.
24//! Invariants: none beyond the types.
25
26#![allow(unused_imports, clippy::too_many_arguments)]
27
28use anyhow::Result;
29use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
30use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
31
32/// 2026-09-25: Batched BF16-KV paged prefill attention, 32 query rows per block.
33pub fn prefill_attention_paged_batched(
34    gpu: &dyn GpuBackend,
35    kernel: KernelHandle,
36    q: DevicePtr,
37    k_cache: DevicePtr,
38    v_cache: DevicePtr,
39    output: DevicePtr,
40    block_table_ptrs: DevicePtr,
41    batch_size: u32,
42    cu_seqlens: DevicePtr,
43    kv_lens: DevicePtr,
44    q_len: u32,
45    kv_len: u32,
46    q_offset: u32,
47    num_q_heads: u32,
48    num_kv_heads: u32,
49    head_dim: u32,
50    cache_block_size: u32,
51    sliding_window: u32,
52    inv_sqrt_d: f32,
53    stream: u64,
54) -> Result<()> {
55    let br = 32u32;
56    KernelLaunch::new(gpu, kernel)
57        .grid([num_q_heads, div_ceil(q_len, br), batch_size])
58        .block([128, 1, 1])
59        .arg_ptr(q)
60        .arg_ptr(k_cache)
61        .arg_ptr(v_cache)
62        .arg_ptr(output)
63        .arg_ptr(block_table_ptrs)
64        .arg_u32(batch_size)
65        .arg_ptr(cu_seqlens)
66        .arg_ptr(kv_lens)
67        .arg_u32(q_len)
68        .arg_u32(kv_len)
69        .arg_u32(q_offset)
70        .arg_u32(num_q_heads)
71        .arg_u32(num_kv_heads)
72        .arg_u32(head_dim)
73        .arg_u32(cache_block_size)
74        .arg_u32(sliding_window)
75        .arg_u32(1u32)
76        .arg_f32(inv_sqrt_d)
77        .launch(stream)
78}
79
80/// 2026-09-25: Batched BF16-KV paged prefill attention, 64 query rows per block.
81pub fn prefill_attention_paged_batched_64(
82    gpu: &dyn GpuBackend,
83    kernel: KernelHandle,
84    q: DevicePtr,
85    k_cache: DevicePtr,
86    v_cache: DevicePtr,
87    output: DevicePtr,
88    block_table_ptrs: DevicePtr,
89    batch_size: u32,
90    cu_seqlens: DevicePtr,
91    kv_lens: DevicePtr,
92    q_len: u32,
93    kv_len: u32,
94    q_offset: u32,
95    num_q_heads: u32,
96    num_kv_heads: u32,
97    head_dim: u32,
98    cache_block_size: u32,
99    sliding_window: u32,
100    inv_sqrt_d: f32,
101    stream: u64,
102) -> Result<()> {
103    let br = 64u32;
104    KernelLaunch::new(gpu, kernel)
105        .grid([num_q_heads, div_ceil(q_len, br), batch_size])
106        .block([256, 1, 1])
107        .arg_ptr(q)
108        .arg_ptr(k_cache)
109        .arg_ptr(v_cache)
110        .arg_ptr(output)
111        .arg_ptr(block_table_ptrs)
112        .arg_u32(batch_size)
113        .arg_ptr(cu_seqlens)
114        .arg_ptr(kv_lens)
115        .arg_u32(q_len)
116        .arg_u32(kv_len)
117        .arg_u32(q_offset)
118        .arg_u32(num_q_heads)
119        .arg_u32(num_kv_heads)
120        .arg_u32(head_dim)
121        .arg_u32(cache_block_size)
122        .arg_u32(sliding_window)
123        .arg_u32(1u32)
124        .arg_f32(inv_sqrt_d)
125        .launch(stream)
126}
127
128/// 2026-09-25: Batched FP8-KV paged prefill attention, 32 query rows per block.
129pub fn prefill_attention_paged_fp8_batched(
130    gpu: &dyn GpuBackend,
131    kernel: KernelHandle,
132    q: DevicePtr,
133    k_cache: DevicePtr,
134    v_cache: DevicePtr,
135    output: DevicePtr,
136    block_table_ptrs: DevicePtr,
137    batch_size: u32,
138    cu_seqlens: DevicePtr,
139    kv_lens: DevicePtr,
140    q_len: u32,
141    kv_len: u32,
142    q_offset: u32,
143    num_q_heads: u32,
144    num_kv_heads: u32,
145    head_dim: u32,
146    cache_block_size: u32,
147    sliding_window: u32,
148    inv_sqrt_d: f32,
149    k_scale: f32,
150    v_scale: f32,
151    cache_stride: u64,
152    stream: u64,
153) -> Result<()> {
154    let br = 32u32;
155    KernelLaunch::new(gpu, kernel)
156        .grid([num_q_heads, div_ceil(q_len, br), batch_size])
157        .block([128, 1, 1])
158        .arg_ptr(q)
159        .arg_ptr(k_cache)
160        .arg_ptr(v_cache)
161        .arg_ptr(output)
162        .arg_ptr(block_table_ptrs)
163        .arg_u32(batch_size)
164        .arg_ptr(cu_seqlens)
165        .arg_ptr(kv_lens)
166        .arg_u32(q_len)
167        .arg_u32(kv_len)
168        .arg_u32(q_offset)
169        .arg_u32(num_q_heads)
170        .arg_u32(num_kv_heads)
171        .arg_u32(head_dim)
172        .arg_u32(cache_block_size)
173        .arg_u32(sliding_window)
174        .arg_u32(1u32)
175        .arg_f32(inv_sqrt_d)
176        .arg_f32(k_scale)
177        .arg_f32(v_scale)
178        .arg_u64(cache_stride)
179        .launch(stream)
180}
181
182/// 2026-09-25: Batched FP8-KV paged prefill attention, 64 query rows per block.
183pub fn prefill_attention_paged_fp8_batched_64(
184    gpu: &dyn GpuBackend,
185    kernel: KernelHandle,
186    q: DevicePtr,
187    k_cache: DevicePtr,
188    v_cache: DevicePtr,
189    output: DevicePtr,
190    block_table_ptrs: DevicePtr,
191    batch_size: u32,
192    cu_seqlens: DevicePtr,
193    kv_lens: DevicePtr,
194    q_len: u32,
195    kv_len: u32,
196    q_offset: u32,
197    num_q_heads: u32,
198    num_kv_heads: u32,
199    head_dim: u32,
200    cache_block_size: u32,
201    sliding_window: u32,
202    inv_sqrt_d: f32,
203    k_scale: f32,
204    v_scale: f32,
205    cache_stride: u64,
206    stream: u64,
207) -> Result<()> {
208    let br = 64u32;
209    KernelLaunch::new(gpu, kernel)
210        .grid([num_q_heads, div_ceil(q_len, br), batch_size])
211        .block([256, 1, 1])
212        .arg_ptr(q)
213        .arg_ptr(k_cache)
214        .arg_ptr(v_cache)
215        .arg_ptr(output)
216        .arg_ptr(block_table_ptrs)
217        .arg_u32(batch_size)
218        .arg_ptr(cu_seqlens)
219        .arg_ptr(kv_lens)
220        .arg_u32(q_len)
221        .arg_u32(kv_len)
222        .arg_u32(q_offset)
223        .arg_u32(num_q_heads)
224        .arg_u32(num_kv_heads)
225        .arg_u32(head_dim)
226        .arg_u32(cache_block_size)
227        .arg_u32(sliding_window)
228        .arg_u32(1u32)
229        .arg_f32(inv_sqrt_d)
230        .arg_f32(k_scale)
231        .arg_f32(v_scale)
232        .arg_u64(cache_stride)
233        .launch(stream)
234}
235
236/// 2026-09-25: Batched NVFP4-KV paged prefill attention, 32 query rows per block.
237pub fn prefill_attention_paged_nvfp4_batched(
238    gpu: &dyn GpuBackend,
239    kernel: KernelHandle,
240    q: DevicePtr,
241    k_cache: DevicePtr,
242    v_cache: DevicePtr,
243    output: DevicePtr,
244    block_table_ptrs: DevicePtr,
245    batch_size: u32,
246    cu_seqlens: DevicePtr,
247    kv_lens: DevicePtr,
248    q_len: u32,
249    kv_len: u32,
250    q_offset: u32,
251    num_q_heads: u32,
252    num_kv_heads: u32,
253    head_dim: u32,
254    cache_block_size: u32,
255    sliding_window: u32,
256    inv_sqrt_d: f32,
257    block_stride_bytes: u64,
258    data_section_bytes: u64,
259    stream: u64,
260) -> Result<()> {
261    let br = 32u32;
262    KernelLaunch::new(gpu, kernel)
263        .grid([num_q_heads, div_ceil(q_len, br), batch_size])
264        .block([128, 1, 1])
265        .arg_ptr(q)
266        .arg_ptr(k_cache)
267        .arg_ptr(v_cache)
268        .arg_ptr(output)
269        .arg_ptr(block_table_ptrs)
270        .arg_u32(batch_size)
271        .arg_ptr(cu_seqlens)
272        .arg_ptr(kv_lens)
273        .arg_u32(q_len)
274        .arg_u32(kv_len)
275        .arg_u32(q_offset)
276        .arg_u32(num_q_heads)
277        .arg_u32(num_kv_heads)
278        .arg_u32(head_dim)
279        .arg_u32(cache_block_size)
280        .arg_u32(sliding_window)
281        .arg_u32(1u32)
282        .arg_f32(inv_sqrt_d)
283        .arg_u64(block_stride_bytes)
284        .arg_u64(data_section_bytes)
285        .launch(stream)
286}
287
288/// 2026-09-25: Batched NVFP4-KV paged prefill attention, 64 query rows per block,
289/// with the arguments of [`prefill_attention_paged_nvfp4_batched`]. The caller
290/// (`qwen3_attention/prefill/paged_attn_batched.rs`) picks the 64-row kernels
291/// for every KV type when `chunk_len >= 256`.
292pub fn prefill_attention_paged_nvfp4_batched_64(
293    gpu: &dyn GpuBackend,
294    kernel: KernelHandle,
295    q: DevicePtr,
296    k_cache: DevicePtr,
297    v_cache: DevicePtr,
298    output: DevicePtr,
299    block_table_ptrs: DevicePtr,
300    batch_size: u32,
301    cu_seqlens: DevicePtr,
302    kv_lens: DevicePtr,
303    q_len: u32,
304    kv_len: u32,
305    q_offset: u32,
306    num_q_heads: u32,
307    num_kv_heads: u32,
308    head_dim: u32,
309    cache_block_size: u32,
310    sliding_window: u32,
311    inv_sqrt_d: f32,
312    block_stride_bytes: u64,
313    data_section_bytes: u64,
314    stream: u64,
315) -> Result<()> {
316    let br = 64u32;
317    KernelLaunch::new(gpu, kernel)
318        .grid([num_q_heads, div_ceil(q_len, br), batch_size])
319        .block([256, 1, 1])
320        .arg_ptr(q)
321        .arg_ptr(k_cache)
322        .arg_ptr(v_cache)
323        .arg_ptr(output)
324        .arg_ptr(block_table_ptrs)
325        .arg_u32(batch_size)
326        .arg_ptr(cu_seqlens)
327        .arg_ptr(kv_lens)
328        .arg_u32(q_len)
329        .arg_u32(kv_len)
330        .arg_u32(q_offset)
331        .arg_u32(num_q_heads)
332        .arg_u32(num_kv_heads)
333        .arg_u32(head_dim)
334        .arg_u32(cache_block_size)
335        .arg_u32(sliding_window)
336        .arg_u32(1u32)
337        .arg_f32(inv_sqrt_d)
338        .arg_u64(block_stride_bytes)
339        .arg_u64(data_section_bytes)
340        .launch(stream)
341}
metrale_model_layers/layers/ops/prefill_attn_batched.rs

metrale_model_layers/layers/ops/
prefill_attn_batched.rs
