metrale_model_layers/layers/ops/
kv_cache_bf16k_turbov.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Cache-write and paged-decode launchers for the asymmetric KV cache with BF16 K and TurboQuant V (Turbo2V, Turbo3V, Turbo4V).
4//!
5//! Owner: model-layers ops.
6//! Invariants:
7//! - K is stored as raw BF16 in the `[num_blocks, block_size, num_kv_heads, head_dim]` layout; V
8//!   as Lloyd-Max codes (2, 3 or 4 bits per element) followed by one FP8 scale per 16 elements
9//!   (kernels/gb10/common/reshape_and_cache_turbo.cu). The two pools therefore have separate
10//!   block strides, passed in bytes.
11//! - Cache writes run one 256-thread block per token; decode runs one 256-thread block per
12//!   (query head, sequence). Each launcher passes its kernel's parameters in declaration order.
13
14use super::*;
15
16/// 2026-09-25: Write K and V to the BF16-K / Turbo3V cache (kernel
17/// `reshape_and_cache_flash_bf16k_turbo3v`): K as raw BF16, V as 3-bit Lloyd-Max codes (3/8 byte
18/// per element) plus one FP8 scale per 16 elements. `v_data_section_bytes` is the size of a V
19/// block's code section, after which its scales start.
20#[allow(clippy::too_many_arguments)]
21pub fn reshape_and_cache_bf16k_turbo3v(
22    gpu: &dyn GpuBackend,
23    kernel: KernelHandle,
24    key: DevicePtr,
25    value: DevicePtr,
26    k_cache: DevicePtr,
27    v_cache: DevicePtr,
28    slot_mapping: DevicePtr,
29    num_tokens: u32,
30    num_kv_heads: u32,
31    head_dim: u32,
32    block_size: u32,
33    key_stride: u32,
34    value_stride: u32,
35    k_block_stride_bytes: u64,
36    v_block_stride_bytes: u64,
37    v_data_section_bytes: u64,
38    stream: u64,
39) -> Result<()> {
40    KernelLaunch::new(gpu, kernel)
41        .grid([num_tokens, 1, 1])
42        .block([256, 1, 1])
43        .arg_ptr(key)
44        .arg_ptr(value)
45        .arg_ptr(k_cache)
46        .arg_ptr(v_cache)
47        .arg_ptr(slot_mapping)
48        .arg_u32(num_kv_heads)
49        .arg_u32(head_dim)
50        .arg_u32(block_size)
51        .arg_u32(key_stride)
52        .arg_u32(value_stride)
53        .arg_u64(k_block_stride_bytes)
54        .arg_u64(v_block_stride_bytes)
55        .arg_u64(v_data_section_bytes)
56        .launch(stream)
57}
58
59/// 2026-09-25: Paged decode attention over the BF16-K / Turbo3V cache (kernel
60/// `paged_decode_attn_bf16k_turbo3v`). The kernel skips a position's V load and dequant when its
61/// attention weight is at most `TQ_PLUS_SPARSE_V_THRESHOLD` (1e-3 by default).
62#[allow(clippy::too_many_arguments)]
63pub fn paged_decode_attn_bf16k_turbo3v(
64    gpu: &dyn GpuBackend,
65    kernel: KernelHandle,
66    q: DevicePtr,
67    k_cache: DevicePtr,
68    v_cache: DevicePtr,
69    output: DevicePtr,
70    block_tables: DevicePtr,
71    seq_lens: DevicePtr,
72    max_blocks_per_seq: u32,
73    num_seqs: u32,
74    num_q_heads: u32,
75    num_kv_heads: u32,
76    head_dim: u32,
77    block_size: u32,
78    inv_sqrt_d: f32,
79    q_stride: u32,
80    v_block_stride_bytes: u64,
81    v_data_section_bytes: u64,
82    sliding_window: u32,
83    stream: u64,
84) -> Result<()> {
85    KernelLaunch::new(gpu, kernel)
86        .grid([num_q_heads, num_seqs, 1])
87        .block([256, 1, 1])
88        .arg_ptr(q)
89        .arg_ptr(k_cache)
90        .arg_ptr(v_cache)
91        .arg_ptr(output)
92        .arg_ptr(block_tables)
93        .arg_ptr(seq_lens)
94        .arg_u32(max_blocks_per_seq)
95        .arg_u32(num_q_heads)
96        .arg_u32(num_kv_heads)
97        .arg_u32(head_dim)
98        .arg_u32(block_size)
99        .arg_f32(inv_sqrt_d)
100        .arg_u32(q_stride)
101        .arg_u64(v_block_stride_bytes)
102        .arg_u64(v_data_section_bytes)
103        .arg_u32(sliding_window)
104        .launch(stream)
105}
106
107/// 2026-09-25: Write K and V to the BF16-K / Turbo4V cache (kernel
108/// `reshape_and_cache_flash_bf16k_turbo4v`): K as raw BF16, V as 4-bit Lloyd-Max codes plus one
109/// FP8 scale per 16 elements.
110#[allow(clippy::too_many_arguments)]
111pub fn reshape_and_cache_bf16k_turbo4v(
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_block_stride_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(v_block_stride_bytes)
145        .arg_u64(v_data_section_bytes)
146        .launch(stream)
147}
148
149/// 2026-09-25: Write K and V to the BF16-K / Turbo2V cache (kernel
150/// `reshape_and_cache_flash_bf16k_turbo2v`): K as raw BF16, V as 2-bit Lloyd-Max codes plus one
151/// FP8 scale per 16 elements, 2.5 bits per V element.
152#[allow(clippy::too_many_arguments)]
153pub fn reshape_and_cache_bf16k_turbo2v(
154    gpu: &dyn GpuBackend,
155    kernel: KernelHandle,
156    key: DevicePtr,
157    value: DevicePtr,
158    k_cache: DevicePtr,
159    v_cache: DevicePtr,
160    slot_mapping: DevicePtr,
161    num_tokens: u32,
162    num_kv_heads: u32,
163    head_dim: u32,
164    block_size: u32,
165    key_stride: u32,
166    value_stride: u32,
167    k_block_stride_bytes: u64,
168    v_block_stride_bytes: u64,
169    v_data_section_bytes: u64,
170    stream: u64,
171) -> Result<()> {
172    KernelLaunch::new(gpu, kernel)
173        .grid([num_tokens, 1, 1])
174        .block([256, 1, 1])
175        .arg_ptr(key)
176        .arg_ptr(value)
177        .arg_ptr(k_cache)
178        .arg_ptr(v_cache)
179        .arg_ptr(slot_mapping)
180        .arg_u32(num_kv_heads)
181        .arg_u32(head_dim)
182        .arg_u32(block_size)
183        .arg_u32(key_stride)
184        .arg_u32(value_stride)
185        .arg_u64(k_block_stride_bytes)
186        .arg_u64(v_block_stride_bytes)
187        .arg_u64(v_data_section_bytes)
188        .launch(stream)
189}
190
191/// 2026-09-25: Paged decode attention over the BF16-K / Turbo4V cache (kernel
192/// `paged_decode_attn_bf16k_turbo4v`), with the same sparse-V skip as the Turbo3V kernel.
193#[allow(clippy::too_many_arguments)]
194pub fn paged_decode_attn_bf16k_turbo4v(
195    gpu: &dyn GpuBackend,
196    kernel: KernelHandle,
197    q: DevicePtr,
198    k_cache: DevicePtr,
199    v_cache: DevicePtr,
200    output: DevicePtr,
201    block_tables: DevicePtr,
202    seq_lens: DevicePtr,
203    max_blocks_per_seq: u32,
204    num_seqs: u32,
205    num_q_heads: u32,
206    num_kv_heads: u32,
207    head_dim: u32,
208    block_size: u32,
209    inv_sqrt_d: f32,
210    q_stride: u32,
211    v_block_stride_bytes: u64,
212    v_data_section_bytes: u64,
213    sliding_window: u32,
214    stream: u64,
215) -> Result<()> {
216    KernelLaunch::new(gpu, kernel)
217        .grid([num_q_heads, num_seqs, 1])
218        .block([256, 1, 1])
219        .arg_ptr(q)
220        .arg_ptr(k_cache)
221        .arg_ptr(v_cache)
222        .arg_ptr(output)
223        .arg_ptr(block_tables)
224        .arg_ptr(seq_lens)
225        .arg_u32(max_blocks_per_seq)
226        .arg_u32(num_q_heads)
227        .arg_u32(num_kv_heads)
228        .arg_u32(head_dim)
229        .arg_u32(block_size)
230        .arg_f32(inv_sqrt_d)
231        .arg_u32(q_stride)
232        .arg_u64(v_block_stride_bytes)
233        .arg_u64(v_data_section_bytes)
234        .arg_u32(sliding_window)
235        .launch(stream)
236}
237
238/// 2026-09-25: Paged decode attention over the BF16-K / Turbo2V cache (kernel
239/// `paged_decode_attn_bf16k_turbo2v`), with the same sparse-V skip as the Turbo3V kernel.
240#[allow(clippy::too_many_arguments)]
241pub fn paged_decode_attn_bf16k_turbo2v(
242    gpu: &dyn GpuBackend,
243    kernel: KernelHandle,
244    q: DevicePtr,
245    k_cache: DevicePtr,
246    v_cache: DevicePtr,
247    output: DevicePtr,
248    block_tables: DevicePtr,
249    seq_lens: DevicePtr,
250    max_blocks_per_seq: u32,
251    num_seqs: u32,
252    num_q_heads: u32,
253    num_kv_heads: u32,
254    head_dim: u32,
255    block_size: u32,
256    inv_sqrt_d: f32,
257    q_stride: u32,
258    v_block_stride_bytes: u64,
259    v_data_section_bytes: u64,
260    sliding_window: u32,
261    stream: u64,
262) -> Result<()> {
263    KernelLaunch::new(gpu, kernel)
264        .grid([num_q_heads, num_seqs, 1])
265        .block([256, 1, 1])
266        .arg_ptr(q)
267        .arg_ptr(k_cache)
268        .arg_ptr(v_cache)
269        .arg_ptr(output)
270        .arg_ptr(block_tables)
271        .arg_ptr(seq_lens)
272        .arg_u32(max_blocks_per_seq)
273        .arg_u32(num_q_heads)
274        .arg_u32(num_kv_heads)
275        .arg_u32(head_dim)
276        .arg_u32(block_size)
277        .arg_f32(inv_sqrt_d)
278        .arg_u32(q_stride)
279        .arg_u64(v_block_stride_bytes)
280        .arg_u64(v_data_section_bytes)
281        .arg_u32(sliding_window)
282        .launch(stream)
283}
metrale_model_layers/layers/ops/kv_cache_bf16k_turbov.rs

metrale_model_layers/layers/ops/
kv_cache_bf16k_turbov.rs
