Struct Glm5NextLayer 

pub struct Glm5NextLayer {

    pub layer_idx: usize,
    pub mixer: Glm5NextMixer,
    pub mlp: Glm5NextMlpSite,
    pub mlp_cfg: Glm5NextMlpConfig,
    pub mlp_kernels: Glm5NextMlpKernels,
    pub mlp_ws: Glm5NextMlpWorkspace,
    pub mhc: Option<Glm5NextMhc>,
    pub input_norm: DevicePtr,
    pub post_attn_norm: DevicePtr,
    pub rms_norm_k: KernelHandle,
    pub add_k: KernelHandle,
    pub rms_eps: f32,
    pub hidden: usize,
    pub mixer_all_reduce: bool,
    pub is_first: bool,
    pub is_last: bool,
}

One bound GLM-5.3 decoder layer.

Fields

layer_idx: usize

mixer: Glm5NextMixer

mlp: Glm5NextMlpSite

mlp_cfg: Glm5NextMlpConfig

mlp_kernels: Glm5NextMlpKernels

mlp_ws: Glm5NextMlpWorkspace

mhc: Option<Glm5NextMhc>

None only for a layer with no hyper-connection — i.e. the MTP layer, which carries zero hc_* tensors. Every text layer has one.

input_norm: DevicePtr

input_layernorm.weight / post_attention_layernorm.weight, both plain RMSNorm.

post_attn_norm: DevicePtr

rms_norm_k: KernelHandle

🪤 rms_norm_vanilla, never rms_norm. See the module header.

add_k: KernelHandle

bf16_add_inplace, the residual add the MTP layer needs and the text layers do not: a text layer’s residual lives in the mHC highway and hc_post folds the block output into it. 0 on a target without the kernel, which the MTP path refuses.

rms_eps: f32

hidden: usize

mixer_all_reduce: bool

🔴 Whether the MIXER output is a partial sum. Both mixers end in a row-parallel o_proj (KdaShard::ChannelCols / DsaShard::HeadCols), so at TP>1 each rank holds only part of the attention output and it must be all-reduced before hc_post folds it into the highway. Reducing after would mix a half-answer into every later layer’s residual stream; not reducing at all is a plausible, wrong output with no shape error.

is_first: bool

Expand the highway here. True for layer 0 only.

is_last: bool

Collapse the highway here. True for the last TEXT layer only.

Implementations

impl Glm5NextLayer

