Struct Glm5NextDsaLayer 

pub struct Glm5NextDsaLayer {
    pub cfg: Glm5NextDsaConfig,
    pub weights: Glm5NextDsaWeights,
    pub kernels: Glm5NextDsaLayerKernels,
    pub select_kernels: Glm5NextDsaKernels,
    pub decode_kernel: Glm5NextDsaDecodeKernel,
    pub workspace: Glm5NextDsaWorkspace,
    pub layer_idx: usize,
    pub attn_layer_idx: usize,
    pub rms_eps: f32,
    pub kv_scale: f32,
    pub persist_bt: bool,
}

Fields

cfg: Glm5NextDsaConfig

weights: Glm5NextDsaWeights

kernels: Glm5NextDsaLayerKernels

select_kernels: Glm5NextDsaKernels

decode_kernel: Glm5NextDsaDecodeKernel

workspace: Glm5NextDsaWorkspace

layer_idx: usize

2026-09-25: Index in the model’s layer stack; used only in error messages.

attn_layer_idx: usize

2026-09-25: Index into the KV pool (kv_cache.k_pool_ptr): the loader counts DSA layers only, so it runs 0..11 on GLM-5.3 while layer_idx runs to 45. The MTP drafter uses 0.

rms_eps: f32

kv_scale: f32

2026-09-25: FP8 latent-cache scale. The decode reads with it and the latent write takes 1/scale.

persist_bt: bool

2026-09-25: Use the workspace’s bt/sl rather than allocating them per call. The loaders set it true unless METRALE_GLM_DSA_ALLOC_PER_STEP=1.

Implementations

impl Glm5NextDsaLayer

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

2026-09-25: k consecutive tokens of one sequence, from position seq_len.

q_a_proj, q_absorb, kv_a_proj and the o_absorb output projection each run once for all k rows. Per row, in order: the latent write to the row’s paged slot, the indexer write, and (unless the batched selector runs) that row’s selection. After the rows: the batched selection when batch_select_enabled holds, then one attend_rows launch.

Row r’s position, KV slot, seq_len and block table come from element r of ctx.attn_metadata on a decode step, or when k > 1 and the metadata has exactly k rows; otherwise they are computed here and copied to the device. On a decode step with k > 1, metadata with another row count is an error.

impl Glm5NextDsaLayer

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

2026-09-25: One MTP drafter context row: writes the KV latent and the indexer row, with no query, selection or attention.

The row goes to KV slot seq_len, and the indexer writes at state.len(). An indexer cache ahead of seq_len is rewound; one behind is an error, so rows must be written densely from 0.

impl Glm5NextDsaLayer

pub fn indexer_forward( &self, gpu: &dyn GpuBackend, hidden: DevicePtr, state: &mut Glm5NextDsaState, pos_dev: Option<DevicePtr>, stream: u64, ) -> Result<()>

2026-09-25: Project hidden into indexer cache row state.len(), then advance by one.

With pos_dev, k_normed and gate go to the workspace staging rows and dsa_indexer_store places them at the device-side position; without it they are written straight into the cache row. Also leaves this row’s selector head weights in the workspace. Fails before any write when the cache is full.

Trait Implementations

impl CircuitBindings for Glm5NextDsaLayer

fn circuit_prepare( &self, _gpu: &dyn GpuBackend, _config: &ModelConfig, _levers: &ModelLevers, _stream: u64, ) -> Result<(), Error>

2026-09-28: Build the lazily built weights the bound kernels read (the MMQ repacks a first wide batch would otherwise build), so a binding can hand them out. Queued on stream.

fn circuit_layer( &self, _config: &ModelConfig, _levers: &ModelLevers, ) -> Option<CircuitLayer>

2026-09-28: This layer for the circuit executor; None for a layer type it does not bind. config and levers are the model’s, as its decode reads them.

impl LayerAuxState for Glm5NextDsaLayer

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

2026-09-25: Host-serialized per-sequence state that must travel with an SSM snapshot for a prefix-cache hit to be complete, such as the PLE and QSA carries. Device-to-host copies inside must be ordered on stream. Default None.

