metrale_gpu_runtime/
flashinfer.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Host-side FFI to the FlashInfer ragged (varlen) BF16 prefill
4//! attention wrapper in `cuda/flashinfer_ragged_prefill.cu`: several requests'
5//! attention in one launch, delimited by `qo_indptr`/`kv_indptr`. `hd128`
6//! holds the head_dim-128 entry.
7//!
8//! Owner: gpu-runtime.
9//! Invariants:
10//! - The FFI and workspaces exist only under `cfg(metrale_flashinfer)`, which
11//!   `build.rs` sets only when `FLASHINFER_HOME` is set at build time. Without
12//!   it `available()` is false and every entry point returns an error.
13//! - The only caller, the batched first-chunk prefill in metrale-model-layers
14//!   (`qwen3_attention/prefill/paged.rs`), also requires
15//!   `METRALE_FLASHINFER_PREFILL=1`.
16
17use anyhow::{Result, bail};
18
19mod hd128;
20pub use hd128::ragged_prefill_bf16_hd128;
21
22#[cfg(metrale_flashinfer)]
23use std::ffi::c_void;
24#[cfg(metrale_flashinfer)]
25use std::sync::OnceLock;
26
27#[cfg(metrale_flashinfer)]
28unsafe extern "C" {
29    // 2026-09-25: Declared but not called: the workspaces use fixed budgets.
30    #[allow(dead_code)]
31    fn metrale_fi_ragged_prefill_workspace_sizes(
32        max_batch: u32,
33        max_total_qo_rows: u32,
34        num_qo_heads: u32,
35        num_kv_heads: u32,
36        head_dim: u32,
37        float_ws_bytes_out: *mut usize,
38        int_ws_bytes_out: *mut usize,
39        pinned_int_ws_bytes_out: *mut usize,
40    ) -> i32;
41
42    #[allow(clippy::too_many_arguments)]
43    fn metrale_fi_ragged_prefill_bf16_hd256(
44        q: *const c_void,
45        k: *const c_void,
46        v: *const c_void,
47        o: *mut c_void,
48        qo_indptr_h: *const i32,
49        kv_indptr_h: *const i32,
50        qo_indptr_d: *const i32,
51        kv_indptr_d: *const i32,
52        batch: u32,
53        total_qo_rows: u32,
54        total_kv_rows: u32,
55        num_qo_heads: u32,
56        num_kv_heads: u32,
57        head_dim: u32,
58        sm_scale: f32,
59        causal: i32,
60        float_ws: *mut c_void,
61        float_ws_bytes: usize,
62        int_ws: *mut c_void,
63        int_ws_bytes: usize,
64        pinned_int_ws: *mut c_void,
65        pinned_int_ws_bytes: usize,
66        stream: *mut c_void,
67    ) -> i32;
68
69    #[cfg(metrale_flashinfer)]
70    fn cuMemAlloc_v2(dptr: *mut u64, bytesize: usize) -> i32;
71    #[cfg(metrale_flashinfer)]
72    fn cudaHostAlloc(ptr: *mut *mut c_void, size: usize, flags: u32) -> i32;
73}
74
75/// 2026-09-25: Whether the FlashInfer wrapper was compiled in (`FLASHINFER_HOME`
76/// was set at build time).
77pub fn available() -> bool {
78    cfg!(metrale_flashinfer)
79}
80
81// 2026-09-25: Workspaces allocated on first use and reused; each call plans
82// into them afresh (`PrefillPlan` in the wrapper), so no state carries over.
83#[cfg(metrale_flashinfer)]
84struct Workspaces {
85    float_ws: u64,
86    int_ws: u64,
87    pinned_int_ws: u64,
88    float_sz: usize,
89    int_sz: usize,
90    pinned_sz: usize,
91}
92#[cfg(metrale_flashinfer)]
93unsafe impl Send for Workspaces {}
94#[cfg(metrale_flashinfer)]
95unsafe impl Sync for Workspaces {}
96#[cfg(metrale_flashinfer)]
97/// 2026-09-25: Never freed. Static like the process CUDA context
98/// (`crate::cuda_host`): the sizes are the fixed `*_WS_BYTES` budgets, not taken
99/// from any model, so the workspaces are kept across model loads.
100static WS: OnceLock<Workspaces> = OnceLock::new();
101
102#[cfg(metrale_flashinfer)]
103const MAX_BATCH: u32 = 16;
104#[cfg(metrale_flashinfer)]
105const MAX_TOTAL_QO_ROWS: u32 = 16 * 16384;
106#[cfg(metrale_flashinfer)]
107const N_QO_HEADS: u32 = 16;
108#[cfg(metrale_flashinfer)]
109const N_KV_HEADS: u32 = 2;
110#[cfg(metrale_flashinfer)]
111const HEAD_DIM: u32 = 256;
112
113// 2026-09-25: Workspace budgets passed to `PrefillPlan`: the float workspace
114// holds split-KV partial results, the int and pinned ones its scheduler metadata.
115#[cfg(metrale_flashinfer)]
116const FLOAT_WS_BYTES: usize = 256 << 20;
117#[cfg(metrale_flashinfer)]
118const INT_WS_BYTES: usize = 64 << 20;
119#[cfg(metrale_flashinfer)]
120const PINNED_WS_BYTES: usize = 64 << 20;
121
122#[cfg(metrale_flashinfer)]
123fn workspaces() -> Result<&'static Workspaces> {
124    if let Some(w) = WS.get() {
125        return Ok(w);
126    }
127    let _ = (
128        MAX_BATCH,
129        MAX_TOTAL_QO_ROWS,
130        N_QO_HEADS,
131        N_KV_HEADS,
132        HEAD_DIM,
133    );
134    let (fsz, isz, psz) = (FLOAT_WS_BYTES, INT_WS_BYTES, PINNED_WS_BYTES);
135    let mut float_ws = 0u64;
136    let mut int_ws = 0u64;
137    let mut pinned = std::ptr::null_mut::<c_void>();
138    unsafe {
139        let s1 = cuMemAlloc_v2(&mut float_ws, fsz.max(1));
140        if s1 != 0 {
141            bail!("cuMemAlloc FlashInfer float ws ({fsz}B) failed: {s1}");
142        }
143        let s2 = cuMemAlloc_v2(&mut int_ws, isz.max(1));
144        if s2 != 0 {
145            bail!("cuMemAlloc FlashInfer int ws ({isz}B) failed: {s2}");
146        }
147        let s3 = cudaHostAlloc(&mut pinned, psz.max(1), 0);
148        if s3 != 0 {
149            bail!("cudaHostAlloc FlashInfer pinned int ws ({psz}B) failed: {s3}");
150        }
151    }
152    let _ = WS.set(Workspaces {
153        float_ws,
154        int_ws,
155        pinned_int_ws: pinned as u64,
156        float_sz: fsz,
157        int_sz: isz,
158        pinned_sz: psz,
159    });
160    Ok(WS.get().unwrap())
161}
162
163/// 2026-09-25: Ragged batched prefill attention, BF16, head_dim 256, GQA,
164/// optionally causal.
165///
166/// `q`/`o`: `[total_qo_rows, num_qo_heads, 256]` BF16 on the device; `k`/`v`:
167/// `[total_kv_rows, num_kv_heads, 256]`. `qo_indptr`/`kv_indptr` are `[batch+1]`
168/// int32 prefix sums, on the host (`*_h`, read by `PrefillPlan`) and on the
169/// device (`*_d`, read by the kernel). Errors when `head_dim != 256`, when a host
170/// indptr slice is not `batch + 1` long, or on a non-zero status.
171#[allow(clippy::too_many_arguments)]
172pub fn ragged_prefill_bf16_hd256(
173    q: u64,
174    k: u64,
175    v: u64,
176    o: u64,
177    qo_indptr_h: &[i32],
178    kv_indptr_h: &[i32],
179    qo_indptr_d: u64,
180    kv_indptr_d: u64,
181    batch: u32,
182    total_qo_rows: u32,
183    total_kv_rows: u32,
184    num_qo_heads: u32,
185    num_kv_heads: u32,
186    head_dim: u32,
187    sm_scale: f32,
188    causal: bool,
189    stream: u64,
190) -> Result<()> {
191    #[cfg(metrale_flashinfer)]
192    {
193        if head_dim != HEAD_DIM {
194            bail!("FlashInfer wrapper is head_dim=256 only (got {head_dim})");
195        }
196        if qo_indptr_h.len() != (batch + 1) as usize || kv_indptr_h.len() != (batch + 1) as usize {
197            bail!("indptr host slices must be batch+1 long");
198        }
199        let ws = workspaces()?;
200        let st = unsafe {
201            metrale_fi_ragged_prefill_bf16_hd256(
202                q as *const c_void,
203                k as *const c_void,
204                v as *const c_void,
205                o as *mut c_void,
206                qo_indptr_h.as_ptr(),
207                kv_indptr_h.as_ptr(),
208                qo_indptr_d as *const i32,
209                kv_indptr_d as *const i32,
210                batch,
211                total_qo_rows,
212                total_kv_rows,
213                num_qo_heads,
214                num_kv_heads,
215                head_dim,
216                sm_scale,
217                if causal { 1 } else { 0 },
218                ws.float_ws as *mut c_void,
219                ws.float_sz,
220                ws.int_ws as *mut c_void,
221                ws.int_sz,
222                ws.pinned_int_ws as *mut c_void,
223                ws.pinned_sz,
224                stream as *mut c_void,
225            )
226        };
227        if st != 0 {
228            bail!(
229                "FlashInfer ragged prefill failed: status {st} (batch={batch}, qo={total_qo_rows})"
230            );
231        }
232        Ok(())
233    }
234    #[cfg(not(metrale_flashinfer))]
235    {
236        let _ = (
237            q,
238            k,
239            v,
240            o,
241            qo_indptr_h,
242            kv_indptr_h,
243            qo_indptr_d,
244            kv_indptr_d,
245            batch,
246            total_qo_rows,
247            total_kv_rows,
248            num_qo_heads,
249            num_kv_heads,
250            head_dim,
251            sm_scale,
252            causal,
253            stream,
254        );
255        bail!("FlashInfer support was not built; set FLASHINFER_HOME when building")
256    }
257}
258
259#[cfg(all(test, metrale_flashinfer))]
260mod tests {
261    use super::*;
262    use std::ffi::c_void;
263
264    const H2D: i32 = 1;
265    const D2H: i32 = 2;
266    unsafe extern "C" {
267        fn cudaMalloc(p: *mut *mut c_void, n: usize) -> i32;
268        fn cudaFree(p: *mut c_void) -> i32;
269        fn cudaMemcpy(d: *mut c_void, s: *const c_void, n: usize, k: i32) -> i32;
270        fn cudaDeviceSynchronize() -> i32;
271    }
272    fn f32_to_bf16(x: f32) -> u16 {
273        let b = x.to_bits();
274        ((b + 0x7fff + ((b >> 16) & 1)) >> 16) as u16
275    }
276    fn bf16_to_f32(x: u16) -> f32 {
277        f32::from_bits((x as u32) << 16)
278    }
279    unsafe fn dev<T>(data: &[T]) -> u64 {
280        let bytes = std::mem::size_of_val(data);
281        let mut p = std::ptr::null_mut();
282        assert_eq!(unsafe { cudaMalloc(&mut p, bytes.max(1)) }, 0);
283        assert_eq!(
284            unsafe { cudaMemcpy(p, data.as_ptr() as *const c_void, bytes, H2D) },
285            0
286        );
287        p as u64
288    }
289
290    #[test]
291    #[ignore = "requires a free CUDA device + FLASHINFER_HOME build"]
292    #[allow(clippy::needless_range_loop)] // 2026-09-25: index loops mirror the reference math.
293    fn flashinfer_ragged_prefill_matches_cpu_reference() {
294        const HD: usize = 256;
295        const NQO: usize = 4;
296        const NKV: usize = 2;
297        let lens = [6usize, 10usize];
298        let qo_indptr: Vec<i32> = {
299            let mut v = vec![0i32];
300            for &l in &lens {
301                v.push(v.last().unwrap() + l as i32);
302            }
303            v
304        };
305        let kv_indptr = qo_indptr.clone();
306        let total: usize = lens.iter().sum();
307        let sm_scale = 1.0f32 / (HD as f32).sqrt();
308
309        let rnd = |seed: u64| -> f32 {
310            let mut x = seed.wrapping_mul(0x9E3779B97F4A7C15).wrapping_add(1);
311            x ^= x >> 31;
312            x = x.wrapping_mul(0xBF58476D1CE4E5B9);
313            ((x >> 40) as f32 / (1u64 << 24) as f32 - 0.5) * 0.5
314        };
315        let q: Vec<u16> = (0..total * NQO * HD)
316            .map(|i| f32_to_bf16(rnd(i as u64)))
317            .collect();
318        let k: Vec<u16> = (0..total * NKV * HD)
319            .map(|i| f32_to_bf16(rnd(i as u64 ^ 0x1111)))
320            .collect();
321        let v: Vec<u16> = (0..total * NKV * HD)
322            .map(|i| f32_to_bf16(rnd(i as u64 ^ 0x2222)))
323            .collect();
324        let mut o = vec![0u16; total * NQO * HD];
325
326        let (q_d, k_d, v_d, o_d, qo_d, kv_d);
327        unsafe {
328            q_d = dev(&q);
329            k_d = dev(&k);
330            v_d = dev(&v);
331            o_d = dev(&o);
332            qo_d = dev(&qo_indptr);
333            kv_d = dev(&kv_indptr);
334        }
335
336        ragged_prefill_bf16_hd256(
337            q_d,
338            k_d,
339            v_d,
340            o_d,
341            &qo_indptr,
342            &kv_indptr,
343            qo_d,
344            kv_d,
345            lens.len() as u32,
346            total as u32,
347            total as u32,
348            NQO as u32,
349            NKV as u32,
350            HD as u32,
351            sm_scale,
352            true,
353            0,
354        )
355        .unwrap();
356        unsafe {
357            assert_eq!(cudaDeviceSynchronize(), 0);
358            assert_eq!(
359                cudaMemcpy(
360                    o.as_mut_ptr() as *mut c_void,
361                    o_d as *const c_void,
362                    o.len() * 2,
363                    D2H
364                ),
365                0
366            );
367        }
368
369        // 2026-09-25: CPU reference: per-request causal GQA attention.
370        let group = NQO / NKV;
371        let qf = |r: usize, h: usize, d: usize| bf16_to_f32(q[(r * NQO + h) * HD + d]);
372        let kf = |r: usize, kh: usize, d: usize| bf16_to_f32(k[(r * NKV + kh) * HD + d]);
373        let vf = |r: usize, kh: usize, d: usize| bf16_to_f32(v[(r * NKV + kh) * HD + d]);
374        let mut max_rel = 0.0f64;
375        let mut worst_cos = 1.0f64;
376        for (b, &len) in lens.iter().enumerate() {
377            let start = qo_indptr[b] as usize;
378            for qi in 0..len {
379                for h in 0..NQO {
380                    let kh = h / group;
381                    let mut scores = vec![0f32; qi + 1];
382                    for j in 0..=qi {
383                        let mut s = 0.0f32;
384                        for d in 0..HD {
385                            s += qf(start + qi, h, d) * kf(start + j, kh, d);
386                        }
387                        scores[j] = s * sm_scale;
388                    }
389                    let mx = scores.iter().cloned().fold(f32::MIN, f32::max);
390                    let mut den = 0.0f32;
391                    for s in &mut scores {
392                        *s = (*s - mx).exp();
393                        den += *s;
394                    }
395                    let mut out_ref = vec![0f32; HD];
396                    for (j, &p) in scores.iter().enumerate() {
397                        let w = p / den;
398                        for d in 0..HD {
399                            out_ref[d] += w * vf(start + j, kh, d);
400                        }
401                    }
402                    let mut dot = 0.0f64;
403                    let mut na = 0.0f64;
404                    let mut nb = 0.0f64;
405                    for d in 0..HD {
406                        let g = bf16_to_f32(o[(start + qi) * NQO * HD + h * HD + d]) as f64;
407                        let r = out_ref[d] as f64;
408                        dot += g * r;
409                        na += g * g;
410                        nb += r * r;
411                        max_rel = max_rel.max((g - r).abs() / (r.abs() + 1e-3));
412                    }
413                    let cos = dot / (na.sqrt() * nb.sqrt() + 1e-12);
414                    worst_cos = worst_cos.min(cos);
415                }
416            }
417        }
418        unsafe {
419            for p in [q_d, k_d, v_d, o_d, qo_d, kv_d] {
420                cudaFree(p as *mut c_void);
421            }
422        }
423        tracing::debug!("FLASHINFER_RAGGED worst_cos={worst_cos:.6} max_rel={max_rel:.4}");
424        assert!(
425            worst_cos > 0.99,
426            "FlashInfer ragged prefill diverges from CPU ref: cos {worst_cos}"
427        );
428    }
429}
metrale_gpu_runtime/flashinfer.rs

metrale_gpu_runtime/
flashinfer.rs
