metrale_model_layers/layers/ops/
ple.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the PLE kernels in
4//! `kernels/gb10/qwen3.8-flash-next/nvfp4/ple.cu`: the gate, the dilated
5//! depthwise conv, and the highway add.
6//!
7//! Owner: model-layers ops.
8//! Invariants: none beyond the types.
9
10use anyhow::Result;
11use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
12use metrale_gpu_runtime::kernel_args::KernelLaunch;
13
14/// 2026-09-25: Gate the n-gram value by the highway, and write both the gated
15/// value and its `norm_conv` output. `hidden` is the FP32 highway `[T, hc*H]`;
16/// `key` and `value` are BF16 projection outputs. Both outputs are FP32, because
17/// the chain ends on the FP32 highway (the precision note in `ple.cu`). One
18/// block per token.
19#[allow(clippy::too_many_arguments)]
20pub fn ple_gate(
21    gpu: &dyn GpuBackend,
22    kernel: KernelHandle,
23    hidden: DevicePtr,
24    key: DevicePtr,
25    value: DevicePtr,
26    norm_query_w: DevicePtr,
27    norm_key_w: DevicePtr,
28    norm_conv_w: DevicePtr,
29    gated_out: DevicePtr,
30    gated_normed: DevicePtr,
31    num_tokens: u32,
32    hidden_size: u32,
33    hc_mult: u32,
34    norm_eps: f32,
35    stream: u64,
36) -> Result<()> {
37    KernelLaunch::new(gpu, kernel)
38        .grid([num_tokens, 1, 1])
39        .block([256, 1, 1])
40        .arg_ptr(hidden)
41        .arg_ptr(key)
42        .arg_ptr(value)
43        .arg_ptr(norm_query_w)
44        .arg_ptr(norm_key_w)
45        .arg_ptr(norm_conv_w)
46        .arg_ptr(gated_out)
47        .arg_ptr(gated_normed)
48        .arg_u32(hidden_size)
49        .arg_u32(hc_mult)
50        .arg_f32(norm_eps)
51        .launch(stream)
52}
53
54/// 2026-09-25: Depthwise causal conv with kernel `k_size` and dilation
55/// `dilation`, then `out = gated + silu(conv)`. All buffers but `weight` (BF16)
56/// are FP32. `state` is `[(k_size - 1) * dilation, channels]` and is rolled in
57/// place, so prefill and decode use the same launch.
58#[allow(clippy::too_many_arguments)]
59pub fn ple_conv(
60    gpu: &dyn GpuBackend,
61    kernel: KernelHandle,
62    x: DevicePtr,
63    gated: DevicePtr,
64    weight: DevicePtr,
65    state: DevicePtr,
66    out: DevicePtr,
67    num_tokens: u32,
68    channels: u32,
69    k_size: u32,
70    dilation: u32,
71    stream: u64,
72) -> Result<()> {
73    let threads = 256u32;
74    KernelLaunch::new(gpu, kernel)
75        .grid([channels.div_ceil(threads), 1, 1])
76        .block([threads, 1, 1])
77        .arg_ptr(x)
78        .arg_ptr(gated)
79        .arg_ptr(weight)
80        .arg_ptr(state)
81        .arg_ptr(out)
82        .arg_u32(num_tokens)
83        .arg_u32(channels)
84        .arg_u32(k_size)
85        .arg_u32(dilation)
86        .launch(stream)
87}
88
89/// 2026-09-25: `highway += ple_out`, in FP32. The reference model
90/// (`bench/qwen4_exp/ref/modeling_qwen4_exp.py`) adds PLE's output to the
91/// hidden state 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}
metrale_model_layers/layers/ops/ple.rs

metrale_model_layers/layers/ops/
ple.rs
