Trait TransformerLayer 

pub trait TransformerLayer: Send + Sync {
    // Required methods
    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<()>;
    fn alloc_state(&self, gpu: &dyn GpuBackend) -> Result<Box<dyn LayerState>>;

    // Provided methods
    fn uses_local_mla_prefill(&self) -> bool { ... }
    fn as_any_mut(&mut self) -> Option<&mut dyn Any> { ... }
    fn fp8_calibration_frozen(&self) -> Option<bool> { ... }
    fn decode_prestage(
        &self,
        _token: u32,
        _state: &mut dyn LayerState,
        _gpu: &dyn GpuBackend,
        _stream: u64,
    ) -> Result<()> { ... }
    fn decode_prestage_rearm(&self, _state: &mut dyn LayerState) { ... }
    fn decode_graph_unsupported(&self) -> bool { ... }
    fn decode_multi_seq_unsupported(&self) -> bool { ... }
    fn decode_verify_multi_unsupported(&self) -> bool { ... }
    fn snapshot_aux(
        &self,
        _state: &dyn LayerState,
        _gpu: &dyn GpuBackend,
        _stream: u64,
    ) -> Result<Option<Vec<u8>>> { ... }
    fn has_aux_state(&self) -> bool { ... }
    fn restore_aux(
        &self,
        _state: &mut dyn LayerState,
        _blob: &[u8],
        _gpu: &dyn GpuBackend,
        _stream: u64,
    ) -> Result<()> { ... }
    fn graph_stale_on_new_sequence(&self) -> bool { ... }
    fn sync_replayed_step(
        &self,
        _state: &mut dyn LayerState,
        _seq_len: usize,
        _k: usize,
    ) -> Result<()> { ... }
    fn check_replay_room(
        &self,
        _state: &dyn LayerState,
        _seq_len: usize,
        _k: usize,
    ) -> Result<()> { ... }
    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 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<()> { ... }
    fn prefill_phase1_proj_batched(
        &self,
        hidden_stacked: DevicePtr,
        residual_stacked: DevicePtr,
        total_tokens: usize,
        gdn_bufs: &GdnPrefillBuffers,
        ctx: &ForwardContext<'_>,
        stream: u64,
    ) -> Result<()> { ... }
    fn prefill_phase1_conv1d_one(
        &self,
        state: &mut dyn LayerState,
        token_offset: usize,
        len: usize,
        gdn_bufs: &GdnPrefillBuffers,
        ctx: &ForwardContext<'_>,
        stream: u64,
    ) -> Result<()> { ... }
    fn prefill_phase1_l2_batched(
        &self,
        total_tokens: usize,
        gdn_bufs: &GdnPrefillBuffers,
        ctx: &ForwardContext<'_>,
        stream: u64,
    ) -> Result<()> { ... }
    fn prefill_gdn_full(
        &self,
        _state: &mut dyn LayerState,
        _gdn_bufs: &GdnPrefillBuffers,
        _ctx: &ForwardContext<'_>,
        _stream: u64,
    ) -> Result<()> { ... }
    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<()> { ... }
    fn prefill_gdn_full_batched(
        &self,
        _h_state_ptrs: DevicePtr,
        _gdn_bufs: &GdnPrefillBuffers,
        _batch_size: u32,
        _chunk_len: u32,
        _ctx: &ForwardContext<'_>,
        _stream: u64,
    ) -> Result<()> { ... }
    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> { ... }
    fn prefill_phase3(
        &self,
        _hidden: DevicePtr,
        _residual: DevicePtr,
        _num_tokens: usize,
        _gdn_bufs: &GdnPrefillBuffers,
        _token_offset: usize,
        _ctx: &ForwardContext<'_>,
        _stream: u64,
    ) -> Result<()> { ... }
    fn is_ssm_layer(&self) -> bool { ... }
    fn transpose_moe_for_prefill(
        &mut self,
        _gpu: &dyn GpuBackend,
        _config: &ModelConfig,
    ) -> Result<()> { ... }
    fn transpose_moe_gate_up_for_prefill(
        &mut self,
        _gpu: &dyn GpuBackend,
        _config: &ModelConfig,
    ) -> Result<()> { ... }
    fn set_moe_down_transpose_scratch(
        &mut self,
        _scratch_packed: DevicePtr,
        _scratch_scale: DevicePtr,
        _packed_ptrs_t: DevicePtr,
        _scale_ptrs_t: DevicePtr,
    ) { ... }
    fn transpose_moe_for_prefill_unified(
        &mut self,
        _gpu: &dyn GpuBackend,
        _config: &ModelConfig,
    ) -> Result<()> { ... }
    fn transpose_moe_for_prefill_hybrid(
        &mut self,
        _gpu: &dyn GpuBackend,
        _config: &ModelConfig,
    ) -> Result<()> { ... }
    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<()> { ... }
    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<()> { ... }
    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<()> { ... }
    fn release_state(
        &self,
        _state: &mut dyn LayerState,
        _gpu: &dyn GpuBackend,
    ) -> Result<()> { ... }
    fn uses_ssm_pool(&self) -> bool { ... }
}

