Struct Qwen3SsmLayer 

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

Qwen3-Next SSM/GDN layer (36 of 48 layers).

Supports two QKVZ projection modes:

• Interleaved (80B): w4a16_gemv_qkvz or GEMV + deinterleave_qkvz
• Sequential (3.5-35B): plain GEMV → [Q|K|V|Z] already in 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>

Construct an SSM layer where QKVZ projection output is already sequential.

Used by Qwen3.5 where separate QKV and Z weights are concatenated at load time into [Q|K|V|Z] row order. The deinterleave_qkvz kernel is skipped and plain w4a16_gemv writes directly to the deinterleaved buffer.

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

Install native FP8 block-scaled weights for the decode GEMV path.

Inputs MUST be tagged WeightQuantFormat::Fp8BlockScaled — that is the canonical input format for the w8a16_gemv kernel (out[n] = sum_k A[k] * E4M3_LUT[B[n,k]] * block_scale[n/BS, k/BS], see kernels/gb10/common/w8a16_gemv.cu). The kernel reads the scale buffer at [N/BS, K/BS] BF16 — exactly the shape produced by load_fp8_block_scaled_as_fp8weight.

This setter does NOT install the raw FP8 DevicePtr fields used by the prefill fp8_gemm_n128 kernel — that kernel takes no scale argument and assumes single-scale FP8 (baked-in scale) produced by bf16_to_fp8. Block-scaled bytes would silently produce wrong outputs there. For prefill, call set_fp8_prefill_only_weights separately with single-scale FP8 derived from a BF16 dequant.

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

Install PER-ROW FP8 weights for the row-wise cuBLASLt PREFILL arm (ATLAS_FP8_ROWWISE=1, mixed-precision compressed-tensors checkpoints). Decode is untouched and keeps the NVFP4 copy.

The Fp8PerRow assertion is the mirror of set_fp8_decode_weights’s Fp8BlockScaled one: each setter refuses the other’s layout, so the two FP8 shapes cannot cross into each other’s kernels. That crossing does not fault — the smaller buffer is read in-bounds — so an assert is the only thing that catches it.

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

Set raw FP8 DevicePtrs for the prefill GEMM path ONLY (no decode GEMV scale fields). Used by the Qwen3.6-27B-FP8 native-FP8 SSM prefill path: the FP8 buffer here is a single-scale FP8 (BF16 → FP8 truncation; values already in FP8 range) suitable for fp8_gemm_n128. Decode falls back to the NVFP4/BF16 paths via the existing qkvz_nvfp4* fields. PCND: caller decides whether to install — never set implicitly.

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

Pre-dequant NVFP4 → FP8 for QKVZ and out_proj transposed weights. Eliminates per-inference dequant overhead in prefill GEMMs.

impl Qwen3SsmLayer

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

Install the Tier-1c keep-packed ternary Q2_0 fused in_proj_qkvz (ATLAS_GGUF_NATIVE_Q2). Decode dispatches q2_0_gemv_vec; prefill transient-dequants via Self::qkvz_q2_prefill_gemm. out_proj is unaffected (stays NVFP4). Requires sequential_qkvz (Bonsai concats [Q|K|V|Z] at load).

The keep-packed MMQ kernels are resolved HERE, not in the constructor: they ship only in GGUF-serving targets, and probing them on models that never install packed-Q2 weights fails the fail-closed boot audit on every other GDN target.

impl Qwen3SsmLayer

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

Install this GDN layer’s MoE router + routed-expert LoRA onto its FfnComponent::Moe. Hard-rejects (never silently drops) when the layer’s FFN is not MoE — an expert/router delta on a dense-FFN GDN layer is a loader/adapter mismatch.

impl Qwen3SsmLayer

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

Install this linear-attention layer’s DENSE-FFN LoRA onto its FfnComponent::Dense.

The mirror of set_moe_lora_weights for dense-FFN hybrids. A linear-attention layer carries no attention projections, but on Qwen3.8-27B it does carry the SwiGLU FFN — all 64 layers do, only 16 of which are full attention — and real adapters for that architecture ship gate/up/down for every one of them. Rejecting those rejected three quarters of the adapter, and the old message could only suggest retraining with layers_to_transform.

The component is the same DenseFfnLayer the full-attention layers hold, so the delta path, its pinned dispatch arms and its refusals are identical here — this only hands it the weights.

Hard-rejects a non-dense FFN rather than dropping the pairs: a dense delta arriving at a MoE or absent FFN is a loader/adapter mismatch, and silently ignoring it would be an adapter that reports success and does nothing — the exact failure this whole change removes.

impl Qwen3SsmLayer

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

Install this layer’s GDN out_proj delta.

Separate from set_ffn_lora_weights: that one targets the block’s FFN, this one the linear-attention block’s own output projection.

impl Qwen3SsmLayer

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

Attach mHC weights. Both concrete layer types carry them on this model: the 12 full-attention layers are Qwen3AttentionLayer, the 36 GDN layers are this one.

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

Attach the PLE n-gram injection to this layer. Exactly one model layer carries it.

Trait Implementations

impl TransformerLayer for Qwen3SsmLayer

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

Downcast hook so the LoRA install walk can reach this layer’s MoE FFN (Feature-1: routed-expert/router deltas exist on GDN layers too).

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

PLE’s host half (hash + NVMe fault-in + slot upload), hoisted before graph replay/capture. No-op on the 47 layers without a PLE site.

fn decode_graph_unsupported(&self) -> bool

PLE’s per-seq host hash on the hc multi-seq decode path is capture-illegal (pageable reads); the single-decode path prestages around it, the batched path does not — veto batched graphs.

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

Free the PLE carry this sequence lazily attached.

Only the ple field — the h/conv state in SsmLayerState is pooled and released by slot in free_sequence_dispatch, so freeing it here would be a double free. The PLE conv buffer is the one piece that is allocated per sequence and owned by nothing.

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

Decode K tokens through this layer using GEMM-batched projections.

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.

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

fn is_ssm_layer(&self) -> bool

Returns true if this layer is an SSM layer (supports two-phase prefill).

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_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 alloc_state(&self, gpu: &dyn GpuBackend) -> Result<Box<dyn LayerState>>

Allocate per-sequence state for this layer.

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 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_multi_seq_unsupported(&self) -> bool

True when this layer cannot serve a BATCHED multi-sequence decode step — i.e. decode_multi_seq’s shared-ForwardContext loop would alias per-sequence state across rows rather than merely run slowly.

fn decode_verify_multi_unsupported(&self) -> bool

True when this layer cannot serve a BATCHED multi-sequence VERIFY sweep (decode_verify_multi). Consumed by can_batch_verify_dispatch; a true layer falls back to the per-sequence verify loop, which is the sealed single-sequence path.

fn graph_stale_on_new_sequence(&self) -> bool

Prefill N tokens through this layer using GEMM-batched projections.

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

Reconcile whatever HOST-side per-sequence bookkeeping a step would have done, when that step was served by a replayed CUDA graph instead of being run. seq_len is the sequence length BEFORE this step’s k rows.

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

Refuse a step whose writes would land past a host-tracked cache — BEFORE the graph that performs them is replayed.

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 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 uses_ssm_pool(&self) -> bool

Does this layer’s recurrent state live in the shared SSM pool?