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):
- 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).
- 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.
- 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.