Function ragged_prefill_bf16_hd128 

pub fn ragged_prefill_bf16_hd128(
    q: u64,
    k: u64,
    v: u64,
    o: u64,
    qo_indptr_h: &[i32],
    kv_indptr_h: &[i32],
    qo_indptr_d: u64,
    kv_indptr_d: u64,
    batch: u32,
    total_qo_rows: u32,
    total_kv_rows: u32,
    num_qo_heads: u32,
    num_kv_heads: u32,
    head_dim: u32,
    sm_scale: f32,
    causal: bool,
    sliding_window: Option<u32>,
    stream: u64,
) -> Result<()>

Ragged batched prefill attention, BF16, head_dim=128 (Laguna), GQA.

Same contract as super::ragged_prefill_bf16_hd256 but with a sliding-window bound. sliding_window is the Atlas convention (mask when q - k >= w, see kernels/gb10/common/inferspark_prefill.cu); pass None for the full-attention layers. It is converted to FlashInfer’s window_left (= w - 1) internally.

Passing the window is not just a correctness requirement for Laguna’s 36 sliding-window-512 layers — FlashInfer’s scheduler uses it to skip out-of-window KV tiles entirely, so it is also where the speedup comes from.