metrale_model_layers/layers/ops/
ssm_gdn_woa.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the GDN write-on-accept K=4 verify
4//! (`kernels/gb10/common/gated_delta_rule_wy4_woa.cu`): the verify kernel, its
5//! post-verdict fold, and the engaged-word clear.
6//!
7//! Owner: model-layers ops (GDN).
8//! Invariants: none beyond the types.
9//!
10//! provenance-id: 526f6e616c6420522e205374657369616b
11
12use anyhow::Result;
13use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
14use metrale_gpu_runtime::kernel_args::KernelLaunch;
15
16/// 2026-09-25: Write-on-accept K=4 verify (`gated_delta_rule_wy4_woa`). It
17/// writes the output but no state, and stashes the per-row update terms for
18/// [`gdn_wy4_fold`]. The state arguments are pointer tables.
19#[allow(clippy::too_many_arguments)]
20pub fn gdn_decode_wy4_woa(
21    gpu: &dyn GpuBackend,
22    kernel: KernelHandle,
23    h_table: DevicePtr,
24    query: DevicePtr,
25    key: DevicePtr,
26    value: DevicePtr,
27    gate: DevicePtr,
28    beta: DevicePtr,
29    output: DevicePtr,
30    stash: DevicePtr,
31    batch_size: u32,
32    num_k_heads: u32,
33    num_v_heads: u32,
34    k_dim: u32,
35    v_dim: u32,
36    qk_stride: u32,
37    v_stride: u32,
38    gb_stride: u32,
39    stash_seq_floats: u32,
40    engaged_flag: DevicePtr,
41    stream: u64,
42) -> Result<()> {
43    anyhow::ensure!(
44        !h_table.is_null() && !stash.is_null() && !engaged_flag.is_null(),
45        "gdn_decode_wy4_woa: null table/stash/flag"
46    );
47    KernelLaunch::new(gpu, kernel)
48        .grid([num_v_heads, batch_size, 1])
49        .block([128, 1, 1])
50        .shared_mem(k_dim * v_dim * 4)
51        .arg_ptr(h_table)
52        .arg_ptr(query)
53        .arg_ptr(key)
54        .arg_ptr(value)
55        .arg_ptr(gate)
56        .arg_ptr(beta)
57        .arg_ptr(output)
58        .arg_ptr(stash)
59        .arg_u32(batch_size)
60        .arg_u32(num_k_heads)
61        .arg_u32(num_v_heads)
62        .arg_u32(k_dim)
63        .arg_u32(v_dim)
64        .arg_u32(qk_stride)
65        .arg_u32(v_stride)
66        .arg_u32(gb_stride)
67        .arg_u32(stash_seq_floats)
68        .arg_ptr(engaged_flag)
69        .launch(stream)
70}
71
72/// 2026-09-25: Post-verdict fold (`gated_delta_rule_wy4_fold`). When the
73/// engaged word is set, it applies stashed rows `0..na_tab[b]` to H, reading
74/// and writing H once. Otherwise the parent kernel ran, and a partial accept
75/// (`na < k_rows`) restores `Hi(na - 1)` from `hi_tables`.
76#[allow(clippy::too_many_arguments)]
77pub fn gdn_wy4_fold(
78    gpu: &dyn GpuBackend,
79    kernel: KernelHandle,
80    h_table: DevicePtr,
81    stash: DevicePtr,
82    na_tab: DevicePtr,
83    hi_tables: DevicePtr,
84    slab_entries: u32,
85    engaged_flag: DevicePtr,
86    k_rows: u32,
87    batch_size: u32,
88    num_k_heads: u32,
89    num_v_heads: u32,
90    k_dim: u32,
91    v_dim: u32,
92    stash_seq_floats: u32,
93    stream: u64,
94) -> Result<()> {
95    KernelLaunch::new(gpu, kernel)
96        .grid([num_v_heads, batch_size, 1])
97        .block([128, 1, 1])
98        .arg_ptr(h_table)
99        .arg_ptr(stash)
100        .arg_ptr(na_tab)
101        .arg_ptr(hi_tables)
102        .arg_u32(slab_entries)
103        .arg_ptr(engaged_flag)
104        .arg_u32(k_rows)
105        .arg_u32(batch_size)
106        .arg_u32(num_k_heads)
107        .arg_u32(num_v_heads)
108        .arg_u32(k_dim)
109        .arg_u32(v_dim)
110        .arg_u32(stash_seq_floats)
111        .launch(stream)
112}
113
114/// 2026-09-25: Clear a layer's write-on-accept engaged word. The caller issues
115/// it at the start of a batched verify that requested write-on-accept, on the
116/// stream of the launches that follow.
117pub fn gdn_wy4_flag_clear(
118    gpu: &dyn GpuBackend,
119    kernel: KernelHandle,
120    flag: DevicePtr,
121    stream: u64,
122) -> Result<()> {
123    KernelLaunch::new(gpu, kernel)
124        .grid([1, 1, 1])
125        .block([32, 1, 1])
126        .arg_ptr(flag)
127        .launch(stream)
128}
metrale_model_layers/layers/ops/ssm_gdn_woa.rs

metrale_model_layers/layers/ops/
ssm_gdn_woa.rs
