Struct Qwen3SsmLayer 

pub struct Qwen3SsmLayer {
    pub out_proj_dense: Option<DenseWeight>,
    /* private fields */
}

2026-09-25: The Gated DeltaNet layer. Its QKVZ projection comes in two layouts: interleaved per key-head group, which deinterleave_qkvz (or the fused w4a16_gemv_qkvz) turns into [Q | K | V | Z], and sequential (sequential_qkvz), already in that order.

Fields

out_proj_dense: Option<DenseWeight>

Implementations

impl Qwen3SsmLayer

pub fn new_sequential( input_norm: DenseWeight, ssm: SsmWeights, post_attn_norm: DenseWeight, ffn: FfnComponent, qkvz_nvfp4: Option<QuantizedWeight>, qkvz_nvfp4_t: Option<QuantizedWeight>, out_proj_nvfp4_t: Option<QuantizedWeight>, config: &ModelConfig, gpu: &dyn GpuBackend, ) -> Result<Self>

2026-09-25: new plus sequential_qkvz = true and the two transposed NVFP4 prefill weights. With sequential_qkvz the QKVZ projection writes straight into the deinterleaved buffer and deinterleave_qkvz is not launched. Called by the qwen35 and qwen35_dense weight loaders.

impl Qwen3SsmLayer

pub fn new( input_norm: DenseWeight, ssm: SsmWeights, post_attn_norm: DenseWeight, ffn: FfnComponent, qkvz_nvfp4: Option<QuantizedWeight>, config: &ModelConfig, gpu: &dyn GpuBackend, ) -> Result<Self>

impl Qwen3SsmLayer

pub fn set_fp8_decode_weights( &mut self, qkvz: Option<Fp8Weight>, out_proj: Option<Fp8Weight>, )

2026-09-25: Install block-scaled FP8 weights as qkvz_fp8w / out_proj_fp8w. Panics unless each is tagged Fp8BlockScaled: their readers include w8a16_gemv, which computes sum_k A[k] * E4M3_LUT[B[n,k]] * block_scale[n/128, k/128] with FP32 scales (kernels/gb10/common/w8a16_gemv.cu).

It does not touch qkvz_fp8 / out_proj_fp8: those feed fp8_gemm_n128, which takes no scale argument, so block-scaled bytes there would be read unscaled. set_fp8_prefill_only_weights installs them.

pub fn set_fp8_rowwise_prefill_weights( &mut self, qkvz: Option<Fp8Weight>, out_proj: Option<Fp8Weight>, )

2026-09-25: Install per-row FP8 weights as qkvz_fp8w_rowwise / out_proj_fp8w_rowwise. Panics unless each is tagged Fp8PerRow. Only the row-wise prefill arms read them (trait_prefill_proj.rs, trait_prefill_helper.rs); decode never does.

pub fn set_fp8_prefill_only_weights( &mut self, qkvz_fp8: Option<DevicePtr>, out_proj_fp8: Option<DevicePtr>, )

2026-09-25: Install unscaled FP8 weights as qkvz_fp8 / out_proj_fp8, the operands of the fp8_gemm_n128 arms in prefill and in the batched decode/verify projections (trait_decode_batched.rs). A None argument leaves that field as it was.

pub fn predequant_for_prefill( &mut self, gpu: &dyn GpuBackend, config: &ModelConfig, stream: u64, ) -> Result<()>

2026-09-25: When the layer holds out_proj_nvfp4_t, dequantize the NVFP4 out_proj to FP8 into a new allocation and store it as out_proj_fp8. QKVZ is not converted. Synchronizes stream.

impl Qwen3SsmLayer

pub fn set_packed_q2_qkvz(&mut self, qkvz: PackedQ2Weight, gpu: &dyn GpuBackend)

2026-09-25: Install the packed Q2_0 fused in_proj_qkvz. Decode runs q2_0_gemv_vec on it and prefill runs qkvz_q2_prefill_gemm. The layer must be sequential_qkvz: both paths write the projection straight into the deinterleaved buffer, and decode reads qkvz_q2 only on the sequential branch.

The MMQ kernels are looked up here rather than in new, so a layer that never installs a packed weight issues no lookup and leaves no failed row in the boot kernel audit.

impl Qwen3SsmLayer

pub fn set_moe_lora_weights( &mut self, router: Option<LoraPair>, experts: ExpertLoraLayer, kernels: LoraKernels, gpu: &dyn GpuBackend, ) -> Result<()>

2026-09-25: Install router and routed-expert LoRA on this layer’s MoE FFN (MoeLayer::set_lora_weights). Errors when the FFN is not MoE.

impl Qwen3SsmLayer

pub fn set_ffn_lora_weights(&mut self, ffn: LoraFfnWeights) -> Result<()>

2026-09-25: Install dense-FFN LoRA on this layer’s FfnComponent::Dense (DenseFfnLayer::set_lora_weights). Errors when the FFN is MoE or absent, rather than dropping the deltas.

impl Qwen3SsmLayer

pub fn set_out_proj_lora(&mut self, pair: LoraPair, kernels: LoraKernels)

2026-09-25: Install this layer’s GDN out_proj delta, which apply_lora_out_proj adds.

impl Qwen3SsmLayer

