spark_model/layers/ops/
prefill_attn_batched.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Q12 Phase 3: same-chunk-len batched paged-prefill attention ops.
4//!
5//! Wraps the `inferspark_prefill_paged{,_fp8,_nvfp4}_batched` kernels
6//! introduced in commit 4ec2cf2 (kernel commit) and the BR=64 siblings
7//! generated from the shared `prefill_paged_compute.cuh` header.
8//!
9//! Caller is responsible for:
10//!   1. Uploading a device array `block_table_ptrs: [int*; batch_size]`
11//!      that holds the per-stream paged-KV block-table device pointers.
12//!   2. Stacking `Q` and `O` for all batched streams contiguously.
13//!      UNIFORM (`cu_seqlens` NULL): `[batch_size, q_len, num_q_heads,
14//!      head_dim]` BF16, stream `b` at `b * q_len * num_q_heads * head_dim`.
15//!      VARLEN (`cu_seqlens` non-NULL): the buffer is PACKED by the prefix
16//!      sum, stream `b` at `cu_seqlens[b] * num_q_heads * head_dim` with
17//!      length `cu_seqlens[b+1] - cu_seqlens[b]` and KV extent `kv_lens[b]`.
18//!   3. UNIFORM only: ensuring all batched streams share the same `q_len`,
19//!      `kv_len` and `q_offset`. Under VARLEN those come from `cu_seqlens` /
20//!      `kv_lens` per stream and `q_len` is passed as the MAX (it bounds the
21//!      grid's Q-tile dimension only). `q_offset`, `sliding_window` and (for
22//!      FP8/NVFP4) quantisation scales are still shared.
23//!
24//! Validation status: kernels unvalidated against hardware.
25
26#![allow(unused_imports, clippy::too_many_arguments)]
27
28use anyhow::Result;
29use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
30use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
31
32/// Batched BF16-KV paged prefill attention (BR=32).
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/// Batched BF16-KV paged prefill attention (BR=64, 256-thread variant).
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/// Batched FP8-KV paged prefill attention (BR=32).
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/// Batched FP8-KV paged prefill attention (BR=64, 256-thread variant).
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/// Batched NVFP4-KV paged prefill attention (BR=32).
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/// Batched NVFP4-KV paged prefill attention (BR=64, chunk_len >= 256).
289///
290/// Same argument contract as [`prefill_attention_paged_nvfp4_batched`];
291/// the `_64` kernel sibling is generated from the shared
292/// `prefill_paged_compute.cuh` (8 warps, 64 Q rows per CTA).
293pub fn prefill_attention_paged_nvfp4_batched_64(
294    gpu: &dyn GpuBackend,
295    kernel: KernelHandle,
296    q: DevicePtr,
297    k_cache: DevicePtr,
298    v_cache: DevicePtr,
299    output: DevicePtr,
300    block_table_ptrs: DevicePtr,
301    batch_size: u32,
302    cu_seqlens: DevicePtr,
303    kv_lens: DevicePtr,
304    q_len: u32,
305    kv_len: u32,
306    q_offset: u32,
307    num_q_heads: u32,
308    num_kv_heads: u32,
309    head_dim: u32,
310    cache_block_size: u32,
311    sliding_window: u32,
312    inv_sqrt_d: f32,
313    block_stride_bytes: u64,
314    data_section_bytes: u64,
315    stream: u64,
316) -> Result<()> {
317    let br = 64u32;
318    KernelLaunch::new(gpu, kernel)
319        .grid([num_q_heads, div_ceil(q_len, br), batch_size])
320        .block([256, 1, 1])
321        .arg_ptr(q)
322        .arg_ptr(k_cache)
323        .arg_ptr(v_cache)
324        .arg_ptr(output)
325        .arg_ptr(block_table_ptrs)
326        .arg_u32(batch_size)
327        .arg_ptr(cu_seqlens)
328        .arg_ptr(kv_lens)
329        .arg_u32(q_len)
330        .arg_u32(kv_len)
331        .arg_u32(q_offset)
332        .arg_u32(num_q_heads)
333        .arg_u32(num_kv_heads)
334        .arg_u32(head_dim)
335        .arg_u32(cache_block_size)
336        .arg_u32(sliding_window)
337        .arg_u32(1u32)
338        .arg_f32(inv_sqrt_d)
339        .arg_u64(block_stride_bytes)
340        .arg_u64(data_section_bytes)
341        .launch(stream)
342}
spark_model/layers/ops/prefill_attn_batched.rs

spark_model/layers/ops/
prefill_attn_batched.rs
