Trait TransformerLayer 

pub trait TransformerLayer:
    Send
    + Sync
    + LayerCapabilities
    + LayerWeightSetup
    + LayerWriteOnAccept
    + LayerGraphHooks
    + LayerAuxState
    + LayerSplitPrefill
    + CircuitBindings {
    // 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 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 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<()> { ... }
}

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

2026-09-25: Decode one token through this layer, updating hidden in place.

• hidden - [1, hidden_size] BF16, read and written
• residual - [1, hidden_size] BF16 residual stream
• state - this layer’s state for the sequence
• seq_len - the token’s position

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.

Provided Methods

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

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

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

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

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

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.

The teardown chokepoint free_sequence_dispatch calls it for every layer, logs an error and continues. An implementation must be idempotent, null what it frees, never free pool addresses, and decide by the state’s type rather than rely on a filter at the call site. The default does nothing, which suits state that owns no device memory and pool-backed SSM state, released through the pool slot.

Implementors

impl TransformerLayer for Qwen3AttentionLayer

impl TransformerLayer for Qwen3SsmLayer