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/// Write K/V to paged Bf16K + Turbo3V (TurboQuant+ safer-asym) cache.
246///
247/// K is written as raw BF16 (NHD contiguous), V as 3-bit Lloyd-Max + FP8
248/// per-group scale with matched-norm correction. K and V pools have separate
249/// strides because K is 2 b/elem and V is ~0.5 b/elem + scale.
250///
251/// Kernel: `reshape_and_cache_flash_bf16k_turbo3v(key, value, k_cache, v_cache,
252///          slot_mapping, num_kv_heads, head_dim, block_size,
253///          key_stride, value_stride, k_block_stride_bytes,
254///          v_block_stride_bytes, v_data_section_bytes)`
255/// Grid: (num_tokens, 1, 1)  Block: (256, 1, 1)
256#[allow(clippy::too_many_arguments)]
257pub fn reshape_and_cache_bf16k_turbo3v(
258    gpu: &dyn GpuBackend,
259    kernel: KernelHandle,
260    key: DevicePtr,
261    value: DevicePtr,
262    k_cache: DevicePtr,
263    v_cache: DevicePtr,
264    slot_mapping: DevicePtr,
265    num_tokens: u32,
266    num_kv_heads: u32,
267    head_dim: u32,
268    block_size: u32,
269    key_stride: u32,
270    value_stride: u32,
271    k_block_stride_bytes: u64,
272    v_block_stride_bytes: u64,
273    v_data_section_bytes: u64,
274    stream: u64,
275) -> Result<()> {
276    KernelLaunch::new(gpu, kernel)
277        .grid([num_tokens, 1, 1])
278        .block([256, 1, 1])
279        .arg_ptr(key)
280        .arg_ptr(value)
281        .arg_ptr(k_cache)
282        .arg_ptr(v_cache)
283        .arg_ptr(slot_mapping)
284        .arg_u32(num_kv_heads)
285        .arg_u32(head_dim)
286        .arg_u32(block_size)
287        .arg_u32(key_stride)
288        .arg_u32(value_stride)
289        .arg_u64(k_block_stride_bytes)
290        .arg_u64(v_block_stride_bytes)
291        .arg_u64(v_data_section_bytes)
292        .launch(stream)
293}
294
295/// Paged decode attention for Bf16K + Turbo3V asymmetric KV cache.
296///
297/// K is read as BF16 NHD (vector loads), V as 3-bit Lloyd-Max packed bytes
298/// with FP8 per-group scale (sparse V on batched AND remainder paths).
299///
300/// Kernel: `paged_decode_attn_bf16k_turbo3v(Q, K_cache, V_cache, O,
301///          block_tables, seq_lens, max_blocks_per_seq, num_q_heads,
302///          num_kv_heads, head_dim, block_size, inv_sqrt_d, q_stride,
303///          v_block_stride_bytes, v_data_section_bytes, sliding_window)`
304/// Grid: (num_q_heads, num_seqs, 1)  Block: (256, 1, 1)
305#[allow(clippy::too_many_arguments)]
306pub fn paged_decode_attn_bf16k_turbo3v(
307    gpu: &dyn GpuBackend,
308    kernel: KernelHandle,
309    q: DevicePtr,
310    k_cache: DevicePtr,
311    v_cache: DevicePtr,
312    output: DevicePtr,
313    block_tables: DevicePtr,
314    seq_lens: DevicePtr,
315    max_blocks_per_seq: u32,
316    num_seqs: u32,
317    num_q_heads: u32,
318    num_kv_heads: u32,
319    head_dim: u32,
320    block_size: u32,
321    inv_sqrt_d: f32,
322    q_stride: u32,
323    v_block_stride_bytes: u64,
324    v_data_section_bytes: u64,
325    sliding_window: u32,
326    stream: u64,
327) -> Result<()> {
328    KernelLaunch::new(gpu, kernel)
329        .grid([num_q_heads, num_seqs, 1])
330        .block([256, 1, 1])
331        .arg_ptr(q)
332        .arg_ptr(k_cache)
333        .arg_ptr(v_cache)
334        .arg_ptr(output)
335        .arg_ptr(block_tables)
336        .arg_ptr(seq_lens)
337        .arg_u32(max_blocks_per_seq)
338        .arg_u32(num_q_heads)
339        .arg_u32(num_kv_heads)
340        .arg_u32(head_dim)
341        .arg_u32(block_size)
342        .arg_f32(inv_sqrt_d)
343        .arg_u32(q_stride)
344        .arg_u64(v_block_stride_bytes)
345        .arg_u64(v_data_section_bytes)
346        .arg_u32(sliding_window)
347        .launch(stream)
348}
349
350/// Write K/V to paged Bf16K + Turbo4V (TurboQuant+ safer-asym) cache.
351///
352/// K written as raw BF16 NHD, V as 4-bit Lloyd-Max + FP8 per-group scale with
353/// matched-norm correction. K and V pools have separate strides.
354#[allow(clippy::too_many_arguments)]
355pub fn reshape_and_cache_bf16k_turbo4v(
356    gpu: &dyn GpuBackend,
357    kernel: KernelHandle,
358    key: DevicePtr,
359    value: DevicePtr,
360    k_cache: DevicePtr,
361    v_cache: DevicePtr,
362    slot_mapping: DevicePtr,
363    num_tokens: u32,
364    num_kv_heads: u32,
365    head_dim: u32,
366    block_size: u32,
367    key_stride: u32,
368    value_stride: u32,
369    k_block_stride_bytes: u64,
370    v_block_stride_bytes: u64,
371    v_data_section_bytes: u64,
372    stream: u64,
373) -> Result<()> {
374    KernelLaunch::new(gpu, kernel)
375        .grid([num_tokens, 1, 1])
376        .block([256, 1, 1])
377        .arg_ptr(key)
378        .arg_ptr(value)
379        .arg_ptr(k_cache)
380        .arg_ptr(v_cache)
381        .arg_ptr(slot_mapping)
382        .arg_u32(num_kv_heads)
383        .arg_u32(head_dim)
384        .arg_u32(block_size)
385        .arg_u32(key_stride)
386        .arg_u32(value_stride)
387        .arg_u64(k_block_stride_bytes)
388        .arg_u64(v_block_stride_bytes)
389        .arg_u64(v_data_section_bytes)
390        .launch(stream)
391}
392
393/// Write K/V to paged Bf16K + Turbo2V (TurboQuant+ safer-asym) cache (6.4x V comp).
394#[allow(clippy::too_many_arguments)]
395pub fn reshape_and_cache_bf16k_turbo2v(
396    gpu: &dyn GpuBackend,
397    kernel: KernelHandle,
398    key: DevicePtr,
399    value: DevicePtr,
400    k_cache: DevicePtr,
401    v_cache: DevicePtr,
402    slot_mapping: DevicePtr,
403    num_tokens: u32,
404    num_kv_heads: u32,
405    head_dim: u32,
406    block_size: u32,
407    key_stride: u32,
408    value_stride: u32,
409    k_block_stride_bytes: u64,
410    v_block_stride_bytes: u64,
411    v_data_section_bytes: u64,
412    stream: u64,
413) -> Result<()> {
414    KernelLaunch::new(gpu, kernel)
415        .grid([num_tokens, 1, 1])
416        .block([256, 1, 1])
417        .arg_ptr(key)
418        .arg_ptr(value)
419        .arg_ptr(k_cache)
420        .arg_ptr(v_cache)
421        .arg_ptr(slot_mapping)
422        .arg_u32(num_kv_heads)
423        .arg_u32(head_dim)
424        .arg_u32(block_size)
425        .arg_u32(key_stride)
426        .arg_u32(value_stride)
427        .arg_u64(k_block_stride_bytes)
428        .arg_u64(v_block_stride_bytes)
429        .arg_u64(v_data_section_bytes)
430        .launch(stream)
431}
432
433/// Paged decode attention for Bf16K + Turbo4V asymmetric KV cache.
434///
435/// K read as BF16 NHD, V as 4-bit Lloyd-Max packed bytes + FP8 per-group scale
436/// (sparse V on batched + remainder paths).
437#[allow(clippy::too_many_arguments)]
438pub fn paged_decode_attn_bf16k_turbo4v(
439    gpu: &dyn GpuBackend,
440    kernel: KernelHandle,
441    q: DevicePtr,
442    k_cache: DevicePtr,
443    v_cache: DevicePtr,
444    output: DevicePtr,
445    block_tables: DevicePtr,
446    seq_lens: DevicePtr,
447    max_blocks_per_seq: u32,
448    num_seqs: u32,
449    num_q_heads: u32,
450    num_kv_heads: u32,
451    head_dim: u32,
452    block_size: u32,
453    inv_sqrt_d: f32,
454    q_stride: u32,
455    v_block_stride_bytes: u64,
456    v_data_section_bytes: u64,
457    sliding_window: u32,
458    stream: u64,
459) -> Result<()> {
460    KernelLaunch::new(gpu, kernel)
461        .grid([num_q_heads, num_seqs, 1])
462        .block([256, 1, 1])
463        .arg_ptr(q)
464        .arg_ptr(k_cache)
465        .arg_ptr(v_cache)
466        .arg_ptr(output)
467        .arg_ptr(block_tables)
468        .arg_ptr(seq_lens)
469        .arg_u32(max_blocks_per_seq)
470        .arg_u32(num_q_heads)
471        .arg_u32(num_kv_heads)
472        .arg_u32(head_dim)
473        .arg_u32(block_size)
474        .arg_f32(inv_sqrt_d)
475        .arg_u32(q_stride)
476        .arg_u64(v_block_stride_bytes)
477        .arg_u64(v_data_section_bytes)
478        .arg_u32(sliding_window)
479        .launch(stream)
480}
481
482/// Paged decode attention for Bf16K + Turbo2V asymmetric KV cache (6.4x V comp).
483#[allow(clippy::too_many_arguments)]
484pub fn paged_decode_attn_bf16k_turbo2v(
485    gpu: &dyn GpuBackend,
486    kernel: KernelHandle,
487    q: DevicePtr,
488    k_cache: DevicePtr,
489    v_cache: DevicePtr,
490    output: DevicePtr,
491    block_tables: DevicePtr,
492    seq_lens: DevicePtr,
493    max_blocks_per_seq: u32,
494    num_seqs: u32,
495    num_q_heads: u32,
496    num_kv_heads: u32,
497    head_dim: u32,
498    block_size: u32,
499    inv_sqrt_d: f32,
500    q_stride: u32,
501    v_block_stride_bytes: u64,
502    v_data_section_bytes: u64,
503    sliding_window: u32,
504    stream: u64,
505) -> Result<()> {
506    KernelLaunch::new(gpu, kernel)
507        .grid([num_q_heads, num_seqs, 1])
508        .block([256, 1, 1])
509        .arg_ptr(q)
510        .arg_ptr(k_cache)
511        .arg_ptr(v_cache)
512        .arg_ptr(output)
513        .arg_ptr(block_tables)
514        .arg_ptr(seq_lens)
515        .arg_u32(max_blocks_per_seq)
516        .arg_u32(num_q_heads)
517        .arg_u32(num_kv_heads)
518        .arg_u32(head_dim)
519        .arg_u32(block_size)
520        .arg_f32(inv_sqrt_d)
521        .arg_u32(q_stride)
522        .arg_u64(v_block_stride_bytes)
523        .arg_u64(v_data_section_bytes)
524        .arg_u32(sliding_window)
525        .launch(stream)
526}
527
528/// `k_cache`/`v_cache` are the full pool base pointers.
529/// `cache_stride` is in elements (block_size * num_kv_heads * head_dim).
530pub fn reshape_and_cache_fp8(
531    gpu: &dyn GpuBackend,
532    kernel: KernelHandle,
533    key: DevicePtr,
534    value: DevicePtr,
535    k_cache: DevicePtr,
536    v_cache: DevicePtr,
537    slot_mapping: DevicePtr,
538    num_tokens: u32,
539    num_kv_heads: u32,
540    head_dim: u32,
541    block_size: u32,
542    k_scale: f32,
543    v_scale: f32,
544    key_stride: u32,
545    value_stride: u32,
546    cache_stride: u64,
547    stream: u64,
548) -> Result<()> {
549    KernelLaunch::new(gpu, kernel)
550        .grid([num_tokens, 1, 1])
551        .block([256, 1, 1])
552        .arg_ptr(key)
553        .arg_ptr(value)
554        .arg_ptr(k_cache)
555        .arg_ptr(v_cache)
556        .arg_ptr(slot_mapping)
557        .arg_u32(num_kv_heads)
558        .arg_u32(head_dim)
559        .arg_u32(block_size)
560        .arg_f32(k_scale)
561        .arg_f32(v_scale)
562        .arg_u32(key_stride)
563        .arg_u32(value_stride)
564        .arg_u64(cache_stride)
565        .launch(stream)
566}
567
568/// Paged decode attention (FP8 KV cache, single/multi sequence).
569///
570/// Kernel: `paged_decode_attn_fp8(Q, K_cache, V_cache, O, block_tables,
571///          seq_lens, max_blocks_per_seq, num_q_heads, num_kv_heads,
572///          head_dim, block_size, inv_sqrt_d, k_scale, v_scale,
573///          q_stride, cache_stride)`
574/// Grid: (num_q_heads, num_seqs, 1)  Block: (256, 1, 1)
575///
576/// `block_tables`: device ptr to `i32[num_seqs * max_blocks_per_seq]`
577/// `seq_lens`: device ptr to `i32[num_seqs]`
578/// `cache_stride` is in elements (u64).
579/// BF16 paged decode attention — no FP8 quantization, direct BF16 KV cache.
580/// MLA batched GEMV: output[head, n] = sum_k(weight[head, n, k] * input[head, k])
581/// Replaces 32 sequential dense_gemv calls with a single kernel launch.
582pub fn mla_batched_gemv(
583    gpu: &dyn GpuBackend,
584    kernel: KernelHandle,
585    input: DevicePtr,
586    weight: DevicePtr,
587    output: DevicePtr,
588    n_out: u32,
589    k: u32,
590    num_heads: u32,
591    input_stride: u32,
592    output_stride: u32,
593    stream: u64,
594) -> Result<()> {
595    KernelLaunch::new(gpu, kernel)
596        .grid([div_ceil(n_out, 8), num_heads, 1]) // N_PER_BLOCK*2=8
597        .block([256, 1, 1])
598        .arg_ptr(input)
599        .arg_ptr(weight)
600        .arg_ptr(output)
601        .arg_u32(n_out)
602        .arg_u32(k)
603        .arg_u32(input_stride)
604        .arg_u32(output_stride)
605        .launch(stream)
606}
607
608/// MLA Q_rope scatter: copy rope portion from q_full to strided q_absorbed_buf. 1 kernel replaces 32 D2D copies.
609#[allow(clippy::too_many_arguments)]
610pub fn mla_q_rope_scatter(
611    gpu: &dyn GpuBackend,
612    kernel: KernelHandle,
613    q_full: DevicePtr,
614    q_absorbed_buf: DevicePtr,
615    q_rope_contiguous: DevicePtr,
616    nq: u32,
617    hd: u32,
618    nope: u32,
619    rope: u32,
620    kv_lora: u32,
621    mla_cache_dim: u32,
622    stream: u64,
623) -> Result<()> {
624    KernelLaunch::new(gpu, kernel)
625        .grid([1, 1, 1])
626        .block([256, 1, 1])
627        .arg_ptr(q_full)
628        .arg_ptr(q_absorbed_buf)
629        .arg_ptr(q_rope_contiguous)
630        .arg_u32(nq)
631        .arg_u32(hd)
632        .arg_u32(nope)
633        .arg_u32(rope)
634        .arg_u32(kv_lora)
635        .arg_u32(mla_cache_dim)
636        .launch(stream)
637}
638
639/// MLA Q_rope writeback: scatter RoPE'd rope portions to strided layout. 1 kernel replaces 32 D2D copies.
640pub fn mla_q_rope_writeback(
641    gpu: &dyn GpuBackend,
642    kernel: KernelHandle,
643    q_rope_direct: DevicePtr,
644    q_absorbed_buf: DevicePtr,
645    nq: u32,
646    rope: u32,
647    kv_lora: u32,
648    mla_cache_dim: u32,
649    stream: u64,
650) -> Result<()> {
651    KernelLaunch::new(gpu, kernel)
652        .grid([1, 1, 1])
653        .block([256, 1, 1])
654        .arg_ptr(q_rope_direct)
655        .arg_ptr(q_absorbed_buf)
656        .arg_u32(nq)
657        .arg_u32(rope)
658        .arg_u32(kv_lora)
659        .arg_u32(mla_cache_dim)
660        .launch(stream)
661}
662
663/// MLA cache assembly: fuse [kv_latent|k_rope]→K and [kv_latent|zeros]→V into 1 kernel.
664pub fn mla_cache_assemble(
665    gpu: &dyn GpuBackend,
666    kernel: KernelHandle,
667    kv_latent: DevicePtr,
668    k_rope: DevicePtr,
669    k_cache: DevicePtr,
670    v_cache: DevicePtr,
671    kv_lora: u32,
672    rope: u32,
673    mla_cache_dim: u32,
674    stream: u64,
675) -> Result<()> {
676    KernelLaunch::new(gpu, kernel)
677        .grid([1, 1, 1])
678        .block([mla_cache_dim.max(256), 1, 1])
679        .arg_ptr(kv_latent)
680        .arg_ptr(k_rope)
681        .arg_ptr(k_cache)
682        .arg_ptr(v_cache)
683        .arg_u32(kv_lora)
684        .arg_u32(rope)
685        .arg_u32(mla_cache_dim)
686        .launch(stream)
687}
688
689/// MLA Paged Decode — NVFP4 variant for DeepSeek-V4-Flash.
690///
691/// Kernel: `mla_paged_decode_nvfp4(Q, K_cache, V_cache, O, block_tables,
692///          seq_lens, max_blocks_per_seq, num_q_heads, num_kv_heads,
693///          q_head_dim, kv_cache_dim, block_size, inv_sqrt_d,
694///          block_stride_bytes, data_section_bytes)`
695/// Grid: (num_q_heads, num_seqs, 1)  Block: (256, 1, 1)
696///
697/// V4-Flash uses compressed KV cache (576 dims: 512 latent + 64 rope)
698/// and flattened Q layout (32768 dims = 64 heads × 512).
699#[allow(clippy::too_many_arguments)]
700pub fn mla_paged_decode_nvfp4(
701    gpu: &dyn GpuBackend,
702    kernel: KernelHandle,
703    q: DevicePtr,
704    k_cache: DevicePtr,
705    v_cache: DevicePtr,
706    o: DevicePtr,
707    block_tables: DevicePtr,
708    seq_lens: DevicePtr,
709    max_blocks_per_seq: u32,
710    num_q_heads: u32,
711    num_kv_heads: u32,
712    q_head_dim: u32,
713    kv_cache_dim: u32,
714    block_size: u32,
715    inv_sqrt_d: f32,
716    block_stride_bytes: u64,
717    data_section_bytes: u64,
718    num_seqs: u32,
719    stream: u64,
720) -> Result<()> {
721    KernelLaunch::new(gpu, kernel)
722        .grid([num_q_heads, num_seqs, 1])
723        .block([256, 1, 1])
724        .arg_ptr(q)
725        .arg_ptr(k_cache)
726        .arg_ptr(v_cache)
727        .arg_ptr(o)
728        .arg_ptr(block_tables)
729        .arg_ptr(seq_lens)
730        .arg_u32(max_blocks_per_seq)
731        .arg_u32(num_q_heads)
732        .arg_u32(num_kv_heads)
733        .arg_u32(q_head_dim)
734        .arg_u32(kv_cache_dim)
735        .arg_u32(block_size)
736        .arg_f32(inv_sqrt_d)
737        .arg_u64(block_stride_bytes)
738        .arg_u64(data_section_bytes)
739        .launch(stream)
740}
741
742/// MLA Paged Decode — FP8 variant for DeepSeek-V4-Flash with FP8 KV cache.
743///
744/// Kernel: `mla_paged_decode_fp8(Q, K_cache, V_cache, O, block_tables,
745///          seq_lens, max_blocks_per_seq, num_q_heads, num_kv_heads,
746///          q_head_dim, kv_cache_dim, block_size, inv_sqrt_d,
747///          k_scale, v_scale, cache_stride)`
748/// Grid: (num_q_heads, num_seqs, 1)  Block: (256, 1, 1)
749///
750/// V4-Flash uses compressed KV cache (576 dims: 512 latent + 64 rope)
751/// and FP8 quantization with per-layer scales.
752#[allow(clippy::too_many_arguments)]
753pub fn mla_paged_decode_fp8(
754    gpu: &dyn GpuBackend,
755    kernel: KernelHandle,
756    q: DevicePtr,
757    k_cache: DevicePtr,
758    v_cache: DevicePtr,
759    o: DevicePtr,
760    block_tables: DevicePtr,
761    seq_lens: DevicePtr,
762    max_blocks_per_seq: u32,
763    num_q_heads: u32,
764    num_kv_heads: u32,
765    q_head_dim: u32,
766    kv_cache_dim: u32,
767    block_size: u32,
768    inv_sqrt_d: f32,
769    k_scale: f32,
770    v_scale: f32,
771    cache_stride: u64,
772    num_seqs: u32,
773    sliding_window: u32,
774    sinks: DevicePtr,
775    comp_pool: DevicePtr, // 4b: flat FP8 compressed-KV pool (NULL = no compressed arm)
776    comp_block_count: u32, // 4b: # compressed blocks to attend (0 = no-op)
777    stream: u64,
778) -> Result<()> {
779    KernelLaunch::new(gpu, kernel)
780        .grid([num_q_heads, num_seqs, 1])
781        .block([256, 1, 1])
782        .arg_ptr(q)
783        .arg_ptr(k_cache)
784        .arg_ptr(v_cache)
785        .arg_ptr(o)
786        .arg_ptr(block_tables)
787        .arg_ptr(seq_lens)
788        .arg_u32(max_blocks_per_seq)
789        .arg_u32(num_q_heads)
790        .arg_u32(num_kv_heads)
791        .arg_u32(q_head_dim)
792        .arg_u32(kv_cache_dim)
793        .arg_u32(block_size)
794        .arg_f32(inv_sqrt_d)
795        .arg_f32(k_scale)
796        .arg_f32(v_scale)
797        .arg_u64(cache_stride)
798        .arg_u32(sliding_window)
799        .arg_ptr(sinks)
800        .arg_ptr(comp_pool)
801        .arg_u32(comp_block_count)
802        .launch(stream)
803}
804
805// ── Batched prefill variants (N tokens) ──
spark_model/layers/ops/kv_cache.rs

spark_model/layers/ops/
kv_cache.rs
