metrale_model_layers/layers/ops/
prefill_attn_a.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the MLA prefill helpers (rope extract and writeback,
4//! K/V and cache assembly, fused prefill, grouped GEMM, absorbed attention) and
5//! the paged decode attention kernels, unpacked and GQA-packed.
6//!
7//! Owner: model-layers ops.
8//! Invariants: none beyond the types.
9
10#![allow(unused_imports)]
11
12use anyhow::Result;
13use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
14use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
15
16use crate::layers::moe;
17use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
18
19use super::*;
20
21/// 2026-09-25: Copy each head's rope slice (offset `nope` in the head's `hd`) out
22/// of Q, rows `q_dim` apart, into contiguous `[num_tokens, nq, rope]`.
23#[allow(clippy::too_many_arguments)]
24pub fn mla_q_rope_extract_batched(
25    gpu: &dyn GpuBackend,
26    kernel: KernelHandle,
27    q_full: DevicePtr,
28    q_rope_out: DevicePtr,
29    num_tokens: u32,
30    nq: u32,
31    hd: u32,
32    nope: u32,
33    rope: u32,
34    q_dim: u32,
35    stream: u64,
36) -> Result<()> {
37    let total = num_tokens * nq * rope;
38    KernelLaunch::new(gpu, kernel)
39        .grid([div_ceil(total, 256), 1, 1])
40        .block([256, 1, 1])
41        .arg_ptr(q_full)
42        .arg_ptr(q_rope_out)
43        .arg_u32(num_tokens)
44        .arg_u32(nq)
45        .arg_u32(hd)
46        .arg_u32(nope)
47        .arg_u32(rope)
48        .arg_u32(q_dim)
49        .launch(stream)
50}
51
52/// 2026-09-25: Inverse of [`mla_q_rope_extract_batched`]: write `[num_tokens, nq,
53/// rope]` back into each head's rope slice of Q.
54#[allow(clippy::too_many_arguments)]
55pub fn mla_q_rope_writeback_batched(
56    gpu: &dyn GpuBackend,
57    kernel: KernelHandle,
58    q_rope_in: DevicePtr,
59    q_full: DevicePtr,
60    num_tokens: u32,
61    nq: u32,
62    hd: u32,
63    nope: u32,
64    rope: u32,
65    q_dim: u32,
66    stream: u64,
67) -> Result<()> {
68    let total = num_tokens * nq * rope;
69    KernelLaunch::new(gpu, kernel)
70        .grid([div_ceil(total, 256), 1, 1])
71        .block([256, 1, 1])
72        .arg_ptr(q_rope_in)
73        .arg_ptr(q_full)
74        .arg_u32(num_tokens)
75        .arg_u32(nq)
76        .arg_u32(hd)
77        .arg_u32(nope)
78        .arg_u32(rope)
79        .arg_u32(q_dim)
80        .launch(stream)
81}
82
83/// 2026-09-25: Per token, build K as `[k_nope | k_rope]` for each KV head (k_rope
84/// shared by all heads) and extract V from `kv_expanded`. `blockIdx.y` 0 builds
85/// K, 1 extracts V.
86#[allow(clippy::too_many_arguments)]
87pub fn mla_kv_assemble_batched(
88    gpu: &dyn GpuBackend,
89    kernel: KernelHandle,
90    kv_expanded: DevicePtr,
91    k_rope_buf: DevicePtr,
92    k_out: DevicePtr,
93    v_out: DevicePtr,
94    num_tokens: u32,
95    nkv: u32,
96    nope: u32,
97    v_dim: u32,
98    rope: u32,
99    hd: u32,
100    kv_expanded_stride: u32,
101    stream: u64,
102) -> Result<()> {
103    KernelLaunch::new(gpu, kernel)
104        .grid([num_tokens, 2, 1])
105        .block([256, 1, 1])
106        .arg_ptr(kv_expanded)
107        .arg_ptr(k_rope_buf)
108        .arg_ptr(k_out)
109        .arg_ptr(v_out)
110        .arg_u32(nkv)
111        .arg_u32(nope)
112        .arg_u32(v_dim)
113        .arg_u32(rope)
114        .arg_u32(hd)
115        .arg_u32(kv_expanded_stride)
116        .launch(stream)
117}
118
119/// 2026-09-26: Per token, write the compressed MLA cache rows: K = `[kv_latent |
120/// k_rope]`, and V = `[kv_latent | zeros]` in the mistral-small-4 kernel or
121/// `[kv_latent | k_rope]` in the deepseek-v4-flash kernel.
122#[allow(clippy::too_many_arguments)]
123pub fn mla_cache_assemble_batched(
124    gpu: &dyn GpuBackend,
125    kernel: KernelHandle,
126    kv_latent: DevicePtr,
127    k_rope: DevicePtr,
128    k_cache: DevicePtr,
129    v_cache: DevicePtr,
130    num_tokens: u32,
131    kv_lora: u32,
132    rope: u32,
133    mla_cache_dim: u32,
134    stream: u64,
135) -> Result<()> {
136    KernelLaunch::new(gpu, kernel)
137        .grid([num_tokens, 1, 1])
138        .block([mla_cache_dim.max(256), 1, 1])
139        .arg_ptr(kv_latent)
140        .arg_ptr(k_rope)
141        .arg_ptr(k_cache)
142        .arg_ptr(v_cache)
143        .arg_u32(kv_lora)
144        .arg_u32(rope)
145        .arg_u32(mla_cache_dim)
146        .launch(stream)
147}
148
149/// 2026-09-25: MLA prefill in one kernel: Q absorption through `w_uk`, attention
150/// over the latent cache, and V extraction through `w_uv`. One block per
151/// `(head, query token)`.
152#[allow(clippy::too_many_arguments)]
153pub fn mla_fused_prefill(
154    gpu: &dyn GpuBackend,
155    kernel: KernelHandle,
156    q_full: DevicePtr,
157    q_rope: DevicePtr,
158    kv_latent: DevicePtr,
159    k_rope: DevicePtr,
160    w_uk: DevicePtr,
161    w_uv: DevicePtr,
162    v_out: DevicePtr,
163    k_cache_out: DevicePtr,
164    v_cache_out: DevicePtr,
165    seq_len: u32,
166    nq: u32,
167    nope: u32,
168    rope: u32,
169    kv_lora: u32,
170    v_dim: u32,
171    hd: u32,
172    num_kv_heads: u32,
173    inv_sqrt_d: f32,
174    stream: u64,
175) -> Result<()> {
176    KernelLaunch::new(gpu, kernel)
177        .grid([nq, seq_len, 1])
178        .block([256, 1, 1])
179        .arg_ptr(q_full)
180        .arg_ptr(q_rope)
181        .arg_ptr(kv_latent)
182        .arg_ptr(k_rope)
183        .arg_ptr(w_uk)
184        .arg_ptr(w_uv)
185        .arg_ptr(v_out)
186        .arg_ptr(k_cache_out)
187        .arg_ptr(v_cache_out)
188        .arg_u32(seq_len)
189        .arg_u32(nq)
190        .arg_u32(nope)
191        .arg_u32(rope)
192        .arg_u32(kv_lora)
193        .arg_u32(v_dim)
194        .arg_u32(hd)
195        .arg_u32(num_kv_heads)
196        .arg_f32(inv_sqrt_d)
197        .launch(stream)
198}
199
200/// 2026-09-25: Build Q as `[q_absorbed (kv_lora) | q_rope (rope)]` per head and
201/// token, `mla_cache_dim = kv_lora + rope` wide.
202#[allow(clippy::too_many_arguments)]
203pub fn mla_q_final_assemble_batched(
204    gpu: &dyn GpuBackend,
205    kernel: KernelHandle,
206    q_absorbed: DevicePtr,
207    q_rope: DevicePtr,
208    q_final: DevicePtr,
209    num_tokens: u32,
210    nq: u32,
211    kv_lora: u32,
212    rope: u32,
213    mla_cache_dim: u32,
214    stream: u64,
215) -> Result<()> {
216    let total = num_tokens * nq * mla_cache_dim;
217    KernelLaunch::new(gpu, kernel)
218        .grid([div_ceil(total, 256), 1, 1])
219        .block([256, 1, 1])
220        .arg_ptr(q_absorbed)
221        .arg_ptr(q_rope)
222        .arg_ptr(q_final)
223        .arg_u32(num_tokens)
224        .arg_u32(nq)
225        .arg_u32(kv_lora)
226        .arg_u32(rope)
227        .arg_u32(mla_cache_dim)
228        .launch(stream)
229}
230
231/// 2026-09-25: `g` independent GEMMs `C_g[m, n_g] = A_g[m, k_g] @ B_g[n_g, k_g]^T`
232/// in one launch. `A_g` sits at column `g * k_g` of A, `B_g` at row `g * n_g` of
233/// B, and `C_g` at column `g * n_g` of C; one block per `(row, group)` computes
234/// 4 outputs.
235#[allow(clippy::too_many_arguments)]
236pub fn grouped_gemm_mla(
237    gpu: &dyn GpuBackend,
238    kernel: KernelHandle,
239    a: DevicePtr,
240    b: DevicePtr,
241    c: DevicePtr,
242    m: u32,
243    g: u32,
244    k_g: u32,
245    n_g: u32,
246    a_stride: u32,
247    c_stride: u32,
248    stream: u64,
249) -> Result<()> {
250    KernelLaunch::new(gpu, kernel)
251        .grid([m * g, div_ceil(n_g, 4), 1])
252        .block([256, 1, 1])
253        .arg_ptr(a)
254        .arg_ptr(b)
255        .arg_ptr(c)
256        .arg_u32(m)
257        .arg_u32(g)
258        .arg_u32(k_g)
259        .arg_u32(n_g)
260        .arg_u32(a_stride)
261        .arg_u32(c_stride)
262        .launch(stream)
263}
264
265/// 2026-09-25: Absorbed MLA prefill attention with scalar BF16 dot products and a
266/// single KV head. The head dim is the kernel's `MLA_HDIM`: 320 in the
267/// mistral-small-4 fork, 576 in the deepseek-v4-flash one; both keep the entry
268/// name `mla_prefill_attn_320`.
269#[allow(clippy::too_many_arguments)]
270pub fn mla_prefill_attention_320(
271    gpu: &dyn GpuBackend,
272    kernel: KernelHandle,
273    q: DevicePtr,
274    k: DevicePtr,
275    v: DevicePtr,
276    output: DevicePtr,
277    seq_len: u32,
278    batch: u32,
279    num_q_heads: u32,
280    num_kv_heads: u32,
281    head_dim: u32,
282    inv_sqrt_d: f32,
283    causal: bool,
284    stream: u64,
285) -> Result<()> {
286    let br = 16u32; // 2026-09-25: `MLA_BR` in both `mla_prefill_attn.cu` forks.
287    KernelLaunch::new(gpu, kernel)
288        .grid([num_q_heads, div_ceil(seq_len, br), batch])
289        .block([256, 1, 1])
290        .arg_ptr(q)
291        .arg_ptr(k)
292        .arg_ptr(v)
293        .arg_ptr(output)
294        .arg_u32(seq_len)
295        .arg_u32(num_q_heads)
296        .arg_u32(num_kv_heads)
297        .arg_u32(head_dim)
298        .arg_f32(inv_sqrt_d)
299        .arg_u32(if causal { 1 } else { 0 })
300        .launch(stream)
301}
302
303/// 2026-09-25: Paged BF16-KV decode attention, one block per `(q_head, seq)`.
304/// `sliding_window == 0` attends to the whole sequence; otherwise only the last
305/// `sliding_window` positions.
306pub fn paged_decode_attn_bf16(
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    sliding_window: u32,
324    stream: u64,
325) -> Result<()> {
326    KernelLaunch::new(gpu, kernel)
327        .grid([num_q_heads, num_seqs, 1])
328        .block([256, 1, 1])
329        .arg_ptr(q)
330        .arg_ptr(k_cache)
331        .arg_ptr(v_cache)
332        .arg_ptr(output)
333        .arg_ptr(block_tables)
334        .arg_ptr(seq_lens)
335        .arg_u32(max_blocks_per_seq)
336        .arg_u32(num_q_heads)
337        .arg_u32(num_kv_heads)
338        .arg_u32(head_dim)
339        .arg_u32(block_size)
340        .arg_f32(inv_sqrt_d)
341        .arg_u32(q_stride)
342        .arg_u32(sliding_window)
343        .launch(stream)
344}
345
346pub fn paged_decode_attn_fp8(
347    gpu: &dyn GpuBackend,
348    kernel: KernelHandle,
349    q: DevicePtr,
350    k_cache: DevicePtr,
351    v_cache: DevicePtr,
352    output: DevicePtr,
353    block_tables: DevicePtr,
354    seq_lens: DevicePtr,
355    max_blocks_per_seq: u32,
356    num_seqs: u32,
357    num_q_heads: u32,
358    num_kv_heads: u32,
359    head_dim: u32,
360    block_size: u32,
361    inv_sqrt_d: f32,
362    k_scale: f32,
363    v_scale: f32,
364    q_stride: u32,
365    cache_stride: u64,
366    sliding_window: u32,
367    stream: u64,
368) -> Result<()> {
369    KernelLaunch::new(gpu, kernel)
370        .grid([num_q_heads, num_seqs, 1])
371        .block([256, 1, 1])
372        .arg_ptr(q)
373        .arg_ptr(k_cache)
374        .arg_ptr(v_cache)
375        .arg_ptr(output)
376        .arg_ptr(block_tables)
377        .arg_ptr(seq_lens)
378        .arg_u32(max_blocks_per_seq)
379        .arg_u32(num_q_heads)
380        .arg_u32(num_kv_heads)
381        .arg_u32(head_dim)
382        .arg_u32(block_size)
383        .arg_f32(inv_sqrt_d)
384        .arg_f32(k_scale)
385        .arg_f32(v_scale)
386        .arg_u32(q_stride)
387        .arg_u64(cache_stride)
388        .arg_u32(sliding_window)
389        .launch(stream)
390}
391
392/// 2026-09-25: GQA-packed BF16-KV paged decode, one block per `(kv_head, seq)`.
393/// Same arguments, argument order and block as [`paged_decode_attn_bf16`]; only
394/// the grid's x extent changes, to `num_kv_heads`, because the kernel
395/// (`kernels/gb10/common/paged_decode_attn_bf16_gqa.cu`) keeps a KV head's
396/// whole query group in registers.
397///
398/// Preconditions: `num_q_heads == num_kv_heads * DECODE_GQA_PACK_WIDTH` and
399/// `head_dim == DECODE_GQA_PACK_HEAD_DIM`. The kernel derives its heads as
400/// `kv_head * PD_GQA + h` and sizes its register arrays from `PD_GQA`. The
401/// caller checks both with `metrale_kernels::attn_splitk::gqa_pack_shape_ok`
402/// and uses the unpacked kernel when they fail.
403#[allow(clippy::too_many_arguments)]
404pub fn paged_decode_attn_bf16_gqa(
405    gpu: &dyn GpuBackend,
406    kernel: KernelHandle,
407    q: DevicePtr,
408    k_cache: DevicePtr,
409    v_cache: DevicePtr,
410    output: DevicePtr,
411    block_tables: DevicePtr,
412    seq_lens: DevicePtr,
413    max_blocks_per_seq: u32,
414    num_seqs: u32,
415    num_q_heads: u32,
416    num_kv_heads: u32,
417    head_dim: u32,
418    block_size: u32,
419    inv_sqrt_d: f32,
420    q_stride: u32,
421    sliding_window: u32,
422    stream: u64,
423) -> Result<()> {
424    KernelLaunch::new(gpu, kernel)
425        .grid([num_kv_heads, num_seqs, 1])
426        .block([256, 1, 1])
427        .arg_ptr(q)
428        .arg_ptr(k_cache)
429        .arg_ptr(v_cache)
430        .arg_ptr(output)
431        .arg_ptr(block_tables)
432        .arg_ptr(seq_lens)
433        .arg_u32(max_blocks_per_seq)
434        .arg_u32(num_q_heads)
435        .arg_u32(num_kv_heads)
436        .arg_u32(head_dim)
437        .arg_u32(block_size)
438        .arg_f32(inv_sqrt_d)
439        .arg_u32(q_stride)
440        .arg_u32(sliding_window)
441        .launch(stream)
442}
443
444/// 2026-09-25: GQA-packed FP8-KV paged decode, one block per `(kv_head, seq)`:
445/// the arguments and order of [`paged_decode_attn_fp8`], grid x extent
446/// `num_kv_heads`, and the preconditions of [`paged_decode_attn_bf16_gqa`].
447#[allow(clippy::too_many_arguments)]
448pub fn paged_decode_attn_fp8_gqa(
449    gpu: &dyn GpuBackend,
450    kernel: KernelHandle,
451    q: DevicePtr,
452    k_cache: DevicePtr,
453    v_cache: DevicePtr,
454    output: DevicePtr,
455    block_tables: DevicePtr,
456    seq_lens: DevicePtr,
457    max_blocks_per_seq: u32,
458    num_seqs: u32,
459    num_q_heads: u32,
460    num_kv_heads: u32,
461    head_dim: u32,
462    block_size: u32,
463    inv_sqrt_d: f32,
464    k_scale: f32,
465    v_scale: f32,
466    q_stride: u32,
467    cache_stride: u64,
468    sliding_window: u32,
469    stream: u64,
470) -> Result<()> {
471    KernelLaunch::new(gpu, kernel)
472        .grid([num_kv_heads, num_seqs, 1])
473        .block([256, 1, 1])
474        .arg_ptr(q)
475        .arg_ptr(k_cache)
476        .arg_ptr(v_cache)
477        .arg_ptr(output)
478        .arg_ptr(block_tables)
479        .arg_ptr(seq_lens)
480        .arg_u32(max_blocks_per_seq)
481        .arg_u32(num_q_heads)
482        .arg_u32(num_kv_heads)
483        .arg_u32(head_dim)
484        .arg_u32(block_size)
485        .arg_f32(inv_sqrt_d)
486        .arg_f32(k_scale)
487        .arg_f32(v_scale)
488        .arg_u32(q_stride)
489        .arg_u64(cache_stride)
490        .arg_u32(sliding_window)
491        .launch(stream)
492}
metrale_model_layers/layers/ops/prefill_attn_a.rs

metrale_model_layers/layers/ops/
prefill_attn_a.rs
