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
Ragged batched prefill attention (BF16, head_dim=256, GQA, causal selectable).
q/o: [total_qo_rows, num_qo_heads, 256] BF16 device; k/v:
[total_kv_rows, num_kv_heads, 256] BF16 device. qo_indptr/kv_indptr are
[batch+1] int32 prefix-sum offsets — provided both on host (*_h, for the
scheduler plan) and as device copies (*_d, read by the kernel).