spark_model/layers/ops/
prefill_attn_fp8k.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! TurboQuant+ asymmetric Fp8K + TurboNV prefill (BR=64) ops wrappers.
4//!
5//! Sibling of `prefill_attn_main_b.rs`: keeps the bf16k_* prefill wrappers
6//! there from growing past the 500-LoC cap. Each wrapper mirrors the
7//! corresponding bf16k_turbo*v_64 function plus a `k_scale` parameter
8//! threaded into the kernel call (FP8 K dequant scale).
9
10#![allow(unused_imports)]
11
12use anyhow::Result;
13use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
14use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
15
16/// Prefill paged attention — TurboQuant+ asym Fp8K + Turbo3V (BR=64).
17///
18/// K is read as FP8 (per-tensor `k_scale` dequant in LOAD_K_TILE),
19/// V as turbo3 (3-bit packed + FP8 group scale).
20///
21/// Kernel: `inferspark_prefill_paged_fp8k_turbo3v_64(Q, K_cache, V_cache,
22///          O, block_table, q_len, kv_len, q_offset, num_q_heads,
23///          num_kv_heads, head_dim, cache_block_size, sliding_window,
24///          causal_mask_enabled, inv_sqrt_d, k_scale,
25///          v_block_stride_bytes, v_data_section_bytes)`
26/// Grid: (num_q_heads, div_ceil(q_len, BR), 1)  Block: (256, 1, 1)
27#[allow(clippy::too_many_arguments)]
28pub fn prefill_attention_paged_fp8k_turbo3v_64(
29    gpu: &dyn GpuBackend,
30    kernel: KernelHandle,
31    q: DevicePtr,
32    k_cache: DevicePtr,
33    v_cache: DevicePtr,
34    output: DevicePtr,
35    block_table: DevicePtr,
36    q_len: u32,
37    kv_len: u32,
38    q_offset: u32,
39    num_q_heads: u32,
40    num_kv_heads: u32,
41    head_dim: u32,
42    cache_block_size: u32,
43    sliding_window: u32,
44    inv_sqrt_d: f32,
45    k_scale: f32,
46    v_block_stride_bytes: u64,
47    v_data_section_bytes: u64,
48    stream: u64,
49) -> Result<()> {
50    let br = 64u32;
51    KernelLaunch::new(gpu, kernel)
52        .grid([num_q_heads, div_ceil(q_len, br), 1])
53        .block([256, 1, 1])
54        .arg_ptr(q)
55        .arg_ptr(k_cache)
56        .arg_ptr(v_cache)
57        .arg_ptr(output)
58        .arg_ptr(block_table)
59        .arg_u32(q_len)
60        .arg_u32(kv_len)
61        .arg_u32(q_offset)
62        .arg_u32(num_q_heads)
63        .arg_u32(num_kv_heads)
64        .arg_u32(head_dim)
65        .arg_u32(cache_block_size)
66        .arg_u32(sliding_window)
67        .arg_u32(1u32)
68        .arg_f32(inv_sqrt_d)
69        .arg_f32(k_scale)
70        .arg_u64(v_block_stride_bytes)
71        .arg_u64(v_data_section_bytes)
72        .launch(stream)
73}
74
75/// Prefill paged attention — TurboQuant+ asym Fp8K + Turbo4V (BR=64).
76#[allow(clippy::too_many_arguments)]
77pub fn prefill_attention_paged_fp8k_turbo4v_64(
78    gpu: &dyn GpuBackend,
79    kernel: KernelHandle,
80    q: DevicePtr,
81    k_cache: DevicePtr,
82    v_cache: DevicePtr,
83    output: DevicePtr,
84    block_table: DevicePtr,
85    q_len: u32,
86    kv_len: u32,
87    q_offset: u32,
88    num_q_heads: u32,
89    num_kv_heads: u32,
90    head_dim: u32,
91    cache_block_size: u32,
92    sliding_window: u32,
93    inv_sqrt_d: f32,
94    k_scale: f32,
95    v_block_stride_bytes: u64,
96    v_data_section_bytes: u64,
97    stream: u64,
98) -> Result<()> {
99    let br = 64u32;
100    KernelLaunch::new(gpu, kernel)
101        .grid([num_q_heads, div_ceil(q_len, br), 1])
102        .block([256, 1, 1])
103        .arg_ptr(q)
104        .arg_ptr(k_cache)
105        .arg_ptr(v_cache)
106        .arg_ptr(output)
107        .arg_ptr(block_table)
108        .arg_u32(q_len)
109        .arg_u32(kv_len)
110        .arg_u32(q_offset)
111        .arg_u32(num_q_heads)
112        .arg_u32(num_kv_heads)
113        .arg_u32(head_dim)
114        .arg_u32(cache_block_size)
115        .arg_u32(sliding_window)
116        .arg_u32(1u32)
117        .arg_f32(inv_sqrt_d)
118        .arg_f32(k_scale)
119        .arg_u64(v_block_stride_bytes)
120        .arg_u64(v_data_section_bytes)
121        .launch(stream)
122}
123
124/// Prefill paged attention — TurboQuant+ asym Fp8K + Turbo2V (BR=64, 6.4x V comp).
125#[allow(clippy::too_many_arguments)]
126pub fn prefill_attention_paged_fp8k_turbo2v_64(
127    gpu: &dyn GpuBackend,
128    kernel: KernelHandle,
129    q: DevicePtr,
130    k_cache: DevicePtr,
131    v_cache: DevicePtr,
132    output: DevicePtr,
133    block_table: DevicePtr,
134    q_len: u32,
135    kv_len: u32,
136    q_offset: u32,
137    num_q_heads: u32,
138    num_kv_heads: u32,
139    head_dim: u32,
140    cache_block_size: u32,
141    sliding_window: u32,
142    inv_sqrt_d: f32,
143    k_scale: f32,
144    v_block_stride_bytes: u64,
145    v_data_section_bytes: u64,
146    stream: u64,
147) -> Result<()> {
148    let br = 64u32;
149    KernelLaunch::new(gpu, kernel)
150        .grid([num_q_heads, div_ceil(q_len, br), 1])
151        .block([256, 1, 1])
152        .arg_ptr(q)
153        .arg_ptr(k_cache)
154        .arg_ptr(v_cache)
155        .arg_ptr(output)
156        .arg_ptr(block_table)
157        .arg_u32(q_len)
158        .arg_u32(kv_len)
159        .arg_u32(q_offset)
160        .arg_u32(num_q_heads)
161        .arg_u32(num_kv_heads)
162        .arg_u32(head_dim)
163        .arg_u32(cache_block_size)
164        .arg_u32(sliding_window)
165        .arg_u32(1u32)
166        .arg_f32(inv_sqrt_d)
167        .arg_f32(k_scale)
168        .arg_u64(v_block_stride_bytes)
169        .arg_u64(v_data_section_bytes)
170        .launch(stream)
171}
spark_model/layers/ops/prefill_attn_fp8k.rs

spark_model/layers/ops/
prefill_attn_fp8k.rs
