spark_runtime/flashinfer/
hd128.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//! head_dim=128 (Laguna) FlashInfer ragged-prefill wrapper. Split from
3//! `flashinfer.rs` (500-LoC cap); a child module so it shares the private
4//! workspace singleton.
5
6use anyhow::{Result, bail};
7
8#[cfg(atlas_flashinfer)]
9use std::ffi::c_void;
10
11#[cfg(atlas_flashinfer)]
12use super::workspaces;
13
14#[cfg(atlas_flashinfer)]
15unsafe extern "C" {
16    // head_dim=128 (Laguna). Same ABI as hd256 plus `window_left` after
17    // `causal`: -1 = full causal attention, otherwise sliding_window - 1.
18    #[allow(clippy::too_many_arguments)]
19    fn atlas_fi_ragged_prefill_bf16_hd128(
20        q: *const c_void,
21        k: *const c_void,
22        v: *const c_void,
23        o: *mut c_void,
24        qo_indptr_h: *const i32,
25        kv_indptr_h: *const i32,
26        qo_indptr_d: *const i32,
27        kv_indptr_d: *const i32,
28        batch: u32,
29        total_qo_rows: u32,
30        total_kv_rows: u32,
31        num_qo_heads: u32,
32        num_kv_heads: u32,
33        head_dim: u32,
34        sm_scale: f32,
35        causal: i32,
36        window_left: i32,
37        float_ws: *mut c_void,
38        float_ws_bytes: usize,
39        int_ws: *mut c_void,
40        int_ws_bytes: usize,
41        pinned_int_ws: *mut c_void,
42        pinned_int_ws_bytes: usize,
43        stream: *mut c_void,
44    ) -> i32;
45}
46
47/// Ragged batched prefill attention, BF16, **head_dim=128** (Laguna), GQA.
48///
49/// Same contract as [`super::ragged_prefill_bf16_hd256`] but with a sliding-window
50/// bound. `sliding_window` is the Atlas convention (mask when `q - k >= w`,
51/// see `kernels/gb10/common/inferspark_prefill.cu`); pass `None` for the
52/// full-attention layers. It is converted to FlashInfer's `window_left`
53/// (= `w - 1`) internally.
54///
55/// Passing the window is not just a correctness requirement for Laguna's 36
56/// sliding-window-512 layers — FlashInfer's scheduler uses it to skip
57/// out-of-window KV tiles entirely, so it is also where the speedup comes from.
58#[allow(clippy::too_many_arguments)]
59pub fn ragged_prefill_bf16_hd128(
60    q: u64,
61    k: u64,
62    v: u64,
63    o: u64,
64    qo_indptr_h: &[i32],
65    kv_indptr_h: &[i32],
66    qo_indptr_d: u64,
67    kv_indptr_d: u64,
68    batch: u32,
69    total_qo_rows: u32,
70    total_kv_rows: u32,
71    num_qo_heads: u32,
72    num_kv_heads: u32,
73    head_dim: u32,
74    sm_scale: f32,
75    causal: bool,
76    sliding_window: Option<u32>,
77    stream: u64,
78) -> Result<()> {
79    #[cfg(atlas_flashinfer)]
80    {
81        if head_dim != 128 {
82            bail!("ragged_prefill_bf16_hd128 is head_dim=128 only (got {head_dim})");
83        }
84        if qo_indptr_h.len() != (batch + 1) as usize || kv_indptr_h.len() != (batch + 1) as usize {
85            bail!("indptr host slices must be batch+1 long");
86        }
87        // Atlas masks at (q - k) >= w; FlashInfer keeps kv within window_left of
88        // the query, so window_left = w - 1. w == 0 means "no window".
89        let window_left: i32 = match sliding_window {
90            Some(w) if w > 0 => (w - 1) as i32,
91            _ => -1,
92        };
93        let ws = workspaces()?;
94        let st = unsafe {
95            atlas_fi_ragged_prefill_bf16_hd128(
96                q as *const c_void,
97                k as *const c_void,
98                v as *const c_void,
99                o as *mut c_void,
100                qo_indptr_h.as_ptr(),
101                kv_indptr_h.as_ptr(),
102                qo_indptr_d as *const i32,
103                kv_indptr_d as *const i32,
104                batch,
105                total_qo_rows,
106                total_kv_rows,
107                num_qo_heads,
108                num_kv_heads,
109                head_dim,
110                sm_scale,
111                if causal { 1 } else { 0 },
112                window_left,
113                ws.float_ws as *mut c_void,
114                ws.float_sz,
115                ws.int_ws as *mut c_void,
116                ws.int_sz,
117                ws.pinned_int_ws as *mut c_void,
118                ws.pinned_sz,
119                stream as *mut c_void,
120            )
121        };
122        if st != 0 {
123            bail!(
124                "FlashInfer ragged prefill hd128 failed: status {st} \
125                 (batch={batch}, qo={total_qo_rows}, window_left={window_left})"
126            );
127        }
128        Ok(())
129    }
130    #[cfg(not(atlas_flashinfer))]
131    {
132        let _ = (
133            q,
134            k,
135            v,
136            o,
137            qo_indptr_h,
138            kv_indptr_h,
139            qo_indptr_d,
140            kv_indptr_d,
141            batch,
142            total_qo_rows,
143            total_kv_rows,
144            num_qo_heads,
145            num_kv_heads,
146            head_dim,
147            sm_scale,
148            causal,
149            sliding_window,
150            stream,
151        );
152        bail!("FlashInfer support was not built; set FLASHINFER_HOME when building")
153    }
154}
spark_runtime/flashinfer/hd128.rs

spark_runtime/flashinfer/
hd128.rs
