spark_model/layers/glm5next_layer/
state.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! One GLM-5.3 decoder layer's per-sequence state.
4//!
5//! A GLM layer is one of two mixers and the two need *different kinds* of state: KDA carries a
6//! recurrent hidden state plus a causal-conv window and touches no KV cache at all, while DSA
7//! carries an indexer key cache alongside paged KV blocks. Each mixer therefore returns its own
8//! concrete `LayerState` and [`crate::layers::glm5next_layer::Glm5NextLayer`] downcasts to the
9//! one its mixer expects.
10//!
11//! 🔴 **KDA's state is the pool's [`SsmLayerState`], not a GLM-private type.**
12//! `rollback_ssm_states_dispatch` walks every `LayerType::LinearAttention` layer and downcasts
13//! to exactly that type to rewind a rejected speculative draft. GLM's KDA blocks *are*
14//! `linear_attention` in `layer_types`, so a GLM-private state means the first rejected draft
15//! is a hard error — and the shapes line up byte-for-byte anyway:
16//!
17//! | | pool (`config` fields, TP-local) | GLM (`Glm5NextKdaConfig`) |
18//! |---|---|---|
19//! | h    | `nv · vd · kd · 4`               | `heads · head_dim² · 4` |
20//! | conv | `(nk · kd · 2 + nv · vd) · d_conv · 4` | `3 · heads · head_dim · conv_kernel · 4` |
21//!
22//! The parser fills `linear_num_{key,value}_heads` / `linear_{key,value}_head_dim` /
23//! `linear_conv_kernel_dim` from `linear_attn_config`, already divided by TP, so
24//! `ModelConfig::ssm_h_state_bytes()` and `ssm_conv_state_bytes()` return GLM's own numbers.
25//!
26//! 🪤 The two state kinds are NOT interchangeable and admission needs both kinds satisfied — a
27//! KDA slot is not a KV block. That is [`crate::layers::glm5next_skeleton::StateKind`], made
28//! real.
29
30use anyhow::Result;
31use spark_runtime::gpu::GpuBackend;
32
33use crate::layer::SsmLayerState;
34use crate::layers::glm5next_kda::Glm5NextKdaConfig;
35
36/// Allocate and **zero** a KDA layer's recurrent + conv state, pool-free.
37///
38/// Used only where a state is built outside the SSM pool; the serving path takes pool slots
39/// (`Glm5NextLayer::uses_ssm_pool()`), which is what carries the checkpoints and per-token
40/// intermediates a speculative rollback needs.
41///
42/// 🪤 **FP32 is not negotiable.** HF casts the recurrent state to float32 and vLLM hardcodes
43/// `kda_state_dtype`, so a BF16/FP16 state is a deviation from the reference, not a memory
44/// setting. `h_is_f16: false` here, and `Glm5NextLayer::kda_state` refuses a narrowed slot.
45pub fn alloc_kda_ssm_state(gpu: &dyn GpuBackend, cfg: &Glm5NextKdaConfig) -> Result<SsmLayerState> {
46    let h_bytes = cfg.recurrent_state_elems() * 4;
47    let conv_bytes = cfg.conv_state_elems() * 4;
48    let h_state = gpu.alloc(h_bytes)?;
49    let conv_state = gpu.alloc(conv_bytes)?;
50    gpu.memset_async(h_state, 0, h_bytes, 0)?;
51    gpu.memset_async(conv_state, 0, conv_bytes, 0)?;
52    gpu.synchronize(0)?;
53    Ok(SsmLayerState {
54        h_state,
55        conv_state,
56        h_state_checkpoint: None,
57        conv_state_checkpoint: None,
58        h_state_intermediates: Vec::new(),
59        conv_state_intermediates: Vec::new(),
60        h_is_f16: false,
61        h_prefill_stage: None,
62        // GLM-5.3 hosts no `PleLayer` (its linear-attention block is KDA), so
63        // there is no PLE per-sequence carry to hold. Upstream #753 item B
64        // added this field; `None` is the correct answer, not a placeholder.
65        ple: None,
66    })
67}
spark_model/layers/glm5next_layer/state.rs

spark_model/layers/glm5next_layer/
state.rs
