compute_gdn_gates

Function compute_gdn_gates 

Source
pub fn compute_gdn_gates(
    gpu: &dyn GpuBackend,
    kernel: KernelHandle,
    ba_interleaved: DevicePtr,
    a_log: DevicePtr,
    dt_bias: DevicePtr,
    gate_out: DevicePtr,
    beta_out: DevicePtr,
    num_tokens: u32,
    num_v_heads: u32,
    num_groups: u32,
    vheads_per_group: u32,
    ba_stride: u32,
    stream: u64,
) -> Result<()>
Expand description

2026-09-25: GDN gates from the interleaved BA projection and the learned A_log / dt_bias: FP32 gate (decay) and beta (write gate) per value head and token.

Input rows are ba_stride BF16 elements apart. Output rows of both gate_out and beta_out are 2 * num_v_heads floats apart.

Kernel: compute_gdn_gates(ba_interleaved, A_log, dt_bias, gate_out, beta_out, num_v_heads, num_groups, vheads_per_group, ba_stride) Grid: (num_tokens, 1, 1) Block: (num_v_heads, 1, 1)