metrale_model_arch/glm5next_layer/
types.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The data of a bound GLM-5.3 layer: `Glm5NextLayer` and its mixer, MLP site and
4//! hyper-connection types.
5//!
6//! Owner: model-arch (GLM-5.3).
7//! Invariants: none beyond the types.
8
9use super::*;
10
11/// 2026-09-25: The layer's mixer. KDA's `decode` is an inherent method that takes a
12/// `KdaSeqState` and a workspace and leaves its output in `ws.final_out`; DSA's `decode` is the
13/// `TransformerLayer` method and writes its output over the buffer it was given.
14/// `Glm5NextLayer::mixer_forward` handles both.
15pub enum Glm5NextMixer {
16    Kda {
17        layer: Box<Glm5NextKdaLayer>,
18        /// 2026-09-25: One workspace shared by every KDA layer of the stack: the loader builds
19        /// one and clones the `Arc`.
20        ws: Arc<Glm5NextKdaWorkspace>,
21        cfg: Glm5NextKdaConfig,
22    },
23    Dsa(Box<Glm5NextDsaLayer>),
24}
25
26/// 2026-09-25: The layer's MLP: dense for the layers the config lists in `mlp_only_layers`,
27/// routed MoE for the rest.
28pub enum Glm5NextMlpSite {
29    Dense(Glm5NextDenseMlpWeights),
30    Moe(Box<Glm5NextMoeWeights>),
31}
32
33/// 2026-09-25: The layer's hyper-connection: the kernels, each site's weights, and the config's
34/// `hc_mult`, `hc_sinkhorn_iters` and `hc_eps`.
35pub struct Glm5NextMhc {
36    pub kernels: Glm5NextMhcKernels,
37    pub attn: Glm5NextMhcSiteWeights,
38    pub ffn: Glm5NextMhcSiteWeights,
39    pub hc_mult: usize,
40    pub sinkhorn_iters: usize,
41    pub hc_eps: f32,
42}
43
44/// 2026-09-25: One bound GLM-5.3 decoder layer.
45pub struct Glm5NextLayer {
46    pub layer_idx: usize,
47    pub mixer: Glm5NextMixer,
48    pub mlp: Glm5NextMlpSite,
49    pub mlp_cfg: Glm5NextMlpConfig,
50    pub mlp_kernels: Glm5NextMlpKernels,
51    /// 2026-09-25: MLP scratch. The text layers share one workspace unless
52    /// `METRALE_GLM_MLP_WS_SHARED=0` (`mlp_ws_shared`); the MTP block builds its own.
53    pub mlp_ws: Arc<Glm5NextMlpWorkspace>,
54    /// 2026-09-25: `None` for the MTP block, which runs `forward_one_plain`. Every text layer has
55    /// one: the skeleton marks every text layer `hyper_connection`.
56    pub mhc: Option<Glm5NextMhc>,
57    /// 2026-09-25: `input_layernorm.weight`; `post_attn_norm` is `post_attention_layernorm.weight`.
58    pub input_norm: DevicePtr,
59    pub post_attn_norm: DevicePtr,
60    /// 2026-09-25: `rms_norm_vanilla`; see the module invariants.
61    pub rms_norm_k: KernelHandle,
62    /// 2026-09-25: `bf16_add_inplace`, the residual add of the MTP block, whose residual is not
63    /// in a highway. `0` on a target without the kernel; `add_inplace` then returns an error.
64    pub add_k: KernelHandle,
65    pub rms_eps: f32,
66    pub hidden: usize,
67    /// 2026-09-25: The mixer output is a partial sum across TP ranks, all-reduced before
68    /// `hc_post` (or, in the MTP block, before the residual add). Set from the mixer's TP plan
69    /// (`needs_output_all_reduce`).
70    pub mixer_all_reduce: bool,
71    /// 2026-09-25: Expand the highway in this layer; true for layer 0 only.
72    pub is_first: bool,
73    /// 2026-09-25: Collapse the highway in this layer; true for the last text layer only.
74    pub is_last: bool,
75}
metrale_model_arch/glm5next_layer/types.rs

metrale_model_arch/glm5next_layer/
types.rs