Required Methods

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.

Arguments

• hidden - [1, hidden_size] BF16, read and written
• residual - [1, hidden_size] BF16, scratch for residual stream
• state - Per-layer state (empty for attention, SSM state for recurrent)
• kv_cache - Paged KV cache (may be mutated for block allocation)
• seq_len - Current sequence length (for position encoding + cache)
• block_table - Sequence’s block table (may grow if new blocks needed)
• ctx - Shared forward context (buffers, gpu, config)
• stream - CUDA stream handle

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

Allocate per-sequence state for this layer.

Called once when a new sequence is created. Returns:

• EmptyLayerState for pure attention layers
• SsmLayerState for SSM/recurrent layers

Provided Methods

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

True when this layer’s decode can NEVER be captured into a CUDA graph — e.g. the QSA indexer’s host top-k round trip, whose captured dense fallback would silently replay WRONG attention once selection activates. The scheduler ORs this across layers once and keeps the whole model eager.

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.

Mirrors Self::decode_graph_unsupported exactly: layer-level statement, default false, ORed across layers by the caller and consumed at the DISPATCH site. A true layer is NOT refused concurrency — it is routed onto the per-sequence highway loop that #753 item B already built for mHC models, so C>1 keeps serving.

Wired at BOTH multi-seq callers (decode_a2’s hc_perseq and decode_b’s hc_qsa_perseq), because decode_b is the single-GPU fused decode+prefill path and a decision made only in decode_a2 leaves it exposed.

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.

Separate from Self::decode_multi_seq_unsupported because the two answers can differ: verify carries its rows on the k axis with its own R-row metadata block, decode carries them on the sequence axis.

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

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

Used during prompt processing: reads weight matrices once for all N tokens (GEMM M=N) instead of N separate GEMV calls. Attention uses Flash Attention on contiguous Q/K/V. SSM/GDN recurrence remains sequential per-token.

Arguments

• hidden - [N, hidden_size] BF16, read and written
• residual - [N, hidden_size] BF16, scratch for residual stream
• num_tokens - Number of tokens (N)
• state - Per-layer state (SSM state updated sequentially)
• kv_cache - Paged KV cache (attention layers write K/V for all N)
• seq_len_start - Sequence position of first token (usually 0)
• block_table - Block table for KV cache (pre-allocated for N tokens)
• ctx - Shared forward context (buffers, gpu, config)
• stream - CUDA stream handle

Default: falls back to sequential single-token decode calls.

kv_write_start: number of tokens whose KV cache entries are already populated (prefix caching). Attention layers skip KV writes for positions < kv_write_start. SSM layers ignore this (recurrent). Does a captured decode graph go STALE when a new sequence takes this slot?

🔴 decode_graph is keyed by slot_idx on the premise that the only per-sequence addresses a capture bakes live in the SSM pool, which is slot-addressed and stable. A layer that allocates its own per-sequence state (GLM-5.3 allocates a fresh indexer cache and KDA state per sequence) breaks that premise: the next sequence gets new buffers and the old graph still reads and writes the freed ones — the second request continues the first one’s text. Such a layer says so here and free_sequence drops the slot’s graph, costing one re-capture per request.

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.

