Struct Qwen3AttentionLayer 

pub struct Qwen3AttentionLayer { /* private fields */ }

Qwen3-Next full attention layer (12 of 48 layers).

Implementations

impl Qwen3AttentionLayer

pub fn set_mla_weights(&mut self, mla: MlaWeights)

Set MLA weights for 2-step latent decode. When set, decode uses latent→norm→expand instead of single-step GEMV.

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

Set per-block Manifold-Constrained Hyper-Connection weights (DeepSeek-V4). When set, the attn/ffn residual sites route through hc_pre/hc_post against the model-level hc_streams buffer.

pub fn set_qsa(&mut self, qsa: QsaIndexer)

Attach the QSA indexer (Qwen3.8-Flash-Next full-attention layers).

pub fn set_dimension_overrides( &mut self, head_dim: usize, num_q_heads: usize, num_kv_heads: usize, )

Set per-layer dimension overrides for heterogeneous models (Gemma-4). Full-attention layers have different Q/KV head counts and head_dim than sliding layers.

pub fn set_sliding_window(&mut self, window: Option<u32>)

Set per-layer sliding-window size (Gemma-4 hybrid attention). Call with Some(window_size) on sliding layers, None on full-attention layers. Non-Gemma-4 models never call this.

pub fn set_rope_overrides(&mut self, theta: f32, rotary_dim: u32)

Set per-layer RoPE overrides (theta, rotary_dim) for dual-RoPE models (Gemma-4).

pub fn set_rope_proportional(&mut self, enable: bool)

Enable proportional RoPE (Gemma-4 full-attention layers). Must be called AFTER set_rope_overrides; the rotary_dim set there is reinterpreted as the number of non-zero rotation pairs.

pub fn set_attn_scale_override(&mut self, scale: f32)

Set per-layer attention scale override. Gemma-4 uses QK-norm, so attention scale should be 1.0 (not 1/sqrt(head_dim)).

pub fn set_k_eq_v(&mut self, v_norm_weight: DenseWeight)

Set K=V mode (Gemma-4 full-attention layers).

v_norm_weight is a BF16 weight buffer of size [head_dim]. For Gemma-4 it’s ones-filled because Gemma-4’s rms_norm kernel uses the absolute convention out = x * rms * weight, and weight = 1.0 gives pure RMSNorm (matching HF Gemma4RMSNorm(with_scale=False)).

pub fn set_v_norm(&mut self, v_norm_weight: DenseWeight)

Install a pure-RMSNorm v_norm WITHOUT enabling K=V aliasing. Used for Gemma-4 sliding-attention layers where V_proj exists on disk but HF Gemma4TextAttention.forward() still applies value_states = self.v_norm(value_states) with Gemma4RMSNorm(with_scale=False) — pure x * rms.

pub fn set_o_dense_bf16(&mut self, o_dense: DenseWeight)

Install a BF16 dense fallback for the output projection. When set, decode + prefill skip the NVFP4 attn.o_proj path and use this BF16 dense_gemv / dense_gemm instead. Required for Gemma-4 dense (Nvidia ModelOpt’s official ignore list keeps ALL self_attn projections at BF16).

pub fn set_post_sublayer_norms( &mut self, post_attn_out: DenseWeight, post_ffn_out: DenseWeight, )

Set post-sublayer norms (Gemma-4: 4-norm residual structure).

pub fn set_layer_scalar(&mut self, scalar: f32)

Set per-layer scalar (Gemma-4: hidden_states *= scalar at end of layer).

pub fn set_moe_ffn( &mut self, ffn: FfnComponent, pre_norm: DenseWeight, post_norm: DenseWeight, post_dense_norm: DenseWeight, )

Set secondary MoE FFN (Gemma-4 26B dual-FFN: dense + MoE per layer).

pub fn set_shortcut_moe( &mut self, moe: FfnComponent, carry: DevicePtr, carry_tokens: usize, )

LongCat: install the shortcut MoE on the FIRST sublayer of a dual-sublayer block. The MoE runs on this sublayer’s post-attention normed input; its output is stashed into carry (capacity carry_tokens tokens) and added by the SECOND sublayer via Self::set_shortcut_carry_in.

pub fn set_shortcut_carry_in(&mut self, carry: DevicePtr, carry_tokens: usize)

LongCat: the SECOND sublayer of a dual-sublayer block adds the paired first sublayer’s stashed shortcut-MoE output at its end.

impl Qwen3AttentionLayer

pub fn new( input_norm: DenseWeight, attn: AttentionWeights, post_attn_norm: DenseWeight, ffn: FfnComponent, attn_layer_idx: usize, q_nvfp4: Option<QuantizedWeight>, k_nvfp4: Option<QuantizedWeight>, v_nvfp4: Option<QuantizedWeight>, gpu: &dyn GpuBackend, kv_dtype: KvCacheDtype, fp8_calibration_tokens: usize, config: &ModelConfig, ) -> Result<Self>

