Struct DeepSeekV41Layer 

pub struct DeepSeekV41Layer {

    pub idx: usize,
    pub role: LayerRole,
    pub rt: Arc<V41Runtime>,
    pub attn_w: AttnV41LayerWeights,
    pub moe_w: MoeV41LayerWeights,
    pub next_router: Option<RouterWeights>,
    pub engram_index: Option<usize>,
    pub hc_attn: HcSiteWeights,
    pub hc_ffn: HcSiteWeights,
    pub attn_norm: DenseWeight,
    pub ffn_norm: DenseWeight,
    pub k_hc_expand: KernelHandle,
    pub k_hc_post: KernelHandle,
    pub k_mixes_dot: KernelHandle,
    pub k_mixes_finish: KernelHandle,
    pub k_collapse: KernelHandle,
    pub k_finish_collapse: KernelHandle,
    pub k_collapse_wide: KernelHandle,
    pub k_post_wide: KernelHandle,
    pub k_rms_norm: KernelHandle,
}

Fields

idx: usize

role: LayerRole

rt: Arc<V41Runtime>

attn_w: AttnV41LayerWeights

moe_w: MoeV41LayerWeights

next_router: Option<RouterWeights>

The next layer’s router, run on this layer’s MoE input to predict (and start reading) the experts the next layer will ask for.

engram_index: Option<usize>

Some(index into engram_layer_ids) on engram layers.

hc_attn: HcSiteWeights

hc_ffn: HcSiteWeights

attn_norm: DenseWeight

ffn_norm: DenseWeight

k_hc_expand: KernelHandle

k_hc_post: KernelHandle

k_mixes_dot: KernelHandle

k_mixes_finish: KernelHandle

k_collapse: KernelHandle

k_finish_collapse: KernelHandle

the site’s finish and the block’s collapse in one wide launch

k_collapse_wide: KernelHandle

hc_v41_collapse / hc_post over ceil(H/256) blocks a token

k_post_wide: KernelHandle

k_rms_norm: KernelHandle

Trait Implementations

impl TransformerLayer for DeepSeekV41Layer

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

The buffers go with the sequence. alloc_state allocates the window ring, the compressor’s partial group, the latent cache and the index-key cache per sequence with gpu.alloc, so a release that did nothing left all four on the device for the life of the process. The latent cache alone is max_seq / ratio rows of 512 bf16 (4 MiB at 4096 tokens, ratio 1), so every request left on the order of 15 to 25 MiB across the forty layers, and it never came back. The captured segments bake this sequence’s buffers, so they go first.

fn decode_rollback_unsupported(&self) -> bool

Lowering seq_len rewinds none of this layer’s state: the shared compressed-row count only grows (SharedV41::compress_len), the ratio-2 compressor holds a partial group of the dropped tokens, and the engram hasher’s n-gram history ends in them. A watchdog rollback then re-steers from a corrupted context (09-16: the second MinHeap request spliced _wHere at the boundary and degenerated from there). Declined until a rewind exists; the loop watchdog still ends the response.

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

Allocate per-sequence state for this layer.

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.

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 gdn_woa_stash_seq_floats(&self) -> Option<usize>

Write-on-accept capability (GDN layers only): the per-sequence stash size in f32 elements when this layer can run the K=4 write-on-accept twin (kernels linked, write-on-accept not disabled), None otherwise. The model sizes the stash from it and binds with Self::gdn_woa_bind.

fn gdn_woa_bind(&self, _flag: DevicePtr, _stash: DevicePtr, _seqs: usize)

Bind this layer’s write-on-accept engaged word and stash slab (seqs sequences of Self::gdn_woa_stash_seq_floats f32 each). Called once by the model, pre-capture, on the first write-on-accept request; the addresses are baked into verify graphs afterwards and never move.

fn gdn_fold_accepted( &self, _gpu: &dyn GpuBackend, _h_table: DevicePtr, _na_tab: DevicePtr, _k_rows: usize, _n: usize, _stream: u64, ) -> Result<bool>

Write-on-accept fold (GDN layers only): apply the accepted rows of the last batched K=4 verify to this layer’s h states. h_table is the layer’s slab-0 pointer table from the verify, na_tab a device u32[n] of accepted row counts in batch order. Ok(false) when the layer has nothing bound (not a GDN layer, or write-on-accept is off). When the parent kernel ran for the last verify (engaged word 0) the fold performs the parent’s partial-accept restore instead.

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

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