pub fn ple_gate(
gpu: &dyn GpuBackend,
kernel: KernelHandle,
hidden: DevicePtr,
key: DevicePtr,
value: DevicePtr,
norm_query_w: DevicePtr,
norm_key_w: DevicePtr,
norm_conv_w: DevicePtr,
gated_out: DevicePtr,
gated_normed: DevicePtr,
num_tokens: u32,
hidden_size: u32,
hc_mult: u32,
norm_eps: f32,
stream: u64,
) -> Result<()>Expand description
Gate the n-gram value by the highway, and emit both the gated value and
its norm_conv’d twin.
hidden is the FP32 mHC highway [T, hc*H]; key/value are the BF16
projection outputs. Both outputs are FP32 — the whole PLE chain is,
because its result lands on the FP32 highway; see the PRECISION NOTE in
ple.cu. One block per token.