spark_model/layers/ops/
kv_cache_fp8k.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! TurboQuant+ asymmetric Fp8K + TurboNV KV-cache ops wrappers (write + decode).
4//!
5//! Sibling of `kv_cache.rs`: keeps the bf16k_* paths there from growing past
6//! the 500-LoC cap. Each wrapper mirrors the corresponding bf16k_turbo*v
7//! function plus a `k_scale` parameter threaded into the kernel call.
8
9#![allow(unused_imports)]
10
11use anyhow::Result;
12use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
13use spark_runtime::kernel_args::KernelLaunch;
14
15/// Write K/V to paged Fp8K + Turbo3V (TurboQuant+ asym) cache.
16///
17/// K is written as FP8 E4M3 (per-tensor `k_scale`, NHD contiguous),
18/// V as 3-bit Lloyd-Max + FP8 per-group scale with matched-norm correction.
19/// K and V pools have separate strides (K 1 b/elem; V ~0.5 b/elem + scale).
20///
21/// Kernel: `reshape_and_cache_flash_fp8k_turbo3v(key, value, k_cache, v_cache,
22///          slot_mapping, num_kv_heads, head_dim, block_size,
23///          key_stride, value_stride, k_scale, k_block_stride_bytes,
24///          v_block_stride_bytes, v_data_section_bytes)`
25/// Grid: (num_tokens, 1, 1)  Block: (256, 1, 1)
26#[allow(clippy::too_many_arguments)]
27pub fn reshape_and_cache_fp8k_turbo3v(
28    gpu: &dyn GpuBackend,
29    kernel: KernelHandle,
30    key: DevicePtr,
31    value: DevicePtr,
32    k_cache: DevicePtr,
33    v_cache: DevicePtr,
34    slot_mapping: DevicePtr,
35    num_tokens: u32,
36    num_kv_heads: u32,
37    head_dim: u32,
38    block_size: u32,
39    key_stride: u32,
40    value_stride: u32,
41    k_scale: f32,
42    k_block_stride_bytes: u64,
43    v_block_stride_bytes: u64,
44    v_data_section_bytes: u64,
45    stream: u64,
46) -> Result<()> {
47    KernelLaunch::new(gpu, kernel)
48        .grid([num_tokens, 1, 1])
49        .block([256, 1, 1])
50        .arg_ptr(key)
51        .arg_ptr(value)
52        .arg_ptr(k_cache)
53        .arg_ptr(v_cache)
54        .arg_ptr(slot_mapping)
55        .arg_u32(num_kv_heads)
56        .arg_u32(head_dim)
57        .arg_u32(block_size)
58        .arg_u32(key_stride)
59        .arg_u32(value_stride)
60        .arg_f32(k_scale)
61        .arg_u64(k_block_stride_bytes)
62        .arg_u64(v_block_stride_bytes)
63        .arg_u64(v_data_section_bytes)
64        .launch(stream)
65}
66
67/// Write K/V to paged Fp8K + Turbo4V (TurboQuant+ asym) cache.
68#[allow(clippy::too_many_arguments)]
69pub fn reshape_and_cache_fp8k_turbo4v(
70    gpu: &dyn GpuBackend,
71    kernel: KernelHandle,
72    key: DevicePtr,
73    value: DevicePtr,
74    k_cache: DevicePtr,
75    v_cache: DevicePtr,
76    slot_mapping: DevicePtr,
77    num_tokens: u32,
78    num_kv_heads: u32,
79    head_dim: u32,
80    block_size: u32,
81    key_stride: u32,
82    value_stride: u32,
83    k_scale: f32,
84    k_block_stride_bytes: u64,
85    v_block_stride_bytes: u64,
86    v_data_section_bytes: u64,
87    stream: u64,
88) -> Result<()> {
89    KernelLaunch::new(gpu, kernel)
90        .grid([num_tokens, 1, 1])
91        .block([256, 1, 1])
92        .arg_ptr(key)
93        .arg_ptr(value)
94        .arg_ptr(k_cache)
95        .arg_ptr(v_cache)
96        .arg_ptr(slot_mapping)
97        .arg_u32(num_kv_heads)
98        .arg_u32(head_dim)
99        .arg_u32(block_size)
100        .arg_u32(key_stride)
101        .arg_u32(value_stride)
102        .arg_f32(k_scale)
103        .arg_u64(k_block_stride_bytes)
104        .arg_u64(v_block_stride_bytes)
105        .arg_u64(v_data_section_bytes)
106        .launch(stream)
107}
108
109/// Write K/V to paged Fp8K + Turbo2V (TurboQuant+ asym) cache (6.4x V comp).
110#[allow(clippy::too_many_arguments)]
111pub fn reshape_and_cache_fp8k_turbo2v(
112    gpu: &dyn GpuBackend,
113    kernel: KernelHandle,
114    key: DevicePtr,
115    value: DevicePtr,
116    k_cache: DevicePtr,
117    v_cache: DevicePtr,
118    slot_mapping: DevicePtr,
119    num_tokens: u32,
120    num_kv_heads: u32,
121    head_dim: u32,
122    block_size: u32,
123    key_stride: u32,
124    value_stride: u32,
125    k_scale: f32,
126    k_block_stride_bytes: u64,
127    v_block_stride_bytes: u64,
128    v_data_section_bytes: u64,
129    stream: u64,
130) -> Result<()> {
131    KernelLaunch::new(gpu, kernel)
132        .grid([num_tokens, 1, 1])
133        .block([256, 1, 1])
134        .arg_ptr(key)
135        .arg_ptr(value)
136        .arg_ptr(k_cache)
137        .arg_ptr(v_cache)
138        .arg_ptr(slot_mapping)
139        .arg_u32(num_kv_heads)
140        .arg_u32(head_dim)
141        .arg_u32(block_size)
142        .arg_u32(key_stride)
143        .arg_u32(value_stride)
144        .arg_f32(k_scale)
145        .arg_u64(k_block_stride_bytes)
146        .arg_u64(v_block_stride_bytes)
147        .arg_u64(v_data_section_bytes)
148        .launch(stream)
149}
150
151/// Paged decode attention for Fp8K + Turbo3V asymmetric KV cache.
152///
153/// K is read as FP8 NHD with per-tensor `k_scale` dequant, V as 3-bit
154/// Lloyd-Max packed bytes + FP8 per-group scale (sparse-V threshold on
155/// batched + remainder paths).
156///
157/// Kernel: `paged_decode_attn_fp8k_turbo3v(Q, K_cache, V_cache, O,
158///          block_tables, seq_lens, max_blocks_per_seq, num_q_heads,
159///          num_kv_heads, head_dim, block_size, inv_sqrt_d, k_scale,
160///          q_stride, v_block_stride_bytes, v_data_section_bytes,
161///          sliding_window)`
162/// Grid: (num_q_heads, num_seqs, 1)  Block: (256, 1, 1)
163#[allow(clippy::too_many_arguments)]
164pub fn paged_decode_attn_fp8k_turbo3v(
165    gpu: &dyn GpuBackend,
166    kernel: KernelHandle,
167    q: DevicePtr,
168    k_cache: DevicePtr,
169    v_cache: DevicePtr,
170    output: DevicePtr,
171    block_tables: DevicePtr,
172    seq_lens: DevicePtr,
173    max_blocks_per_seq: u32,
174    num_seqs: u32,
175    num_q_heads: u32,
176    num_kv_heads: u32,
177    head_dim: u32,
178    block_size: u32,
179    inv_sqrt_d: f32,
180    k_scale: f32,
181    q_stride: u32,
182    v_block_stride_bytes: u64,
183    v_data_section_bytes: u64,
184    sliding_window: u32,
185    stream: u64,
186) -> Result<()> {
187    KernelLaunch::new(gpu, kernel)
188        .grid([num_q_heads, num_seqs, 1])
189        .block([256, 1, 1])
190        .arg_ptr(q)
191        .arg_ptr(k_cache)
192        .arg_ptr(v_cache)
193        .arg_ptr(output)
194        .arg_ptr(block_tables)
195        .arg_ptr(seq_lens)
196        .arg_u32(max_blocks_per_seq)
197        .arg_u32(num_q_heads)
198        .arg_u32(num_kv_heads)
199        .arg_u32(head_dim)
200        .arg_u32(block_size)
201        .arg_f32(inv_sqrt_d)
202        .arg_f32(k_scale)
203        .arg_u32(q_stride)
204        .arg_u64(v_block_stride_bytes)
205        .arg_u64(v_data_section_bytes)
206        .arg_u32(sliding_window)
207        .launch(stream)
208}
209
210/// Paged decode attention for Fp8K + Turbo4V asymmetric KV cache.
211#[allow(clippy::too_many_arguments)]
212pub fn paged_decode_attn_fp8k_turbo4v(
213    gpu: &dyn GpuBackend,
214    kernel: KernelHandle,
215    q: DevicePtr,
216    k_cache: DevicePtr,
217    v_cache: DevicePtr,
218    output: DevicePtr,
219    block_tables: DevicePtr,
220    seq_lens: DevicePtr,
221    max_blocks_per_seq: u32,
222    num_seqs: u32,
223    num_q_heads: u32,
224    num_kv_heads: u32,
225    head_dim: u32,
226    block_size: u32,
227    inv_sqrt_d: f32,
228    k_scale: f32,
229    q_stride: u32,
230    v_block_stride_bytes: u64,
231    v_data_section_bytes: u64,
232    sliding_window: u32,
233    stream: u64,
234) -> Result<()> {
235    KernelLaunch::new(gpu, kernel)
236        .grid([num_q_heads, num_seqs, 1])
237        .block([256, 1, 1])
238        .arg_ptr(q)
239        .arg_ptr(k_cache)
240        .arg_ptr(v_cache)
241        .arg_ptr(output)
242        .arg_ptr(block_tables)
243        .arg_ptr(seq_lens)
244        .arg_u32(max_blocks_per_seq)
245        .arg_u32(num_q_heads)
246        .arg_u32(num_kv_heads)
247        .arg_u32(head_dim)
248        .arg_u32(block_size)
249        .arg_f32(inv_sqrt_d)
250        .arg_f32(k_scale)
251        .arg_u32(q_stride)
252        .arg_u64(v_block_stride_bytes)
253        .arg_u64(v_data_section_bytes)
254        .arg_u32(sliding_window)
255        .launch(stream)
256}
257
258/// Paged decode attention for Fp8K + Turbo2V asymmetric KV cache (6.4x V comp).
259#[allow(clippy::too_many_arguments)]
260pub fn paged_decode_attn_fp8k_turbo2v(
261    gpu: &dyn GpuBackend,
262    kernel: KernelHandle,
263    q: DevicePtr,
264    k_cache: DevicePtr,
265    v_cache: DevicePtr,
266    output: DevicePtr,
267    block_tables: DevicePtr,
268    seq_lens: DevicePtr,
269    max_blocks_per_seq: u32,
270    num_seqs: u32,
271    num_q_heads: u32,
272    num_kv_heads: u32,
273    head_dim: u32,
274    block_size: u32,
275    inv_sqrt_d: f32,
276    k_scale: f32,
277    q_stride: u32,
278    v_block_stride_bytes: u64,
279    v_data_section_bytes: u64,
280    sliding_window: u32,
281    stream: u64,
282) -> Result<()> {
283    KernelLaunch::new(gpu, kernel)
284        .grid([num_q_heads, num_seqs, 1])
285        .block([256, 1, 1])
286        .arg_ptr(q)
287        .arg_ptr(k_cache)
288        .arg_ptr(v_cache)
289        .arg_ptr(output)
290        .arg_ptr(block_tables)
291        .arg_ptr(seq_lens)
292        .arg_u32(max_blocks_per_seq)
293        .arg_u32(num_q_heads)
294        .arg_u32(num_kv_heads)
295        .arg_u32(head_dim)
296        .arg_u32(block_size)
297        .arg_f32(inv_sqrt_d)
298        .arg_f32(k_scale)
299        .arg_u32(q_stride)
300        .arg_u64(v_block_stride_bytes)
301        .arg_u64(v_data_section_bytes)
302        .arg_u32(sliding_window)
303        .launch(stream)
304}
spark_model/layers/ops/kv_cache_fp8k.rs

spark_model/layers/ops/
kv_cache_fp8k.rs
