Struct NemotronMamba2Layer 

pub struct NemotronMamba2Layer { /* private fields */ }

Implementations

impl NemotronMamba2Layer

pub fn new( input_norm: DenseWeight, ssm: NemotronSsmWeights, config: &ModelConfig, gpu: &dyn GpuBackend, layer_idx: usize, ) -> Result<Self>

pub fn set_fp8_weights( &mut self, in_proj: Option<Fp8Weight>, out_proj: Option<Fp8Weight>, prefill: bool, ) -> Result<()>

2026-09-25: Install the checkpoint’s native FP8 projections. Decode then uses w8a16_gemv, and when prefill is true the prefill GEMMs use w8a16_gemm_pipelined (or w8a16_gemm) ahead of the NVFP4 / W4A4 arms.

Inputs must be tagged WeightQuantFormat::Fp8BlockScaled: the w8a16 kernels index block_scale[n_block * k_blocks + k_block] with 128-wide blocks, so a per-row [N] scale or a scalar weight_scale would be read past its end. A wrong tag panics (expect); a missing w8a16_gemv, or with prefill a missing prefill GEMM, is an error here at load.

When prefill is false (METRALE_NEMOTRON_NATIVE_FP8_SSM=decode) only w8a16_gemv reads these weights and prefill uses the NVFP4 / pre-dequantized copies, which the loader builds in that mode.

pub fn ssm_weights(&self) -> &NemotronSsmWeights

2026-09-25: The layer’s Mamba-2 mixer weights.

pub fn set_prefill_weights( &mut self, in_proj_t: Option<QuantizedWeight>, out_proj_t: Option<QuantizedWeight>, )

2026-09-25: Install transposed NVFP4 projections: prefill then uses w4a16_gemm_t (or, above 128 tokens and when it resolved, w4a16_gemm_t_m128) instead of the plain w4a16_gemm, unless an earlier arm in prefill_proj.rs applies.

pub fn set_bf16_weights(&mut self, in_proj: DenseWeight, out_proj: DenseWeight)

2026-09-25: Install the checkpoint’s own BF16 projections; decode and prefill then use the dense BF16 kernels ahead of every other arm. Valid only when both projections are BF16 in the checkpoint and the dense kernels resolved; the caller checks that (bf16_native_ready).

pub fn bf16_native_ready(&self) -> bool

2026-09-25: Whether this layer can run natively BF16 (weights installed and both dense kernels present).

pub fn set_fp8_prefill_weights( &mut self, in_proj: DevicePtr, out_proj: DevicePtr, )

2026-09-25: Install pre-dequantized FP8 E4M3 copies of in_proj / out_proj for prefill; their arms come after the native and W4A4 arms and before the transposed NVFP4 ones (prefill_proj.rs).

Trait Implementations

impl CircuitBindings for NemotronMamba2Layer

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 NemotronMamba2Layer

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 NemotronMamba2Layer

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 NemotronMamba2Layer

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.

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

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

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 NemotronMamba2Layer

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 NemotronMamba2Layer

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 NemotronMamba2Layer

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 NemotronMamba2Layer

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

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

fn release_state( &self, _state: &mut dyn LayerState, _gpu: &dyn GpuBackend, ) -> Result<(), Error>

2026-09-25: Free the device memory this layer allocated for one sequence, in alloc_state or attached to the state later. LayerState holds bare DevicePtrs, so dropping it frees only the host struct.