spark_model/layers/ops/
ple.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! PLE kernel dispatch — the gate, the dilated depthwise conv, and the
4//! highway add. See `kernels/gb10/qwen3.8-flash-next/nvfp4/ple.cu`.
5
6use anyhow::Result;
7use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
8use spark_runtime::kernel_args::KernelLaunch;
9
10/// Gate the n-gram value by the highway, and emit both the gated value and
11/// its `norm_conv`'d twin.
12///
13/// `hidden` is the FP32 mHC highway `[T, hc*H]`; `key`/`value` are the BF16
14/// projection outputs. **Both outputs are FP32** — the whole PLE chain is,
15/// because its result lands on the FP32 highway; see the PRECISION NOTE in
16/// `ple.cu`. One block per token.
17#[allow(clippy::too_many_arguments)]
18pub fn ple_gate(
19    gpu: &dyn GpuBackend,
20    kernel: KernelHandle,
21    hidden: DevicePtr,
22    key: DevicePtr,
23    value: DevicePtr,
24    norm_query_w: DevicePtr,
25    norm_key_w: DevicePtr,
26    norm_conv_w: DevicePtr,
27    gated_out: DevicePtr,
28    gated_normed: DevicePtr,
29    num_tokens: u32,
30    hidden_size: u32,
31    hc_mult: u32,
32    norm_eps: f32,
33    stream: u64,
34) -> Result<()> {
35    KernelLaunch::new(gpu, kernel)
36        .grid([num_tokens, 1, 1])
37        .block([256, 1, 1])
38        .arg_ptr(hidden)
39        .arg_ptr(key)
40        .arg_ptr(value)
41        .arg_ptr(norm_query_w)
42        .arg_ptr(norm_key_w)
43        .arg_ptr(norm_conv_w)
44        .arg_ptr(gated_out)
45        .arg_ptr(gated_normed)
46        .arg_u32(hidden_size)
47        .arg_u32(hc_mult)
48        .arg_f32(norm_eps)
49        .launch(stream)
50}
51
52/// Depthwise causal conv, kernel `k_size`, **dilation `dilation`**, plus the
53/// SiLU and the residual add against the un-normalized gated value.
54///
55/// Everything but `weight` is FP32 — see the PRECISION NOTE in `ple.cu`.
56///
57/// `state` is `[(k_size-1)*dilation, channels]` and is rolled in place, so
58/// prefill and decode share one launch — there is no decode twin to drift.
59#[allow(clippy::too_many_arguments)]
60pub fn ple_conv(
61    gpu: &dyn GpuBackend,
62    kernel: KernelHandle,
63    x: DevicePtr,
64    gated: DevicePtr,
65    weight: DevicePtr,
66    state: DevicePtr,
67    out: DevicePtr,
68    num_tokens: u32,
69    channels: u32,
70    k_size: u32,
71    dilation: u32,
72    stream: u64,
73) -> Result<()> {
74    let threads = 256u32;
75    KernelLaunch::new(gpu, kernel)
76        .grid([channels.div_ceil(threads), 1, 1])
77        .block([threads, 1, 1])
78        .arg_ptr(x)
79        .arg_ptr(gated)
80        .arg_ptr(weight)
81        .arg_ptr(state)
82        .arg_ptr(out)
83        .arg_u32(num_tokens)
84        .arg_u32(channels)
85        .arg_u32(k_size)
86        .arg_u32(dilation)
87        .launch(stream)
88}
89
90/// `highway += ple_out`, in FP32. The reference adds PLE's output to the
91/// residual before that layer's attention hyper-connection.
92pub fn ple_add_highway(
93    gpu: &dyn GpuBackend,
94    kernel: KernelHandle,
95    ple_out: DevicePtr,
96    hidden: DevicePtr,
97    n: u32,
98    stream: u64,
99) -> Result<()> {
100    let threads = 256u32;
101    KernelLaunch::new(gpu, kernel)
102        .grid([n.div_ceil(threads), 1, 1])
103        .block([threads, 1, 1])
104        .arg_ptr(ple_out)
105        .arg_ptr(hidden)
106        .arg_u32(n)
107        .launch(stream)
108}
spark_model/layers/ops/ple.rs

spark_model/layers/ops/
ple.rs
