spark_model/layers/ops/
hyper_connection.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Manifold-Constrained Hyper-Connections (mHC) kernel dispatch (DeepSeek-V4).
4//!
5//! Wraps the `hyper_connection` module kernels (`hc_pre`, `hc_post`,
6//! `hc_head`). The hidden state is stored BF16 as `[T, hc_mult, H]`
7//! (stream-major per token). HC parameters are float32 device buffers.
8
9use anyhow::Result;
10use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
11use spark_runtime::kernel_args::KernelLaunch;
12
13/// Broadcast a single hidden state into `hc_mult` identical streams:
14/// `streams[t, i, d] = hidden[t, d]`. One block per token.
15pub fn hc_expand(
16    gpu: &dyn GpuBackend,
17    kernel: KernelHandle,
18    hidden: DevicePtr,
19    streams: DevicePtr,
20    num_tokens: u32,
21    hidden_size: u32,
22    hc_mult: u32,
23    stream: u64,
24) -> Result<()> {
25    KernelLaunch::new(gpu, kernel)
26        .grid([num_tokens, 1, 1])
27        .block([256, 1, 1])
28        .arg_ptr(hidden)
29        .arg_ptr(streams)
30        .arg_u32(hidden_size)
31        .arg_u32(hc_mult)
32        .launch(stream)
33}
34
35/// Collapse `hc_mult` streams to one (RMS-rescaled mix → sigmoid `pre`
36/// weighted sum) and emit `post` / `comb` (Sinkhorn) for the matching
37/// `hc_post`. One block per token.
38#[allow(clippy::too_many_arguments)]
39pub fn hc_pre(
40    gpu: &dyn GpuBackend,
41    kernel: KernelHandle,
42    streams: DevicePtr,
43    hc_fn: DevicePtr,
44    hc_scale: DevicePtr,
45    hc_base: DevicePtr,
46    y_out: DevicePtr,
47    post_out: DevicePtr,
48    comb_out: DevicePtr,
49    num_tokens: u32,
50    hidden_size: u32,
51    hc_mult: u32,
52    sinkhorn_iters: u32,
53    norm_eps: f32,
54    hc_eps: f32,
55    stream: u64,
56) -> Result<()> {
57    KernelLaunch::new(gpu, kernel)
58        .grid([num_tokens, 1, 1])
59        .block([256, 1, 1])
60        .arg_ptr(streams)
61        .arg_ptr(hc_fn)
62        .arg_ptr(hc_scale)
63        .arg_ptr(hc_base)
64        .arg_ptr(y_out)
65        .arg_ptr(post_out)
66        .arg_ptr(comb_out)
67        .arg_u32(hidden_size)
68        .arg_u32(hc_mult)
69        .arg_u32(sinkhorn_iters)
70        .arg_f32(norm_eps)
71        .arg_f32(hc_eps)
72        .launch(stream)
73}
74
75/// Expand the sublayer output back into `hc_mult` streams, mixing the saved
76/// residual streams through the doubly-stochastic `comb`. `out` may alias
77/// `residual`. One block per token.
78#[allow(clippy::too_many_arguments)]
79pub fn hc_post(
80    gpu: &dyn GpuBackend,
81    kernel: KernelHandle,
82    block_out: DevicePtr,
83    residual: DevicePtr,
84    post: DevicePtr,
85    comb: DevicePtr,
86    out: DevicePtr,
87    num_tokens: u32,
88    hidden_size: u32,
89    hc_mult: u32,
90    stream: u64,
91) -> Result<()> {
92    KernelLaunch::new(gpu, kernel)
93        .grid([num_tokens, 1, 1])
94        .block([256, 1, 1])
95        .arg_ptr(block_out)
96        .arg_ptr(residual)
97        .arg_ptr(post)
98        .arg_ptr(comb)
99        .arg_ptr(out)
100        .arg_u32(hidden_size)
101        .arg_u32(hc_mult)
102        .launch(stream)
103}
104
105/// Final collapse before the LM head: a single learned sigmoid-weighted sum
106/// over the `hc_mult` streams. One block per token.
107#[allow(clippy::too_many_arguments)]
108pub fn hc_head(
109    gpu: &dyn GpuBackend,
110    kernel: KernelHandle,
111    streams: DevicePtr,
112    head_fn: DevicePtr,
113    head_scale: DevicePtr,
114    head_base: DevicePtr,
115    y_out: DevicePtr,
116    num_tokens: u32,
117    hidden_size: u32,
118    hc_mult: u32,
119    norm_eps: f32,
120    hc_eps: f32,
121    stream: u64,
122) -> Result<()> {
123    KernelLaunch::new(gpu, kernel)
124        .grid([num_tokens, 1, 1])
125        .block([256, 1, 1])
126        .arg_ptr(streams)
127        .arg_ptr(head_fn)
128        .arg_ptr(head_scale)
129        .arg_ptr(head_base)
130        .arg_ptr(y_out)
131        .arg_u32(hidden_size)
132        .arg_u32(hc_mult)
133        .arg_f32(norm_eps)
134        .arg_f32(hc_eps)
135        .launch(stream)
136}
spark_model/layers/ops/hyper_connection.rs

spark_model/layers/ops/
hyper_connection.rs
