pub fn mla_fused_prefill(
gpu: &dyn GpuBackend,
kernel: KernelHandle,
q_full: DevicePtr,
q_rope: DevicePtr,
kv_latent: DevicePtr,
k_rope: DevicePtr,
w_uk: DevicePtr,
w_uv: DevicePtr,
v_out: DevicePtr,
k_cache_out: DevicePtr,
v_cache_out: DevicePtr,
seq_len: u32,
nq: u32,
nope: u32,
rope: u32,
kv_lora: u32,
v_dim: u32,
hd: u32,
num_kv_heads: u32,
inv_sqrt_d: f32,
stream: u64,
) -> Result<()>Expand description
Fused MLA prefill: Q_absorption + attention + V_extraction in one kernel. Grid: (num_heads, seq_len, 1) Block: (256, 1, 1)