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:
recompute_wu, or its Hopper twin: per chunk, W and U (bf16) and the cumulative gategc(f32);- a state spine, chosen below: walks the chunks in order, writes each
chunk’s entry state
Sanduc(bf16), and updatesh_statein place; 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.