spark_model/layers/ops/
prefill_attn_turbok.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! TurboQuant+ both-sides-quantized Turbo*K + Turbo*V prefill (BR=64) ops wrappers.
4//!
5//! Sibling of `prefill_attn_main_b.rs` / `prefill_attn_fp8k.rs`. Each wrapper
6//! takes BOTH K-side and V-side (block_stride_bytes, data_section_bytes) pairs
7//! since the two pools have independent byte layouts.
8
9#![allow(unused_imports)]
10
11use anyhow::Result;
12use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
13use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
14
15/// Prefill paged attention — TurboQuant+ asym Turbo4K + Turbo3V (BR=64).
16#[allow(clippy::too_many_arguments)]
17pub fn prefill_attention_paged_turbo4k_turbo3v_64(
18    gpu: &dyn GpuBackend,
19    kernel: KernelHandle,
20    q: DevicePtr,
21    k_cache: DevicePtr,
22    v_cache: DevicePtr,
23    output: DevicePtr,
24    block_table: DevicePtr,
25    q_len: u32,
26    kv_len: u32,
27    q_offset: u32,
28    num_q_heads: u32,
29    num_kv_heads: u32,
30    head_dim: u32,
31    cache_block_size: u32,
32    sliding_window: u32,
33    inv_sqrt_d: f32,
34    k_block_stride_bytes: u64,
35    k_data_section_bytes: u64,
36    v_block_stride_bytes: u64,
37    v_data_section_bytes: u64,
38    stream: u64,
39) -> Result<()> {
40    let br = 64u32;
41    KernelLaunch::new(gpu, kernel)
42        .grid([num_q_heads, div_ceil(q_len, br), 1])
43        .block([256, 1, 1])
44        .arg_ptr(q)
45        .arg_ptr(k_cache)
46        .arg_ptr(v_cache)
47        .arg_ptr(output)
48        .arg_ptr(block_table)
49        .arg_u32(q_len)
50        .arg_u32(kv_len)
51        .arg_u32(q_offset)
52        .arg_u32(num_q_heads)
53        .arg_u32(num_kv_heads)
54        .arg_u32(head_dim)
55        .arg_u32(cache_block_size)
56        .arg_u32(sliding_window)
57        .arg_u32(1u32)
58        .arg_f32(inv_sqrt_d)
59        .arg_u64(k_block_stride_bytes)
60        .arg_u64(k_data_section_bytes)
61        .arg_u64(v_block_stride_bytes)
62        .arg_u64(v_data_section_bytes)
63        .launch(stream)
64}
65
66/// Prefill paged attention — TurboQuant+ asym Turbo4K + Turbo8V (BR=64).
67#[allow(clippy::too_many_arguments)]
68pub fn prefill_attention_paged_turbo4k_turbo8v_64(
69    gpu: &dyn GpuBackend,
70    kernel: KernelHandle,
71    q: DevicePtr,
72    k_cache: DevicePtr,
73    v_cache: DevicePtr,
74    output: DevicePtr,
75    block_table: DevicePtr,
76    q_len: u32,
77    kv_len: u32,
78    q_offset: u32,
79    num_q_heads: u32,
80    num_kv_heads: u32,
81    head_dim: u32,
82    cache_block_size: u32,
83    sliding_window: u32,
84    inv_sqrt_d: f32,
85    k_block_stride_bytes: u64,
86    k_data_section_bytes: u64,
87    v_block_stride_bytes: u64,
88    v_data_section_bytes: u64,
89    stream: u64,
90) -> Result<()> {
91    let br = 64u32;
92    KernelLaunch::new(gpu, kernel)
93        .grid([num_q_heads, div_ceil(q_len, br), 1])
94        .block([256, 1, 1])
95        .arg_ptr(q)
96        .arg_ptr(k_cache)
97        .arg_ptr(v_cache)
98        .arg_ptr(output)
99        .arg_ptr(block_table)
100        .arg_u32(q_len)
101        .arg_u32(kv_len)
102        .arg_u32(q_offset)
103        .arg_u32(num_q_heads)
104        .arg_u32(num_kv_heads)
105        .arg_u32(head_dim)
106        .arg_u32(cache_block_size)
107        .arg_u32(sliding_window)
108        .arg_u32(1u32)
109        .arg_f32(inv_sqrt_d)
110        .arg_u64(k_block_stride_bytes)
111        .arg_u64(k_data_section_bytes)
112        .arg_u64(v_block_stride_bytes)
113        .arg_u64(v_data_section_bytes)
114        .launch(stream)
115}
116
117/// Prefill paged attention — TurboQuant+ asym Turbo3K + Turbo8V (BR=64).
118#[allow(clippy::too_many_arguments)]
119pub fn prefill_attention_paged_turbo3k_turbo8v_64(
120    gpu: &dyn GpuBackend,
121    kernel: KernelHandle,
122    q: DevicePtr,
123    k_cache: DevicePtr,
124    v_cache: DevicePtr,
125    output: DevicePtr,
126    block_table: DevicePtr,
127    q_len: u32,
128    kv_len: u32,
129    q_offset: u32,
130    num_q_heads: u32,
131    num_kv_heads: u32,
132    head_dim: u32,
133    cache_block_size: u32,
134    sliding_window: u32,
135    inv_sqrt_d: f32,
136    k_block_stride_bytes: u64,
137    k_data_section_bytes: u64,
138    v_block_stride_bytes: u64,
139    v_data_section_bytes: u64,
140    stream: u64,
141) -> Result<()> {
142    let br = 64u32;
143    KernelLaunch::new(gpu, kernel)
144        .grid([num_q_heads, div_ceil(q_len, br), 1])
145        .block([256, 1, 1])
146        .arg_ptr(q)
147        .arg_ptr(k_cache)
148        .arg_ptr(v_cache)
149        .arg_ptr(output)
150        .arg_ptr(block_table)
151        .arg_u32(q_len)
152        .arg_u32(kv_len)
153        .arg_u32(q_offset)
154        .arg_u32(num_q_heads)
155        .arg_u32(num_kv_heads)
156        .arg_u32(head_dim)
157        .arg_u32(cache_block_size)
158        .arg_u32(sliding_window)
159        .arg_u32(1u32)
160        .arg_f32(inv_sqrt_d)
161        .arg_u64(k_block_stride_bytes)
162        .arg_u64(k_data_section_bytes)
163        .arg_u64(v_block_stride_bytes)
164        .arg_u64(v_data_section_bytes)
165        .launch(stream)
166}
spark_model/layers/ops/prefill_attn_turbok.rs

spark_model/layers/ops/
prefill_attn_turbok.rs
