spark_model/layers/ops/
ssm_gdn_woa.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Launchers for the GDN write-on-accept K=4 verify pair
4//! (`kernels/gb10/common/gated_delta_rule_wy4_woa.cu`): the twin, its
5//! post-verdict fold, and the engaged-word clear. Split from `ssm_gdn_b.rs`
6//! to keep that file under the 500-LoC cap.
7//!
8//! provenance-id: 526f6e616c6420522e205374657369616b
9
10use anyhow::Result;
11use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
12use spark_runtime::kernel_args::KernelLaunch;
13
14/// Write-on-accept K=4 verify (`gated_delta_rule_wy4_woa`): byte-identical
15/// `output` to `gdn_decode_wy4`, writes NO state, stashes the per-row update
16/// terms for `gdn_wy4_fold`. Table form only. Grid (num_v_heads, batch), 128.
17#[allow(clippy::too_many_arguments)]
18pub fn gdn_decode_wy4_woa(
19    gpu: &dyn GpuBackend,
20    kernel: KernelHandle,
21    h_table: DevicePtr,
22    query: DevicePtr,
23    key: DevicePtr,
24    value: DevicePtr,
25    gate: DevicePtr,
26    beta: DevicePtr,
27    output: DevicePtr,
28    stash: DevicePtr,
29    batch_size: u32,
30    num_k_heads: u32,
31    num_v_heads: u32,
32    k_dim: u32,
33    v_dim: u32,
34    qk_stride: u32,
35    v_stride: u32,
36    gb_stride: u32,
37    stash_seq_floats: u32,
38    engaged_flag: DevicePtr,
39    stream: u64,
40) -> Result<()> {
41    anyhow::ensure!(
42        !h_table.is_null() && !stash.is_null() && !engaged_flag.is_null(),
43        "gdn_decode_wy4_woa: null table/stash/flag"
44    );
45    KernelLaunch::new(gpu, kernel)
46        .grid([num_v_heads, batch_size, 1])
47        .block([128, 1, 1])
48        .shared_mem(k_dim * v_dim * 4)
49        .arg_ptr(h_table)
50        .arg_ptr(query)
51        .arg_ptr(key)
52        .arg_ptr(value)
53        .arg_ptr(gate)
54        .arg_ptr(beta)
55        .arg_ptr(output)
56        .arg_ptr(stash)
57        .arg_u32(batch_size)
58        .arg_u32(num_k_heads)
59        .arg_u32(num_v_heads)
60        .arg_u32(k_dim)
61        .arg_u32(v_dim)
62        .arg_u32(qk_stride)
63        .arg_u32(v_stride)
64        .arg_u32(gb_stride)
65        .arg_u32(stash_seq_floats)
66        .arg_ptr(engaged_flag)
67        .launch(stream)
68}
69
70/// Post-verdict fold (`gated_delta_rule_wy4_fold`): applies rows `0..na_tab[b]`
71/// of the stashed updates to H, one read + one write per state.
72#[allow(clippy::too_many_arguments)]
73pub fn gdn_wy4_fold(
74    gpu: &dyn GpuBackend,
75    kernel: KernelHandle,
76    h_table: DevicePtr,
77    stash: DevicePtr,
78    na_tab: DevicePtr,
79    hi_tables: DevicePtr,
80    slab_entries: u32,
81    engaged_flag: DevicePtr,
82    k_rows: u32,
83    batch_size: u32,
84    num_k_heads: u32,
85    num_v_heads: u32,
86    k_dim: u32,
87    v_dim: u32,
88    stash_seq_floats: u32,
89    stream: u64,
90) -> Result<()> {
91    KernelLaunch::new(gpu, kernel)
92        .grid([num_v_heads, batch_size, 1])
93        .block([128, 1, 1])
94        .arg_ptr(h_table)
95        .arg_ptr(stash)
96        .arg_ptr(na_tab)
97        .arg_ptr(hi_tables)
98        .arg_u32(slab_entries)
99        .arg_ptr(engaged_flag)
100        .arg_u32(k_rows)
101        .arg_u32(batch_size)
102        .arg_u32(num_k_heads)
103        .arg_u32(num_v_heads)
104        .arg_u32(k_dim)
105        .arg_u32(v_dim)
106        .arg_u32(stash_seq_floats)
107        .launch(stream)
108}
109
110/// Reset a layer's write-on-accept engaged word (at the top of a requesting
111/// batched verify, same stream and capture as the launches that follow).
112pub fn gdn_wy4_flag_clear(
113    gpu: &dyn GpuBackend,
114    kernel: KernelHandle,
115    flag: DevicePtr,
116    stream: u64,
117) -> Result<()> {
118    KernelLaunch::new(gpu, kernel)
119        .grid([1, 1, 1])
120        .block([32, 1, 1])
121        .arg_ptr(flag)
122        .launch(stream)
123}
spark_model/layers/ops/ssm_gdn_woa.rs

spark_model/layers/ops/
ssm_gdn_woa.rs
