spark_model/layers/ops/
hyper_connection_dispatch.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! One call site per mHC entry point, for both variants.
4//!
5//! Atlas now runs two different hyper-connection families over the same
6//! `[T, hc_mult, H]` FP32 highway:
7//!
8//! * DeepSeek-V4's — a Sinkhorn-normalized mix over `hc_fn`/`hc_scale`/
9//!   `hc_base`, emitting a `[hc, hc]` combine matrix that `hc_post` mixes the
10//!   saved streams through.
11//! * Qwen3.8-Flash-Next's — a low-rank pair (rank 320) with a grouped
12//!   RMSNorm, emitting one scalar per stream that `hc_post` scales by.
13//!
14//! They share the four kernel NAMES (a model shadow overrides the whole
15//! `hyper_connection.cu` file) but take different argument lists. Rather than
16//! branch at each of the 23 call sites across `prefill_inner`, `decode_inner`
17//! and `multi_seq`, the branch lives here once, behind a signature that is
18//! the union of both.
19//!
20//! **`post_out` doubles as the low-rank injection vector.** Both are `[T, hc]`
21//! f32 and both are consumed by the matching `hc_post`, so the existing
22//! `ctx.buffers.hc_post()` serves each variant without a second allocation.
23//! `comb_out` is untouched on the low-rank path — Qwen scales each stream by
24//! one scalar where DeepSeek mixes them through a full matrix.
25//!
26//! Selection is by WEIGHTS (`lowrank.is_some()`), never by model name. A site
27//! carrying both would be a load-time bug rather than a silent coin-flip.
28
29use anyhow::Result;
30use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
31
32use super::hyper_connection as sinkhorn;
33use super::hyper_connection_lowrank as lowrank;
34use crate::layers::qwen3_attention::{HcHeadWeights, HcSiteWeights, HcWeights};
35
36/// Which family a site's weights select.
37///
38/// Exposed so a caller can answer questions the launch itself cannot — most
39/// importantly whether to apply the layer's `input_norm` after `hc_pre`.
40/// Qwen has no per-layer `input_layernorm`; `hc_norm` inside `hc_pre` plays
41/// that role, and the loader's ones-filled placeholder does NOT make a second
42/// RMS pass an identity.
43#[derive(Clone, Copy, PartialEq, Eq, Debug)]
44pub enum HcVariant {
45    /// DeepSeek-V4: Sinkhorn mix, `[hc, hc]` combine matrix.
46    Sinkhorn,
47    /// Qwen3.8-Flash-Next: low-rank pair, per-stream scalar injection.
48    LowRank,
49}
50
51impl HcVariant {
52    pub fn of_site(site: &HcSiteWeights) -> Self {
53        if site.lowrank.is_some() {
54            Self::LowRank
55        } else {
56            Self::Sinkhorn
57        }
58    }
59
60    pub fn of(hc: &HcWeights) -> Self {
61        Self::of_site(&hc.attn)
62    }
63
64    /// Whether the block's own `input_norm` should run on `hc_pre`'s output.
65    /// False for Qwen — see the type's doc comment.
66    pub fn applies_block_input_norm(self) -> bool {
67        self == Self::Sinkhorn
68    }
69}
70
71/// Collapse the streams to one and emit whatever the matching `hc_post`
72/// needs — `post`+`comb` for Sinkhorn, the injection vector in `post_out` for
73/// low-rank.
74#[allow(clippy::too_many_arguments)]
75pub fn hc_pre_site(
76    gpu: &dyn GpuBackend,
77    kernel: KernelHandle,
78    streams: DevicePtr,
79    site: &HcSiteWeights,
80    hc: &HcWeights,
81    y_out: DevicePtr,
82    post_out: DevicePtr,
83    comb_out: DevicePtr,
84    scratch: DevicePtr,
85    num_tokens: u32,
86    hidden_size: u32,
87    norm_eps: f32,
88    stream: u64,
89) -> Result<()> {
90    match &site.lowrank {
91        Some(w) => lowrank::hc_pre_lowrank(
92            gpu,
93            kernel,
94            streams,
95            w,
96            y_out,
97            post_out,
98            scratch,
99            num_tokens,
100            hidden_size,
101            hc.hc_mult as u32,
102            norm_eps,
103            stream,
104        ),
105        None => sinkhorn::hc_pre(
106            gpu,
107            kernel,
108            streams,
109            site.hc_fn,
110            site.hc_scale,
111            site.hc_base,
112            y_out,
113            post_out,
114            comb_out,
115            num_tokens,
116            hidden_size,
117            hc.hc_mult as u32,
118            hc.sinkhorn_iters as u32,
119            norm_eps,
120            hc.hc_eps,
121            stream,
122        ),
123    }
124}
125
126/// Inject the block output back into every stream. `out` may alias
127/// `residual`.
128///
129/// Takes the whole [`HcWeights`] rather than a site, because NEITHER variant's
130/// `hc_post` reads site weights — Sinkhorn consumes the `comb` its `hc_pre`
131/// emitted, low-rank the injection vector — so the layer's variant is the
132/// only thing being selected on.
133#[allow(clippy::too_many_arguments)]
134pub fn hc_post_site(
135    gpu: &dyn GpuBackend,
136    kernel: KernelHandle,
137    hc: &HcWeights,
138    block_out: DevicePtr,
139    residual: DevicePtr,
140    post: DevicePtr,
141    comb: DevicePtr,
142    out: DevicePtr,
143    num_tokens: u32,
144    hidden_size: u32,
145    stream: u64,
146) -> Result<()> {
147    match HcVariant::of(hc) {
148        // `post` IS the injection vector here; `comb` is not read.
149        HcVariant::LowRank => lowrank::hc_post_lowrank(
150            gpu,
151            kernel,
152            block_out,
153            residual,
154            post,
155            out,
156            num_tokens,
157            hidden_size,
158            hc.hc_mult as u32,
159            stream,
160        ),
161        HcVariant::Sinkhorn => sinkhorn::hc_post(
162            gpu,
163            kernel,
164            block_out,
165            residual,
166            post,
167            comb,
168            out,
169            num_tokens,
170            hidden_size,
171            hc.hc_mult as u32,
172            stream,
173        ),
174    }
175}
176
177/// The model-level final collapse before the LM head.
178///
179/// On Qwen this is ALSO the model's final normalization — the checkpoint
180/// ships no `model.norm.weight` because `hyper_connection_mixer`'s `hc_norm`
181/// plays that role.
182#[allow(clippy::too_many_arguments)]
183pub fn hc_head_site(
184    gpu: &dyn GpuBackend,
185    kernel: KernelHandle,
186    streams: DevicePtr,
187    head: &HcHeadWeights,
188    hc: &HcWeights,
189    y_out: DevicePtr,
190    scratch: DevicePtr,
191    num_tokens: u32,
192    hidden_size: u32,
193    norm_eps: f32,
194    stream: u64,
195) -> Result<()> {
196    match &head.lowrank {
197        Some(w) => lowrank::hc_head_lowrank(
198            gpu,
199            kernel,
200            streams,
201            w,
202            y_out,
203            scratch,
204            num_tokens,
205            hidden_size,
206            hc.hc_mult as u32,
207            norm_eps,
208            stream,
209        ),
210        None => sinkhorn::hc_head(
211            gpu,
212            kernel,
213            streams,
214            head.hc_fn,
215            head.hc_scale,
216            head.hc_base,
217            y_out,
218            num_tokens,
219            hidden_size,
220            hc.hc_mult as u32,
221            norm_eps,
222            hc.hc_eps,
223            stream,
224        ),
225    }
226}
spark_model/layers/ops/hyper_connection_dispatch.rs

spark_model/layers/ops/
hyper_connection_dispatch.rs
