spark_model/layers/ops/
prefill_attn_b.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/// Write K/V to paged NVFP4 cache (E2M1 data + per-group FP8 scales).
17///
18/// Kernel: `reshape_and_cache_flash_nvfp4(key, value, k_cache, v_cache,
19///          slot_mapping, num_kv_heads, head_dim, block_size,
20///          key_stride, value_stride, block_stride_bytes, data_section_bytes)`
21/// Grid: (num_tokens, 1, 1)  Block: (256, 1, 1)
22pub fn reshape_and_cache_nvfp4(
23    gpu: &dyn GpuBackend,
24    kernel: KernelHandle,
25    key: DevicePtr,
26    value: DevicePtr,
27    k_cache: DevicePtr,
28    v_cache: DevicePtr,
29    slot_mapping: DevicePtr,
30    num_tokens: u32,
31    num_kv_heads: u32,
32    head_dim: u32,
33    block_size: u32,
34    key_stride: u32,
35    value_stride: u32,
36    block_stride_bytes: u64,
37    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(block_stride_bytes)
54        .arg_u64(data_section_bytes)
55        .launch(stream)
56}
57
58/// Compute max absolute value of a BF16 buffer into a device-side f32.
59///
60/// Used for FP8 KV cache online scale calibration: accumulates max |K| and
61/// max |V| during warmup tokens. The output f32 is updated via atomicMax,
62/// so the caller must initialize it to 0.0 before the first call.
63///
64/// Kernel: `bf16_absmax(data, out_max, n_elems)`
65/// Grid: (ceil(n_elems / (256*2)), 1, 1)  Block: (256, 1, 1)
66pub fn bf16_absmax(
67    gpu: &dyn GpuBackend,
68    kernel: KernelHandle,
69    data: DevicePtr,
70    out_max: DevicePtr,
71    n_elems: u32,
72    stream: u64,
73) -> Result<()> {
74    // Each thread handles multiple pairs; use enough blocks to cover the buffer.
75    // 256 threads per block, each reads ~8 pairs in the inner loop.
76    let grid_x = (n_elems as u64).div_ceil(256 * 2).min(256) as u32;
77    KernelLaunch::new(gpu, kernel)
78        .grid([grid_x, 1, 1])
79        .block([256, 1, 1])
80        .arg_ptr(data)
81        .arg_ptr(out_max)
82        .arg_u32(n_elems)
83        .launch(stream)
84}
85
86/// Paged decode attention (NVFP4 KV cache, single/multi sequence).
87///
88/// Kernel: `paged_decode_attn_nvfp4(Q, K_cache, V_cache, O, block_tables,
89///          seq_lens, max_blocks_per_seq, num_q_heads, num_kv_heads,
90///          head_dim, block_size, inv_sqrt_d, q_stride,
91///          block_stride_bytes, data_section_bytes)`
92/// Grid: (num_q_heads, num_seqs, 1)  Block: (256, 1, 1)
93pub fn paged_decode_attn_nvfp4(
94    gpu: &dyn GpuBackend,
95    kernel: KernelHandle,
96    q: DevicePtr,
97    k_cache: DevicePtr,
98    v_cache: DevicePtr,
99    output: DevicePtr,
100    block_tables: DevicePtr,
101    seq_lens: DevicePtr,
102    max_blocks_per_seq: u32,
103    num_seqs: u32,
104    num_q_heads: u32,
105    num_kv_heads: u32,
106    head_dim: u32,
107    block_size: u32,
108    inv_sqrt_d: f32,
109    q_stride: u32,
110    block_stride_bytes: u64,
111    data_section_bytes: u64,
112    stream: u64,
113) -> Result<()> {
114    KernelLaunch::new(gpu, kernel)
115        .grid([num_q_heads, num_seqs, 1])
116        .block([256, 1, 1])
117        .arg_ptr(q)
118        .arg_ptr(k_cache)
119        .arg_ptr(v_cache)
120        .arg_ptr(output)
121        .arg_ptr(block_tables)
122        .arg_ptr(seq_lens)
123        .arg_u32(max_blocks_per_seq)
124        .arg_u32(num_q_heads)
125        .arg_u32(num_kv_heads)
126        .arg_u32(head_dim)
127        .arg_u32(block_size)
128        .arg_f32(inv_sqrt_d)
129        .arg_u32(q_stride)
130        .arg_u64(block_stride_bytes)
131        .arg_u64(data_section_bytes)
132        .launch(stream)
133}
134
135/// Split-K paged decode attention (NVFP4 KV cache).
136///
137/// Partitions the KV sequence across `num_splits` CTAs per (q_head, seq).
138/// Each CTA computes partial softmax + weighted output, written to `workspace`.
139///
140/// Grid: (num_q_heads, num_splits, num_seqs)  Block: (256, 1, 1)
141#[allow(clippy::too_many_arguments)]
142pub fn paged_decode_attn_splitk_nvfp4(
143    gpu: &dyn GpuBackend,
144    kernel: KernelHandle,
145    q: DevicePtr,
146    k_cache: DevicePtr,
147    v_cache: DevicePtr,
148    workspace: DevicePtr,
149    block_tables: DevicePtr,
150    seq_lens: DevicePtr,
151    max_blocks_per_seq: u32,
152    num_q_heads: u32,
153    num_kv_heads: u32,
154    head_dim: u32,
155    block_size: u32,
156    inv_sqrt_d: f32,
157    num_splits: u32,
158    q_stride: u32,
159    block_stride_bytes: u64,
160    data_section_bytes: u64,
161    num_seqs: u32,
162    stream: u64,
163) -> Result<()> {
164    KernelLaunch::new(gpu, kernel)
165        .grid([num_q_heads, num_splits, num_seqs])
166        .block([256, 1, 1])
167        .arg_ptr(q)
168        .arg_ptr(k_cache)
169        .arg_ptr(v_cache)
170        .arg_ptr(workspace)
171        .arg_ptr(block_tables)
172        .arg_ptr(seq_lens)
173        .arg_u32(max_blocks_per_seq)
174        .arg_u32(num_q_heads)
175        .arg_u32(num_kv_heads)
176        .arg_u32(head_dim)
177        .arg_u32(block_size)
178        .arg_f32(inv_sqrt_d)
179        .arg_u32(num_splits)
180        .arg_u32(q_stride)
181        .arg_u64(block_stride_bytes)
182        .arg_u64(data_section_bytes)
183        .launch(stream)
184}
185
186/// Reduce split-K partials into final BF16 output.
187///
188/// Grid: (num_q_heads, num_seqs, 1)  Block: (32, 1, 1)
189#[allow(clippy::too_many_arguments)]
190pub fn paged_decode_attn_reduce_nvfp4(
191    gpu: &dyn GpuBackend,
192    kernel: KernelHandle,
193    workspace: DevicePtr,
194    output: DevicePtr,
195    seq_lens: DevicePtr,
196    num_q_heads: u32,
197    head_dim: u32,
198    num_splits: u32,
199    num_seqs: u32,
200    stream: u64,
201) -> Result<()> {
202    KernelLaunch::new(gpu, kernel)
203        .grid([num_q_heads, num_seqs, 1])
204        .block([32, 1, 1])
205        .arg_ptr(workspace)
206        .arg_ptr(output)
207        .arg_ptr(seq_lens)
208        .arg_u32(num_q_heads)
209        .arg_u32(head_dim)
210        .arg_u32(num_splits)
211        .launch(stream)
212}
213
214/// Split-K paged decode attention (FP8 KV cache).
215///
216/// Partitions the KV sequence across `num_splits` CTAs per (q_head, seq).
217/// Each CTA computes partial softmax + weighted output, written to `workspace`.
218///
219/// Grid: (num_q_heads, num_splits, num_seqs)  Block: (256, 1, 1)
220#[allow(clippy::too_many_arguments)]
221pub fn paged_decode_attn_splitk_fp8(
222    gpu: &dyn GpuBackend,
223    kernel: KernelHandle,
224    q: DevicePtr,
225    k_cache: DevicePtr,
226    v_cache: DevicePtr,
227    workspace: DevicePtr,
228    block_tables: DevicePtr,
229    seq_lens: DevicePtr,
230    max_blocks_per_seq: u32,
231    num_q_heads: u32,
232    num_kv_heads: u32,
233    head_dim: u32,
234    block_size: u32,
235    inv_sqrt_d: f32,
236    num_splits: u32,
237    k_scale: f32,
238    v_scale: f32,
239    q_stride: u32,
240    cache_stride: u64,
241    num_seqs: u32,
242    sliding_window: u32,
243    stream: u64,
244) -> Result<()> {
245    KernelLaunch::new(gpu, kernel)
246        .grid([num_q_heads, num_splits, num_seqs])
247        .block([256, 1, 1])
248        .arg_ptr(q)
249        .arg_ptr(k_cache)
250        .arg_ptr(v_cache)
251        .arg_ptr(workspace)
252        .arg_ptr(block_tables)
253        .arg_ptr(seq_lens)
254        .arg_u32(max_blocks_per_seq)
255        .arg_u32(num_q_heads)
256        .arg_u32(num_kv_heads)
257        .arg_u32(head_dim)
258        .arg_u32(block_size)
259        .arg_f32(inv_sqrt_d)
260        .arg_u32(num_splits)
261        .arg_f32(k_scale)
262        .arg_f32(v_scale)
263        .arg_u32(q_stride)
264        .arg_u64(cache_stride)
265        .arg_u32(sliding_window)
266        .launch(stream)
267}
268
269/// Reduce split-K partials into final BF16 output (FP8 variant).
270///
271/// Grid: (num_q_heads, num_seqs, 1)  Block: (32, 1, 1)
272#[allow(clippy::too_many_arguments)]
273pub fn paged_decode_attn_reduce_fp8(
274    gpu: &dyn GpuBackend,
275    kernel: KernelHandle,
276    workspace: DevicePtr,
277    output: DevicePtr,
278    seq_lens: DevicePtr,
279    num_q_heads: u32,
280    head_dim: u32,
281    num_splits: u32,
282    num_seqs: u32,
283    stream: u64,
284) -> Result<()> {
285    KernelLaunch::new(gpu, kernel)
286        .grid([num_q_heads, num_seqs, 1])
287        .block([32, 1, 1])
288        .arg_ptr(workspace)
289        .arg_ptr(output)
290        .arg_ptr(seq_lens)
291        .arg_u32(num_q_heads)
292        .arg_u32(head_dim)
293        .arg_u32(num_splits)
294        .launch(stream)
295}
296
297// ── SSM / Convolution ──────────────────────────────────────────────
spark_model/layers/ops/prefill_attn_b.rs

spark_model/layers/ops/
prefill_attn_b.rs