fn has_aux_state(&self) -> bool

2026-09-25: True when this layer produces aux state. When any layer does, restore sites decline a snapshot that has no aux blobs (requires_aux_state).

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

2026-09-25: Apply a blob from Self::snapshot_aux to this layer’s state (apply_aux_states). The default returns an error.

impl LayerCapabilities for Glm5NextDsaLayer

fn uses_local_mla_prefill(&self) -> bool

2026-09-25: True when this layer’s prefill attends only over the tokens it is handed, so skipping a cached prefix would hide it from attention. The Qwen3 attention layer returns true when it runs MLA; the prefix lookup then matches nothing unless METRALE_MLA_PREFIX_SKIP=1 (mla_prefill_needs_full_recompute).

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

2026-09-25: Whether this layer’s online FP8-KV calibration has frozen its scale; None when the layer runs no online calibration (the default). The model lifts its CUDA-graph suppression once no layer reports Some(false) (graphs_ready_after_fp8_kv_cal).

fn decode_graph_unsupported(&self) -> bool

2026-09-25: True when this layer’s decode cannot run inside a CUDA graph, such as the QSA indexer’s host top-k. The model ORs it across layers into decode_graph_veto when it is built.

fn decode_multi_seq_unsupported(&self) -> bool

2026-09-25: True when this layer cannot serve a batched multi-sequence decode step. The batched decode ORs it across layers into hc_perseq (decode_a2.rs) and then runs each sequence through decode; the single-GPU fused decode+prefill ORs it into hc_qsa_perseq (decode_b.rs) and then runs the batched decode and the prefill separately.

fn decode_rollback_unsupported(&self) -> bool

2026-09-25: True when this layer keeps per-sequence state that lowering the sequence’s KV cursor does not rewind, such as a monotonic cache count or an n-gram history. The model ORs it across layers, and the scheduler’s rollback_to_boundary then declines the rollback (LayerStateNotRewindable).

fn decode_verify_multi_unsupported(&self) -> bool

2026-09-25: True when this layer cannot serve a batched multi-sequence verify (decode_verify_multi); can_batch_verify_dispatch then refuses the batch. It is separate from Self::decode_multi_seq_unsupported because the answers can differ.

fn graph_stale_on_new_sequence(&self) -> bool

2026-09-25: True when a captured decode graph goes stale once a new sequence takes this slot. Decode graphs are keyed by slot_idx, which is safe only while every per-sequence address a capture bakes lives in the slot-addressed SSM pool. A layer that allocates its own per-sequence buffers (GLM-5.3’s DSA indexer cache) returns true, and free_sequence_dispatch then drops that slot’s graphs, except for a slot index past the SSM pool’s max_slots.

fn is_ssm_layer(&self) -> bool

2026-09-25: True for an SSM layer. The model’s split prefill then runs prefill_phase1, prefill_gdn_full and prefill_phase3 instead of prefill.

fn uses_ssm_pool(&self) -> bool

2026-09-25: Whether this layer’s recurrent state lives in the shared SSM pool. When true (the default), sequence setup hands a linear-attention layer an SsmLayerState with pool addresses and never calls its alloc_state. A linear-attention layer with its own state type returns false; the Kimi K3 layer (kimi_k3/bound.rs) does, and downcasts its state to K3CpuFallbackState.

fn supports_replay_tail_split(&self) -> bool

2026-10-01: True when this layer’s prefill honours a replay-tail capture (ForwardContext::midchunk_capture with replay_tail) on every build: the rows from cap_local on get the bits a prefill pass starting at cap_local gives them (an SSM layer splits its conv and recurrence there and copies its state into the plan’s slots; an attention layer runs those rows as their own paged call). The model plans a replay-tail capture only when every layer says so (TransformerModel::inpass_cut_capture_supported); otherwise it keeps the two-pass tail split.

impl LayerGraphHooks for Glm5NextDsaLayer

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

2026-09-25: Fails when a replay writing indexer rows up to seq_len + k would pass the cache capacity; the error carries the context “DSA replay pre-check”.

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

