metrale_model_layers/layers/ops/
kv_cache_fp8k.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 FP8 K and TurboQuant V (Turbo2V, Turbo3V, Turbo4V).
4//!
5//! Owner: model-layers ops.
6//! Invariants:
7//! - Each launcher takes the arguments of its `*_bf16k_turbo*v` counterpart in
8//!   `kv_cache_bf16k_turbov.rs` plus `k_scale`, the per-tensor FP8 K scale: the write kernel
9//!   stores `k / k_scale` as E4M3 and the decode kernel multiplies by it.
10//! - K is FP8 E4M3 in the `[num_blocks, block_size, num_kv_heads, head_dim]` layout; V is the
11//!   same TurboQuant layout as the BF16-K cache (kernels/gb10/common/reshape_and_cache_turbo.cu).
12
13#![allow(unused_imports)]
14
15use anyhow::Result;
16use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
17use metrale_gpu_runtime::kernel_args::KernelLaunch;
18
19/// 2026-09-25: Write K and V to the FP8-K / Turbo3V cache (kernel
20/// `reshape_and_cache_flash_fp8k_turbo3v`): K as FP8 E4M3, V as 3-bit Lloyd-Max codes plus one
21/// FP8 scale per 16 elements. One 256-thread block per token.
22#[allow(clippy::too_many_arguments)]
23pub fn reshape_and_cache_fp8k_turbo3v(
24    gpu: &dyn GpuBackend,
25    kernel: KernelHandle,
26    key: DevicePtr,
27    value: DevicePtr,
28    k_cache: DevicePtr,
29    v_cache: DevicePtr,
30    slot_mapping: DevicePtr,
31    num_tokens: u32,
32    num_kv_heads: u32,
33    head_dim: u32,
34    block_size: u32,
35    key_stride: u32,
36    value_stride: u32,
37    k_scale: f32,
38    k_block_stride_bytes: u64,
39    v_block_stride_bytes: u64,
40    v_data_section_bytes: u64,
41    stream: u64,
42) -> Result<()> {
43    KernelLaunch::new(gpu, kernel)
44        .grid([num_tokens, 1, 1])
45        .block([256, 1, 1])
46        .arg_ptr(key)
47        .arg_ptr(value)
48        .arg_ptr(k_cache)
49        .arg_ptr(v_cache)
50        .arg_ptr(slot_mapping)
51        .arg_u32(num_kv_heads)
52        .arg_u32(head_dim)
53        .arg_u32(block_size)
54        .arg_u32(key_stride)
55        .arg_u32(value_stride)
56        .arg_f32(k_scale)
57        .arg_u64(k_block_stride_bytes)
58        .arg_u64(v_block_stride_bytes)
59        .arg_u64(v_data_section_bytes)
60        .launch(stream)
61}
62
63/// 2026-09-25: Write K and V to the FP8-K / Turbo4V cache (kernel
64/// `reshape_and_cache_flash_fp8k_turbo4v`).
65#[allow(clippy::too_many_arguments)]
66pub fn reshape_and_cache_fp8k_turbo4v(
67    gpu: &dyn GpuBackend,
68    kernel: KernelHandle,
69    key: DevicePtr,
70    value: DevicePtr,
71    k_cache: DevicePtr,
72    v_cache: DevicePtr,
73    slot_mapping: DevicePtr,
74    num_tokens: u32,
75    num_kv_heads: u32,
76    head_dim: u32,
77    block_size: u32,
78    key_stride: u32,
79    value_stride: u32,
80    k_scale: f32,
81    k_block_stride_bytes: u64,
82    v_block_stride_bytes: u64,
83    v_data_section_bytes: u64,
84    stream: u64,
85) -> Result<()> {
86    KernelLaunch::new(gpu, kernel)
87        .grid([num_tokens, 1, 1])
88        .block([256, 1, 1])
89        .arg_ptr(key)
90        .arg_ptr(value)
91        .arg_ptr(k_cache)
92        .arg_ptr(v_cache)
93        .arg_ptr(slot_mapping)
94        .arg_u32(num_kv_heads)
95        .arg_u32(head_dim)
96        .arg_u32(block_size)
97        .arg_u32(key_stride)
98        .arg_u32(value_stride)
99        .arg_f32(k_scale)
100        .arg_u64(k_block_stride_bytes)
101        .arg_u64(v_block_stride_bytes)
102        .arg_u64(v_data_section_bytes)
103        .launch(stream)
104}
105
106/// 2026-09-25: Write K and V to the FP8-K / Turbo2V cache (kernel
107/// `reshape_and_cache_flash_fp8k_turbo2v`).
108#[allow(clippy::too_many_arguments)]
109pub fn reshape_and_cache_fp8k_turbo2v(
110    gpu: &dyn GpuBackend,
111    kernel: KernelHandle,
112    key: DevicePtr,
113    value: DevicePtr,
114    k_cache: DevicePtr,
115    v_cache: DevicePtr,
116    slot_mapping: DevicePtr,
117    num_tokens: u32,
118    num_kv_heads: u32,
119    head_dim: u32,
120    block_size: u32,
121    key_stride: u32,
122    value_stride: u32,
123    k_scale: f32,
124    k_block_stride_bytes: u64,
125    v_block_stride_bytes: u64,
126    v_data_section_bytes: u64,
127    stream: u64,
128) -> Result<()> {
129    KernelLaunch::new(gpu, kernel)
130        .grid([num_tokens, 1, 1])
131        .block([256, 1, 1])
132        .arg_ptr(key)
133        .arg_ptr(value)
134        .arg_ptr(k_cache)
135        .arg_ptr(v_cache)
136        .arg_ptr(slot_mapping)
137        .arg_u32(num_kv_heads)
138        .arg_u32(head_dim)
139        .arg_u32(block_size)
140        .arg_u32(key_stride)
141        .arg_u32(value_stride)
142        .arg_f32(k_scale)
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: Paged decode attention over the FP8-K / Turbo3V cache (kernel
150/// `paged_decode_attn_fp8k_turbo3v`), one 256-thread block per (query head, sequence). The
151/// kernel skips a position's V load and dequant when its attention weight is at most
152/// `TQ_PLUS_SPARSE_V_THRESHOLD`.
153#[allow(clippy::too_many_arguments)]
154pub fn paged_decode_attn_fp8k_turbo3v(
155    gpu: &dyn GpuBackend,
156    kernel: KernelHandle,
157    q: DevicePtr,
158    k_cache: DevicePtr,
159    v_cache: DevicePtr,
160    output: DevicePtr,
161    block_tables: DevicePtr,
162    seq_lens: DevicePtr,
163    max_blocks_per_seq: u32,
164    num_seqs: u32,
165    num_q_heads: u32,
166    num_kv_heads: u32,
167    head_dim: u32,
168    block_size: u32,
169    inv_sqrt_d: f32,
170    k_scale: f32,
171    q_stride: u32,
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_f32(k_scale)
193        .arg_u32(q_stride)
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/// 2026-09-25: Paged decode attention over the FP8-K / Turbo4V cache (kernel
201/// `paged_decode_attn_fp8k_turbo4v`).
202#[allow(clippy::too_many_arguments)]
203pub fn paged_decode_attn_fp8k_turbo4v(
204    gpu: &dyn GpuBackend,
205    kernel: KernelHandle,
206    q: DevicePtr,
207    k_cache: DevicePtr,
208    v_cache: DevicePtr,
209    output: DevicePtr,
210    block_tables: DevicePtr,
211    seq_lens: DevicePtr,
212    max_blocks_per_seq: u32,
213    num_seqs: u32,
214    num_q_heads: u32,
215    num_kv_heads: u32,
216    head_dim: u32,
217    block_size: u32,
218    inv_sqrt_d: f32,
219    k_scale: f32,
220    q_stride: u32,
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_f32(k_scale)
242        .arg_u32(q_stride)
243        .arg_u64(v_block_stride_bytes)
244        .arg_u64(v_data_section_bytes)
245        .arg_u32(sliding_window)
246        .launch(stream)
247}
248
249/// 2026-09-25: Paged decode attention over the FP8-K / Turbo2V cache (kernel
250/// `paged_decode_attn_fp8k_turbo2v`).
251#[allow(clippy::too_many_arguments)]
252pub fn paged_decode_attn_fp8k_turbo2v(
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    k_scale: f32,
269    q_stride: u32,
270    v_block_stride_bytes: u64,
271    v_data_section_bytes: u64,
272    sliding_window: u32,
273    stream: u64,
274) -> Result<()> {
275    KernelLaunch::new(gpu, kernel)
276        .grid([num_q_heads, num_seqs, 1])
277        .block([256, 1, 1])
278        .arg_ptr(q)
279        .arg_ptr(k_cache)
280        .arg_ptr(v_cache)
281        .arg_ptr(output)
282        .arg_ptr(block_tables)
283        .arg_ptr(seq_lens)
284        .arg_u32(max_blocks_per_seq)
285        .arg_u32(num_q_heads)
286        .arg_u32(num_kv_heads)
287        .arg_u32(head_dim)
288        .arg_u32(block_size)
289        .arg_f32(inv_sqrt_d)
290        .arg_f32(k_scale)
291        .arg_u32(q_stride)
292        .arg_u64(v_block_stride_bytes)
293        .arg_u64(v_data_section_bytes)
294        .arg_u32(sliding_window)
295        .launch(stream)
296}
metrale_model_layers/layers/ops/kv_cache_fp8k.rs

metrale_model_layers/layers/ops/
kv_cache_fp8k.rs
