metrale_model_arch/glm5next_mlp/forward/
dense.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: The GLM-5.3 dense SwiGLU MLP forward, used by the dense layers and the shared
4//! expert.
5//!
6//! Owner: model-arch (GLM-5.3).
7//! Invariants:
8//! - `forward_dense` leaves a partial sum in `out` when `inter` is a TP shard; the caller
9//!   all-reduces.
10
11use anyhow::{Result, bail};
12use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
13
14use super::Glm5NextMlpWorkspace;
15use super::launch::{gemm, swiglu};
16use crate::glm5next_mlp::weights::Glm5NextDenseMlpWeights;
17use crate::glm5next_mlp::{Glm5NextMlpConfig, Glm5NextMlpKernels};
18
19/// 2026-09-25: A BF16 SwiGLU MLP of width `inter`, `down(clamped_swiglu(gate(x), up(x)))`, for a
20/// dense layer or the shared expert. Errors when `inter` or `m` does not fit the workspace.
21#[allow(clippy::too_many_arguments)]
22pub fn forward_dense(
23    gpu: &dyn GpuBackend,
24    k: &Glm5NextMlpKernels,
25    cfg: &Glm5NextMlpConfig,
26    w: &Glm5NextDenseMlpWeights,
27    inter: usize,
28    x: DevicePtr,
29    out: DevicePtr,
30    m: usize,
31    ws: &Glm5NextMlpWorkspace,
32    stream: u64,
33) -> Result<()> {
34    if inter == 0 || inter > ws.max_inter {
35        bail!(
36            "GLM dense MLP: width {inter} does not fit a workspace built for {}",
37            ws.max_inter
38        );
39    }
40    if m == 0 || m > ws.max_rows {
41        bail!(
42            "GLM dense MLP: {m} rows do not fit a workspace built for {}",
43            ws.max_rows
44        );
45    }
46    gemm(
47        gpu,
48        k.gemm,
49        k.gemv,
50        k.gemv_batchm,
51        x,
52        w.gate_proj,
53        ws.a_gate,
54        m,
55        inter,
56        cfg.hidden,
57        stream,
58    )?;
59    gemm(
60        gpu,
61        k.gemm,
62        k.gemv,
63        k.gemv_batchm,
64        x,
65        w.up_proj,
66        ws.a_up,
67        m,
68        inter,
69        cfg.hidden,
70        stream,
71    )?;
72    swiglu(
73        gpu,
74        k.swiglu,
75        ws.a_gate,
76        ws.a_up,
77        ws.a_act,
78        // 2026-09-25: Elementwise over all `m * inter` values.
79        m * inter,
80        cfg.swiglu_limit,
81        stream,
82    )?;
83    gemm(
84        gpu,
85        k.gemm,
86        k.gemv,
87        k.gemv_batchm,
88        ws.a_act,
89        w.down_proj,
90        out,
91        m,
92        cfg.hidden,
93        inter,
94        stream,
95    )
96}
metrale_model_arch/glm5next_mlp/forward/dense.rs

metrale_model_arch/glm5next_mlp/forward/
dense.rs