🔴 A graph replay executes kernels and nothing else: the layer’s decode never runs, so a layer that tracks its own cache length on the host silently stops advancing and every replayed step overwrites the same row.

🔴 It is a RECONCILE, not an advance. A K-row verify writes K rows and the scheduler then keeps only the accepted prefix, so the counter has to be rewound to seq_len first — exactly what decode_k’s own lockstep check does on the eager path. Advancing blindly leaves the counter (k - accepted) ahead of the sequence on every rejected draft, and that drift is ANOMALIES A56: the DRAFTER writes its indexer rows at state.len(), so a counter running ahead lands them on rows the target then selects over. Default is a no-op — only a layer with host-side state (GLM-5.3’s DSA indexer cache) needs this.

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.

🔴 sync_replayed_step above is the RECONCILE and it deliberately runs AFTER launch_graph, which is too late to prevent a write. A replayed dsa_indexer_store places its row from a DEVICE position with no host code in the loop, so at the DSA ceiling it writes one row past the buffer and the refusal arrives afterwards. The resulting CUDA_ERROR_ILLEGAL_ADDRESS (700) is STICKY: it fails every later CUDA call in the context, so one over-long sequence takes the serve down for every subsequent request while the health endpoints keep answering 200. ANOMALIES A62.

seq_len is the length BEFORE this step’s k rows, so the step ends at seq_len + k — the same post-condition sync_replayed_step reconciles to. Default is a no-op: only a layer with host-side cache bookkeeping needs it.

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

Runs RMS norm, QKVZ projection, BA+gates, conv1d, and L2 norm for a chunk of num_tokens tokens, then copies the GDN inputs (packed QKV, gate/beta, Z) into the full-sequence gdn_bufs at token_offset.

Does NOT run the GDN recurrence — that happens in prefill_gdn_full after all chunks have staged their inputs.

Attention layers: default falls back to full prefill (no phasing).

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.

Runs the WY4-persistent GDN kernel over all total_len tokens in gdn_bufs in a single launch. The kernel reads packed QKV and gate/beta from the full-sequence buffers and writes the GDN output.

Only meaningful for SSM layers. Attention layers return Ok(()).

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.

Runs the full attention-layer prefill (rms_norm + residual, QKV proj, RoPE, KV-write, batched attention compute, O proj, post-attn norm, FFN, final residual) over num_tokens = batch_size * chunk_len stacked tokens, using batched_meta for per-stream metadata resolution.

Default impl returns Err — only Qwen3AttentionLayer overrides. SSM/dense layers don’t override (they have their own batched paths or work without batched metadata).

Caller (model-level prefill_attn_batched_layer) is responsible for ensuring all streams share the same chunk_len, seq_len_start (q_offset), and that the layer is not MLA / not HDIM=512 / not HSS- engaged. The override bails Err if any unsupported case is detected.

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.

Runs the same WY32 / persistent / split4 GDN kernel as prefill_gdn_full but with batch_size = batch_size and h_state_ptrs pointing to a device array of N per-stream h_state pointers (staged by TransformerModel::stage_h_state_ptrs). gdn_bufs.qkv / gate_beta / output are stacked across N streams contiguously: each stream’s data lives at b * chunk_len * conv_dim (BF16) within the buffer.

Default impl returns Err — the SSM layer override implements the actual batched dispatch using the kernel handles loaded in commit 8d07ca4. Attention layers don’t override (they don’t have a GDN step).

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.

Reads GDN output and Z gate from gdn_bufs at token_offset, then runs gated RMS norm, output projection, residual add, and MoE for the chunk of num_tokens tokens.

Only meaningful for SSM layers. Attention layers return Ok(()).

fn is_ssm_layer(&self) -> bool

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

When true, the model loop can use prefill_phase1 / prefill_gdn_full / prefill_phase3 instead of the monolithic prefill.

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.

Default: no-op (non-MoE layers, and MoE layers whose loader already called MoeLayer::transpose_for_prefill inline during construction).

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