pub fn new_ungated( input_norm: DenseWeight, attn: AttentionWeights, post_attn_norm: DenseWeight, ffn: FfnComponent, attn_layer_idx: usize, q_nvfp4: Option<QuantizedWeight>, k_nvfp4: Option<QuantizedWeight>, v_nvfp4: Option<QuantizedWeight>, gpu: &dyn GpuBackend, kv_dtype: KvCacheDtype, fp8_calibration_tokens: usize, config: &ModelConfig, ) -> Result<Self>

impl Qwen3AttentionLayer

pub fn set_prefill_weights( &mut self, q_nvfp4_t: Option<QuantizedWeight>, k_nvfp4_t: Option<QuantizedWeight>, v_nvfp4_t: Option<QuantizedWeight>, o_nvfp4_t: Option<QuantizedWeight>, )

Set transposed NVFP4 weight copies for prefill GEMM (w4a16_gemm_t, N_TILE=128).

pub fn set_packed_q2_weights( &mut self, q: PackedQ2Weight, k: PackedQ2Weight, v: PackedQ2Weight, o: PackedQ2Weight, gpu: &dyn GpuBackend, )

Install keep-packed ternary Q2_0 q/k/v/o weights (Tier-1c, ATLAS_GGUF_NATIVE_Q2=1). Decode dispatches q2_0_gemv_vec (2-bit resident, no NVFP4); prefill transient-dequants each to BF16 via Self::q2_prefill_gemm. Replaces the NVFP4 decode weights (which are NULL on this path — no NVFP4 was allocated).

pub fn set_fused_qkv_prefill_weight( &mut self, qkv_nvfp4_t: Option<QuantizedWeight>, )

Install the fused [q|k|v] transposed twin. Separate from set_prefill_weights so the fused path is opt-in per loader and the separate twins stay available as the fallback.

pub fn set_fp8_weights( &mut self, q: Option<Fp8Weight>, k: Option<Fp8Weight>, v: Option<Fp8Weight>, o: Option<Fp8Weight>, )

Set native FP8 checkpoint weights for the w8a16_gemv decode path.

The block-scaled FP8 weights stored here (weight + per-128 row_scale) are ALSO consumed by block-scaled prefill: fp8_gemm_t_blockscaled folds both the per-token activation scale and the per-block weight scale in an FP32 epilogue. (Historical note: the older single-scale fp8_gemm_t/fp8_gemm_n128 prefill could not apply block scales, so prefill used to fall through to the NVFP4/BF16 dequant path — that is no longer the case; block-scaled prefill is the default, see ops::fp8_blockscaled_prefill_enabled.)

pub fn set_lora_weights( &mut self, attn: LoraAttnWeights, ffn: Option<LoraFfnWeights>, ) -> Result<()>

Install the startup-static LoRA adapter overlay (post-construction, mirroring Self::set_fp8_weights). attn carries the K/V/O pairs; ffn (when Some) is routed into this layer’s dense FFN component — it lives here rather than on the model because self.ffn is pub(super). M0: weights are stored only; compute reads land in M1.

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

Feature-1: install this layer’s MoE router + routed-expert LoRA onto its FfnComponent::Moe. The MoE FFN lives in self.ffn or (some loaders) self.moe_ffn — try both, else the adapter targeted experts on a layer with no MoE FFN (hard reject). Scratch is allocated inside crate::layers::MoeLayer::set_lora_weights.

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

Transpose FP8 weights for fast prefill (w8a16_gemm_t: coalesced reads). Must be called after Self::set_fp8_weights. Allocates new GPU buffers.

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

Pre-dequant NVFP4 → FP8 for Q/K/V/O transposed weights.

Trait Implementations

impl TransformerLayer for Qwen3AttentionLayer

fn decode_graph_unsupported(&self) -> bool

QSA selection does a host top-k per step — never capturable, and a graph captured on the dense path would replay wrong attention once selection activates.

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-mode attention prefill via prefill_inner with batched_meta = Some. The model-level prefill_attn_batched_layer calls this method. Per-stream block_table is unused under batched mode (block_table_ptrs from batched_meta carries them); we still pass an empty Vec to satisfy the signature.

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

Free the QSA indexer carry this sequence lazily attached.

alloc_state hands back an EMPTY AttnLayerState; the buffers appear later, on first use, via qsa_seq_state. So the thing to release is not what alloc_state returned — it is whatever the sequence grew. take() makes this idempotent and leaves the state in the same shape alloc_state produced.

A layer with no QSA indexer (plain attention) never populates the field, so the take() yields None and this costs nothing.

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

Allocate per-sequence state for this layer.

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

Returns true if this layer is an SSM layer (supports two-phase 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<()>

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

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