metrale_model_layers/layers/ops/
hyper_connection.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the DeepSeek-V4 Sinkhorn hyper-connection (mHC) kernels `hc_expand`, `hc_pre`, `hc_post` and `hc_head`.
4//!
5//! Owner: model-layers ops.
6//! Invariants:
7//! - The stream state is FP32 `[T, hc_mult, H]`, stream-major per token; the collapsed hidden
8//!   state and the sublayer output are BF16 `[T, H]`; the HC parameters are FP32
9//!   (kernels/gb10/deepseek-v4-flash/nvfp4/hyper_connection.cu).
10//! - Every launch is one 256-thread block per token.
11//! - The kernels support `hc_mult <= 4` (`HC_MAX_MULT`); these launchers do not check it.
12
13use anyhow::Result;
14use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
15use metrale_gpu_runtime::kernel_args::KernelLaunch;
16
17/// 2026-09-25: Broadcast the BF16 hidden state into `hc_mult` identical FP32 streams:
18/// `streams[t, i, d] = hidden[t, d]`.
19pub fn hc_expand(
20    gpu: &dyn GpuBackend,
21    kernel: KernelHandle,
22    hidden: DevicePtr,
23    streams: DevicePtr,
24    num_tokens: u32,
25    hidden_size: u32,
26    hc_mult: u32,
27    stream: u64,
28) -> Result<()> {
29    KernelLaunch::new(gpu, kernel)
30        .grid([num_tokens, 1, 1])
31        .block([256, 1, 1])
32        .arg_ptr(hidden)
33        .arg_ptr(streams)
34        .arg_u32(hidden_size)
35        .arg_u32(hc_mult)
36        .launch(stream)
37}
38
39/// 2026-09-25: Collapse the `hc_mult` streams to one BF16 row (an RMS-rescaled mix gives sigmoid
40/// `pre` weights for the weighted sum), and write `post_out [T, hc]` and the Sinkhorn-normalized
41/// `comb_out [T, hc, hc]` for the matching [`hc_post`].
42#[allow(clippy::too_many_arguments)]
43pub fn hc_pre(
44    gpu: &dyn GpuBackend,
45    kernel: KernelHandle,
46    streams: DevicePtr,
47    hc_fn: DevicePtr,
48    hc_scale: DevicePtr,
49    hc_base: DevicePtr,
50    y_out: DevicePtr,
51    post_out: DevicePtr,
52    comb_out: DevicePtr,
53    num_tokens: u32,
54    hidden_size: u32,
55    hc_mult: u32,
56    sinkhorn_iters: u32,
57    norm_eps: f32,
58    hc_eps: f32,
59    stream: u64,
60) -> Result<()> {
61    KernelLaunch::new(gpu, kernel)
62        .grid([num_tokens, 1, 1])
63        .block([256, 1, 1])
64        .arg_ptr(streams)
65        .arg_ptr(hc_fn)
66        .arg_ptr(hc_scale)
67        .arg_ptr(hc_base)
68        .arg_ptr(y_out)
69        .arg_ptr(post_out)
70        .arg_ptr(comb_out)
71        .arg_u32(hidden_size)
72        .arg_u32(hc_mult)
73        .arg_u32(sinkhorn_iters)
74        .arg_f32(norm_eps)
75        .arg_f32(hc_eps)
76        .launch(stream)
77}
78
79/// 2026-09-25: Expand the sublayer output back into `hc_mult` streams:
80/// `out[t,j,d] = post[t,j] * block_out[t,d] + sum_i comb[t,i,j] * residual[t,i,d]`. `out` may
81/// alias `residual`: each element's residual values are read before its outputs are written.
82#[allow(clippy::too_many_arguments)]
83pub fn hc_post(
84    gpu: &dyn GpuBackend,
85    kernel: KernelHandle,
86    block_out: DevicePtr,
87    residual: DevicePtr,
88    post: DevicePtr,
89    comb: DevicePtr,
90    out: DevicePtr,
91    num_tokens: u32,
92    hidden_size: u32,
93    hc_mult: u32,
94    stream: u64,
95) -> Result<()> {
96    KernelLaunch::new(gpu, kernel)
97        .grid([num_tokens, 1, 1])
98        .block([256, 1, 1])
99        .arg_ptr(block_out)
100        .arg_ptr(residual)
101        .arg_ptr(post)
102        .arg_ptr(comb)
103        .arg_ptr(out)
104        .arg_u32(hidden_size)
105        .arg_u32(hc_mult)
106        .launch(stream)
107}
108
109/// 2026-09-25: Final collapse before the LM head: one learned sigmoid-weighted sum over the
110/// `hc_mult` streams, written as BF16 `[T, H]`.
111#[allow(clippy::too_many_arguments)]
112pub fn hc_head(
113    gpu: &dyn GpuBackend,
114    kernel: KernelHandle,
115    streams: DevicePtr,
116    head_fn: DevicePtr,
117    head_scale: DevicePtr,
118    head_base: DevicePtr,
119    y_out: DevicePtr,
120    num_tokens: u32,
121    hidden_size: u32,
122    hc_mult: u32,
123    norm_eps: f32,
124    hc_eps: f32,
125    stream: u64,
126) -> Result<()> {
127    KernelLaunch::new(gpu, kernel)
128        .grid([num_tokens, 1, 1])
129        .block([256, 1, 1])
130        .arg_ptr(streams)
131        .arg_ptr(head_fn)
132        .arg_ptr(head_scale)
133        .arg_ptr(head_base)
134        .arg_ptr(y_out)
135        .arg_u32(hidden_size)
136        .arg_u32(hc_mult)
137        .arg_f32(norm_eps)
138        .arg_f32(hc_eps)
139        .launch(stream)
140}
metrale_model_layers/layers/ops/hyper_connection.rs

metrale_model_layers/layers/ops/
hyper_connection.rs
