metrale_model_layers/layers/ops/
hyper_connection_lowrank.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the Qwen3.8-Flash-Next low-rank hyper-connection (mHC) kernels.
4//!
5//! Owner: model-layers ops.
6//! Invariants:
7//! - The kernels come from kernels/gb10/qwen3.8-flash-next/nvfp4/hyper_connection.cu, which
8//!   uses the same kernel names as the DeepSeek-V4 file launched by `hyper_connection.rs` but
9//!   different argument lists; the two sets of launchers are kept apart for that reason.
10//! - `hc_expand` behaves the same in both kernel files, so it has no launcher here.
11//! - `hyper_connection_dispatch.rs` routes a site here when `HcSiteWeights::lowrank` is `Some`.
12//! - With `T <= 64` and a non-null scratch, [`hc_pre_lowrank`] and [`hc_head_lowrank`] run the
13//!   cuBLASLt GEMM path (or the three-launch split path when `METRALE_HC_DECODE_SPLIT=1`); with
14//!   larger `T` and a scratch they run the `dense_gemm_bf16_pipelined` GEMM path unless
15//!   `METRALE_QWEN4EXP_NO_HC_GEMM=1`; otherwise the fused kernel, one 1024-thread block per token.
16
17use anyhow::Result;
18#[path = "hyper_connection_lowrank_gemm.rs"]
19mod gemm;
20pub(crate) use gemm::hc_pre_gemm;
21
22use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
23use metrale_gpu_runtime::kernel_args::KernelLaunch;
24
25use crate::layers::qwen3_attention::HcLowRank;
26
27/// 2026-09-25: `METRALE_QWEN4EXP_NO_HC_GEMM` equal to `"1"`, read once per process: the large-T
28/// collapse runs the fused FP32 kernel instead of the GEMM path, which rounds `normed` to BF16
29/// before the projections (`hc_pre_stage_bf16`).
30fn hc_gemm_disabled() -> bool {
31    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
32    *ON.get_or_init(|| std::env::var("METRALE_QWEN4EXP_NO_HC_GEMM").as_deref() == Ok("1"))
33}
34
35/// 2026-09-25: `METRALE_HC_DECODE_SPLIT` equal to `"1"`, read once per process: at `T <= 64` run
36/// the three-launch split path instead of the cuBLASLt GEMM path.
37fn hc_decode_split_forced() -> bool {
38    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
39    *V.get_or_init(|| std::env::var("METRALE_HC_DECODE_SPLIT").as_deref() == Ok("1"))
40}
41
42/// 2026-09-25: Collapse the `hc_mult` streams to one, and emit the per-stream injection weights
43/// the matching [`hc_post_lowrank`] needs: `streams [T, hc, H] -> y_out [T, H]` BF16,
44/// `inj_out [T, hc]` FP32. Fails when `w.inject_w` is null (a site without an injection weight
45/// is the model-level mixer, which uses [`hc_head_lowrank`]).
46#[allow(clippy::too_many_arguments)]
47pub fn hc_pre_lowrank(
48    gpu: &dyn GpuBackend,
49    kernel: KernelHandle,
50    streams: DevicePtr,
51    w: &HcLowRank,
52    y_out: DevicePtr,
53    inj_out: DevicePtr,
54    scratch: DevicePtr,
55    num_tokens: u32,
56    hidden_size: u32,
57    hc_mult: u32,
58    norm_eps: f32,
59    stream: u64,
60) -> Result<()> {
61    anyhow::ensure!(
62        !w.inject_w.is_null(),
63        "hc_pre_lowrank needs block_inject_weight; a site loaded without one \
64         is the model-level mixer and must use hc_head_lowrank"
65    );
66    // 2026-09-25: Small T: the fused kernel's grid is [T], a single block at T = 1, so small T
67    // uses multi-block launches through the scratch instead.
68    if num_tokens <= 64 && !scratch.is_null() {
69        if !hc_decode_split_forced() {
70            return hc_pre_gemm(
71                gpu,
72                streams,
73                w,
74                y_out,
75                inj_out,
76                scratch,
77                num_tokens,
78                hidden_size,
79                hc_mult,
80                norm_eps,
81                true,
82                true,
83                stream,
84            );
85        }
86        return hc_pre_split(
87            gpu,
88            streams,
89            w,
90            y_out,
91            inj_out,
92            scratch,
93            num_tokens,
94            hidden_size,
95            hc_mult,
96            norm_eps,
97            true,
98            stream,
99        );
100    }
101    // 2026-09-25: Large T: the GEMM formulation on `dense_gemm_bf16_pipelined`, unless the kill
102    // switch sends it to the fused kernel below.
103    if !scratch.is_null() && !hc_gemm_disabled() {
104        return hc_pre_gemm(
105            gpu,
106            streams,
107            w,
108            y_out,
109            inj_out,
110            scratch,
111            num_tokens,
112            hidden_size,
113            hc_mult,
114            norm_eps,
115            true,
116            false,
117            stream,
118        );
119    }
120    // 2026-09-25: 1024 threads, and dynamic shared memory for the staged `normed` vector
121    // `[hc*H]` and the rank vector, both FP32.
122    let smem = (hc_mult * hidden_size + w.rank as u32) * 4;
123    KernelLaunch::new(gpu, kernel)
124        .grid([num_tokens, 1, 1])
125        .block([1024, 1, 1])
126        .shared_mem(smem)
127        .arg_ptr(streams)
128        .arg_ptr(w.norm_w)
129        .arg_ptr(w.down_w)
130        .arg_ptr(w.up_w)
131        .arg_ptr(w.inject_w)
132        .arg_ptr(y_out)
133        .arg_ptr(inj_out)
134        .arg_u32(hidden_size)
135        .arg_u32(hc_mult)
136        .arg_u32(w.rank as u32)
137        .arg_f32(norm_eps)
138        .launch(stream)
139}
140
141/// 2026-09-25: The model-level mixer (`use_combine=False` in the reference
142/// bench/qwen4_exp/ref/modeling_qwen4_exp.py): the same collapse with no injection vector, on
143/// the same path choice as [`hc_pre_lowrank`]. It is also the model's final normalization: the
144/// config parser sets `final_norm_identity` (config/src/parsers/qwen4_exp.rs).
145#[allow(clippy::too_many_arguments)]
146pub fn hc_head_lowrank(
147    gpu: &dyn GpuBackend,
148    kernel: KernelHandle,
149    streams: DevicePtr,
150    w: &HcLowRank,
151    y_out: DevicePtr,
152    scratch: DevicePtr,
153    num_tokens: u32,
154    hidden_size: u32,
155    hc_mult: u32,
156    norm_eps: f32,
157    stream: u64,
158) -> Result<()> {
159    if num_tokens <= 64 && !scratch.is_null() {
160        if !hc_decode_split_forced() {
161            return hc_pre_gemm(
162                gpu,
163                streams,
164                w,
165                y_out,
166                DevicePtr::NULL,
167                scratch,
168                num_tokens,
169                hidden_size,
170                hc_mult,
171                norm_eps,
172                false,
173                true,
174                stream,
175            );
176        }
177        return hc_pre_split(
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            false,
189            stream,
190        );
191    }
192    // 2026-09-25: The GEMM path without the injection GEMM; `hc_pre_mix` skips the injection
193    // when `inj_pre` is null.
194    if !scratch.is_null() && !hc_gemm_disabled() {
195        return hc_pre_gemm(
196            gpu,
197            streams,
198            w,
199            y_out,
200            DevicePtr::NULL,
201            scratch,
202            num_tokens,
203            hidden_size,
204            hc_mult,
205            norm_eps,
206            false,
207            false,
208            stream,
209        );
210    }
211    let smem = (hc_mult * hidden_size + w.rank as u32) * 4;
212    KernelLaunch::new(gpu, kernel)
213        .grid([num_tokens, 1, 1])
214        .block([1024, 1, 1])
215        .shared_mem(smem)
216        .arg_ptr(streams)
217        .arg_ptr(w.norm_w)
218        .arg_ptr(w.down_w)
219        .arg_ptr(w.up_w)
220        .arg_ptr(y_out)
221        .arg_u32(hidden_size)
222        .arg_u32(hc_mult)
223        .arg_u32(w.rank as u32)
224        .arg_f32(norm_eps)
225        .launch(stream)
226}
227
228/// 2026-09-25: Inject the block output back into every stream:
229/// `out[t, s*H + d] = residual[t, s*H + d] + block_out[t, d] * inj[t, s]`. There is no `comb`
230/// argument: this family scales each stream by one scalar.
231#[allow(clippy::too_many_arguments)]
232pub fn hc_post_lowrank(
233    gpu: &dyn GpuBackend,
234    kernel: KernelHandle,
235    block_out: DevicePtr,
236    residual: DevicePtr,
237    inj: DevicePtr,
238    out: DevicePtr,
239    num_tokens: u32,
240    hidden_size: u32,
241    hc_mult: u32,
242    stream: u64,
243) -> Result<()> {
244    KernelLaunch::new(gpu, kernel)
245        .grid([num_tokens, 1, 1])
246        .block([256, 1, 1])
247        .arg_ptr(block_out)
248        .arg_ptr(residual)
249        .arg_ptr(inj)
250        .arg_ptr(out)
251        .arg_u32(hidden_size)
252        .arg_u32(hc_mult)
253        .launch(stream)
254}
255
256/// 2026-09-25: The three-launch collapse (`hc_pre_stage`, `hc_pre_down`, `hc_pre_finish`) for
257/// `T <= 64`, the row count its scratch layout reserves. `hyper_connection_lowrank_tests.rs`
258/// calls it directly.
259#[allow(clippy::too_many_arguments)]
260pub(crate) fn hc_pre_split(
261    gpu: &dyn GpuBackend,
262    streams: DevicePtr,
263    w: &HcLowRank,
264    y_out: DevicePtr,
265    inj_out: DevicePtr,
266    scratch: DevicePtr,
267    num_tokens: u32,
268    hidden_size: u32,
269    hc_mult: u32,
270    norm_eps: f32,
271    inject: bool,
272    stream: u64,
273) -> Result<()> {
274    let hc_dim = hc_mult * hidden_size;
275    // 2026-09-25: Scratch layout: normed [64, hc_dim] then low [64, rank], FP32.
276    let normed = scratch;
277    let low = scratch.offset(64 * hc_dim as usize * 4);
278
279    let k_stage = gpu.kernel("hyper_connection", "hc_pre_stage")?;
280    let k_down = gpu.kernel("hyper_connection", "hc_pre_down")?;
281    let k_fin = gpu.kernel("hyper_connection", "hc_pre_finish")?;
282
283    KernelLaunch::new(gpu, k_stage)
284        .grid([num_tokens, 1, 1])
285        .block([1024, 1, 1])
286        .arg_ptr(streams)
287        .arg_ptr(w.norm_w)
288        .arg_ptr(normed)
289        .arg_u32(hidden_size)
290        .arg_u32(hc_mult)
291        .arg_f32(norm_eps)
292        .launch(stream)?;
293
294    // 2026-09-25: Grid y splits the rank rows so that small T still launches several blocks.
295    let dsplit = (48 / num_tokens.max(1)).clamp(1, 10);
296    KernelLaunch::new(gpu, k_down)
297        .grid([num_tokens, dsplit, 1])
298        .block([1024, 1, 1])
299        .arg_ptr(normed)
300        .arg_ptr(w.down_w)
301        .arg_ptr(low)
302        .arg_u32(hidden_size)
303        .arg_u32(hc_mult)
304        .arg_u32(w.rank as u32)
305        .launch(stream)?;
306
307    let fsplit = (48 / num_tokens.max(1)).clamp(1, 10);
308    KernelLaunch::new(gpu, k_fin)
309        .grid([num_tokens, fsplit, 1])
310        .block([256, 1, 1])
311        .shared_mem(w.rank as u32 * 4)
312        .arg_ptr(normed)
313        .arg_ptr(low)
314        .arg_ptr(w.up_w)
315        .arg_ptr(if inject { w.inject_w } else { DevicePtr::NULL })
316        .arg_ptr(y_out)
317        .arg_ptr(inj_out)
318        .arg_u32(hidden_size)
319        .arg_u32(hc_mult)
320        .arg_u32(w.rank as u32)
321        .launch(stream)
322}
metrale_model_layers/layers/ops/hyper_connection_lowrank.rs

metrale_model_layers/layers/ops/
hyper_connection_lowrank.rs
