metrale_model_layers/layers/ops/
kv_cache_turbok.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Cache-write and paged-decode launchers for the KV cache with TurboQuant on both sides: Turbo4K/Turbo3V, Turbo4K/Turbo8V and Turbo3K/Turbo8V.
4//!
5//! Owner: model-layers ops.
6//! Invariants:
7//! - Both pools are byte-addressed blocks with a data section followed by a scale section, each
8//!   with its own layout, so every launcher takes `k_block_stride_bytes`, `k_data_section_bytes`,
9//!   `v_block_stride_bytes` and `v_data_section_bytes`.
10//! - Turbo3 and Turbo4 store Lloyd-Max codes with one FP8 scale per 16 elements; Turbo8 stores
11//!   FP8 E4M3 with a BF16 group scale (kernels/gb10/common/reshape_and_cache_turbo.cu).
12//! - Cache writes run one 256-thread block per token; decode runs one 256-thread block per
13//!   (query head, sequence).
14
15#![allow(unused_imports)]
16
17use anyhow::Result;
18use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
19use metrale_gpu_runtime::kernel_args::KernelLaunch;
20
21/// 2026-09-25: Write K and V to the Turbo4K / Turbo3V cache (kernel
22/// `reshape_and_cache_flash_turbo4k_turbo3v`).
23#[allow(clippy::too_many_arguments)]
24pub fn reshape_and_cache_turbo4k_turbo3v(
25    gpu: &dyn GpuBackend,
26    kernel: KernelHandle,
27    key: DevicePtr,
28    value: DevicePtr,
29    k_cache: DevicePtr,
30    v_cache: DevicePtr,
31    slot_mapping: DevicePtr,
32    num_tokens: u32,
33    num_kv_heads: u32,
34    head_dim: u32,
35    block_size: u32,
36    key_stride: u32,
37    value_stride: u32,
38    k_block_stride_bytes: u64,
39    k_data_section_bytes: u64,
40    v_block_stride_bytes: u64,
41    v_data_section_bytes: u64,
42    stream: u64,
43) -> Result<()> {
44    KernelLaunch::new(gpu, kernel)
45        .grid([num_tokens, 1, 1])
46        .block([256, 1, 1])
47        .arg_ptr(key)
48        .arg_ptr(value)
49        .arg_ptr(k_cache)
50        .arg_ptr(v_cache)
51        .arg_ptr(slot_mapping)
52        .arg_u32(num_kv_heads)
53        .arg_u32(head_dim)
54        .arg_u32(block_size)
55        .arg_u32(key_stride)
56        .arg_u32(value_stride)
57        .arg_u64(k_block_stride_bytes)
58        .arg_u64(k_data_section_bytes)
59        .arg_u64(v_block_stride_bytes)
60        .arg_u64(v_data_section_bytes)
61        .launch(stream)
62}
63
64/// 2026-09-25: Write K and V to the Turbo4K / Turbo8V cache (kernel
65/// `reshape_and_cache_flash_turbo4k_turbo8v`).
66#[allow(clippy::too_many_arguments)]
67pub fn reshape_and_cache_turbo4k_turbo8v(
68    gpu: &dyn GpuBackend,
69    kernel: KernelHandle,
70    key: DevicePtr,
71    value: DevicePtr,
72    k_cache: DevicePtr,
73    v_cache: DevicePtr,
74    slot_mapping: DevicePtr,
75    num_tokens: u32,
76    num_kv_heads: u32,
77    head_dim: u32,
78    block_size: u32,
79    key_stride: u32,
80    value_stride: u32,
81    k_block_stride_bytes: u64,
82    k_data_section_bytes: u64,
83    v_block_stride_bytes: u64,
84    v_data_section_bytes: u64,
85    stream: u64,
86) -> Result<()> {
87    KernelLaunch::new(gpu, kernel)
88        .grid([num_tokens, 1, 1])
89        .block([256, 1, 1])
90        .arg_ptr(key)
91        .arg_ptr(value)
92        .arg_ptr(k_cache)
93        .arg_ptr(v_cache)
94        .arg_ptr(slot_mapping)
95        .arg_u32(num_kv_heads)
96        .arg_u32(head_dim)
97        .arg_u32(block_size)
98        .arg_u32(key_stride)
99        .arg_u32(value_stride)
100        .arg_u64(k_block_stride_bytes)
101        .arg_u64(k_data_section_bytes)
102        .arg_u64(v_block_stride_bytes)
103        .arg_u64(v_data_section_bytes)
104        .launch(stream)
105}
106
107/// 2026-09-25: Write K and V to the Turbo3K / Turbo8V cache (kernel
108/// `reshape_and_cache_flash_turbo3k_turbo8v`).
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/// 2026-09-25: Paged decode attention over the Turbo4K / Turbo3V cache (kernel
151/// `paged_decode_attn_turbo4k_turbo3v`).
152#[allow(clippy::too_many_arguments)]
153pub fn paged_decode_attn_turbo4k_turbo3v(
154    gpu: &dyn GpuBackend,
155    kernel: KernelHandle,
156    q: DevicePtr,
157    k_cache: DevicePtr,
158    v_cache: DevicePtr,
159    output: DevicePtr,
160    block_tables: DevicePtr,
161    seq_lens: DevicePtr,
162    max_blocks_per_seq: u32,
163    num_seqs: u32,
164    num_q_heads: u32,
165    num_kv_heads: u32,
166    head_dim: u32,
167    block_size: u32,
168    inv_sqrt_d: f32,
169    q_stride: u32,
170    k_block_stride_bytes: u64,
171    k_data_section_bytes: u64,
172    v_block_stride_bytes: u64,
173    v_data_section_bytes: u64,
174    sliding_window: u32,
175    stream: u64,
176) -> Result<()> {
177    KernelLaunch::new(gpu, kernel)
178        .grid([num_q_heads, num_seqs, 1])
179        .block([256, 1, 1])
180        .arg_ptr(q)
181        .arg_ptr(k_cache)
182        .arg_ptr(v_cache)
183        .arg_ptr(output)
184        .arg_ptr(block_tables)
185        .arg_ptr(seq_lens)
186        .arg_u32(max_blocks_per_seq)
187        .arg_u32(num_q_heads)
188        .arg_u32(num_kv_heads)
189        .arg_u32(head_dim)
190        .arg_u32(block_size)
191        .arg_f32(inv_sqrt_d)
192        .arg_u32(q_stride)
193        .arg_u64(k_block_stride_bytes)
194        .arg_u64(k_data_section_bytes)
195        .arg_u64(v_block_stride_bytes)
196        .arg_u64(v_data_section_bytes)
197        .arg_u32(sliding_window)
198        .launch(stream)
199}
200
201/// 2026-09-25: Paged decode attention over the Turbo4K / Turbo8V cache (kernel
202/// `paged_decode_attn_turbo4k_turbo8v`).
203#[allow(clippy::too_many_arguments)]
204pub fn paged_decode_attn_turbo4k_turbo8v(
205    gpu: &dyn GpuBackend,
206    kernel: KernelHandle,
207    q: DevicePtr,
208    k_cache: DevicePtr,
209    v_cache: DevicePtr,
210    output: DevicePtr,
211    block_tables: DevicePtr,
212    seq_lens: DevicePtr,
213    max_blocks_per_seq: u32,
214    num_seqs: u32,
215    num_q_heads: u32,
216    num_kv_heads: u32,
217    head_dim: u32,
218    block_size: u32,
219    inv_sqrt_d: f32,
220    q_stride: u32,
221    k_block_stride_bytes: u64,
222    k_data_section_bytes: u64,
223    v_block_stride_bytes: u64,
224    v_data_section_bytes: u64,
225    sliding_window: u32,
226    stream: u64,
227) -> Result<()> {
228    KernelLaunch::new(gpu, kernel)
229        .grid([num_q_heads, num_seqs, 1])
230        .block([256, 1, 1])
231        .arg_ptr(q)
232        .arg_ptr(k_cache)
233        .arg_ptr(v_cache)
234        .arg_ptr(output)
235        .arg_ptr(block_tables)
236        .arg_ptr(seq_lens)
237        .arg_u32(max_blocks_per_seq)
238        .arg_u32(num_q_heads)
239        .arg_u32(num_kv_heads)
240        .arg_u32(head_dim)
241        .arg_u32(block_size)
242        .arg_f32(inv_sqrt_d)
243        .arg_u32(q_stride)
244        .arg_u64(k_block_stride_bytes)
245        .arg_u64(k_data_section_bytes)
246        .arg_u64(v_block_stride_bytes)
247        .arg_u64(v_data_section_bytes)
248        .arg_u32(sliding_window)
249        .launch(stream)
250}
251
252/// 2026-09-25: Paged decode attention over the Turbo3K / Turbo8V cache (kernel
253/// `paged_decode_attn_turbo3k_turbo8v`).
254#[allow(clippy::too_many_arguments)]
255pub fn paged_decode_attn_turbo3k_turbo8v(
256    gpu: &dyn GpuBackend,
257    kernel: KernelHandle,
258    q: DevicePtr,
259    k_cache: DevicePtr,
260    v_cache: DevicePtr,
261    output: DevicePtr,
262    block_tables: DevicePtr,
263    seq_lens: DevicePtr,
264    max_blocks_per_seq: u32,
265    num_seqs: u32,
266    num_q_heads: u32,
267    num_kv_heads: u32,
268    head_dim: u32,
269    block_size: u32,
270    inv_sqrt_d: f32,
271    q_stride: u32,
272    k_block_stride_bytes: u64,
273    k_data_section_bytes: u64,
274    v_block_stride_bytes: u64,
275    v_data_section_bytes: u64,
276    sliding_window: u32,
277    stream: u64,
278) -> Result<()> {
279    KernelLaunch::new(gpu, kernel)
280        .grid([num_q_heads, num_seqs, 1])
281        .block([256, 1, 1])
282        .arg_ptr(q)
283        .arg_ptr(k_cache)
284        .arg_ptr(v_cache)
285        .arg_ptr(output)
286        .arg_ptr(block_tables)
287        .arg_ptr(seq_lens)
288        .arg_u32(max_blocks_per_seq)
289        .arg_u32(num_q_heads)
290        .arg_u32(num_kv_heads)
291        .arg_u32(head_dim)
292        .arg_u32(block_size)
293        .arg_f32(inv_sqrt_d)
294        .arg_u32(q_stride)
295        .arg_u64(k_block_stride_bytes)
296        .arg_u64(k_data_section_bytes)
297        .arg_u64(v_block_stride_bytes)
298        .arg_u64(v_data_section_bytes)
299        .arg_u32(sliding_window)
300        .launch(stream)
301}
metrale_model_layers/layers/ops/kv_cache_turbok.rs

metrale_model_layers/layers/ops/
kv_cache_turbok.rs
