ple_conv

Function ple_conv 

Source
pub fn ple_conv(
    gpu: &dyn GpuBackend,
    kernel: KernelHandle,
    x: DevicePtr,
    gated: DevicePtr,
    weight: DevicePtr,
    state: DevicePtr,
    out: DevicePtr,
    num_tokens: u32,
    channels: u32,
    k_size: u32,
    dilation: u32,
    stream: u64,
) -> Result<()>
Expand description

Depthwise causal conv, kernel k_size, dilation dilation, plus the SiLU and the residual add against the un-normalized gated value.

Everything but weight is FP32 — see the PRECISION NOTE in ple.cu.

state is [(k_size-1)*dilation, channels] and is rolled in place, so prefill and decode share one launch — there is no decode twin to drift.