gdn_prefill_fla

Function gdn_prefill_fla 

Source
pub fn gdn_prefill_fla(
    gpu: &dyn GpuBackend,
    k_recompute_wu: KernelHandle,
    k_chunk_delta_h: KernelHandle,
    k_chunk_delta_h_tc_vblock: 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

FLA multi-kernel chunked GDN prefill (ATLAS_GDN_FLA=1).

Three sequential launches on stream (CPU-serialized → no GPU sync needed):

  1. recompute_wu (grid [num_chunks, nv, batch], 128 thr): solve (I+L)U=βV, (I+L)W=β·exp(gc)·K → W_out, U_out (bf16), gc_out (f32).
  2. chunk_delta_h_ksplit (grid [nv, batch], 256 thr): serial f32 state spine, 2 threads/v-column for occupancy → S_out (per-chunk entry states bf16), uc_out (bf16); updates h_state in-place.
  3. chunk_fwd_o (grid [num_chunks, nv, batch], 128 thr): O = Q̃·S_c + tril(decay·Q̃·Kᵀ)·uc → output (bf16, same layout as wy4).

W_out/U_out/S_out/uc_out are the caller’s pre-sized scratch (BufferArena gdn_fla_scratch, sub-divided). Strides match the packed conv layout (qk_stride=v_stride=conv_dim, gb_stride=2*nv) exactly like the wy4/chunk64 path.