pub fn decode_one_for_drafter( &self, hidden: DevicePtr, state: &mut dyn LayerState, kv_cache: &mut PagedKvCache, seq_len: usize, block_table: &mut Vec<u32>, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

One token through this layer for the MTP drafter.

Only valid on a mhc: None block — the drafter’s layer. hidden is read and written in place (the plain residual path accumulates into it), and there are no disk tiers because the drafter’s KV pool is small, private and fully resident.

pub fn drafter_write_kv_row( &self, x: DevicePtr, state: &mut dyn LayerState, kv_cache: &mut PagedKvCache, seq_len: usize, block_table: &mut Vec<u32>, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

One drafter CONTEXT row: input_norm then the DSA caches only.

x is the block input (post eh_proj) for a row whose OUTPUT is discarded — a prompt or catch-up row. See Glm5NextDsaLayer::write_kv_row for why that is enough.

Trait Implementations

impl TransformerLayer for Glm5NextLayer

fn decode_multi_seq_unsupported(&self) -> bool

🔴 GLM-5.3 CANNOT serve a batched multi-sequence decode step. Two independent row-0 aliases, both structural, either one sufficient:

• D1 — mHC highway slot. Glm5NextLayer::forward_one pins highway slot 0. decode_multi_seq’s default loop shares one ForwardContext across the batch, so every sequence would write and then read the SAME highway stream, and each layer past the first reads the last sequence’s mHC state for all rows. This is the sequence-axis twin of the token-axis argument already written on Self::decode_batched above — the trait default is WRONG there for exactly the same reason.
• D2 — DSA attn_metadata row. The DSA mixer reads metadata row 0 (glm5next_dsa/layer.rs: slot, positions, block_table, seq_len), so every sequence in the batch would attend with sequence 0’s page table and length.

Answering true does NOT cost concurrency: the caller routes GLM onto #753 item B’s per-sequence highway loop, which serves C>1 correctly at C=1-equivalent per-request throughput.

🔒 This must stay true until the Stage 1 commit that adds a real Glm5NextLayer::decode_multi_seq (per-row forward_one with meta_row_base threading and ctx.hc_row_offset honoured as the highway base) flips it to false IN THE SAME COMMIT.

fn decode_verify_multi_unsupported(&self) -> bool

🔴 GLM-5.3 implements no decode_verify_multi, so the batched verify sweep must not be selected for it. The trait default already bail!s, but that is a mid-request abort; declaring it here makes can_batch_verify_dispatch route around it instead, leaving spec-on C>1 on the per-sequence verify loop — the sealed K=3 path.

🔒 Flipped to false by the PR-3 commit that adds Glm5NextLayer::decode_verify_multi (per-sequence forward_k sweep with slot_base = meta_row_base = row_base).

fn release_state( &self, state: &mut dyn LayerState, gpu: &dyn GpuBackend, ) -> Result<()>

Release what alloc_state allocated — ANOMALIES A76. The DSA indexer cache is sized by --max-seq-len, not by the prompt (513 B/token/layer), so leaking one per request walks a unified-memory host into the ground.

🔴 Type-driven on purpose. A KDA layer’s state on the model path is POOL-owned (uses_ssm_pool() is true for Kda, so alloc_sequence hands it pool addresses and free_sequence skips it entirely) — but refusing by TYPE as well means a pool address can never reach gpu.free even if a future call site forgets the skip. SsmLayerState is therefore left alone here, always.

fn graph_stale_on_new_sequence(&self) -> bool

The DSA mixer allocates its per-sequence state with gpu.alloc in alloc_state, so the addresses a capture bakes belong to THAT sequence, not to the slot.

🪤 Only DSA. The KDA mixer is POOL-backed on the model path — uses_ssm_pool() is true for Kda, so meta.rs hands it pool addresses and never calls its alloc_state. The true below is still correct (one owned mixer is enough); the previous wording claimed both mixers own their state, and that was wrong.

fn sync_replayed_step( &self, state: &mut dyn LayerState, seq_len: usize, k: usize, ) -> Result<()>

🔴 A CUDA-graph replay runs kernels and nothing else, so this layer’s one piece of HOST-side per-sequence bookkeeping — the DSA indexer cache length — has to be advanced here. The inner Glm5NextDsaLayer implements this too, but the model’s layer vec holds the COMPOSITE, so the inner impl is never reached and the default no-op left the counter frozen at its capture-time value.

🪤 That was invisible on the spec-off path: after capture, decode never runs again, so nothing compared the counter to seq_len. The first EAGER step after a run of replays — which is exactly what a speculative verify is — then failed with “indexer cache holds 5 tokens but the sequence is at 12”. The rows were there; only the counter was stale.

KDA keeps nothing on the host: its recurrent and conv state are device-resident and the replayed kernels update them in place.

fn check_replay_room( &self, state: &dyn LayerState, seq_len: usize, k: usize, ) -> Result<()>

The model’s layer vec holds the COMPOSITE, so — exactly as with sync_replayed_step — the inner Glm5NextDsaLayer impl is never reached and this one is what runs. A62.

fn uses_ssm_pool(&self) -> bool

GLM’s KDA blocks are linear_attention in layer_types AND carry the pool’s SsmLayerState, so they take pool slots like any other recurrent layer. That is what buys the speculative-verify checkpoints and per-token intermediates for free — meta.rs only wires h_state_checkpoint / h_state_intermediates for layers that answer true here, and rollback_ssm_states_dispatch needs both.

fn prefill( &self, hidden: DevicePtr, residual: DevicePtr, num_tokens: usize, state: &mut dyn LayerState, kv_cache: &mut PagedKvCache, seq_len_start: usize, block_table: &mut Vec<u32>, disk_block_ids: &mut Vec<u32>, disk_last_offloaded_per_layer: &mut Vec<u32>, _kv_write_start: usize, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

Prefill, one token at a time — but with the highway indexed by token.

🔴 The trait’s default fallback would be WRONG here, not merely slow. It runs every token through this layer before the next layer sees any of them, so a single-slot highway would hold only the last token’s streams by the time layer n+1 reads it. The mHC highway is a per-token activation that must survive across layers, so each token gets its own slot.

Per-token (rather than chunked) is deliberate for this slice: KDA’s recurrence is sequential anyway, and the chunked Glm5NextKdaLayer::prefill needs a workspace sized for the chunk. That is an optimisation, explicitly out of scope.

fn decode_batched( &self, hidden: DevicePtr, _residual: DevicePtr, num_tokens: usize, state: &mut dyn LayerState, kv_cache: &mut PagedKvCache, seq_len: usize, block_table: &mut Vec<u32>, _disk_block_ids: &mut Vec<u32>, _disk_last_offloaded_per_layer: &mut Vec<u32>, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

K tokens of ONE sequence in a single call — the speculative-verify body.

The same per-token walk as Self::prefill — same highway slots, same seq_len + t positions, same KV writes — plus the per-token KDA state snapshots that only a verify needs. Prefill is never rolled back, so it does not pay for them.

🔴 The trait’s default would be WRONG, not merely slow: it calls decode per token, and decode pins highway slot 0. K tokens would then overwrite each other’s mHC streams and every layer past the first would read the last token’s highway for all K rows.

✅ Batched. This delegates to Self::forward_k, which sweeps the weights ONCE for all K rows. (An earlier revision of this comment said “still one forward_one per row”; that was stale — forward_k has been the body since the batched-verify work, and the measured K=3 step of ~101 ms against a ~63 ms single-row step is only explicable by it.)

fn is_ssm_layer(&self) -> bool

KDA layers carry recurrent state; DSA layers do not.

fn alloc_state(&self, gpu: &dyn GpuBackend) -> Result<Box<dyn LayerState>>

Allocate per-sequence state for this layer.

fn decode( &self, hidden: DevicePtr, residual: DevicePtr, state: &mut dyn LayerState, kv_cache: &mut PagedKvCache, seq_len: usize, block_table: &mut Vec<u32>, disk_block_ids: &mut Vec<u32>, disk_last_offloaded_per_layer: &mut Vec<u32>, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

Decode one token through this layer, modifying hidden in-place.

fn uses_local_mla_prefill(&self) -> bool

True when this layer’s PREFILL attends only over the tokens it is handed, so a prefix-cache skip would hide the cached prefix from attention entirely. MLA layers on the paged path do; everything else reads the paged cache and is unaffected.

fn as_any_mut(&mut self) -> Option<&mut dyn Any>

&mut dyn Any downcast hook for post-construction weight overlays (e.g. the LoRA install walk). Default None; overlay-capable layers override.

fn fp8_calibration_frozen(&self) -> Option<bool>

Whether this layer’s ONLINE FP8-KV calibration has frozen its scale. None = this layer runs no online calibration (non-attention layer, static checkpoint scales, or a non-FP8 KV dtype). The scheduler’s graph-suppression gate keys off this rather than a token count: the scale freezes on the FIRST observe, so waiting calibration_tokens tokens would run ~256+ eager steps for a calibration that finished immediately.

fn decode_prestage( &self, _token: u32, _state: &mut dyn LayerState, _gpu: &dyn GpuBackend, _stream: u64, ) -> Result<()>

Hoisted per-step HOST work for layers that do host-side computation at decode (PLE: n-gram hash + NVMe fault-in + slot upload). The scheduler calls this every single-token decode step BEFORE any CUDA graph replay/capture — the same phasing as the token_ids upload — so the captured graph contains only kernels over stable device buffers. Layers with no host-side decode work keep the no-op default.

fn decode_prestage_rearm(&self, _state: &mut dyn LayerState)

Re-arm consumed prestage state so a failed CUDA-graph capture attempt can re-run the SAME step eagerly. Must be idempotent, and must not recompute (PLE’s history already advanced in decode_prestage).

fn decode_graph_unsupported(&self) -> bool

True when this layer’s decode can NEVER be captured into a CUDA graph — e.g. the QSA indexer’s host top-k round trip, whose captured dense fallback would silently replay WRONG attention once selection activates. The scheduler ORs this across layers once and keeps the whole model eager.

fn snapshot_aux( &self, _state: &dyn LayerState, _gpu: &dyn GpuBackend, _stream: u64, ) -> Result<Option<Vec<u8>>>

Marconi aux state: host-serialized per-layer SEQUENCE state that must travel with an SSM snapshot for a prefix-cache hit to be complete — PLE’s n-gram history + conv state, QSA’s ingested indexer keys. Without these a restored prefix would silently serve the PREVIOUS request’s lexical state. Called at chunk-boundary snapshot saves; any D2H inside must be stream-ordered (copy_d2h_on_stream). Default: the layer carries no aux sequence state.

fn has_aux_state(&self) -> bool

True when this layer WOULD produce aux state — restore sites use it to decline snapshots that lack aux rather than restore a stale mix.

fn restore_aux( &self, _state: &mut dyn LayerState, _blob: &[u8], _gpu: &dyn GpuBackend, _stream: u64, ) -> Result<()>

Restore the aux state captured by Self::snapshot_aux on a prefix-cache hit, BEFORE the resumed prefill runs.

fn prefill_phase1( &self, hidden: DevicePtr, residual: DevicePtr, num_tokens: usize, state: &mut dyn LayerState, kv_cache: &mut PagedKvCache, seq_len_start: usize, block_table: &mut Vec<u32>, disk_block_ids: &mut Vec<u32>, disk_last_offloaded_per_layer: &mut Vec<u32>, kv_write_start: usize, gdn_bufs: &GdnPrefillBuffers, token_offset: usize, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

Two-phase SSM prefill — Phase 1: projections and GDN input staging.

fn prefill_phase1_proj_batched( &self, hidden_stacked: DevicePtr, residual_stacked: DevicePtr, total_tokens: usize, gdn_bufs: &GdnPrefillBuffers, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

M1 large-M batched Phase-1: token-parallel projections (RMS/QKVZ/BA-gates) over ALL stacked tokens in one large-M GEMM each. SSM-only; the caller runs prefill_phase1_conv1d_one per request then prefill_phase1_l2_batched.

fn prefill_phase1_conv1d_one( &self, state: &mut dyn LayerState, token_offset: usize, len: usize, gdn_bufs: &GdnPrefillBuffers, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

M1: per-request conv1d tail (advances per-request conv_state), reading the request’s slice of the stacked QKVZ scratch and writing into gdn_bufs.qkv.

fn prefill_phase1_l2_batched( &self, total_tokens: usize, gdn_bufs: &GdnPrefillBuffers, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

M1: batched L2 norm over the full stacked QKV buffer after all per-request conv1d tails have written their slices.

fn prefill_gdn_full( &self, _state: &mut dyn LayerState, _gdn_bufs: &GdnPrefillBuffers, _ctx: &ForwardContext<'_>, _stream: u64, ) -> Result<()>

Two-phase SSM prefill — Phase 2: GDN recurrence on the full sequence.

fn prefill_inner_batched_q12( &self, _hidden_stacked: DevicePtr, _residual_stacked: DevicePtr, _num_tokens: usize, _kv_cache: &mut PagedKvCache, _seq_len_start: usize, _batched_meta: &BatchedAttnMetadata, _ctx: &ForwardContext<'_>, _stream: u64, ) -> Result<()>

Q12 Path B: batched attention prefill across N stacked-input streams.

fn prefill_gdn_full_batched( &self, _h_state_ptrs: DevicePtr, _gdn_bufs: &GdnPrefillBuffers, _batch_size: u32, _chunk_len: u32, _ctx: &ForwardContext<'_>, _stream: u64, ) -> Result<()>

Q12 Path B: batched GDN recurrence across N streams.

fn prefill_gdn_full_batched_fla_varlen( &self, _h_state_ptrs: DevicePtr, _gdn_bufs: &GdnPrefillBuffers, _batch_size: u32, _cu_seqlens: DevicePtr, _max_num_chunks: u32, _total_nt: usize, _max_seqlen: u32, _ctx: &ForwardContext<'_>, _stream: u64, ) -> Result<bool>

VARLEN batched GDN: process ragged co-dispatch lengths via cu_seqlens in ONE gdn_prefill_fla(batch=N, is_varlen) call (replaces the non-uniform per-request loop → fills chunk_delta_h’s 32→32N CTAs). Returns Ok(true) if it ran, Ok(false) if not eligible (caller falls back to the loop). Default (non-SSM layers, or FLA disabled): Ok(false).

fn prefill_phase3( &self, _hidden: DevicePtr, _residual: DevicePtr, _num_tokens: usize, _gdn_bufs: &GdnPrefillBuffers, _token_offset: usize, _ctx: &ForwardContext<'_>, _stream: u64, ) -> Result<()>

Two-phase SSM prefill — Phase 3: post-GDN processing.

fn transpose_moe_for_prefill( &mut self, _gpu: &dyn GpuBackend, _config: &ModelConfig, ) -> Result<()>

Allocate the transposed MoE expert weights used by the coalesced prefill GEMM kernels. Called as a post-load pass from factory::build after LM-head NVFP4 quantization has freed BF16 headroom, so memory-tight EP configurations (e.g. MiniMax M2.7-NVFP4 EP=2) can fit the transpose that layer-0 preflight would otherwise reject.

fn transpose_moe_gate_up_for_prefill( &mut self, _gpu: &dyn GpuBackend, _config: &ModelConfig, ) -> Result<()>

Like transpose_moe_for_prefill but only transposes the gate+up projections (skips the down projection), reducing the transpose cost from 3× to 2× per expert. Used as a memory-tight fallback by the MiniMax loader when full transpose doesn’t fit.

fn set_moe_down_transpose_scratch( &mut self, _scratch_packed: DevicePtr, _scratch_scale: DevicePtr, _packed_ptrs_t: DevicePtr, _scale_ptrs_t: DevicePtr, )

Wire a shared per-prefill down_proj transpose scratch into this layer’s MoE block. Used as a memory-tight alternative to the persistent down transpose: factory allocates one shared scratch, every MoE layer reuses it layer-by-layer during sequential prefill. No-op for non-MoE layers and MoE layers that already have a persistent transposed down.

fn transpose_moe_for_prefill_unified( &mut self, _gpu: &dyn GpuBackend, _config: &ModelConfig, ) -> Result<()>

Phase 8a unified-layout MoE transpose: build persistent transposed gate/up/down for all experts and free the untransposed copies. Phased flow keeps memory budget tight enough for MiniMax M2.7 EP=2. After this call, the untransposed-layout decode kernels can no longer execute correctly — MoeLayer::use_t_layout_for_decode() must gate dispatch to the _t decode kernels. Default no-op.

fn transpose_moe_for_prefill_hybrid( &mut self, _gpu: &dyn GpuBackend, _config: &ModelConfig, ) -> Result<()>

Block C Path 2 hybrid-layout MoE transpose: build persistent transposed gate/up/down alongside the untransposed originals (no frees). Doubles MoE-weight memory but recovers the ~15 % decode regression of pure unified mode — decode + MTP verify dispatch keeps using the warp-reduction kernels on the originals while prefill (forward_batched) routes through transposed kernels. Caller must verify enough free memory before invocation. Default no-op for non-MoE layers.

fn decode_multi_seq<'a, 'b: 'a>( &self, hidden: DevicePtr, residual: DevicePtr, num_seqs: usize, states: &'a mut [&'b mut (dyn LayerState + 'static)], kv_cache: &mut PagedKvCache, seq_lens: &[usize], block_tables: &[Vec<u32>], ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

Decode N sequences through this layer in a single batched call.

fn decode_verify_multi<'a, 'b: 'a>( &self, _hidden: DevicePtr, _residual: DevicePtr, _n_seqs: usize, _ks: &[usize], _states: &'a mut [&'b mut (dyn LayerState + 'static)], _kv_cache: &mut PagedKvCache, _wy_tables: DevicePtr, _ctx: &ForwardContext<'_>, _stream: u64, ) -> Result<()>

Batched MTP verify: n_seqs sequences × k tokens through this layer in ONE weight sweep (rows seq-major, r = i*k + j, contiguous in hidden/residual). Projections/FFN batch across all n_seqs*k rows; the stateful recurrence (conv/GDN) runs per-sequence against states[i] with row-offset buffer bases — per-sequence math is byte-identical to the single-sequence decode_batched K-token body.