metrale_model_layers/layers/ops/
kv_cache_mla.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the MLA decode kernels: per-head batched GEMV, Q RoPE scatter and writeback, cache-entry assembly, and the NVFP4 and FP8 paged decode.
4//!
5//! Owner: model-layers ops.
6//! Invariants:
7//! - The scatter, writeback and assembly launchers run a single block, so they handle one
8//!   token (kernels/gb10/deepseek-v4-flash/nvfp4/mla_absorbed.cu).
9//! - The paged decode launchers run one 256-thread block per (query head, sequence) and pass
10//!   their kernel's parameters in declaration order.
11
12use super::*;
13
14/// 2026-09-25: MLA batched GEMV over all heads in one launch:
15/// `output[head, n] = sum_k(weight[head, n, k] * input[head, k])`, BF16, with the per-head input
16/// and output at `input_stride` / `output_stride` and the weight contiguous per head.
17pub fn mla_batched_gemv(
18    gpu: &dyn GpuBackend,
19    kernel: KernelHandle,
20    input: DevicePtr,
21    weight: DevicePtr,
22    output: DevicePtr,
23    n_out: u32,
24    k: u32,
25    num_heads: u32,
26    input_stride: u32,
27    output_stride: u32,
28    stream: u64,
29) -> Result<()> {
30    KernelLaunch::new(gpu, kernel)
31        .grid([div_ceil(n_out, 8), num_heads, 1]) // 2026-09-25: N_PER_BLOCK * 2 = 8 outputs per block.
32        .block([256, 1, 1])
33        .arg_ptr(input)
34        .arg_ptr(weight)
35        .arg_ptr(output)
36        .arg_u32(n_out)
37        .arg_u32(k)
38        .arg_u32(input_stride)
39        .arg_u32(output_stride)
40        .launch(stream)
41}
42
43/// 2026-09-25: Copy each head's rope part of `q_full` (`[nq, hd]`, at offset `nope`) into
44/// `q_absorbed_buf` (`[nq, mla_cache_dim]`, at offset `kv_lora`) and into the contiguous
45/// `q_rope_contiguous` (`[nq, rope]`) in one pass.
46#[allow(clippy::too_many_arguments)]
47pub fn mla_q_rope_scatter(
48    gpu: &dyn GpuBackend,
49    kernel: KernelHandle,
50    q_full: DevicePtr,
51    q_absorbed_buf: DevicePtr,
52    q_rope_contiguous: DevicePtr,
53    nq: u32,
54    hd: u32,
55    nope: u32,
56    rope: u32,
57    kv_lora: u32,
58    mla_cache_dim: u32,
59    stream: u64,
60) -> Result<()> {
61    KernelLaunch::new(gpu, kernel)
62        .grid([1, 1, 1])
63        .block([256, 1, 1])
64        .arg_ptr(q_full)
65        .arg_ptr(q_absorbed_buf)
66        .arg_ptr(q_rope_contiguous)
67        .arg_u32(nq)
68        .arg_u32(hd)
69        .arg_u32(nope)
70        .arg_u32(rope)
71        .arg_u32(kv_lora)
72        .arg_u32(mla_cache_dim)
73        .launch(stream)
74}
75
76/// 2026-09-25: Write the rotated `[nq, rope]` Q rope parts back into `q_absorbed_buf` at offset
77/// `kv_lora` of each head's `mla_cache_dim` row.
78pub fn mla_q_rope_writeback(
79    gpu: &dyn GpuBackend,
80    kernel: KernelHandle,
81    q_rope_direct: DevicePtr,
82    q_absorbed_buf: DevicePtr,
83    nq: u32,
84    rope: u32,
85    kv_lora: u32,
86    mla_cache_dim: u32,
87    stream: u64,
88) -> Result<()> {
89    KernelLaunch::new(gpu, kernel)
90        .grid([1, 1, 1])
91        .block([256, 1, 1])
92        .arg_ptr(q_rope_direct)
93        .arg_ptr(q_absorbed_buf)
94        .arg_u32(nq)
95        .arg_u32(rope)
96        .arg_u32(kv_lora)
97        .arg_u32(mla_cache_dim)
98        .launch(stream)
99}
100
101/// 2026-09-25: Assemble one token's cache entries: K = `[kv_latent | k_rope]` and
102/// V = `[kv_latent | zeros]`, each `mla_cache_dim` wide. The block has one thread per element
103/// (`max(mla_cache_dim, 256)` threads).
104pub fn mla_cache_assemble(
105    gpu: &dyn GpuBackend,
106    kernel: KernelHandle,
107    kv_latent: DevicePtr,
108    k_rope: DevicePtr,
109    k_cache: DevicePtr,
110    v_cache: DevicePtr,
111    kv_lora: u32,
112    rope: u32,
113    mla_cache_dim: u32,
114    stream: u64,
115) -> Result<()> {
116    KernelLaunch::new(gpu, kernel)
117        .grid([1, 1, 1])
118        .block([mla_cache_dim.max(256), 1, 1])
119        .arg_ptr(kv_latent)
120        .arg_ptr(k_rope)
121        .arg_ptr(k_cache)
122        .arg_ptr(v_cache)
123        .arg_u32(kv_lora)
124        .arg_u32(rope)
125        .arg_u32(mla_cache_dim)
126        .launch(stream)
127}
128
129/// 2026-09-25: MLA paged decode for DeepSeek-V4-Flash over an NVFP4 cache (kernel
130/// `mla_paged_decode_nvfp4`). A cache token is `kv_lora_rank + qk_rope_head_dim` = 512 + 64 = 576
131/// values (`MLA_CACHE_DIM` in the kernel); each block holds packed E2M1 data
132/// (`data_section_bytes`) followed by one FP8 E4M3 scale per 16 values.
133#[allow(clippy::too_many_arguments)]
134pub fn mla_paged_decode_nvfp4(
135    gpu: &dyn GpuBackend,
136    kernel: KernelHandle,
137    q: DevicePtr,
138    k_cache: DevicePtr,
139    v_cache: DevicePtr,
140    o: DevicePtr,
141    block_tables: DevicePtr,
142    seq_lens: DevicePtr,
143    max_blocks_per_seq: u32,
144    num_q_heads: u32,
145    num_kv_heads: u32,
146    q_head_dim: u32,
147    kv_cache_dim: u32,
148    block_size: u32,
149    inv_sqrt_d: f32,
150    block_stride_bytes: u64,
151    data_section_bytes: u64,
152    num_seqs: u32,
153    stream: u64,
154) -> Result<()> {
155    KernelLaunch::new(gpu, kernel)
156        .grid([num_q_heads, num_seqs, 1])
157        .block([256, 1, 1])
158        .arg_ptr(q)
159        .arg_ptr(k_cache)
160        .arg_ptr(v_cache)
161        .arg_ptr(o)
162        .arg_ptr(block_tables)
163        .arg_ptr(seq_lens)
164        .arg_u32(max_blocks_per_seq)
165        .arg_u32(num_q_heads)
166        .arg_u32(num_kv_heads)
167        .arg_u32(q_head_dim)
168        .arg_u32(kv_cache_dim)
169        .arg_u32(block_size)
170        .arg_f32(inv_sqrt_d)
171        .arg_u64(block_stride_bytes)
172        .arg_u64(data_section_bytes)
173        .launch(stream)
174}
175
176/// 2026-09-25: MLA paged decode for DeepSeek-V4-Flash over an FP8 cache (kernel
177/// `mla_paged_decode_fp8`), with scalar dequant scales `k_scale` / `v_scale` and `cache_stride`
178/// in bytes. `sinks` (per-head attention sinks, FP32) may be null. The compressed-KV pool is
179/// attended only when `comp_pool` is non-null and `comp_block_count > 0`.
180#[allow(clippy::too_many_arguments)]
181pub fn mla_paged_decode_fp8(
182    gpu: &dyn GpuBackend,
183    kernel: KernelHandle,
184    q: DevicePtr,
185    k_cache: DevicePtr,
186    v_cache: DevicePtr,
187    o: DevicePtr,
188    block_tables: DevicePtr,
189    seq_lens: DevicePtr,
190    max_blocks_per_seq: u32,
191    num_q_heads: u32,
192    num_kv_heads: u32,
193    q_head_dim: u32,
194    kv_cache_dim: u32,
195    block_size: u32,
196    inv_sqrt_d: f32,
197    k_scale: f32,
198    v_scale: f32,
199    cache_stride: u64,
200    num_seqs: u32,
201    sliding_window: u32,
202    sinks: DevicePtr,
203    comp_pool: DevicePtr,
204    comp_block_count: u32,
205    stream: u64,
206) -> Result<()> {
207    KernelLaunch::new(gpu, kernel)
208        .grid([num_q_heads, num_seqs, 1])
209        .block([256, 1, 1])
210        .arg_ptr(q)
211        .arg_ptr(k_cache)
212        .arg_ptr(v_cache)
213        .arg_ptr(o)
214        .arg_ptr(block_tables)
215        .arg_ptr(seq_lens)
216        .arg_u32(max_blocks_per_seq)
217        .arg_u32(num_q_heads)
218        .arg_u32(num_kv_heads)
219        .arg_u32(q_head_dim)
220        .arg_u32(kv_cache_dim)
221        .arg_u32(block_size)
222        .arg_f32(inv_sqrt_d)
223        .arg_f32(k_scale)
224        .arg_f32(v_scale)
225        .arg_u64(cache_stride)
226        .arg_u32(sliding_window)
227        .arg_ptr(sinks)
228        .arg_ptr(comp_pool)
229        .arg_u32(comp_block_count)
230        .launch(stream)
231}
metrale_model_layers/layers/ops/kv_cache_mla.rs

metrale_model_layers/layers/ops/
kv_cache_mla.rs
