metrale_model_layers/layers/ops/
ssm_gdn_exact_f16.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-01: Launchers for the FP16 h-state exact MTP verify: `gdn_exact_chain_f16_{2,3,4}`
4//! (`gdn_exact_carry.cu` in the qwen3.6-27b directory) and the batched FP32 conv chain
5//! `gdn_conv_chain_f32_batched` (`kernels/gb10/common/gated_delta_rule_carry.cu`).
6//!
7//! Owner: model-layers ops (GDN).
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-10-01: Where the FP16 exact chain finds its states: per-sequence pointer tables (the
15/// verify's WY tables: h, Hi0, Hi1, Hi2), or the bases of a single sequence.
16#[derive(Clone, Copy, Debug, PartialEq, Eq)]
17pub enum F16ChainStates {
18    /// 2026-10-01: Device tables of `batch` per-sequence base pointers.
19    Tables([DevicePtr; 4]),
20    /// 2026-10-01: One sequence's h slot and its intermediates 0..2; batch must be 1.
21    Single([DevicePtr; 4]),
22}
23
24/// 2026-10-01: The FP16 h-state exact verify, K = 2..4 rows per sequence (`kernel` is the
25/// `gdn_exact_chain_f16_{K}` entry): the decode's FP16 fused-norm chain per row with the state
26/// in registers, the FP16 intermediates of rows 0..K-2 and the final state written, and the BF16
27/// normed rows. Rows are `b * K + t`; `strides` is `[qk, v, gb, z, out]` per row.
28#[allow(clippy::too_many_arguments)]
29pub fn gdn_exact_chain_f16(
30    gpu: &dyn GpuBackend,
31    kernel: KernelHandle,
32    states: F16ChainStates,
33    query: DevicePtr,
34    key: DevicePtr,
35    value: DevicePtr,
36    gate: DevicePtr,
37    beta: DevicePtr,
38    z_gate: DevicePtr,
39    norm_weight: DevicePtr,
40    output: DevicePtr,
41    batch: u32,
42    heads: [u32; 3],
43    strides: [u32; 5],
44    eps: f32,
45    stream: u64,
46) -> Result<()> {
47    let (ptrs, table) = match states {
48        F16ChainStates::Tables(p) => (p, 1u32),
49        F16ChainStates::Single(p) => {
50            anyhow::ensure!(
51                batch == 1,
52                "gdn_exact_chain_f16: single-sequence bases with batch {batch}"
53            );
54            (p, 0u32)
55        }
56    };
57    let [num_k_heads, num_v_heads, k_dim] = heads;
58    let [qk_stride, v_stride, gb_stride, z_stride, out_stride] = strides;
59    KernelLaunch::new(gpu, kernel)
60        .grid([num_v_heads, batch, 1])
61        .block([128, 1, 1])
62        .arg_ptr(ptrs[0])
63        .arg_ptr(query)
64        .arg_ptr(key)
65        .arg_ptr(value)
66        .arg_ptr(gate)
67        .arg_ptr(beta)
68        .arg_ptr(z_gate)
69        .arg_ptr(norm_weight)
70        .arg_ptr(output)
71        .arg_ptr(ptrs[1])
72        .arg_ptr(ptrs[2])
73        .arg_ptr(ptrs[3])
74        .arg_u32(batch)
75        .arg_u32(num_k_heads)
76        .arg_u32(num_v_heads)
77        .arg_u32(k_dim)
78        .arg_u32(qk_stride)
79        .arg_u32(v_stride)
80        .arg_u32(gb_stride)
81        .arg_u32(z_stride)
82        .arg_u32(out_stride)
83        .arg_f32(eps)
84        .arg_u32(table)
85        .launch(stream)
86}
87
88/// 2026-10-01: [`super::gdn_conv_chain_f32`] for `batch` sequences on consecutive conv slots
89/// (`dim * d_conv` floats apart): row `b * num_tokens + t` reads `new_input + row *
90/// input_stride` and writes `output + row * output_stride`; the window after position t <
91/// num_tokens - 1 goes to `conv_inter + b * inter_strides[1] + t * inter_strides[0]` floats.
92#[allow(clippy::too_many_arguments)]
93pub fn gdn_conv_chain_f32_batched(
94    gpu: &dyn GpuBackend,
95    kernel: KernelHandle,
96    conv_state: DevicePtr,
97    new_input: DevicePtr,
98    weight: &crate::weight_map::DenseWeight,
99    output: DevicePtr,
100    conv_inter: DevicePtr,
101    batch: u32,
102    num_tokens: u32,
103    dims: [u32; 4],
104    l2_eps: f32,
105    row_strides: [u32; 2],
106    inter_strides: [u64; 2],
107    stream: u64,
108) -> Result<()> {
109    let [dim, d_conv, qk_channels, head_dim] = dims;
110    KernelLaunch::new(gpu, kernel)
111        .grid([dim.div_ceil(256), batch, 1])
112        .block([256, 1, 1])
113        .arg_ptr(conv_state)
114        .arg_ptr(new_input)
115        .arg_ptr(weight.weight)
116        .arg_ptr(output)
117        .arg_ptr(conv_inter)
118        .arg_u32(num_tokens)
119        .arg_u32(dim)
120        .arg_u32(d_conv)
121        .arg_u32(qk_channels)
122        .arg_u32(head_dim)
123        .arg_f32(l2_eps)
124        .arg_u32(row_strides[0])
125        .arg_u32(row_strides[1])
126        .arg_u64(inter_strides[0])
127        .arg_u64(inter_strides[1])
128        .launch(stream)
129}
metrale_model_layers/layers/ops/ssm_gdn_exact_f16.rs

metrale_model_layers/layers/ops/
ssm_gdn_exact_f16.rs
