gdn_prefill_fla

Function gdn_prefill_fla 

Source
pub fn gdn_prefill_fla(
    gpu: &dyn GpuBackend,
    k_recompute_wu: KernelHandle,
    k_recompute_wu_hopper: KernelHandle,
    k_chunk_fwd_o_hopper: KernelHandle,
    k_chunk_delta_h: KernelHandle,
    k_chunk_delta_h_tc_vblock: KernelHandle,
    k_chunk_delta_h_tcfuse: KernelHandle,
    k_chunk_delta_h_fused: KernelHandle,
    k_chunk_delta_h_tma: KernelHandle,
    k_chunk_fwd_o: KernelHandle,
    h_state: DevicePtr,
    query: DevicePtr,
    key: DevicePtr,
    value: DevicePtr,
    gate: DevicePtr,
    beta: DevicePtr,
    output: DevicePtr,
    w_out: DevicePtr,
    u_out: DevicePtr,
    s_out: DevicePtr,
    uc_out: DevicePtr,
    gc_out: DevicePtr,
    batch_size: u32,
    seq_len: u32,
    num_chunks: u32,
    num_k_heads: u32,
    num_v_heads: u32,
    k_dim: u32,
    v_dim: u32,
    qk_stride: u32,
    v_stride: u32,
    gb_stride: u32,
    h_state_is_table: bool,
    cu_seqlens: DevicePtr,
    cu_chunks: DevicePtr,
    is_varlen: bool,
    profile: bool,
    stream: u64,
) -> Result<()>
Expand description

2026-09-25: Chunked GDN prefill in three launches, ordered by stream:

  1. recompute_wu, or its Hopper twin: per chunk, W and U (bf16) and the cumulative gate gc (f32);
  2. a state spine, chosen below: walks the chunks in order, writes each chunk’s entry state S and uc (bf16), and updates h_state in place;
  3. chunk_fwd_o, or its Hopper twin: the bf16 output.

w_out, u_out, s_out, uc_out and gc_out are caller-sized scratch.