2026-09-25: A replayed graph writes the indexer rows (positions from device memory) without running decode, so this moves the host-side row counter to where decode_k would have left it (Glm5NextDsaState::sync_to).

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

2026-09-25: Per-step host work of a layer that computes on the host at decode (the PLE hash and slot upload). The model’s single-token decode calls it before any CUDA graph replay or capture, when the token id is on the host (decode_a.rs), so a graph reads only device buffers refreshed here. Default: nothing.

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

2026-09-25: Re-arm prestaged state that a failed CUDA-graph capture consumed, so the same step can re-run eagerly (decode_a.rs). It must not recompute the prestage work: PLE’s history already advanced in decode_prestage.

impl LayerSplitPrefill for Glm5NextDsaLayer

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

2026-09-25: Phase 1 projections (norm, QKVZ, gates) over all stacked tokens of a batched prefill at once. The caller then runs prefill_phase1_conv1d_one per request and prefill_phase1_l2_batched. The default returns an error.

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

2026-09-25: Conv1d over one request’s tokens of a batched prefill, advancing its conv state and writing its slice of gdn_bufs.qkv. The default returns an error.

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

2026-09-25: L2 norm over the whole stacked gdn_bufs.qkv, after every request’s conv1d. The default returns an error.

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

2026-09-25: Split SSM prefill, phase 2: the GDN recurrence over all gdn_bufs.total_len tokens, reading packed QKV and gate/beta and writing output. Default: nothing.

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<(), Error>

2026-09-25: Prefill of one attention layer over the stacked tokens of several streams, with per-stream metadata from batched_meta. The default returns an error; the Qwen3 attention layer overrides it.

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

2026-09-25: GDN recurrence for batch_size streams of equal length chunk_len in one call. h_state_ptrs is a device array of one h-state pointer per stream (TransformerModel::stage_h_state_ptrs), and the streams’ tokens lie back to back in gdn_bufs. prefill_ssm_batched_layer calls it only when every stream has the same length. The default returns an error.

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, Error>

2026-09-25: GDN recurrence for streams of different lengths in one FLA call, driven by cu_seqlens. Ok(true) if it ran; Ok(false) if not eligible, and the caller then runs prefill_gdn_full per stream. Default Ok(false).

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

2026-09-25: Split SSM prefill, phase 3: for num_tokens tokens at token_offset, gated RMSNorm of the GDN output with Z, the output projection, the residual add and the FFN. Default: nothing.

impl LayerWeightSetup for Glm5NextDsaLayer

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

2026-09-25: &mut dyn Any downcast hook for weight overlays applied after construction, such as a LoRA install. Default None.

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

2026-09-25: Build persistent transposed gate/up/down MoE expert weights for the prefill GEMMs. The factory’s MoE transpose pass (factory/m2_setup.rs) calls it, outside the hybrid and unified layouts, when free memory covers the full transpose plus a safety margin. Default: nothing.

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

2026-09-25: Like transpose_moe_for_prefill, for gate and up only. Outside the hybrid and unified layouts, the transpose pass calls it when the full transpose does not fit but gate and up do. Default: nothing.

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

2026-09-25: Give this layer’s MoE block the per-prefill down_proj transpose scratch that every MoE layer shares, after transpose_moe_gate_up_for_prefill. Default: nothing.

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

2026-09-25: Unified layout (METRALE_UNIFIED_MOE_LAYOUT=1): build transposed gate/up/down and free the untransposed copies. Decode reads the transposed copies when MoeLayer::use_t_layout_for_decode holds. Default: nothing.

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

2026-09-25: Hybrid layout (METRALE_HYBRID_MOE_LAYOUT=1): build transposed gate/up/down beside the originals and free nothing, so decode keeps the originals and prefill uses the transposed copies. The transpose pass calls it only when twice the full transpose fits. Default: nothing.

impl LayerWriteOnAccept for Glm5NextDsaLayer

fn gdn_woa_stash_seq_floats(&self) -> Option<usize>

2026-09-25: Per-sequence write-on-accept stash size in f32 elements when this layer can run write-on-accept, None otherwise (the default). The model sizes the stash from the largest answer and binds each GDN layer with Self::gdn_woa_bind.

