pub fn ragged_prefill_bf16_hd256(
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,
stream: u64,
) -> Result<()>Expand description
2026-09-25: Ragged batched prefill attention, BF16, head_dim 256, GQA, optionally causal.
q/o: [total_qo_rows, num_qo_heads, 256] BF16 on the device; k/v:
[total_kv_rows, num_kv_heads, 256]. qo_indptr/kv_indptr are [batch+1]
int32 prefix sums, on the host (*_h, read by PrefillPlan) and on the
device (*_d, read by the kernel). Errors when head_dim != 256, when a host
indptr slice is not batch + 1 long, or on a non-zero status.