Used for speculative decode verification: processes multiple tokens per layer with GEMM for weight-heavy projections (amortizes bandwidth) and sequential ops for stateful/recurrent components.

Arguments

• hidden - [K, hidden_size] BF16, read and written (K tokens contiguous)
• residual - [K, hidden_size] BF16, scratch for residual stream
• num_tokens - Number of tokens (K)
• state - Per-layer state
• kv_cache - Paged KV cache
• seq_len - Starting sequence length (before these tokens)
• block_table - Block table for KV cache
• ctx - Shared context
• stream - CUDA stream

Default: falls back to sequential single-token decode calls.

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.

Each sequence contributes 1 token. The weight matrices are loaded once and applied to all N sequences (amortizing memory bandwidth).

Arguments

• hidden - [N, hidden_size] BF16, contiguous
• residual - [N, hidden_size] BF16, contiguous
• num_seqs - Number of sequences (N)
• states - N per-layer states (one per sequence)
• kv_cache - Shared paged KV cache
• ctx - Forward context (attn_metadata contains N-sequence metadata)
• stream - CUDA stream

Default: falls back to N sequential single-token decode calls.

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.

Only SSM layers override (attention layers are handled by the caller via decode_multi_seq, which already takes per-row block tables and seq lens). Default: unsupported.

wy_tables: this layer’s slice of the model-staged WY pointer tables (layout above, VERIFY_WY_LAYER_STRIDE_BYTES; refreshed pre-graph every step) enabling the single-launch table-form WY batch. NULL → the layer keeps its per-sequence WY path.

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

Release the per-sequence state alloc_state produced, plus anything the layer attached to it lazily afterwards.

Called once per sequence from the teardown chokepoint (free_sequence_dispatch). The default no-op is correct for layers whose state owns no device memory (EmptyLayerState) and for state that comes from a pool reclaimed by slot (SsmLayerState’s h/conv, released via ssm_pool.release_slot).

It exists because LayerState implementors hold BARE DevicePtrs: dropping the box reclaims the host struct and leaks the device buffer. The QSA indexer carry (~739 MB per request at 200K context across the 12 full-attention layers) and the PLE conv carry both leaked this way. On unified memory such a leak is invisible to RSS and reported as N/A by nvidia-smi, so it surfaces only as the host exhausting RAM with no process to blame.

MUST be idempotent — teardown can run after a partial failure. Callers log errors and continue rather than aborting: a sequence that cannot free its state is still finished, and bailing would strand the rest. Owns every device allocation reachable from this LayerState that the layer obtained from gpu.alloc, whether in alloc_state or attached later. Idempotent; nulls what it frees; never touches pool addresses.

🔴 Refuse by TYPE inside the impl, not by a filter at the call site. A call-site filter is a second spelling of “is this pooled?” that can drift out of agreement with the first; the type check lives where the knowledge is.

🔴 Invariant L2 (slot reuse), NOT a line order. It is tempting to write “the graph drop must come before this call” — that over-states a call order as an invariant. The real requirement is that when a slot is re-occupied, its graphs are destroyed AND its owned pointers are freed and nulled. Nothing between the two blocks replays a graph, and destroy_graph does not dereference baked pointers, so either order satisfies it. ANOMALIES A56 is the history; slot reuse is the invariant.

fn uses_ssm_pool(&self) -> bool

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

true (the default) is the long-standing arrangement: sequence setup sees LayerType::LinearAttention and hands the layer an SsmLayerState pointing at pool-owned addresses, so alloc_state is never consulted.

🪤 A linear-attention mixer with its OWN state type must return false, or it is handed an SsmLayerState and the downcast in its forward path fails at layer 0 on the first request. GLM-5.3’s KDA blocks are the case: they are linear_attention in layer_types but carry Glm5NextLayerState::Kda.

Implementors

impl TransformerLayer for Glm5NextDsaLayer

impl TransformerLayer for Glm5NextLayer

impl TransformerLayer for NemotronMamba2Layer

impl TransformerLayer for NemotronMoeLayer

impl TransformerLayer for Qwen3AttentionLayer

impl TransformerLayer for Qwen3SsmLayer