pub fn gdn_prefill_persistent(
gpu: &dyn GpuBackend,
kernel: KernelHandle,
h_state: DevicePtr,
query: DevicePtr,
key: DevicePtr,
value: DevicePtr,
gate: DevicePtr,
beta: DevicePtr,
output: DevicePtr,
batch_size: u32,
seq_len: u32,
num_k_heads: u32,
num_v_heads: u32,
k_dim: u32,
v_dim: u32,
qk_stride: u32,
v_stride: u32,
gb_stride: u32,
stream: u64,
) -> Result<()>Expand description
Persistent GDN prefill — h_state stays in shared memory for entire sequence.
Same parameters as gdn_prefill_split4 but uses persistent CTAs with 128 threads and 67 KB shared memory. Each CTA processes ALL tokens for one v_head, keeping h_state in shared memory (never written to global until the end). Targets L2 bandwidth (~3 TB/s) instead of LPDDR5X (273 GB/s).
Grid: (num_v_heads, batch, 1) Block: (128, 1, 1) Shared: k_dimv_dim4 + 4k_dim4 bytes