metrale_model_arch/glm5next_layer/
state.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Pool-free allocation of a GLM-5.3 KDA layer's per-sequence state.
4//!
5//! The state is the SSM pool's `SsmLayerState`, not a GLM type:
6//! `rollback_ssm_states_dispatch` downcasts the state of every `LinearAttention` layer to that
7//! type, and every KDA layer is `LinearAttention`.
8//!
9//! Owner: model-arch (GLM-5.3).
10//! Invariants: none beyond the types.
11
12use anyhow::Result;
13use metrale_gpu_runtime::gpu::GpuBackend;
14
15use crate::glm5next_kda::Glm5NextKdaConfig;
16use metrale_model_layers::layer::SsmLayerState;
17
18/// 2026-09-25: Allocate a KDA layer's recurrent and conv state outside the SSM pool, FP32
19/// (`h_is_f16: false`), zeroed, with no checkpoints or intermediates. The model path uses pool
20/// slots instead (`Glm5NextLayer::uses_ssm_pool`).
21pub fn alloc_kda_ssm_state(gpu: &dyn GpuBackend, cfg: &Glm5NextKdaConfig) -> Result<SsmLayerState> {
22    let h_bytes = cfg.recurrent_state_elems() * 4;
23    let conv_bytes = cfg.conv_state_elems() * 4;
24    let h_state = gpu.alloc(h_bytes)?;
25    let conv_state = gpu.alloc(conv_bytes)?;
26    gpu.memset_async(h_state, 0, h_bytes, 0)?;
27    gpu.memset_async(conv_state, 0, conv_bytes, 0)?;
28    gpu.synchronize(0)?;
29    Ok(SsmLayerState {
30        h_state,
31        conv_state,
32        h_state_checkpoint: None,
33        conv_state_checkpoint: None,
34        h_state_intermediates: Vec::new(),
35        conv_state_intermediates: Vec::new(),
36        h_is_f16: false,
37        h_prefill_stage: None,
38        ple: None,
39    })
40}
metrale_model_arch/glm5next_layer/state.rs

metrale_model_arch/glm5next_layer/
state.rs
