spark_model/layers/ops/
kv_cache.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Auto-extracted from `ops.rs` during refactor wave 4a.
4
5#![allow(unused_imports)]
6
7use anyhow::Result;
8use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
9use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
10
11use crate::layers::moe;
12use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
13
14use super::*;
15
16/// Fill paged-KV slot mappings on-device from a persistent block table.
17///
18/// Kernel: `fill_slots_from_block_table(slots, block_table, start_pos, count, block_size)`
19/// Grid: (ceil(count/256), 1, 1)  Block: (256, 1, 1)
20pub fn fill_slots_from_block_table(
21    gpu: &dyn GpuBackend,
22    kernel: KernelHandle,
23    slots: DevicePtr,
24    block_table: DevicePtr,
25    start_pos: u32,
26    count: u32,
27    block_size: u32,
28    stream: u64,
29) -> Result<()> {
30    if count == 0 {
31        return Ok(());
32    }
33    KernelLaunch::new(gpu, kernel)
34        .grid([div_ceil(count, 256), 1, 1])
35        .block([256, 1, 1])
36        .arg_ptr(slots)
37        .arg_ptr(block_table)
38        .arg_u32(start_pos)
39        .arg_u32(count)
40        .arg_u32(block_size)
41        .launch(stream)
42}
43
44// ── KV cache ───────────────────────────────────────────────────────
45
46/// Write K/V to paged FP8 cache using slot_mapping.
47///
48/// Kernel: `reshape_and_cache_flash_fp8(key, value, k_cache, v_cache,
49///          slot_mapping, num_kv_heads, head_dim, block_size,
50///          k_scale, v_scale, key_stride, value_stride, cache_stride)`
51/// Grid: (num_tokens, 1, 1)  Block: (256, 1, 1)
52///
53/// `slot_mapping` is a device pointer to `i64[num_tokens]`.
54/// BF16 reshape and cache — no quantization, direct BF16 copy.
55#[allow(clippy::too_many_arguments)]
56pub fn reshape_and_cache(
57    gpu: &dyn GpuBackend,
58    kernel: KernelHandle,
59    key: DevicePtr,
60    value: DevicePtr,
61    k_cache: DevicePtr,
62    v_cache: DevicePtr,
63    slot_mapping: DevicePtr,
64    num_tokens: u32,
65    num_kv_heads: u32,
66    head_dim: u32,
67    block_size: u32,
68    key_stride: u32,
69    value_stride: u32,
70    _cache_stride: u64,
71    stream: u64,
72) -> Result<()> {
73    KernelLaunch::new(gpu, kernel)
74        .grid([num_tokens, 1, 1])
75        .block([256, 1, 1])
76        .arg_ptr(key)
77        .arg_ptr(value)
78        .arg_ptr(k_cache)
79        .arg_ptr(v_cache)
80        .arg_ptr(slot_mapping)
81        .arg_u32(num_kv_heads)
82        .arg_u32(head_dim)
83        .arg_u32(block_size)
84        .arg_u32(key_stride)
85        .arg_u32(value_stride)
86        .launch(stream)
87}
88
89/// V-only paged cache write — companion to the fused K-path so the
90/// K side of the cache stays exclusively owned by
91/// `fused_k_norm_rope_cache_write_*`. Use this when the fused K kernel
92/// is active to avoid the existing `reshape_and_cache` overwriting
93/// the correct K values with a double-rounded copy.
94#[allow(clippy::too_many_arguments)]
95pub fn reshape_and_cache_flash_v_only(
96    gpu: &dyn GpuBackend,
97    kernel: KernelHandle,
98    value: DevicePtr,
99    v_cache: DevicePtr,
100    slot_mapping: DevicePtr,
101    num_tokens: u32,
102    num_kv_heads: u32,
103    head_dim: u32,
104    block_size: u32,
105    value_stride: u32,
106    stream: u64,
107) -> Result<()> {
108    KernelLaunch::new(gpu, kernel)
109        .grid([num_tokens, 1, 1])
110        .block([256, 1, 1])
111        .arg_ptr(value)
112        .arg_ptr(v_cache)
113        .arg_ptr(slot_mapping)
114        .arg_u32(num_kv_heads)
115        .arg_u32(head_dim)
116        .arg_u32(block_size)
117        .arg_u32(value_stride)
118        .launch(stream)
119}
120
121/// Fused K-path: rms_norm → RoPE → BF16 paged cache write in one kernel.
122///
123/// Replaces the chained `ops::rms_norm + ops::rope + ops::reshape_and_cache`
124/// sequence for the K projection. Keeps K in FP32 between the three
125/// operations and BF16-rounds ONLY at cache write — vLLM-equivalent
126/// precision regime. Eliminates the two intermediate BF16 rounding steps
127/// that previously compounded at deep attention layers (L35-L39) where K
128/// magnitudes peak ~18× vs L0, causing the documented BF16-KV cliff.
129#[allow(clippy::too_many_arguments)]
130pub fn fused_k_norm_rope_cache_write_bf16(
131    gpu: &dyn GpuBackend,
132    kernel: KernelHandle,
133    k_in: DevicePtr,
134    k_norm_weight: DevicePtr,
135    positions: DevicePtr,
136    k_cache: DevicePtr,
137    slot_mapping: DevicePtr,
138    num_tokens: u32,
139    num_kv_heads: u32,
140    head_dim: u32,
141    rotary_dim: u32,
142    block_size: u32,
143    rms_eps: f32,
144    theta: f32,
145    stream: u64,
146) -> Result<()> {
147    KernelLaunch::new(gpu, kernel)
148        .grid([num_tokens, num_kv_heads, 1])
149        .block([head_dim, 1, 1])
150        .arg_ptr(k_in)
151        .arg_ptr(k_norm_weight)
152        .arg_ptr(positions)
153        .arg_ptr(k_cache)
154        .arg_ptr(slot_mapping)
155        .arg_u32(num_kv_heads)
156        .arg_u32(head_dim)
157        .arg_u32(rotary_dim)
158        .arg_u32(block_size)
159        .arg_f32(rms_eps)
160        .arg_f32(theta)
161        .launch(stream)
162}
163
164/// MRoPE-interleaved variant — selects abs position from pos_t/pos_h/pos_w
165/// based on `pair_idx % 3`. For text-only inputs (pos_h == pos_w == pos_t)
166/// the result is bit-identical to the scalar-position variant.
167#[allow(clippy::too_many_arguments)]
168pub fn fused_k_norm_rope_cache_write_bf16_mrope(
169    gpu: &dyn GpuBackend,
170    kernel: KernelHandle,
171    k_in: DevicePtr,
172    k_norm_weight: DevicePtr,
173    pos_t: DevicePtr,
174    pos_h: DevicePtr,
175    pos_w: DevicePtr,
176    k_cache: DevicePtr,
177    slot_mapping: DevicePtr,
178    num_tokens: u32,
179    num_kv_heads: u32,
180    head_dim: u32,
181    rotary_dim: u32,
182    block_size: u32,
183    rms_eps: f32,
184    theta: f32,
185    stream: u64,
186) -> Result<()> {
187    KernelLaunch::new(gpu, kernel)
188        .grid([num_tokens, num_kv_heads, 1])
189        .block([head_dim, 1, 1])
190        .arg_ptr(k_in)
191        .arg_ptr(k_norm_weight)
192        .arg_ptr(pos_t)
193        .arg_ptr(pos_h)
194        .arg_ptr(pos_w)
195        .arg_ptr(k_cache)
196        .arg_ptr(slot_mapping)
197        .arg_u32(num_kv_heads)
198        .arg_u32(head_dim)
199        .arg_u32(rotary_dim)
200        .arg_u32(block_size)
201        .arg_f32(rms_eps)
202        .arg_f32(theta)
203        .launch(stream)
204}
205
206/// FP8-output sibling of [`fused_k_norm_rope_cache_write_bf16`]. Same
207/// semantics; one fewer BF16 round before the saturating FP8 cast.
208#[allow(clippy::too_many_arguments)]
209pub fn fused_k_norm_rope_cache_write_fp8(
210    gpu: &dyn GpuBackend,
211    kernel: KernelHandle,
212    k_in: DevicePtr,
213    k_norm_weight: DevicePtr,
214    positions: DevicePtr,
215    k_cache_fp8: DevicePtr,
216    slot_mapping: DevicePtr,
217    num_tokens: u32,
218    num_kv_heads: u32,
219    head_dim: u32,
220    rotary_dim: u32,
221    block_size: u32,
222    rms_eps: f32,
223    theta: f32,
224    inv_scale: f32,
225    stream: u64,
226) -> Result<()> {
227    KernelLaunch::new(gpu, kernel)
228        .grid([num_tokens, num_kv_heads, 1])
229        .block([head_dim, 1, 1])
230        .arg_ptr(k_in)
231        .arg_ptr(k_norm_weight)
232        .arg_ptr(positions)
233        .arg_ptr(k_cache_fp8)
234        .arg_ptr(slot_mapping)
235        .arg_u32(num_kv_heads)
236        .arg_u32(head_dim)
237        .arg_u32(rotary_dim)
238        .arg_u32(block_size)
239        .arg_f32(rms_eps)
240        .arg_f32(theta)
241        .arg_f32(inv_scale)
242        .launch(stream)
243}
244
245/// Decode-path fusion of `rms_norm` + `rope_forward` + `reshape_and_cache_flash_fp8`
246/// for an FP8 KV cache. Writes BOTH K (normed + rotated + quantized) and V
247/// (quantized) in one launch, so the FP8 decode chain drops from
248/// `rms_norm(K)` + `rope(Q,K)` + `reshape_and_cache_fp8(K,V)` to
249/// `rope(Q only)` + this kernel.
250///
251/// BIT-IDENTICAL to the chain it replaces — the kernel reproduces every
252/// intermediate BF16 rounding step rather than keeping intermediates in FP32.
253/// See the header of `reshape_and_cache_fused_k_fp8.cu`; do not "improve" the
254/// precision here, it would change committed cache bytes.
255///
256/// `k_scale`/`v_scale` are the DEQUANT scales; the kernel takes the reciprocal
257/// itself so the lowering matches `reshape_and_cache_flash_fp8` exactly.
258///
259/// Grid: (num_tokens, num_kv_heads, 1)  Block: (head_dim, 1, 1)
260#[allow(clippy::too_many_arguments)]
261pub fn fused_k_norm_rope_cache_write_fp8_kv(
262    gpu: &dyn GpuBackend,
263    kernel: KernelHandle,
264    k_in: DevicePtr,
265    value: DevicePtr,
266    k_norm_weight: DevicePtr,
267    positions: DevicePtr,
268    k_cache: DevicePtr,
269    v_cache: DevicePtr,
270    slot_mapping: DevicePtr,
271    num_tokens: u32,
272    num_kv_heads: u32,
273    head_dim: u32,
274    rotary_dim: u32,
275    block_size: u32,
276    k_scale: f32,
277    v_scale: f32,
278    key_stride: u32,
279    value_stride: u32,
280    cache_stride: u64,
281    rms_eps: f32,
282    theta: f32,
283    stream: u64,
284) -> Result<()> {
285    KernelLaunch::new(gpu, kernel)
286        .grid([num_tokens, num_kv_heads, 1])
287        .block([head_dim, 1, 1])
288        .arg_ptr(k_in)
289        .arg_ptr(value)
290        .arg_ptr(k_norm_weight)
291        .arg_ptr(positions)
292        .arg_ptr(k_cache)
293        .arg_ptr(v_cache)
294        .arg_ptr(slot_mapping)
295        .arg_u32(num_kv_heads)
296        .arg_u32(head_dim)
297        .arg_u32(rotary_dim)
298        .arg_u32(block_size)
299        .arg_f32(k_scale)
300        .arg_f32(v_scale)
301        .arg_u32(key_stride)
302        .arg_u32(value_stride)
303        .arg_u64(cache_stride)
304        .arg_f32(rms_eps)
305        .arg_f32(theta)
306        .launch(stream)
307}
308
309/// Write K/V to paged Bf16K + Turbo3V (TurboQuant+ safer-asym) cache.
310///
311/// K is written as raw BF16 (NHD contiguous), V as 3-bit Lloyd-Max + FP8
312/// per-group scale with matched-norm correction. K and V pools have separate
313/// strides because K is 2 b/elem and V is ~0.5 b/elem + scale.
314///
315/// Kernel: `reshape_and_cache_flash_bf16k_turbo3v(key, value, k_cache, v_cache,
316///          slot_mapping, num_kv_heads, head_dim, block_size,
317///          key_stride, value_stride, k_block_stride_bytes,
318///          v_block_stride_bytes, v_data_section_bytes)`
319/// Grid: (num_tokens, 1, 1)  Block: (256, 1, 1)
320#[allow(clippy::too_many_arguments)]
321pub fn reshape_and_cache_bf16k_turbo3v(
322    gpu: &dyn GpuBackend,
323    kernel: KernelHandle,
324    key: DevicePtr,
325    value: DevicePtr,
326    k_cache: DevicePtr,
327    v_cache: DevicePtr,
328    slot_mapping: DevicePtr,
329    num_tokens: u32,
330    num_kv_heads: u32,
331    head_dim: u32,
332    block_size: u32,
333    key_stride: u32,
334    value_stride: u32,
335    k_block_stride_bytes: u64,
336    v_block_stride_bytes: u64,
337    v_data_section_bytes: u64,
338    stream: u64,
339) -> Result<()> {
340    KernelLaunch::new(gpu, kernel)
341        .grid([num_tokens, 1, 1])
342        .block([256, 1, 1])
343        .arg_ptr(key)
344        .arg_ptr(value)
345        .arg_ptr(k_cache)
346        .arg_ptr(v_cache)
347        .arg_ptr(slot_mapping)
348        .arg_u32(num_kv_heads)
349        .arg_u32(head_dim)
350        .arg_u32(block_size)
351        .arg_u32(key_stride)
352        .arg_u32(value_stride)
353        .arg_u64(k_block_stride_bytes)
354        .arg_u64(v_block_stride_bytes)
355        .arg_u64(v_data_section_bytes)
356        .launch(stream)
357}
358
359/// Paged decode attention for Bf16K + Turbo3V asymmetric KV cache.
360///
361/// K is read as BF16 NHD (vector loads), V as 3-bit Lloyd-Max packed bytes
362/// with FP8 per-group scale (sparse V on batched AND remainder paths).
363///
364/// Kernel: `paged_decode_attn_bf16k_turbo3v(Q, K_cache, V_cache, O,
365///          block_tables, seq_lens, max_blocks_per_seq, num_q_heads,
366///          num_kv_heads, head_dim, block_size, inv_sqrt_d, q_stride,
367///          v_block_stride_bytes, v_data_section_bytes, sliding_window)`
368/// Grid: (num_q_heads, num_seqs, 1)  Block: (256, 1, 1)
369#[allow(clippy::too_many_arguments)]
370pub fn paged_decode_attn_bf16k_turbo3v(
371    gpu: &dyn GpuBackend,
372    kernel: KernelHandle,
373    q: DevicePtr,
374    k_cache: DevicePtr,
375    v_cache: DevicePtr,
376    output: DevicePtr,
377    block_tables: DevicePtr,
378    seq_lens: DevicePtr,
379    max_blocks_per_seq: u32,
380    num_seqs: u32,
381    num_q_heads: u32,
382    num_kv_heads: u32,
383    head_dim: u32,
384    block_size: u32,
385    inv_sqrt_d: f32,
386    q_stride: u32,
387    v_block_stride_bytes: u64,
388    v_data_section_bytes: u64,
389    sliding_window: u32,
390    stream: u64,
391) -> Result<()> {
392    KernelLaunch::new(gpu, kernel)
393        .grid([num_q_heads, num_seqs, 1])
394        .block([256, 1, 1])
395        .arg_ptr(q)
396        .arg_ptr(k_cache)
397        .arg_ptr(v_cache)
398        .arg_ptr(output)
399        .arg_ptr(block_tables)
400        .arg_ptr(seq_lens)
401        .arg_u32(max_blocks_per_seq)
402        .arg_u32(num_q_heads)
403        .arg_u32(num_kv_heads)
404        .arg_u32(head_dim)
405        .arg_u32(block_size)
406        .arg_f32(inv_sqrt_d)
407        .arg_u32(q_stride)
408        .arg_u64(v_block_stride_bytes)
409        .arg_u64(v_data_section_bytes)
410        .arg_u32(sliding_window)
411        .launch(stream)
412}
413
414/// Write K/V to paged Bf16K + Turbo4V (TurboQuant+ safer-asym) cache.
415///
416/// K written as raw BF16 NHD, V as 4-bit Lloyd-Max + FP8 per-group scale with
417/// matched-norm correction. K and V pools have separate strides.
418#[allow(clippy::too_many_arguments)]
419pub fn reshape_and_cache_bf16k_turbo4v(
420    gpu: &dyn GpuBackend,
421    kernel: KernelHandle,
422    key: DevicePtr,
423    value: DevicePtr,
424    k_cache: DevicePtr,
425    v_cache: DevicePtr,
426    slot_mapping: DevicePtr,
427    num_tokens: u32,
428    num_kv_heads: u32,
429    head_dim: u32,
430    block_size: u32,
431    key_stride: u32,
432    value_stride: u32,
433    k_block_stride_bytes: u64,
434    v_block_stride_bytes: u64,
435    v_data_section_bytes: u64,
436    stream: u64,
437) -> Result<()> {
438    KernelLaunch::new(gpu, kernel)
439        .grid([num_tokens, 1, 1])
440        .block([256, 1, 1])
441        .arg_ptr(key)
442        .arg_ptr(value)
443        .arg_ptr(k_cache)
444        .arg_ptr(v_cache)
445        .arg_ptr(slot_mapping)
446        .arg_u32(num_kv_heads)
447        .arg_u32(head_dim)
448        .arg_u32(block_size)
449        .arg_u32(key_stride)
450        .arg_u32(value_stride)
451        .arg_u64(k_block_stride_bytes)
452        .arg_u64(v_block_stride_bytes)
453        .arg_u64(v_data_section_bytes)
454        .launch(stream)
455}
456
457/// Write K/V to paged Bf16K + Turbo2V (TurboQuant+ safer-asym) cache (6.4x V comp).
458#[allow(clippy::too_many_arguments)]
459pub fn reshape_and_cache_bf16k_turbo2v(
460    gpu: &dyn GpuBackend,
461    kernel: KernelHandle,
462    key: DevicePtr,
463    value: DevicePtr,
464    k_cache: DevicePtr,
465    v_cache: DevicePtr,
466    slot_mapping: DevicePtr,
467    num_tokens: u32,
468    num_kv_heads: u32,
469    head_dim: u32,
470    block_size: u32,
471    key_stride: u32,
472    value_stride: u32,
473    k_block_stride_bytes: u64,
474    v_block_stride_bytes: u64,
475    v_data_section_bytes: u64,
476    stream: u64,
477) -> Result<()> {
478    KernelLaunch::new(gpu, kernel)
479        .grid([num_tokens, 1, 1])
480        .block([256, 1, 1])
481        .arg_ptr(key)
482        .arg_ptr(value)
483        .arg_ptr(k_cache)
484        .arg_ptr(v_cache)
485        .arg_ptr(slot_mapping)
486        .arg_u32(num_kv_heads)
487        .arg_u32(head_dim)
488        .arg_u32(block_size)
489        .arg_u32(key_stride)
490        .arg_u32(value_stride)
491        .arg_u64(k_block_stride_bytes)
492        .arg_u64(v_block_stride_bytes)
493        .arg_u64(v_data_section_bytes)
494        .launch(stream)
495}
496
497/// Paged decode attention for Bf16K + Turbo4V asymmetric KV cache.
498///
499/// K read as BF16 NHD, V as 4-bit Lloyd-Max packed bytes + FP8 per-group scale
500/// (sparse V on batched + remainder paths).
501#[allow(clippy::too_many_arguments)]
502pub fn paged_decode_attn_bf16k_turbo4v(
503    gpu: &dyn GpuBackend,
504    kernel: KernelHandle,
505    q: DevicePtr,
506    k_cache: DevicePtr,
507    v_cache: DevicePtr,
508    output: DevicePtr,
509    block_tables: DevicePtr,
510    seq_lens: DevicePtr,
511    max_blocks_per_seq: u32,
512    num_seqs: u32,
513    num_q_heads: u32,
514    num_kv_heads: u32,
515    head_dim: u32,
516    block_size: u32,
517    inv_sqrt_d: f32,
518    q_stride: u32,
519    v_block_stride_bytes: u64,
520    v_data_section_bytes: u64,
521    sliding_window: u32,
522    stream: u64,
523) -> Result<()> {
524    KernelLaunch::new(gpu, kernel)
525        .grid([num_q_heads, num_seqs, 1])
526        .block([256, 1, 1])
527        .arg_ptr(q)
528        .arg_ptr(k_cache)
529        .arg_ptr(v_cache)
530        .arg_ptr(output)
531        .arg_ptr(block_tables)
532        .arg_ptr(seq_lens)
533        .arg_u32(max_blocks_per_seq)
534        .arg_u32(num_q_heads)
535        .arg_u32(num_kv_heads)
536        .arg_u32(head_dim)
537        .arg_u32(block_size)
538        .arg_f32(inv_sqrt_d)
539        .arg_u32(q_stride)
540        .arg_u64(v_block_stride_bytes)
541        .arg_u64(v_data_section_bytes)
542        .arg_u32(sliding_window)
543        .launch(stream)
544}
545
546/// Paged decode attention for Bf16K + Turbo2V asymmetric KV cache (6.4x V comp).
547#[allow(clippy::too_many_arguments)]
548pub fn paged_decode_attn_bf16k_turbo2v(
549    gpu: &dyn GpuBackend,
550    kernel: KernelHandle,
551    q: DevicePtr,
552    k_cache: DevicePtr,
553    v_cache: DevicePtr,
554    output: DevicePtr,
555    block_tables: DevicePtr,
556    seq_lens: DevicePtr,
557    max_blocks_per_seq: u32,
558    num_seqs: u32,
559    num_q_heads: u32,
560    num_kv_heads: u32,
561    head_dim: u32,
562    block_size: u32,
563    inv_sqrt_d: f32,
564    q_stride: u32,
565    v_block_stride_bytes: u64,
566    v_data_section_bytes: u64,
567    sliding_window: u32,
568    stream: u64,
569) -> Result<()> {
570    KernelLaunch::new(gpu, kernel)
571        .grid([num_q_heads, num_seqs, 1])
572        .block([256, 1, 1])
573        .arg_ptr(q)
574        .arg_ptr(k_cache)
575        .arg_ptr(v_cache)
576        .arg_ptr(output)
577        .arg_ptr(block_tables)
578        .arg_ptr(seq_lens)
579        .arg_u32(max_blocks_per_seq)
580        .arg_u32(num_q_heads)
581        .arg_u32(num_kv_heads)
582        .arg_u32(head_dim)
583        .arg_u32(block_size)
584        .arg_f32(inv_sqrt_d)
585        .arg_u32(q_stride)
586        .arg_u64(v_block_stride_bytes)
587        .arg_u64(v_data_section_bytes)
588        .arg_u32(sliding_window)
589        .launch(stream)
590}
591
592/// `k_cache`/`v_cache` are the full pool base pointers.
593/// `cache_stride` is in elements (block_size * num_kv_heads * head_dim).
594pub fn reshape_and_cache_fp8(
595    gpu: &dyn GpuBackend,
596    kernel: KernelHandle,
597    key: DevicePtr,
598    value: DevicePtr,
599    k_cache: DevicePtr,
600    v_cache: DevicePtr,
601    slot_mapping: DevicePtr,
602    num_tokens: u32,
603    num_kv_heads: u32,
604    head_dim: u32,
605    block_size: u32,
606    k_scale: f32,
607    v_scale: f32,
608    key_stride: u32,
609    value_stride: u32,
610    cache_stride: u64,
611    stream: u64,
612) -> Result<()> {
613    KernelLaunch::new(gpu, kernel)
614        .grid([num_tokens, 1, 1])
615        .block([256, 1, 1])
616        .arg_ptr(key)
617        .arg_ptr(value)
618        .arg_ptr(k_cache)
619        .arg_ptr(v_cache)
620        .arg_ptr(slot_mapping)
621        .arg_u32(num_kv_heads)
622        .arg_u32(head_dim)
623        .arg_u32(block_size)
624        .arg_f32(k_scale)
625        .arg_f32(v_scale)
626        .arg_u32(key_stride)
627        .arg_u32(value_stride)
628        .arg_u64(cache_stride)
629        .launch(stream)
630}
631
632/// Paged decode attention (FP8 KV cache, single/multi sequence).
633///
634/// Kernel: `paged_decode_attn_fp8(Q, K_cache, V_cache, O, block_tables,
635///          seq_lens, max_blocks_per_seq, num_q_heads, num_kv_heads,
636///          head_dim, block_size, inv_sqrt_d, k_scale, v_scale,
637///          q_stride, cache_stride)`
638/// Grid: (num_q_heads, num_seqs, 1)  Block: (256, 1, 1)
639///
640/// `block_tables`: device ptr to `i32[num_seqs * max_blocks_per_seq]`
641/// `seq_lens`: device ptr to `i32[num_seqs]`
642/// `cache_stride` is in elements (u64).
643/// BF16 paged decode attention — no FP8 quantization, direct BF16 KV cache.
644/// MLA batched GEMV: output[head, n] = sum_k(weight[head, n, k] * input[head, k])
645/// Replaces 32 sequential dense_gemv calls with a single kernel launch.
646pub fn mla_batched_gemv(
647    gpu: &dyn GpuBackend,
648    kernel: KernelHandle,
649    input: DevicePtr,
650    weight: DevicePtr,
651    output: DevicePtr,
652    n_out: u32,
653    k: u32,
654    num_heads: u32,
655    input_stride: u32,
656    output_stride: u32,
657    stream: u64,
658) -> Result<()> {
659    KernelLaunch::new(gpu, kernel)
660        .grid([div_ceil(n_out, 8), num_heads, 1]) // N_PER_BLOCK*2=8
661        .block([256, 1, 1])
662        .arg_ptr(input)
663        .arg_ptr(weight)
664        .arg_ptr(output)
665        .arg_u32(n_out)
666        .arg_u32(k)
667        .arg_u32(input_stride)
668        .arg_u32(output_stride)
669        .launch(stream)
670}
671
672/// MLA Q_rope scatter: copy rope portion from q_full to strided q_absorbed_buf. 1 kernel replaces 32 D2D copies.
673#[allow(clippy::too_many_arguments)]
674pub fn mla_q_rope_scatter(
675    gpu: &dyn GpuBackend,
676    kernel: KernelHandle,
677    q_full: DevicePtr,
678    q_absorbed_buf: DevicePtr,
679    q_rope_contiguous: DevicePtr,
680    nq: u32,
681    hd: u32,
682    nope: u32,
683    rope: u32,
684    kv_lora: u32,
685    mla_cache_dim: u32,
686    stream: u64,
687) -> Result<()> {
688    KernelLaunch::new(gpu, kernel)
689        .grid([1, 1, 1])
690        .block([256, 1, 1])
691        .arg_ptr(q_full)
692        .arg_ptr(q_absorbed_buf)
693        .arg_ptr(q_rope_contiguous)
694        .arg_u32(nq)
695        .arg_u32(hd)
696        .arg_u32(nope)
697        .arg_u32(rope)
698        .arg_u32(kv_lora)
699        .arg_u32(mla_cache_dim)
700        .launch(stream)
701}
702
703/// MLA Q_rope writeback: scatter RoPE'd rope portions to strided layout. 1 kernel replaces 32 D2D copies.
704pub fn mla_q_rope_writeback(
705    gpu: &dyn GpuBackend,
706    kernel: KernelHandle,
707    q_rope_direct: DevicePtr,
708    q_absorbed_buf: DevicePtr,
709    nq: u32,
710    rope: u32,
711    kv_lora: u32,
712    mla_cache_dim: u32,
713    stream: u64,
714) -> Result<()> {
715    KernelLaunch::new(gpu, kernel)
716        .grid([1, 1, 1])
717        .block([256, 1, 1])
718        .arg_ptr(q_rope_direct)
719        .arg_ptr(q_absorbed_buf)
720        .arg_u32(nq)
721        .arg_u32(rope)
722        .arg_u32(kv_lora)
723        .arg_u32(mla_cache_dim)
724        .launch(stream)
725}
726
727/// MLA cache assembly: fuse [kv_latent|k_rope]→K and [kv_latent|zeros]→V into 1 kernel.
728pub fn mla_cache_assemble(
729    gpu: &dyn GpuBackend,
730    kernel: KernelHandle,
731    kv_latent: DevicePtr,
732    k_rope: DevicePtr,
733    k_cache: DevicePtr,
734    v_cache: DevicePtr,
735    kv_lora: u32,
736    rope: u32,
737    mla_cache_dim: u32,
738    stream: u64,
739) -> Result<()> {
740    KernelLaunch::new(gpu, kernel)
741        .grid([1, 1, 1])
742        .block([mla_cache_dim.max(256), 1, 1])
743        .arg_ptr(kv_latent)
744        .arg_ptr(k_rope)
745        .arg_ptr(k_cache)
746        .arg_ptr(v_cache)
747        .arg_u32(kv_lora)
748        .arg_u32(rope)
749        .arg_u32(mla_cache_dim)
750        .launch(stream)
751}
752
753/// MLA Paged Decode — NVFP4 variant for DeepSeek-V4-Flash.
754///
755/// Kernel: `mla_paged_decode_nvfp4(Q, K_cache, V_cache, O, block_tables,
756///          seq_lens, max_blocks_per_seq, num_q_heads, num_kv_heads,
757///          q_head_dim, kv_cache_dim, block_size, inv_sqrt_d,
758///          block_stride_bytes, data_section_bytes)`
759/// Grid: (num_q_heads, num_seqs, 1)  Block: (256, 1, 1)
760///
761/// V4-Flash uses compressed KV cache (576 dims: 512 latent + 64 rope)
762/// and flattened Q layout (32768 dims = 64 heads × 512).
763#[allow(clippy::too_many_arguments)]
764pub fn mla_paged_decode_nvfp4(
765    gpu: &dyn GpuBackend,
766    kernel: KernelHandle,
767    q: DevicePtr,
768    k_cache: DevicePtr,
769    v_cache: DevicePtr,
770    o: DevicePtr,
771    block_tables: DevicePtr,
772    seq_lens: DevicePtr,
773    max_blocks_per_seq: u32,
774    num_q_heads: u32,
775    num_kv_heads: u32,
776    q_head_dim: u32,
777    kv_cache_dim: u32,
778    block_size: u32,
779    inv_sqrt_d: f32,
780    block_stride_bytes: u64,
781    data_section_bytes: u64,
782    num_seqs: u32,
783    stream: u64,
784) -> Result<()> {
785    KernelLaunch::new(gpu, kernel)
786        .grid([num_q_heads, num_seqs, 1])
787        .block([256, 1, 1])
788        .arg_ptr(q)
789        .arg_ptr(k_cache)
790        .arg_ptr(v_cache)
791        .arg_ptr(o)
792        .arg_ptr(block_tables)
793        .arg_ptr(seq_lens)
794        .arg_u32(max_blocks_per_seq)
795        .arg_u32(num_q_heads)
796        .arg_u32(num_kv_heads)
797        .arg_u32(q_head_dim)
798        .arg_u32(kv_cache_dim)
799        .arg_u32(block_size)
800        .arg_f32(inv_sqrt_d)
801        .arg_u64(block_stride_bytes)
802        .arg_u64(data_section_bytes)
803        .launch(stream)
804}
805
806/// MLA Paged Decode — FP8 variant for DeepSeek-V4-Flash with FP8 KV cache.
807///
808/// Kernel: `mla_paged_decode_fp8(Q, K_cache, V_cache, O, block_tables,
809///          seq_lens, max_blocks_per_seq, num_q_heads, num_kv_heads,
810///          q_head_dim, kv_cache_dim, block_size, inv_sqrt_d,
811///          k_scale, v_scale, cache_stride)`
812/// Grid: (num_q_heads, num_seqs, 1)  Block: (256, 1, 1)
813///
814/// V4-Flash uses compressed KV cache (576 dims: 512 latent + 64 rope)
815/// and FP8 quantization with per-layer scales.
816#[allow(clippy::too_many_arguments)]
817pub fn mla_paged_decode_fp8(
818    gpu: &dyn GpuBackend,
819    kernel: KernelHandle,
820    q: DevicePtr,
821    k_cache: DevicePtr,
822    v_cache: DevicePtr,
823    o: DevicePtr,
824    block_tables: DevicePtr,
825    seq_lens: DevicePtr,
826    max_blocks_per_seq: u32,
827    num_q_heads: u32,
828    num_kv_heads: u32,
829    q_head_dim: u32,
830    kv_cache_dim: u32,
831    block_size: u32,
832    inv_sqrt_d: f32,
833    k_scale: f32,
834    v_scale: f32,
835    cache_stride: u64,
836    num_seqs: u32,
837    sliding_window: u32,
838    sinks: DevicePtr,
839    comp_pool: DevicePtr, // 4b: flat FP8 compressed-KV pool (NULL = no compressed arm)
840    comp_block_count: u32, // 4b: # compressed blocks to attend (0 = no-op)
841    stream: u64,
842) -> Result<()> {
843    KernelLaunch::new(gpu, kernel)
844        .grid([num_q_heads, num_seqs, 1])
845        .block([256, 1, 1])
846        .arg_ptr(q)
847        .arg_ptr(k_cache)
848        .arg_ptr(v_cache)
849        .arg_ptr(o)
850        .arg_ptr(block_tables)
851        .arg_ptr(seq_lens)
852        .arg_u32(max_blocks_per_seq)
853        .arg_u32(num_q_heads)
854        .arg_u32(num_kv_heads)
855        .arg_u32(q_head_dim)
856        .arg_u32(kv_cache_dim)
857        .arg_u32(block_size)
858        .arg_f32(inv_sqrt_d)
859        .arg_f32(k_scale)
860        .arg_f32(v_scale)
861        .arg_u64(cache_stride)
862        .arg_u32(sliding_window)
863        .arg_ptr(sinks)
864        .arg_ptr(comp_pool)
865        .arg_u32(comp_block_count)
866        .launch(stream)
867}
868
869// ── Batched prefill variants (N tokens) ──
spark_model/layers/ops/kv_cache.rs

spark_model/layers/ops/
kv_cache.rs
