spark_model/layers/ops/
kv_cache_turbok.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! TurboQuant+ both-sides-quantized asymmetric Turbo*K + Turbo*V KV-cache
4//! ops wrappers (write + decode). K and V both use turbo (byte-addressed,
5//! data + scale section) layouts but with potentially different K-side and
6//! V-side dtypes — three combos: turbo4k_turbo3v, turbo4k_turbo8v,
7//! turbo3k_turbo8v.
8//!
9//! Sibling of `kv_cache.rs` / `kv_cache_fp8k.rs`. Each wrapper takes BOTH
10//! K-side (k_block_stride_bytes, k_data_section_bytes) AND V-side
11//! (v_block_stride_bytes, v_data_section_bytes) parameters since each pool
12//! has its own byte layout.
13
14#![allow(unused_imports)]
15
16use anyhow::Result;
17use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
18use spark_runtime::kernel_args::KernelLaunch;
19
20/// Write K/V to paged Turbo4K + Turbo3V cache.
21///
22/// K written as turbo4 (4-bit packed + FP8 group scale, matched-norm L2)
23/// and V as turbo3 (3-bit packed + FP8 group scale, matched-norm L2). The
24/// pools have distinct byte strides — passed independently.
25#[allow(clippy::too_many_arguments)]
26pub fn reshape_and_cache_turbo4k_turbo3v(
27    gpu: &dyn GpuBackend,
28    kernel: KernelHandle,
29    key: DevicePtr,
30    value: DevicePtr,
31    k_cache: DevicePtr,
32    v_cache: DevicePtr,
33    slot_mapping: DevicePtr,
34    num_tokens: u32,
35    num_kv_heads: u32,
36    head_dim: u32,
37    block_size: u32,
38    key_stride: u32,
39    value_stride: u32,
40    k_block_stride_bytes: u64,
41    k_data_section_bytes: u64,
42    v_block_stride_bytes: u64,
43    v_data_section_bytes: u64,
44    stream: u64,
45) -> Result<()> {
46    KernelLaunch::new(gpu, kernel)
47        .grid([num_tokens, 1, 1])
48        .block([256, 1, 1])
49        .arg_ptr(key)
50        .arg_ptr(value)
51        .arg_ptr(k_cache)
52        .arg_ptr(v_cache)
53        .arg_ptr(slot_mapping)
54        .arg_u32(num_kv_heads)
55        .arg_u32(head_dim)
56        .arg_u32(block_size)
57        .arg_u32(key_stride)
58        .arg_u32(value_stride)
59        .arg_u64(k_block_stride_bytes)
60        .arg_u64(k_data_section_bytes)
61        .arg_u64(v_block_stride_bytes)
62        .arg_u64(v_data_section_bytes)
63        .launch(stream)
64}
65
66/// Write K/V to paged Turbo4K + Turbo8V cache.
67#[allow(clippy::too_many_arguments)]
68pub fn reshape_and_cache_turbo4k_turbo8v(
69    gpu: &dyn GpuBackend,
70    kernel: KernelHandle,
71    key: DevicePtr,
72    value: DevicePtr,
73    k_cache: DevicePtr,
74    v_cache: DevicePtr,
75    slot_mapping: DevicePtr,
76    num_tokens: u32,
77    num_kv_heads: u32,
78    head_dim: u32,
79    block_size: u32,
80    key_stride: u32,
81    value_stride: u32,
82    k_block_stride_bytes: u64,
83    k_data_section_bytes: u64,
84    v_block_stride_bytes: u64,
85    v_data_section_bytes: u64,
86    stream: u64,
87) -> Result<()> {
88    KernelLaunch::new(gpu, kernel)
89        .grid([num_tokens, 1, 1])
90        .block([256, 1, 1])
91        .arg_ptr(key)
92        .arg_ptr(value)
93        .arg_ptr(k_cache)
94        .arg_ptr(v_cache)
95        .arg_ptr(slot_mapping)
96        .arg_u32(num_kv_heads)
97        .arg_u32(head_dim)
98        .arg_u32(block_size)
99        .arg_u32(key_stride)
100        .arg_u32(value_stride)
101        .arg_u64(k_block_stride_bytes)
102        .arg_u64(k_data_section_bytes)
103        .arg_u64(v_block_stride_bytes)
104        .arg_u64(v_data_section_bytes)
105        .launch(stream)
106}
107
108/// Write K/V to paged Turbo3K + Turbo8V cache.
109#[allow(clippy::too_many_arguments)]
110pub fn reshape_and_cache_turbo3k_turbo8v(
111    gpu: &dyn GpuBackend,
112    kernel: KernelHandle,
113    key: DevicePtr,
114    value: DevicePtr,
115    k_cache: DevicePtr,
116    v_cache: DevicePtr,
117    slot_mapping: DevicePtr,
118    num_tokens: u32,
119    num_kv_heads: u32,
120    head_dim: u32,
121    block_size: u32,
122    key_stride: u32,
123    value_stride: u32,
124    k_block_stride_bytes: u64,
125    k_data_section_bytes: u64,
126    v_block_stride_bytes: u64,
127    v_data_section_bytes: u64,
128    stream: u64,
129) -> Result<()> {
130    KernelLaunch::new(gpu, kernel)
131        .grid([num_tokens, 1, 1])
132        .block([256, 1, 1])
133        .arg_ptr(key)
134        .arg_ptr(value)
135        .arg_ptr(k_cache)
136        .arg_ptr(v_cache)
137        .arg_ptr(slot_mapping)
138        .arg_u32(num_kv_heads)
139        .arg_u32(head_dim)
140        .arg_u32(block_size)
141        .arg_u32(key_stride)
142        .arg_u32(value_stride)
143        .arg_u64(k_block_stride_bytes)
144        .arg_u64(k_data_section_bytes)
145        .arg_u64(v_block_stride_bytes)
146        .arg_u64(v_data_section_bytes)
147        .launch(stream)
148}
149
150/// Paged decode attention for Turbo4K + Turbo3V asymmetric KV cache.
151#[allow(clippy::too_many_arguments)]
152pub fn paged_decode_attn_turbo4k_turbo3v(
153    gpu: &dyn GpuBackend,
154    kernel: KernelHandle,
155    q: DevicePtr,
156    k_cache: DevicePtr,
157    v_cache: DevicePtr,
158    output: DevicePtr,
159    block_tables: DevicePtr,
160    seq_lens: DevicePtr,
161    max_blocks_per_seq: u32,
162    num_seqs: u32,
163    num_q_heads: u32,
164    num_kv_heads: u32,
165    head_dim: u32,
166    block_size: u32,
167    inv_sqrt_d: f32,
168    q_stride: u32,
169    k_block_stride_bytes: u64,
170    k_data_section_bytes: u64,
171    v_block_stride_bytes: u64,
172    v_data_section_bytes: u64,
173    sliding_window: u32,
174    stream: u64,
175) -> Result<()> {
176    KernelLaunch::new(gpu, kernel)
177        .grid([num_q_heads, num_seqs, 1])
178        .block([256, 1, 1])
179        .arg_ptr(q)
180        .arg_ptr(k_cache)
181        .arg_ptr(v_cache)
182        .arg_ptr(output)
183        .arg_ptr(block_tables)
184        .arg_ptr(seq_lens)
185        .arg_u32(max_blocks_per_seq)
186        .arg_u32(num_q_heads)
187        .arg_u32(num_kv_heads)
188        .arg_u32(head_dim)
189        .arg_u32(block_size)
190        .arg_f32(inv_sqrt_d)
191        .arg_u32(q_stride)
192        .arg_u64(k_block_stride_bytes)
193        .arg_u64(k_data_section_bytes)
194        .arg_u64(v_block_stride_bytes)
195        .arg_u64(v_data_section_bytes)
196        .arg_u32(sliding_window)
197        .launch(stream)
198}
199
200/// Paged decode attention for Turbo4K + Turbo8V asymmetric KV cache.
201#[allow(clippy::too_many_arguments)]
202pub fn paged_decode_attn_turbo4k_turbo8v(
203    gpu: &dyn GpuBackend,
204    kernel: KernelHandle,
205    q: DevicePtr,
206    k_cache: DevicePtr,
207    v_cache: DevicePtr,
208    output: DevicePtr,
209    block_tables: DevicePtr,
210    seq_lens: DevicePtr,
211    max_blocks_per_seq: u32,
212    num_seqs: u32,
213    num_q_heads: u32,
214    num_kv_heads: u32,
215    head_dim: u32,
216    block_size: u32,
217    inv_sqrt_d: f32,
218    q_stride: u32,
219    k_block_stride_bytes: u64,
220    k_data_section_bytes: u64,
221    v_block_stride_bytes: u64,
222    v_data_section_bytes: u64,
223    sliding_window: u32,
224    stream: u64,
225) -> Result<()> {
226    KernelLaunch::new(gpu, kernel)
227        .grid([num_q_heads, num_seqs, 1])
228        .block([256, 1, 1])
229        .arg_ptr(q)
230        .arg_ptr(k_cache)
231        .arg_ptr(v_cache)
232        .arg_ptr(output)
233        .arg_ptr(block_tables)
234        .arg_ptr(seq_lens)
235        .arg_u32(max_blocks_per_seq)
236        .arg_u32(num_q_heads)
237        .arg_u32(num_kv_heads)
238        .arg_u32(head_dim)
239        .arg_u32(block_size)
240        .arg_f32(inv_sqrt_d)
241        .arg_u32(q_stride)
242        .arg_u64(k_block_stride_bytes)
243        .arg_u64(k_data_section_bytes)
244        .arg_u64(v_block_stride_bytes)
245        .arg_u64(v_data_section_bytes)
246        .arg_u32(sliding_window)
247        .launch(stream)
248}
249
250/// Paged decode attention for Turbo3K + Turbo8V asymmetric KV cache.
251#[allow(clippy::too_many_arguments)]
252pub fn paged_decode_attn_turbo3k_turbo8v(
253    gpu: &dyn GpuBackend,
254    kernel: KernelHandle,
255    q: DevicePtr,
256    k_cache: DevicePtr,
257    v_cache: DevicePtr,
258    output: DevicePtr,
259    block_tables: DevicePtr,
260    seq_lens: DevicePtr,
261    max_blocks_per_seq: u32,
262    num_seqs: u32,
263    num_q_heads: u32,
264    num_kv_heads: u32,
265    head_dim: u32,
266    block_size: u32,
267    inv_sqrt_d: f32,
268    q_stride: u32,
269    k_block_stride_bytes: u64,
270    k_data_section_bytes: u64,
271    v_block_stride_bytes: u64,
272    v_data_section_bytes: u64,
273    sliding_window: u32,
274    stream: u64,
275) -> Result<()> {
276    KernelLaunch::new(gpu, kernel)
277        .grid([num_q_heads, num_seqs, 1])
278        .block([256, 1, 1])
279        .arg_ptr(q)
280        .arg_ptr(k_cache)
281        .arg_ptr(v_cache)
282        .arg_ptr(output)
283        .arg_ptr(block_tables)
284        .arg_ptr(seq_lens)
285        .arg_u32(max_blocks_per_seq)
286        .arg_u32(num_q_heads)
287        .arg_u32(num_kv_heads)
288        .arg_u32(head_dim)
289        .arg_u32(block_size)
290        .arg_f32(inv_sqrt_d)
291        .arg_u32(q_stride)
292        .arg_u64(k_block_stride_bytes)
293        .arg_u64(k_data_section_bytes)
294        .arg_u64(v_block_stride_bytes)
295        .arg_u64(v_data_section_bytes)
296        .arg_u32(sliding_window)
297        .launch(stream)
298}
spark_model/layers/ops/kv_cache_turbok.rs

spark_model/layers/ops/
kv_cache_turbok.rs
