metrale_model_layers/layers/ops/hyper_connection_dispatch.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: One entry per mHC stage (`hc_pre`, `hc_post`, `hc_head`) that selects between the Sinkhorn and the low-rank hyper-connection launchers.
4//!
5//! Owner: model-layers ops.
6//! Invariants:
7//! - Both families run over the same FP32 `[T, hc_mult, H]` streams and use the same kernel
8//! names, because a model's own `hyper_connection.cu` replaces the whole file
9//! (kernels/gb10/qwen3.8-flash-next/nvfp4/hyper_connection.cu), but their argument lists
10//! differ. The signatures here are the union of both.
11//! - The variant is chosen from the weights (`lowrank.is_some()`), never from the model name;
12//! a site with low-rank weights takes the low-rank launcher.
13//! - `post_out` / `post` carries the low-rank injection vector: both are `[T, hc]` FP32 and read
14//! by the matching `hc_post`, so one buffer serves both variants. The low-rank path neither
15//! writes nor reads `comb`.
16//!
17//! The DeepSeek-V4 mixer is Sinkhorn-normalized over `hc_fn` / `hc_scale` / `hc_base` and emits
18//! a `[hc, hc]` combine matrix; the Qwen3.8-Flash-Next mixer is a low-rank pair with a grouped
19//! RMSNorm and emits one scalar per stream.
20
21use anyhow::Result;
22use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
23
24use super::hyper_connection as sinkhorn;
25use super::hyper_connection_lowrank as lowrank;
26use crate::layers::qwen3_attention::{HcHeadWeights, HcSiteWeights, HcWeights};
27
28/// 2026-09-25: Which family a site's weights select. Callers use it to decide whether the
29/// block's `input_norm` runs on `hc_pre`'s output ([`HcVariant::applies_block_input_norm`]): in
30/// the low-rank family `hc_norm` inside `hc_pre` is the block's input norm, and a second RMS pass
31/// with a ones-filled weight is not an identity.
32#[derive(Clone, Copy, PartialEq, Eq, Debug)]
33pub enum HcVariant {
34 /// 2026-09-25: DeepSeek-V4: Sinkhorn mix, `[hc, hc]` combine matrix.
35 Sinkhorn,
36 /// 2026-09-25: Qwen3.8-Flash-Next: low-rank pair, one injection scalar per stream.
37 LowRank,
38}
39
40impl HcVariant {
41 pub fn of_site(site: &HcSiteWeights) -> Self {
42 if site.lowrank.is_some() {
43 Self::LowRank
44 } else {
45 Self::Sinkhorn
46 }
47 }
48
49 pub fn of(hc: &HcWeights) -> Self {
50 Self::of_site(&hc.attn)
51 }
52
53 /// 2026-09-25: Whether the block's own `input_norm` runs on `hc_pre`'s output: true only for
54 /// [`HcVariant::Sinkhorn`].
55 pub fn applies_block_input_norm(self) -> bool {
56 self == Self::Sinkhorn
57 }
58}
59
60/// 2026-09-25: Collapse the streams to one and emit what the matching `hc_post` needs: `post` and
61/// `comb` for Sinkhorn, the injection vector in `post_out` for low-rank.
62#[allow(clippy::too_many_arguments)]
63pub fn hc_pre_site(
64 gpu: &dyn GpuBackend,
65 kernel: KernelHandle,
66 streams: DevicePtr,
67 site: &HcSiteWeights,
68 hc: &HcWeights,
69 y_out: DevicePtr,
70 post_out: DevicePtr,
71 comb_out: DevicePtr,
72 scratch: DevicePtr,
73 num_tokens: u32,
74 hidden_size: u32,
75 norm_eps: f32,
76 stream: u64,
77) -> Result<()> {
78 match &site.lowrank {
79 Some(w) => lowrank::hc_pre_lowrank(
80 gpu,
81 kernel,
82 streams,
83 w,
84 y_out,
85 post_out,
86 scratch,
87 num_tokens,
88 hidden_size,
89 hc.hc_mult as u32,
90 norm_eps,
91 stream,
92 ),
93 None => sinkhorn::hc_pre(
94 gpu,
95 kernel,
96 streams,
97 site.hc_fn,
98 site.hc_scale,
99 site.hc_base,
100 y_out,
101 post_out,
102 comb_out,
103 num_tokens,
104 hidden_size,
105 hc.hc_mult as u32,
106 hc.sinkhorn_iters as u32,
107 norm_eps,
108 hc.hc_eps,
109 stream,
110 ),
111 }
112}
113
114/// 2026-09-25: Inject the block output back into every stream. `out` may alias `residual`.
115///
116/// Takes the whole [`HcWeights`] rather than a site because neither variant's `hc_post` reads
117/// site weights: Sinkhorn reads the `comb` its `hc_pre` emitted, low-rank the injection vector.
118/// The variant is taken from the layer's attention site.
119#[allow(clippy::too_many_arguments)]
120pub fn hc_post_site(
121 gpu: &dyn GpuBackend,
122 kernel: KernelHandle,
123 hc: &HcWeights,
124 block_out: DevicePtr,
125 residual: DevicePtr,
126 post: DevicePtr,
127 comb: DevicePtr,
128 out: DevicePtr,
129 num_tokens: u32,
130 hidden_size: u32,
131 stream: u64,
132) -> Result<()> {
133 match HcVariant::of(hc) {
134 // 2026-09-25: `post` is the injection vector here; `comb` is not read.
135 HcVariant::LowRank => lowrank::hc_post_lowrank(
136 gpu,
137 kernel,
138 block_out,
139 residual,
140 post,
141 out,
142 num_tokens,
143 hidden_size,
144 hc.hc_mult as u32,
145 stream,
146 ),
147 HcVariant::Sinkhorn => sinkhorn::hc_post(
148 gpu,
149 kernel,
150 block_out,
151 residual,
152 post,
153 comb,
154 out,
155 num_tokens,
156 hidden_size,
157 hc.hc_mult as u32,
158 stream,
159 ),
160 }
161}
162
163/// 2026-09-25: The model-level final collapse before the LM head. In the low-rank family this is
164/// also the model's final normalization: the Qwen3.8-Flash-Next config parser sets
165/// `final_norm_identity` (config/src/parsers/qwen4_exp.rs), and the mixer's `hc_norm` normalizes.
166#[allow(clippy::too_many_arguments)]
167pub fn hc_head_site(
168 gpu: &dyn GpuBackend,
169 kernel: KernelHandle,
170 streams: DevicePtr,
171 head: &HcHeadWeights,
172 hc: &HcWeights,
173 y_out: DevicePtr,
174 scratch: DevicePtr,
175 num_tokens: u32,
176 hidden_size: u32,
177 norm_eps: f32,
178 stream: u64,
179) -> Result<()> {
180 match &head.lowrank {
181 Some(w) => lowrank::hc_head_lowrank(
182 gpu,
183 kernel,
184 streams,
185 w,
186 y_out,
187 scratch,
188 num_tokens,
189 hidden_size,
190 hc.hc_mult as u32,
191 norm_eps,
192 stream,
193 ),
194 None => sinkhorn::hc_head(
195 gpu,
196 kernel,
197 streams,
198 head.hc_fn,
199 head.hc_scale,
200 head.hc_base,
201 y_out,
202 num_tokens,
203 hidden_size,
204 hc.hc_mult as u32,
205 norm_eps,
206 hc.hc_eps,
207 stream,
208 ),
209 }
210}