pub fn set_w8a8_decode_weights( &mut self, w: W8a8Mixer, qkvz_size: u32, hidden: u32, value_dim: u32, ) -> Result<()>

2026-09-28: Install the W8A8 QKV|Z and out_proj. Refuses an interleaved-QKVZ layer or shapes that are not [qkvz, hidden] and [hidden, value_dim].

pub fn set_w8a8_ffn_weights( &mut self, w: W8a8Ffn, hidden: u32, inter: u32, ) -> Result<()>

2026-09-28: Install W8A8 gate/up/down on this layer’s dense FFN (DenseFfnLayer::set_w8a8_decode_weights); refuses a MoE or absent FFN.

pub fn adopt_fp8_block_w8a8( &mut self, ctx: W8a8Ctx, hidden: u32, ) -> Result<bool>

2026-09-28: Run this layer’s block-scaled FP8 [QKV | Z] and out_proj (the decode FP8 weights of a native FP8 checkpoint, set_fp8_decode_weights) W8A8 from now on, with per-(token, 128) activation scales. Ok(false) installs nothing: they are absent or not block-scaled.

impl Qwen3SsmLayer

pub fn set_hc_weights(&mut self, hc: HcWeights)

2026-09-25: Attach the mHC highway weights; the entry points then route as the module header states.

pub fn set_ple(&mut self, ple: PleLayer)

2026-09-25: Attach the PLE n-gram injection. The hc prefill and single-sequence decode paths run it before this layer’s hc_pre_site.

Trait Implementations

impl CircuitBindings for Qwen3SsmLayer

fn circuit_prepare( &self, gpu: &dyn GpuBackend, config: &ModelConfig, levers: &ModelLevers, stream: u64, ) -> Result<()>

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 Qwen3SsmLayer

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 snapshot_aux( &self, state: &dyn LayerState, gpu: &dyn GpuBackend, stream: u64, ) -> Result<Option<Vec<u8>>>

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 restore_aux( &self, state: &mut dyn LayerState, blob: &[u8], gpu: &dyn GpuBackend, stream: u64, ) -> Result<()>

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 Qwen3SsmLayer

fn decode_graph_unsupported(&self) -> bool

2026-09-25: True when this layer carries PLE. The multi-seq decode (forward_row) hashes ids on the host inside the forward, which a graph capture refuses. The model ORs this into decode_graph_veto, which keeps decode eager.

fn supports_replay_tail_split(&self) -> bool

2026-10-01: Both prefill entry paths (prefill_inner, prefill_inner_hc) run prefill_block, whose conv (conv1d_prefill_capture) and recurrence (prefill_gdn_recurrence) split at the capture point on every build, the recurrence’s tail on the exact-replay arm.

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_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_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 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.

impl LayerGraphHooks for Qwen3SsmLayer

fn decode_prestage( &self, token: u32, state: &mut dyn LayerState, gpu: &dyn GpuBackend, stream: u64, ) -> Result<()>

2026-09-25: PLE’s host half (id hash, NVMe fault-in, slot upload), run before graph replay or capture. A no-op on a layer without PLE.

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.

fn sync_replayed_step( &self, _state: &mut dyn LayerState, _seq_len: usize, _k: usize, ) -> Result<()>

2026-09-25: Update host-side per-sequence bookkeeping for a step served by a replayed CUDA graph, in which the layer’s decode did not run. seq_len is the length before the step’s k rows. GLM-5.3’s DSA layer rewinds its indexer-cache length to seq_len when it is ahead, errors when it is behind, and then advances it by k (Glm5NextDsaState::sync_to). The model calls it after a replay in decode_a.rs, verify_b.rs, verify_c.rs and verify_c2.rs. Default: nothing.

fn check_replay_room( &self, _state: &dyn LayerState, _seq_len: usize, _k: usize, ) -> Result<()>

2026-09-25: Return an error for a replayed step whose writes would land past the end of a host-tracked cache. The model calls it before launching the graph; sync_replayed_step runs only after the launch. The step ends at seq_len + k. Default: nothing.

impl LayerSplitPrefill for Qwen3SsmLayer

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

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

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

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

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_gdn_full_batched( &self, h_state_ptrs: DevicePtr, gdn_bufs: &GdnPrefillBuffers, batch_size: u32, chunk_len: u32, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

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>

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

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.

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

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.

impl LayerWeightSetup for Qwen3SsmLayer

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

2026-09-25: Downcast hook: the LoRA install (impl_lora_rotate.rs) and the qwen4_exp loader reach this concrete layer through it.

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

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

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

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

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 Qwen3SsmLayer

fn gdn_verify_run_batched( &self, states: &[&mut (dyn LayerState + 'static)], kk: usize, gdn_wyn: bool, wy_tables: DevicePtr, ) -> Result<Option<bool>>

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_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_fold_accepted( &self, gpu: &dyn GpuBackend, h_table: DevicePtr, na_tab: DevicePtr, k_rows: usize, n: usize, stream: u64, ) -> Result<bool>

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 Qwen3SsmLayer

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

2026-09-25: Free the PLE carry this sequence attached on first use.

Only the ple field: the h/conv state points into the SSM pool, which free_sequence_dispatch releases by slot. The PLE conv buffer is allocated per sequence (new_seq_state), outside that pool.

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

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

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.

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

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

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

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 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.