metrale_gpu_runtime/flashinfer/
hd128.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The head_dim-128 FlashInfer ragged-prefill entry, with an optional
4//! sliding window. A child of `flashinfer`, so it shares that module's
5//! workspaces.
6//!
7//! Owner: gpu-runtime.
8//! Invariants: none beyond the types.
9
10use anyhow::{Result, bail};
11
12#[cfg(metrale_flashinfer)]
13use std::ffi::c_void;
14
15#[cfg(metrale_flashinfer)]
16use super::workspaces;
17
18#[cfg(metrale_flashinfer)]
19unsafe extern "C" {
20    // 2026-09-25: The hd256 ABI plus `window_left` after `causal`: -1 for no
21    // window, otherwise `sliding_window - 1`.
22    #[allow(clippy::too_many_arguments)]
23    fn metrale_fi_ragged_prefill_bf16_hd128(
24        q: *const c_void,
25        k: *const c_void,
26        v: *const c_void,
27        o: *mut c_void,
28        qo_indptr_h: *const i32,
29        kv_indptr_h: *const i32,
30        qo_indptr_d: *const i32,
31        kv_indptr_d: *const i32,
32        batch: u32,
33        total_qo_rows: u32,
34        total_kv_rows: u32,
35        num_qo_heads: u32,
36        num_kv_heads: u32,
37        head_dim: u32,
38        sm_scale: f32,
39        causal: i32,
40        window_left: i32,
41        float_ws: *mut c_void,
42        float_ws_bytes: usize,
43        int_ws: *mut c_void,
44        int_ws_bytes: usize,
45        pinned_int_ws: *mut c_void,
46        pinned_int_ws_bytes: usize,
47        stream: *mut c_void,
48    ) -> i32;
49}
50
51/// 2026-09-25: Ragged batched prefill attention, BF16, head_dim 128, GQA.
52///
53/// The contract of [`super::ragged_prefill_bf16_hd256`], plus `sliding_window`
54/// in the in-tree convention (mask when `q - k >= w`, as in
55/// `kernels/gb10/common/attn_prefill.cu`); `None` or `Some(0)` means no window.
56/// It is passed to FlashInfer as `window_left = w - 1`.
57#[allow(clippy::too_many_arguments)]
58pub fn ragged_prefill_bf16_hd128(
59    q: u64,
60    k: u64,
61    v: u64,
62    o: u64,
63    qo_indptr_h: &[i32],
64    kv_indptr_h: &[i32],
65    qo_indptr_d: u64,
66    kv_indptr_d: u64,
67    batch: u32,
68    total_qo_rows: u32,
69    total_kv_rows: u32,
70    num_qo_heads: u32,
71    num_kv_heads: u32,
72    head_dim: u32,
73    sm_scale: f32,
74    causal: bool,
75    sliding_window: Option<u32>,
76    stream: u64,
77) -> Result<()> {
78    #[cfg(metrale_flashinfer)]
79    {
80        if head_dim != 128 {
81            bail!("ragged_prefill_bf16_hd128 is head_dim=128 only (got {head_dim})");
82        }
83        if qo_indptr_h.len() != (batch + 1) as usize || kv_indptr_h.len() != (batch + 1) as usize {
84            bail!("indptr host slices must be batch+1 long");
85        }
86        let window_left: i32 = match sliding_window {
87            Some(w) if w > 0 => (w - 1) as i32,
88            _ => -1,
89        };
90        let ws = workspaces()?;
91        let st = unsafe {
92            metrale_fi_ragged_prefill_bf16_hd128(
93                q as *const c_void,
94                k as *const c_void,
95                v as *const c_void,
96                o as *mut c_void,
97                qo_indptr_h.as_ptr(),
98                kv_indptr_h.as_ptr(),
99                qo_indptr_d as *const i32,
100                kv_indptr_d as *const i32,
101                batch,
102                total_qo_rows,
103                total_kv_rows,
104                num_qo_heads,
105                num_kv_heads,
106                head_dim,
107                sm_scale,
108                if causal { 1 } else { 0 },
109                window_left,
110                ws.float_ws as *mut c_void,
111                ws.float_sz,
112                ws.int_ws as *mut c_void,
113                ws.int_sz,
114                ws.pinned_int_ws as *mut c_void,
115                ws.pinned_sz,
116                stream as *mut c_void,
117            )
118        };
119        if st != 0 {
120            bail!(
121                "FlashInfer ragged prefill hd128 failed: status {st} \
122                 (batch={batch}, qo={total_qo_rows}, window_left={window_left})"
123            );
124        }
125        Ok(())
126    }
127    #[cfg(not(metrale_flashinfer))]
128    {
129        let _ = (
130            q,
131            k,
132            v,
133            o,
134            qo_indptr_h,
135            kv_indptr_h,
136            qo_indptr_d,
137            kv_indptr_d,
138            batch,
139            total_qo_rows,
140            total_kv_rows,
141            num_qo_heads,
142            num_kv_heads,
143            head_dim,
144            sm_scale,
145            causal,
146            sliding_window,
147            stream,
148        );
149        bail!("FlashInfer support was not built; set FLASHINFER_HOME when building")
150    }
151}
metrale_gpu_runtime/flashinfer/hd128.rs

metrale_gpu_runtime/flashinfer/
hd128.rs
