spark_model/layers/ops/
hyper_connection_lowrank.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Qwen3.8-Flash-Next low-rank mHC dispatch.
4//!
5//! Companion to `hyper_connection.rs`, which drives DeepSeek-V4's Sinkhorn
6//! mixer. Both families share the `[T, hc_mult, H]` FP32 highway and the same
7//! four kernel NAMES — a model shadow overrides the whole
8//! `hyper_connection.cu` file, so `qwen3.8-flash-next` resolves
9//! `hyper_connection::hc_pre` to the low-rank kernel while
10//! `deepseek-v4-flash` resolves it to the Sinkhorn one. The two take
11//! DIFFERENT argument lists, which is why the launches live apart.
12//!
13//! `hc_expand` is byte-identical across both and is not duplicated here.
14//!
15//! Selection is by WEIGHTS, not by model name: `HcSiteWeights::lowrank`
16//! being `Some` is what routes here. A model that somehow carried both would
17//! be a load-time bug, not a silent dispatch coin-flip.
18
19use anyhow::Result;
20#[path = "hyper_connection_lowrank_gemm.rs"]
21mod gemm;
22pub(crate) use gemm::hc_pre_gemm;
23
24use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
25use spark_runtime::kernel_args::KernelLaunch;
26
27use crate::layers::qwen3_attention::HcLowRank;
28
29/// `ATLAS_QWEN4EXP_NO_HC_GEMM=1`: revert the large-T collapse to the fused
30/// FP32 kernel (deploy-time kill switch; the GEMM path rounds `normed` to
31/// BF16 before the projections).
32fn hc_gemm_disabled() -> bool {
33    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
34    *ON.get_or_init(|| std::env::var("ATLAS_QWEN4EXP_NO_HC_GEMM").as_deref() == Ok("1"))
35}
36
37/// `ATLAS_HC_DECODE_SPLIT=1`: keep the pre-cuBLASLt split path for
38/// decode-shaped T (A/B escape hatch, same convention as the GEMM kill
39/// switch above).
40fn hc_decode_split_forced() -> bool {
41    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
42    *V.get_or_init(|| std::env::var("ATLAS_HC_DECODE_SPLIT").as_deref() == Ok("1"))
43}
44
45/// Collapse the `hc_mult` streams to one, and emit the per-stream injection
46/// weights the matching [`hc_post_lowrank`] needs.
47///
48/// `streams [T, hc, H] -> y_out [T, H]`, `inj_out [T, hc]`.
49#[allow(clippy::too_many_arguments)]
50pub fn hc_pre_lowrank(
51    gpu: &dyn GpuBackend,
52    kernel: KernelHandle,
53    streams: DevicePtr,
54    w: &HcLowRank,
55    y_out: DevicePtr,
56    inj_out: DevicePtr,
57    scratch: DevicePtr,
58    num_tokens: u32,
59    hidden_size: u32,
60    hc_mult: u32,
61    norm_eps: f32,
62    stream: u64,
63) -> Result<()> {
64    anyhow::ensure!(
65        !w.inject_w.is_null(),
66        "hc_pre_lowrank needs block_inject_weight; a site loaded without one \
67         is the model-level mixer and must use hc_head_lowrank"
68    );
69    // SMALL T (decode): three multi-block launches instead of the fused
70    // kernel, whose grid=[T] means grid=[1] at decode — one block, one SM,
71    // ~13 MB of weights per call (measured 2.0 ms; the whole token was
72    // 96 x that). The fused kernel stays for prefill, where grid=[T]
73    // already fills the machine and skips the global round trip.
74    if num_tokens <= 64 && !scratch.is_null() {
75        // Decode-shaped T: the GEMM decomposition with cuBLASLt for the
76        // three projections. The split path's hand-rolled k_down/k_fin each
77        // stream ~6.5 MB of low-rank weights well off the bandwidth floor —
78        // the same GEMM-shaped-work-on-hand-rolled-kernels defect class as
79        // the prefill collapse and the batched-decode QKVZ arms, and the
80        // same cure. ATLAS_HC_DECODE_SPLIT=1 keeps the split path (A/B).
81        if !hc_decode_split_forced() {
82            return hc_pre_gemm(
83                gpu,
84                streams,
85                w,
86                y_out,
87                inj_out,
88                scratch,
89                num_tokens,
90                hidden_size,
91                hc_mult,
92                norm_eps,
93                /* inject */ true,
94                /* use_cublas */ true,
95                stream,
96            );
97        }
98        return hc_pre_split(
99            gpu,
100            streams,
101            w,
102            y_out,
103            inj_out,
104            scratch,
105            num_tokens,
106            hidden_size,
107            hc_mult,
108            norm_eps,
109            /* inject */ true,
110            stream,
111        );
112    }
113    // LARGE T (prefill): tensor-core GEMM formulation — 47% of prefill was
114    // this collapse running as FP32 warp loops. Kill switch reverts to the
115    // fused kernel below.
116    if !scratch.is_null() && !hc_gemm_disabled() {
117        return hc_pre_gemm(
118            gpu,
119            streams,
120            w,
121            y_out,
122            inj_out,
123            scratch,
124            num_tokens,
125            hidden_size,
126            hc_mult,
127            norm_eps,
128            /* inject */ true,
129            /* use_cublas */ false,
130            stream,
131        );
132    }
133    // Block 1024 + dynamic shared for the staged normed vector [hc*H] and
134    // the rank vector — the warp-cooperative core. This launch WAS the whole
135    // decode budget at block 256 with per-thread serial rows (4.5 ms/call,
136    // x96 calls/token); see the kernel's PERFORMANCE SHAPE note.
137    let smem = (hc_mult * hidden_size + w.rank as u32) * 4;
138    KernelLaunch::new(gpu, kernel)
139        .grid([num_tokens, 1, 1])
140        .block([1024, 1, 1])
141        .shared_mem(smem)
142        .arg_ptr(streams)
143        .arg_ptr(w.norm_w)
144        .arg_ptr(w.down_w)
145        .arg_ptr(w.up_w)
146        .arg_ptr(w.inject_w)
147        .arg_ptr(y_out)
148        .arg_ptr(inj_out)
149        .arg_u32(hidden_size)
150        .arg_u32(hc_mult)
151        .arg_u32(w.rank as u32)
152        .arg_f32(norm_eps)
153        .launch(stream)
154}
155
156/// The model-level mixer (`use_combine=False`): the same collapse with no
157/// injection vector.
158///
159/// This is also the model's FINAL NORMALIZATION — the checkpoint ships no
160/// `model.norm.weight` because `hc_norm` here plays that role.
161#[allow(clippy::too_many_arguments)]
162pub fn hc_head_lowrank(
163    gpu: &dyn GpuBackend,
164    kernel: KernelHandle,
165    streams: DevicePtr,
166    w: &HcLowRank,
167    y_out: DevicePtr,
168    scratch: DevicePtr,
169    num_tokens: u32,
170    hidden_size: u32,
171    hc_mult: u32,
172    norm_eps: f32,
173    stream: u64,
174) -> Result<()> {
175    if num_tokens <= 64 && !scratch.is_null() {
176        if !hc_decode_split_forced() {
177            return hc_pre_gemm(
178                gpu,
179                streams,
180                w,
181                y_out,
182                DevicePtr::NULL,
183                scratch,
184                num_tokens,
185                hidden_size,
186                hc_mult,
187                norm_eps,
188                /* inject */ false,
189                /* use_cublas */ true,
190                stream,
191            );
192        }
193        return hc_pre_split(
194            gpu,
195            streams,
196            w,
197            y_out,
198            DevicePtr::NULL,
199            scratch,
200            num_tokens,
201            hidden_size,
202            hc_mult,
203            norm_eps,
204            /* inject */ false,
205            stream,
206        );
207    }
208    // Same GEMM formulation as hc_pre — the head is the identical collapse
209    // minus the injection GEMM (hc_pre_mix skips inj on a null inj_pre).
210    if !scratch.is_null() && !hc_gemm_disabled() {
211        return hc_pre_gemm(
212            gpu,
213            streams,
214            w,
215            y_out,
216            DevicePtr::NULL,
217            scratch,
218            num_tokens,
219            hidden_size,
220            hc_mult,
221            norm_eps,
222            /* inject */ false,
223            /* use_cublas */ false,
224            stream,
225        );
226    }
227    let smem = (hc_mult * hidden_size + w.rank as u32) * 4;
228    KernelLaunch::new(gpu, kernel)
229        .grid([num_tokens, 1, 1])
230        .block([1024, 1, 1])
231        .shared_mem(smem)
232        .arg_ptr(streams)
233        .arg_ptr(w.norm_w)
234        .arg_ptr(w.down_w)
235        .arg_ptr(w.up_w)
236        .arg_ptr(y_out)
237        .arg_u32(hidden_size)
238        .arg_u32(hc_mult)
239        .arg_u32(w.rank as u32)
240        .arg_f32(norm_eps)
241        .launch(stream)
242}
243
244/// Inject the block output back into every stream:
245/// `out[t, s*H + d] = residual[t, s*H + d] + block_out[t, d] * inj[t, s]`.
246///
247/// Note there is no `comb` argument: DeepSeek mixes streams with a full
248/// `[hc, hc]` combine matrix on the way back, Qwen scales by one scalar per
249/// stream. Passing a combine matrix here would not type-check, which is the
250/// point of keeping the two launches separate.
251#[allow(clippy::too_many_arguments)]
252pub fn hc_post_lowrank(
253    gpu: &dyn GpuBackend,
254    kernel: KernelHandle,
255    block_out: DevicePtr,
256    residual: DevicePtr,
257    inj: DevicePtr,
258    out: DevicePtr,
259    num_tokens: u32,
260    hidden_size: u32,
261    hc_mult: u32,
262    stream: u64,
263) -> Result<()> {
264    KernelLaunch::new(gpu, kernel)
265        .grid([num_tokens, 1, 1])
266        .block([256, 1, 1])
267        .arg_ptr(block_out)
268        .arg_ptr(residual)
269        .arg_ptr(inj)
270        .arg_ptr(out)
271        .arg_u32(hidden_size)
272        .arg_u32(hc_mult)
273        .launch(stream)
274}
275
276/// The three-launch collapse for small T. Same math as the fused kernel;
277/// the parity probe's T=8 fixture runs THIS path.
278#[allow(clippy::too_many_arguments)]
279pub(crate) fn hc_pre_split(
280    gpu: &dyn GpuBackend,
281    streams: DevicePtr,
282    w: &HcLowRank,
283    y_out: DevicePtr,
284    inj_out: DevicePtr,
285    scratch: DevicePtr,
286    num_tokens: u32,
287    hidden_size: u32,
288    hc_mult: u32,
289    norm_eps: f32,
290    inject: bool,
291    stream: u64,
292) -> Result<()> {
293    let hc_dim = hc_mult * hidden_size;
294    // Scratch layout: normed [T<=64, hc_dim] then low [T<=64, rank], F32.
295    let normed = scratch;
296    let low = scratch.offset(64 * hc_dim as usize * 4);
297
298    let k_stage = gpu.kernel("hyper_connection", "hc_pre_stage")?;
299    let k_down = gpu.kernel("hyper_connection", "hc_pre_down")?;
300    let k_fin = gpu.kernel("hyper_connection", "hc_pre_finish")?;
301
302    KernelLaunch::new(gpu, k_stage)
303        .grid([num_tokens, 1, 1])
304        .block([1024, 1, 1])
305        .arg_ptr(streams)
306        .arg_ptr(w.norm_w)
307        .arg_ptr(normed)
308        .arg_u32(hidden_size)
309        .arg_u32(hc_mult)
310        .arg_f32(norm_eps)
311        .launch(stream)?;
312
313    // Spread rank rows over enough blocks to occupy the part even at T=1.
314    let dsplit = (48 / num_tokens.max(1)).clamp(1, 10);
315    KernelLaunch::new(gpu, k_down)
316        .grid([num_tokens, dsplit, 1])
317        .block([1024, 1, 1])
318        .arg_ptr(normed)
319        .arg_ptr(w.down_w)
320        .arg_ptr(low)
321        .arg_u32(hidden_size)
322        .arg_u32(hc_mult)
323        .arg_u32(w.rank as u32)
324        .launch(stream)?;
325
326    let fsplit = (48 / num_tokens.max(1)).clamp(1, 10);
327    KernelLaunch::new(gpu, k_fin)
328        .grid([num_tokens, fsplit, 1])
329        .block([256, 1, 1])
330        .shared_mem(w.rank as u32 * 4)
331        .arg_ptr(normed)
332        .arg_ptr(low)
333        .arg_ptr(w.up_w)
334        .arg_ptr(if inject { w.inject_w } else { DevicePtr::NULL })
335        .arg_ptr(y_out)
336        .arg_ptr(inj_out)
337        .arg_u32(hidden_size)
338        .arg_u32(hc_mult)
339        .arg_u32(w.rank as u32)
340        .launch(stream)
341}
spark_model/layers/ops/hyper_connection_lowrank.rs

spark_model/layers/ops/
hyper_connection_lowrank.rs