fn gdn_woa_bind(&self, _flag: DevicePtr, _stash: DevicePtr, _seqs: usize)

2026-09-25: Bind this layer’s write-on-accept flag word and stash slab (seqs sequences of Self::gdn_woa_stash_seq_floats f32 each). The model calls it once, on the first write-on-accept request and before its graph decision, and never moves the buffers afterwards (gdn_woa.rs).

fn gdn_carry_seq_floats(&self) -> Option<usize>

2026-09-26: Per-slot carried-state stash size in f32 elements when this layer can run the carried-state verify, None otherwise (the default).

fn gdn_carry_bind(&self, _binding: GdnCarryBinding)

2026-09-26: Bind this layer’s carried-state buffers. Called once, before any capture; the buffers never move afterwards.

fn gdn_verify_run_batched( &self, _states: &[&mut (dyn LayerState + 'static)], _kk: usize, _gdn_wyn: bool, _wy_tables: DevicePtr, ) -> Result<Option<bool>, Error>

2026-09-30: Whether a batched verify runs one run of states (adjacent sequences of width kk, under the WY tables at wy_tables) through the batched conv and the table-form WY, as this layer’s own batched verify decides it; None for a layer without a batched GDN verify (the default). The circuit executor picks each run’s arm with it.

fn gdn_fold_accepted( &self, _gpu: &dyn GpuBackend, _h_table: DevicePtr, _na_tab: DevicePtr, _k_rows: usize, _n: usize, _stream: u64, ) -> Result<bool, Error>

2026-09-25: Apply the accepted rows of the last batched verify to this layer’s h states. h_table is the layer’s WY pointer-table slice, na_tab a device u32[n] of accepted row counts in batch order. Ok(false) when the layer did nothing, as the default does.

impl TransformerLayer for Glm5NextDsaLayer

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

2026-09-25: Frees a Glm5NextDsaState; any other state type is left alone. The composite Glm5NextLayer has its own override that does the same.

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

2026-09-25: Allocate this layer’s per-sequence state. Sequence setup (model-engine trait_impl/meta.rs) calls it for every layer except a linear-attention layer that uses the SSM pool, which gets pool addresses instead.

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<()>

2026-09-25: Decode one token through this layer, updating hidden in place.

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<(), Error>

2026-09-25: Prefill num_tokens tokens starting at position seq_len_start. hidden and residual are [num_tokens, hidden_size] BF16. kv_write_start is how many of these tokens already have their KV in the cache; the Qwen3 attention cache-skip prefill writes KV only for the rest (prefill/cache_skip.rs). Default: decode once per token (default_loops::prefill_default), which ignores kv_write_start.

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<(), Error>

2026-09-25: Split SSM prefill, phase 1: norm, projections, gates, conv1d and L2 norm for num_tokens tokens, then the GDN inputs (packed QKV, gate/beta, Z) are copied into gdn_bufs at token_offset. The recurrence runs later, in prefill_gdn_full. Default: the whole prefill.

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<(), Error>

2026-09-25: Decode num_tokens consecutive tokens of one sequence, starting at position seq_len. hidden and residual are [num_tokens, hidden_size] BF16. Default: decode once per token (default_loops::decode_batched_default).

fn decode_multi_seq<'a, 'b>( &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<(), Error>

where 'b: 'a,

2026-09-25: Decode one token for each of num_seqs rows, each with its own entry in seq_lens and block_tables: one row per sequence in a batched decode, one per verify row in a batched verify (verify_e.rs). hidden and residual are [num_seqs, hidden_size] BF16. Default: decode once per row (default_loops::decode_multi_seq_default).

fn decode_verify_multi<'a, 'b>( &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<(), Error>

where 'b: 'a,

2026-09-25: Batched verify of ks[i] rows for each of n_seqs sequences in one sweep, the rows sequence-major and contiguous in hidden/residual. The verify (verify_e.rs) runs attention layers through decode_multi_seq and every other layer through this. wy_tables is this layer’s slice of the staged WY pointer tables (layout at VERIFY_WY_TABLE_SEQS), or NULL when none were staged or the layer is not linear attention. The default returns an error.