spark_model/layers/ops/
prefill_attn_main_b.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Auto-extracted from `ops.rs` during refactor wave 4a.
4
5#![allow(unused_imports)]
6
7use anyhow::Result;
8use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
9use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
10
11use crate::layers::moe;
12use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
13
14use super::*;
15
16/// Paged prefill Flash Attention — NVFP4 KV cache variant.
17#[allow(clippy::too_many_arguments)]
18pub fn prefill_attention_paged_nvfp4(
19    gpu: &dyn GpuBackend,
20    kernel: KernelHandle,
21    q: DevicePtr,
22    k_cache: DevicePtr,
23    v_cache: DevicePtr,
24    output: DevicePtr,
25    block_table: DevicePtr,
26    q_len: u32,
27    kv_len: u32,
28    q_offset: u32,
29    num_q_heads: u32,
30    num_kv_heads: u32,
31    head_dim: u32,
32    cache_block_size: u32,
33    sliding_window: u32,
34    inv_sqrt_d: f32,
35    block_stride_bytes: u64,
36    data_section_bytes: u64,
37    stream: u64,
38) -> Result<()> {
39    let br = 32u32;
40    KernelLaunch::new(gpu, kernel)
41        .grid([num_q_heads, div_ceil(q_len, br), 1])
42        .block([128, 1, 1])
43        .arg_ptr(q)
44        .arg_ptr(k_cache)
45        .arg_ptr(v_cache)
46        .arg_ptr(output)
47        .arg_ptr(block_table)
48        .arg_u32(q_len)
49        .arg_u32(kv_len)
50        .arg_u32(q_offset)
51        .arg_u32(num_q_heads)
52        .arg_u32(num_kv_heads)
53        .arg_u32(head_dim)
54        .arg_u32(cache_block_size)
55        .arg_u32(sliding_window)
56        // causal_mask_enabled = 1 (default causal). DFlash γ-block kernels
57        // pass 0 via dedicated dispatchers (`prefill_attention_paged_dflash_*`).
58        .arg_u32(1u32)
59        .arg_f32(inv_sqrt_d)
60        .arg_u64(block_stride_bytes)
61        .arg_u64(data_section_bytes)
62        .launch(stream)
63}
64
65/// Paged prefill Flash Attention for HDIM=512 (Gemma-4 full-attention) — BF16 KV.
66///
67/// Uses dynamic shared memory (101,120 B) opt-in. Single-buffered K, 8 warps.
68/// Required for chunked long-context prefill on layers with `head_dim==512`
69/// where the standard 4-warp template doesn't fit GB10's 99 KB smem cap.
70#[allow(clippy::too_many_arguments)]
71pub fn prefill_attention_paged_512(
72    gpu: &dyn GpuBackend,
73    kernel: KernelHandle,
74    q: DevicePtr,
75    k_cache: DevicePtr,
76    v_cache: DevicePtr,
77    output: DevicePtr,
78    block_table: DevicePtr,
79    q_len: u32,
80    kv_len: u32,
81    q_offset: u32,
82    num_q_heads: u32,
83    num_kv_heads: u32,
84    head_dim: u32,
85    cache_block_size: u32,
86    sliding_window: u32,
87    inv_sqrt_d: f32,
88    stream: u64,
89) -> Result<()> {
90    let br = 32u32;
91    KernelLaunch::new(gpu, kernel)
92        .grid([num_q_heads, div_ceil(q_len, br), 1])
93        .block([256, 1, 1])
94        .shared_mem(101_120)
95        .arg_ptr(q)
96        .arg_ptr(k_cache)
97        .arg_ptr(v_cache)
98        .arg_ptr(output)
99        .arg_ptr(block_table)
100        .arg_u32(q_len)
101        .arg_u32(kv_len)
102        .arg_u32(q_offset)
103        .arg_u32(num_q_heads)
104        .arg_u32(num_kv_heads)
105        .arg_u32(head_dim)
106        .arg_u32(cache_block_size)
107        .arg_u32(sliding_window)
108        // causal_mask_enabled = 1 (default causal). DFlash γ-block kernels
109        // pass 0 via dedicated dispatchers (`prefill_attention_paged_dflash_*`).
110        .arg_u32(1u32)
111        .arg_f32(inv_sqrt_d)
112        .launch(stream)
113}
114
115/// Paged prefill Flash Attention — BF16 KV cache, BR=64 (256 threads).
116#[allow(clippy::too_many_arguments)]
117pub fn prefill_attention_paged_64(
118    gpu: &dyn GpuBackend,
119    kernel: KernelHandle,
120    q: DevicePtr,
121    k_cache: DevicePtr,
122    v_cache: DevicePtr,
123    output: DevicePtr,
124    block_table: DevicePtr,
125    q_len: u32,
126    kv_len: u32,
127    q_offset: u32,
128    num_q_heads: u32,
129    num_kv_heads: u32,
130    head_dim: u32,
131    cache_block_size: u32,
132    sliding_window: u32,
133    inv_sqrt_d: f32,
134    stream: u64,
135) -> Result<()> {
136    // Paged prefill kernels clamp BR64 64->32 on AMD (gfx1151 LDS cap;
137    // prefill_paged_compute.cuh). The grid stride must match the kernel's BR64
138    // or query rows 32..63 of every 64-row band are dropped (same class as the
139    // non-paged prefill_attention_64 fix). cfg!(atlas_scale) = strix+strix-hip;
140    // NVIDIA keeps 64 byte-identical.
141    let br = if cfg!(atlas_scale) { 32u32 } else { 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        // causal_mask_enabled = 1 (default causal). DFlash γ-block kernels
159        // pass 0 via dedicated dispatchers (`prefill_attention_paged_dflash_*`).
160        .arg_u32(1u32)
161        .arg_f32(inv_sqrt_d)
162        .launch(stream)
163}
164
165/// Paged prefill Flash Attention — FP8 KV cache, BR=64 (256 threads).
166#[allow(clippy::too_many_arguments)]
167pub fn prefill_attention_paged_fp8_64(
168    gpu: &dyn GpuBackend,
169    kernel: KernelHandle,
170    q: DevicePtr,
171    k_cache: DevicePtr,
172    v_cache: DevicePtr,
173    output: DevicePtr,
174    block_table: DevicePtr,
175    q_len: u32,
176    kv_len: u32,
177    q_offset: u32,
178    num_q_heads: u32,
179    num_kv_heads: u32,
180    head_dim: u32,
181    cache_block_size: u32,
182    sliding_window: u32,
183    inv_sqrt_d: f32,
184    k_scale: f32,
185    v_scale: f32,
186    cache_stride: u64,
187    stream: u64,
188) -> Result<()> {
189    // Paged prefill kernels clamp BR64 64->32 on AMD (gfx1151 LDS cap;
190    // prefill_paged_compute.cuh). The grid stride must match the kernel's BR64
191    // or query rows 32..63 of every 64-row band are dropped (same class as the
192    // non-paged prefill_attention_64 fix). cfg!(atlas_scale) = strix+strix-hip;
193    // NVIDIA keeps 64 byte-identical.
194    let br = if cfg!(atlas_scale) { 32u32 } else { 64u32 };
195    KernelLaunch::new(gpu, kernel)
196        .grid([num_q_heads, div_ceil(q_len, br), 1])
197        .block([256, 1, 1])
198        .arg_ptr(q)
199        .arg_ptr(k_cache)
200        .arg_ptr(v_cache)
201        .arg_ptr(output)
202        .arg_ptr(block_table)
203        .arg_u32(q_len)
204        .arg_u32(kv_len)
205        .arg_u32(q_offset)
206        .arg_u32(num_q_heads)
207        .arg_u32(num_kv_heads)
208        .arg_u32(head_dim)
209        .arg_u32(cache_block_size)
210        .arg_u32(sliding_window)
211        // causal_mask_enabled = 1 (default causal). DFlash γ-block kernels
212        // pass 0 via dedicated dispatchers (`prefill_attention_paged_dflash_*`).
213        .arg_u32(1u32)
214        .arg_f32(inv_sqrt_d)
215        .arg_f32(k_scale)
216        .arg_f32(v_scale)
217        .arg_u64(cache_stride)
218        .launch(stream)
219}
220
221/// Paged prefill Flash Attention — symmetric TurboQuant KV cache, BR=64.
222/// Shared launch wrapper for the turbo8 / turbo4 / turbo3 `_64` kernel
223/// entries: identical ABI, the caller selects the dtype via `kernel` and
224/// passes that pool's block stride + data-section offset.
225#[allow(clippy::too_many_arguments)]
226pub fn prefill_attention_paged_turbo_64(
227    gpu: &dyn GpuBackend,
228    kernel: KernelHandle,
229    q: DevicePtr,
230    k_cache: DevicePtr,
231    v_cache: DevicePtr,
232    output: DevicePtr,
233    block_table: DevicePtr,
234    q_len: u32,
235    kv_len: u32,
236    q_offset: u32,
237    num_q_heads: u32,
238    num_kv_heads: u32,
239    head_dim: u32,
240    cache_block_size: u32,
241    sliding_window: u32,
242    inv_sqrt_d: f32,
243    block_stride_bytes: u64,
244    data_section_bytes: u64,
245    stream: u64,
246) -> Result<()> {
247    // Paged prefill kernels clamp BR64 64->32 on AMD (gfx1151 LDS cap;
248    // prefill_paged_compute.cuh). The grid stride must match the kernel's BR64
249    // or query rows 32..63 of every 64-row band are dropped (same class as the
250    // non-paged prefill_attention_64 fix). cfg!(atlas_scale) = strix+strix-hip;
251    // NVIDIA keeps 64 byte-identical.
252    let br = if cfg!(atlas_scale) { 32u32 } else { 64u32 };
253    KernelLaunch::new(gpu, kernel)
254        .grid([num_q_heads, div_ceil(q_len, br), 1])
255        .block([256, 1, 1])
256        .arg_ptr(q)
257        .arg_ptr(k_cache)
258        .arg_ptr(v_cache)
259        .arg_ptr(output)
260        .arg_ptr(block_table)
261        .arg_u32(q_len)
262        .arg_u32(kv_len)
263        .arg_u32(q_offset)
264        .arg_u32(num_q_heads)
265        .arg_u32(num_kv_heads)
266        .arg_u32(head_dim)
267        .arg_u32(cache_block_size)
268        .arg_u32(sliding_window)
269        .arg_u32(1u32)
270        .arg_f32(inv_sqrt_d)
271        .arg_u64(block_stride_bytes)
272        .arg_u64(data_section_bytes)
273        .launch(stream)
274}
275
276pub fn prefill_attention_paged_turbo2_64(
277    gpu: &dyn GpuBackend,
278    kernel: KernelHandle,
279    q: DevicePtr,
280    k_cache: DevicePtr,
281    v_cache: DevicePtr,
282    output: DevicePtr,
283    block_table: DevicePtr,
284    q_len: u32,
285    kv_len: u32,
286    q_offset: u32,
287    num_q_heads: u32,
288    num_kv_heads: u32,
289    head_dim: u32,
290    cache_block_size: u32,
291    sliding_window: u32,
292    inv_sqrt_d: f32,
293    block_stride_bytes: u64,
294    data_section_bytes: u64,
295    stream: u64,
296) -> Result<()> {
297    let br = 32u32; // try BR=32 entry first while debugging BR=64 OOB
298    // BR=32 entry is sized for 128 threads (4 warps); 256 threads makes
299    // warps 4-7 read past smem_V (OOB shared reads, results discarded).
300    KernelLaunch::new(gpu, kernel)
301        .grid([num_q_heads, div_ceil(q_len, br), 1])
302        .block([128, 1, 1])
303        .arg_ptr(q)
304        .arg_ptr(k_cache)
305        .arg_ptr(v_cache)
306        .arg_ptr(output)
307        .arg_ptr(block_table)
308        .arg_u32(q_len)
309        .arg_u32(kv_len)
310        .arg_u32(q_offset)
311        .arg_u32(num_q_heads)
312        .arg_u32(num_kv_heads)
313        .arg_u32(head_dim)
314        .arg_u32(cache_block_size)
315        .arg_u32(sliding_window)
316        .arg_u32(1u32)
317        .arg_f32(inv_sqrt_d)
318        .arg_u64(block_stride_bytes)
319        .arg_u64(data_section_bytes)
320        .launch(stream)
321}
322
323pub fn prefill_attention_paged_nvfp4_64(
324    gpu: &dyn GpuBackend,
325    kernel: KernelHandle,
326    q: DevicePtr,
327    k_cache: DevicePtr,
328    v_cache: DevicePtr,
329    output: DevicePtr,
330    block_table: DevicePtr,
331    q_len: u32,
332    kv_len: u32,
333    q_offset: u32,
334    num_q_heads: u32,
335    num_kv_heads: u32,
336    head_dim: u32,
337    cache_block_size: u32,
338    sliding_window: u32,
339    inv_sqrt_d: f32,
340    block_stride_bytes: u64,
341    data_section_bytes: u64,
342    stream: u64,
343) -> Result<()> {
344    // Paged prefill kernels clamp BR64 64->32 on AMD (gfx1151 LDS cap;
345    // prefill_paged_compute.cuh). The grid stride must match the kernel's BR64
346    // or query rows 32..63 of every 64-row band are dropped (same class as the
347    // non-paged prefill_attention_64 fix). cfg!(atlas_scale) = strix+strix-hip;
348    // NVIDIA keeps 64 byte-identical.
349    let br = if cfg!(atlas_scale) { 32u32 } else { 64u32 };
350    KernelLaunch::new(gpu, kernel)
351        .grid([num_q_heads, div_ceil(q_len, br), 1])
352        .block([256, 1, 1])
353        .arg_ptr(q)
354        .arg_ptr(k_cache)
355        .arg_ptr(v_cache)
356        .arg_ptr(output)
357        .arg_ptr(block_table)
358        .arg_u32(q_len)
359        .arg_u32(kv_len)
360        .arg_u32(q_offset)
361        .arg_u32(num_q_heads)
362        .arg_u32(num_kv_heads)
363        .arg_u32(head_dim)
364        .arg_u32(cache_block_size)
365        .arg_u32(sliding_window)
366        // causal_mask_enabled = 1 (default causal). DFlash γ-block kernels
367        // pass 0 via dedicated dispatchers (`prefill_attention_paged_dflash_*`).
368        .arg_u32(1u32)
369        .arg_f32(inv_sqrt_d)
370        .arg_u64(block_stride_bytes)
371        .arg_u64(data_section_bytes)
372        .launch(stream)
373}
374
375/// Paged prefill (BR=64) for Bf16K + Turbo3V asymmetric KV cache.
376///
377/// Reads K as BF16 (NHD contiguous) and V as 3-bit Lloyd-Max packed bytes
378/// with FP8 per-group scale. Uses the asym prefill compute template
379/// (prefill_paged_compute_asym.cuh) which takes separate LOAD_K_TILE +
380/// LOAD_V_TILE macros — bf16 cp.async for K, sync dequant for V.
381///
382/// Kernel: `inferspark_prefill_paged_bf16k_turbo3v_64(Q, K_cache, V_cache,
383///          O, block_table, q_len, kv_len, q_offset, num_q_heads,
384///          num_kv_heads, head_dim, cache_block_size, sliding_window,
385///          causal_mask_enabled, inv_sqrt_d, v_block_stride_bytes,
386///          v_data_section_bytes)`
387/// Grid: (num_q_heads, div_ceil(q_len, BR), 1)  Block: (256, 1, 1)
388#[allow(clippy::too_many_arguments)]
389pub fn prefill_attention_paged_bf16k_turbo3v_64(
390    gpu: &dyn GpuBackend,
391    kernel: KernelHandle,
392    q: DevicePtr,
393    k_cache: DevicePtr,
394    v_cache: DevicePtr,
395    output: DevicePtr,
396    block_table: DevicePtr,
397    q_len: u32,
398    kv_len: u32,
399    q_offset: u32,
400    num_q_heads: u32,
401    num_kv_heads: u32,
402    head_dim: u32,
403    cache_block_size: u32,
404    sliding_window: u32,
405    inv_sqrt_d: f32,
406    v_block_stride_bytes: u64,
407    v_data_section_bytes: u64,
408    stream: u64,
409) -> Result<()> {
410    // Paged prefill kernels clamp BR64 64->32 on AMD (gfx1151 LDS cap;
411    // prefill_paged_compute.cuh). The grid stride must match the kernel's BR64
412    // or query rows 32..63 of every 64-row band are dropped (same class as the
413    // non-paged prefill_attention_64 fix). cfg!(atlas_scale) = strix+strix-hip;
414    // NVIDIA keeps 64 byte-identical.
415    let br = if cfg!(atlas_scale) { 32u32 } else { 64u32 };
416    KernelLaunch::new(gpu, kernel)
417        .grid([num_q_heads, div_ceil(q_len, br), 1])
418        .block([256, 1, 1])
419        .arg_ptr(q)
420        .arg_ptr(k_cache)
421        .arg_ptr(v_cache)
422        .arg_ptr(output)
423        .arg_ptr(block_table)
424        .arg_u32(q_len)
425        .arg_u32(kv_len)
426        .arg_u32(q_offset)
427        .arg_u32(num_q_heads)
428        .arg_u32(num_kv_heads)
429        .arg_u32(head_dim)
430        .arg_u32(cache_block_size)
431        .arg_u32(sliding_window)
432        // causal_mask_enabled = 1 (default causal).
433        .arg_u32(1u32)
434        .arg_f32(inv_sqrt_d)
435        .arg_u64(v_block_stride_bytes)
436        .arg_u64(v_data_section_bytes)
437        .launch(stream)
438}
439
440/// Prefill paged attention — TurboQuant+ safer-asym Bf16K + Turbo4V (BR=64).
441///
442/// Same kernel ABI as `prefill_attention_paged_bf16k_turbo3v_64`; the
443/// underlying kernel uses a 4-bit V dequant path in `LOAD_V_TILE`.
444#[allow(clippy::too_many_arguments)]
445pub fn prefill_attention_paged_bf16k_turbo4v_64(
446    gpu: &dyn GpuBackend,
447    kernel: KernelHandle,
448    q: DevicePtr,
449    k_cache: DevicePtr,
450    v_cache: DevicePtr,
451    output: DevicePtr,
452    block_table: DevicePtr,
453    q_len: u32,
454    kv_len: u32,
455    q_offset: u32,
456    num_q_heads: u32,
457    num_kv_heads: u32,
458    head_dim: u32,
459    cache_block_size: u32,
460    sliding_window: u32,
461    inv_sqrt_d: f32,
462    v_block_stride_bytes: u64,
463    v_data_section_bytes: u64,
464    stream: u64,
465) -> Result<()> {
466    // Paged prefill kernels clamp BR64 64->32 on AMD (gfx1151 LDS cap;
467    // prefill_paged_compute.cuh). The grid stride must match the kernel's BR64
468    // or query rows 32..63 of every 64-row band are dropped (same class as the
469    // non-paged prefill_attention_64 fix). cfg!(atlas_scale) = strix+strix-hip;
470    // NVIDIA keeps 64 byte-identical.
471    let br = if cfg!(atlas_scale) { 32u32 } else { 64u32 };
472    KernelLaunch::new(gpu, kernel)
473        .grid([num_q_heads, div_ceil(q_len, br), 1])
474        .block([256, 1, 1])
475        .arg_ptr(q)
476        .arg_ptr(k_cache)
477        .arg_ptr(v_cache)
478        .arg_ptr(output)
479        .arg_ptr(block_table)
480        .arg_u32(q_len)
481        .arg_u32(kv_len)
482        .arg_u32(q_offset)
483        .arg_u32(num_q_heads)
484        .arg_u32(num_kv_heads)
485        .arg_u32(head_dim)
486        .arg_u32(cache_block_size)
487        .arg_u32(sliding_window)
488        .arg_u32(1u32)
489        .arg_f32(inv_sqrt_d)
490        .arg_u64(v_block_stride_bytes)
491        .arg_u64(v_data_section_bytes)
492        .launch(stream)
493}
494
495/// Prefill paged attention — TurboQuant+ safer-asym Bf16K + Turbo2V (BR=64).
496///
497/// 6.4x V compression. Same kernel ABI as `prefill_attention_paged_bf16k_turbo3v_64`;
498/// kernel uses a 2-bit V dequant path in `LOAD_V_TILE`.
499#[allow(clippy::too_many_arguments)]
500pub fn prefill_attention_paged_bf16k_turbo2v_64(
501    gpu: &dyn GpuBackend,
502    kernel: KernelHandle,
503    q: DevicePtr,
504    k_cache: DevicePtr,
505    v_cache: DevicePtr,
506    output: DevicePtr,
507    block_table: DevicePtr,
508    q_len: u32,
509    kv_len: u32,
510    q_offset: u32,
511    num_q_heads: u32,
512    num_kv_heads: u32,
513    head_dim: u32,
514    cache_block_size: u32,
515    sliding_window: u32,
516    inv_sqrt_d: f32,
517    v_block_stride_bytes: u64,
518    v_data_section_bytes: u64,
519    stream: u64,
520) -> Result<()> {
521    // Paged prefill kernels clamp BR64 64->32 on AMD (gfx1151 LDS cap;
522    // prefill_paged_compute.cuh). The grid stride must match the kernel's BR64
523    // or query rows 32..63 of every 64-row band are dropped (same class as the
524    // non-paged prefill_attention_64 fix). cfg!(atlas_scale) = strix+strix-hip;
525    // NVIDIA keeps 64 byte-identical.
526    let br = if cfg!(atlas_scale) { 32u32 } else { 64u32 };
527    KernelLaunch::new(gpu, kernel)
528        .grid([num_q_heads, div_ceil(q_len, br), 1])
529        .block([256, 1, 1])
530        .arg_ptr(q)
531        .arg_ptr(k_cache)
532        .arg_ptr(v_cache)
533        .arg_ptr(output)
534        .arg_ptr(block_table)
535        .arg_u32(q_len)
536        .arg_u32(kv_len)
537        .arg_u32(q_offset)
538        .arg_u32(num_q_heads)
539        .arg_u32(num_kv_heads)
540        .arg_u32(head_dim)
541        .arg_u32(cache_block_size)
542        .arg_u32(sliding_window)
543        .arg_u32(1u32)
544        .arg_f32(inv_sqrt_d)
545        .arg_u64(v_block_stride_bytes)
546        .arg_u64(v_data_section_bytes)
547        .launch(stream)
548}
spark_model/layers/ops/prefill_attn_main_b.rs

spark_model/layers/ops/
prefill_attn_main_b.rs
