gdn_prefill_persistent

Function gdn_prefill_persistent 

Source
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