pub struct Glm5NextLayer {Show 16 fields
pub layer_idx: usize,
pub mixer: Glm5NextMixer,
pub mlp: Glm5NextMlpSite,
pub mlp_cfg: Glm5NextMlpConfig,
pub mlp_kernels: Glm5NextMlpKernels,
pub mlp_ws: Arc<Glm5NextMlpWorkspace>,
pub mhc: Option<Glm5NextMhc>,
pub input_norm: DevicePtr,
pub post_attn_norm: DevicePtr,
pub rms_norm_k: KernelHandle,
pub add_k: KernelHandle,
pub rms_eps: f32,
pub hidden: usize,
pub mixer_all_reduce: bool,
pub is_first: bool,
pub is_last: bool,
}Expand description
2026-09-25: One bound GLM-5.3 decoder layer.
Fields§
§layer_idx: usize§mixer: Glm5NextMixer§mlp: Glm5NextMlpSite§mlp_cfg: Glm5NextMlpConfig§mlp_kernels: Glm5NextMlpKernels§mlp_ws: Arc<Glm5NextMlpWorkspace>2026-09-25: MLP scratch. The text layers share one workspace unless
METRALE_GLM_MLP_WS_SHARED=0 (mlp_ws_shared); the MTP block builds its own.
mhc: Option<Glm5NextMhc>2026-09-25: None for the MTP block, which runs forward_one_plain. Every text layer has
one: the skeleton marks every text layer hyper_connection.
input_norm: DevicePtr2026-09-25: input_layernorm.weight; post_attn_norm is post_attention_layernorm.weight.
post_attn_norm: DevicePtr§rms_norm_k: KernelHandle2026-09-25: rms_norm_vanilla; see the module invariants.
add_k: KernelHandle2026-09-25: bf16_add_inplace, the residual add of the MTP block, whose residual is not
in a highway. 0 on a target without the kernel; add_inplace then returns an error.
rms_eps: f32§mixer_all_reduce: bool2026-09-25: The mixer output is a partial sum across TP ranks, all-reduced before
hc_post (or, in the MTP block, before the residual add). Set from the mixer’s TP plan
(needs_output_all_reduce).
is_first: bool2026-09-25: Expand the highway in this layer; true for layer 0 only.
is_last: bool2026-09-25: Collapse the highway in this layer; true for the last text layer only.
Implementations§
Source§impl Glm5NextLayer
impl Glm5NextLayer
Sourcepub fn decode_one_for_drafter(
&self,
hidden: DevicePtr,
state: &mut dyn LayerState,
kv_cache: &mut PagedKvCache,
seq_len: usize,
block_table: &mut Vec<u32>,
ctx: &ForwardContext<'_>,
stream: u64,
) -> Result<()>
pub fn decode_one_for_drafter( &self, hidden: DevicePtr, state: &mut dyn LayerState, kv_cache: &mut PagedKvCache, seq_len: usize, block_table: &mut Vec<u32>, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>
2026-09-25: One token through this block for the MTP drafter; errors on a block with a
hyper-connection. hidden is both input and residual and is updated in place. No disk
tiers: the disk block lists passed down are empty.
Sourcepub fn drafter_write_kv_row(
&self,
x: DevicePtr,
state: &mut dyn LayerState,
kv_cache: &mut PagedKvCache,
seq_len: usize,
block_table: &mut Vec<u32>,
ctx: &ForwardContext<'_>,
stream: u64,
) -> Result<()>
pub fn drafter_write_kv_row( &self, x: DevicePtr, state: &mut dyn LayerState, kv_cache: &mut PagedKvCache, seq_len: usize, block_table: &mut Vec<u32>, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>
2026-09-25: One drafter context row: input_norm, then Glm5NextDsaLayer::write_kv_row,
which writes the row’s KV latent and indexer entry without computing the block output.
Errors on a block whose mixer is not DSA.
Trait Implementations§
Source§impl CircuitBindings for Glm5NextLayer
impl CircuitBindings for Glm5NextLayer
Source§fn circuit_prepare(
&self,
_gpu: &dyn GpuBackend,
_config: &ModelConfig,
_levers: &ModelLevers,
_stream: u64,
) -> Result<(), Error>
fn circuit_prepare( &self, _gpu: &dyn GpuBackend, _config: &ModelConfig, _levers: &ModelLevers, _stream: u64, ) -> Result<(), Error>
stream.Source§fn circuit_layer(
&self,
_config: &ModelConfig,
_levers: &ModelLevers,
) -> Option<CircuitLayer>
fn circuit_layer( &self, _config: &ModelConfig, _levers: &ModelLevers, ) -> Option<CircuitLayer>
None for a layer type it does not bind.
config and levers are the model’s, as its decode reads them.Source§impl LayerAuxState for Glm5NextLayer
impl LayerAuxState for Glm5NextLayer
Source§fn has_aux_state(&self) -> bool
fn has_aux_state(&self) -> bool
2026-09-25: True for a DSA layer: its indexer cache is the state snapshot_aux and
restore_aux carry. A KDA layer’s state is in the SSM pool (uses_ssm_pool) and is not
carried here.
Source§fn restore_aux(
&self,
state: &mut dyn LayerState,
blob: &[u8],
gpu: &dyn GpuBackend,
stream: u64,
) -> Result<()>
fn restore_aux( &self, state: &mut dyn LayerState, blob: &[u8], gpu: &dyn GpuBackend, stream: u64, ) -> Result<()>
2026-09-25: Errors on a KDA layer, whose state travels with the SSM snapshot. On a DSA
layer it restores the blob through Glm5NextDsaState::restore_blob; apply_aux_states
propagates any error.
Source§fn snapshot_aux(
&self,
state: &dyn LayerState,
gpu: &dyn GpuBackend,
stream: u64,
) -> Result<Option<Vec<u8>>>
fn snapshot_aux( &self, state: &dyn LayerState, gpu: &dyn GpuBackend, stream: u64, ) -> Result<Option<Vec<u8>>>
stream. Default None.Source§impl LayerCapabilities for Glm5NextLayer
impl LayerCapabilities for Glm5NextLayer
Source§fn decode_multi_seq_unsupported(&self) -> bool
fn decode_multi_seq_unsupported(&self) -> bool
2026-09-25: True. decode runs every sequence through highway slot 0, and the trait’s
default decode_multi_seq calls decode once per sequence with one shared
ForwardContext, so every sequence would read and write the same highway slot. The DSA
mixer’s single-row decode also reads attn_metadata row 0 on a decode step.
Source§fn decode_verify_multi_unsupported(&self) -> bool
fn decode_verify_multi_unsupported(&self) -> bool
2026-09-25: True. This layer has no decode_verify_multi, and the trait’s default
returns an error; answering true makes can_batch_verify_dispatch refuse the batched
verify.
Source§fn graph_stale_on_new_sequence(&self) -> bool
fn graph_stale_on_new_sequence(&self) -> bool
2026-09-25: True. A DSA layer’s per-sequence state comes from gpu.alloc in
alloc_state, so the addresses a captured graph holds belong to one sequence.
Source§fn uses_ssm_pool(&self) -> bool
fn uses_ssm_pool(&self) -> bool
2026-09-25: True for a KDA layer. Its recurrent and conv state then come from the SSM
pool, and with MTP on, so do the checkpoint and per-token intermediate pointers that a
speculative rollback restores (meta.rs).
Source§fn is_ssm_layer(&self) -> bool
fn is_ssm_layer(&self) -> bool
2026-09-25: True for a KDA layer, whose mixer carries recurrent state.
Source§fn uses_local_mla_prefill(&self) -> bool
fn uses_local_mla_prefill(&self) -> bool
METRALE_MLA_PREFIX_SKIP=1 (mla_prefill_needs_full_recompute).Source§fn fp8_calibration_frozen(&self) -> Option<bool>
fn fp8_calibration_frozen(&self) -> Option<bool>
None when the layer runs no online calibration (the default). The model lifts its
CUDA-graph suppression once no layer reports Some(false)
(graphs_ready_after_fp8_kv_cal).Source§fn decode_graph_unsupported(&self) -> bool
fn decode_graph_unsupported(&self) -> bool
decode_graph_veto when it is built.Source§fn decode_rollback_unsupported(&self) -> bool
fn decode_rollback_unsupported(&self) -> bool
rollback_to_boundary
then declines the rollback (LayerStateNotRewindable).Source§fn supports_replay_tail_split(&self) -> bool
fn supports_replay_tail_split(&self) -> bool
prefill honours a replay-tail capture
(ForwardContext::midchunk_capture with replay_tail) on every build: the rows
from cap_local on get the bits a prefill pass starting at cap_local gives them
(an SSM layer splits its conv and recurrence there and copies its state into the
plan’s slots; an attention layer runs those rows as their own paged call). The
model plans a replay-tail capture only when every layer says so
(TransformerModel::inpass_cut_capture_supported); otherwise it keeps the
two-pass tail split.Source§impl LayerGraphHooks for Glm5NextLayer
impl LayerGraphHooks for Glm5NextLayer
Source§fn sync_replayed_step(
&self,
state: &mut dyn LayerState,
seq_len: usize,
k: usize,
) -> Result<()>
fn sync_replayed_step( &self, state: &mut dyn LayerState, seq_len: usize, k: usize, ) -> Result<()>
2026-09-25: A graph replay runs only kernels, so the host-side length of the DSA indexer
cache is reconciled here (Glm5NextDsaState::sync_to). The inner Glm5NextDsaLayer
implements this too, but the model’s layer list holds this composite, so this impl is the
one called. KDA keeps no host-side state.
Source§fn check_replay_room(
&self,
state: &dyn LayerState,
seq_len: usize,
k: usize,
) -> Result<()>
fn check_replay_room( &self, state: &dyn LayerState, seq_len: usize, k: usize, ) -> Result<()>
2026-09-25: Checks before a graph replay that the DSA indexer cache can hold seq_len + k
rows. As with sync_replayed_step, this composite impl is the one the model calls.
Source§fn decode_prestage(
&self,
_token: u32,
_state: &mut dyn LayerState,
_gpu: &dyn GpuBackend,
_stream: u64,
) -> Result<(), Error>
fn decode_prestage( &self, _token: u32, _state: &mut dyn LayerState, _gpu: &dyn GpuBackend, _stream: u64, ) -> Result<(), Error>
decode_a.rs), so
a graph reads only device buffers refreshed here. Default: nothing.Source§fn decode_prestage_rearm(&self, _state: &mut dyn LayerState)
fn decode_prestage_rearm(&self, _state: &mut dyn LayerState)
decode_a.rs). It must not recompute the prestage
work: PLE’s history already advanced in decode_prestage.Source§impl LayerSplitPrefill for Glm5NextLayer
impl LayerSplitPrefill for Glm5NextLayer
Source§fn prefill_phase1_proj_batched(
&self,
hidden_stacked: DevicePtr,
residual_stacked: DevicePtr,
total_tokens: usize,
gdn_bufs: &GdnPrefillBuffers,
ctx: &ForwardContext<'_>,
stream: u64,
) -> Result<(), Error>
fn prefill_phase1_proj_batched( &self, hidden_stacked: DevicePtr, residual_stacked: DevicePtr, total_tokens: usize, gdn_bufs: &GdnPrefillBuffers, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<(), Error>
prefill_phase1_conv1d_one per
request and prefill_phase1_l2_batched. The default returns an error.Source§fn prefill_phase1_conv1d_one(
&self,
state: &mut dyn LayerState,
token_offset: usize,
len: usize,
gdn_bufs: &GdnPrefillBuffers,
ctx: &ForwardContext<'_>,
stream: u64,
) -> Result<(), Error>
fn prefill_phase1_conv1d_one( &self, state: &mut dyn LayerState, token_offset: usize, len: usize, gdn_bufs: &GdnPrefillBuffers, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<(), Error>
gdn_bufs.qkv. The default returns an error.Source§fn prefill_phase1_l2_batched(
&self,
total_tokens: usize,
gdn_bufs: &GdnPrefillBuffers,
ctx: &ForwardContext<'_>,
stream: u64,
) -> Result<(), Error>
fn prefill_phase1_l2_batched( &self, total_tokens: usize, gdn_bufs: &GdnPrefillBuffers, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<(), Error>
gdn_bufs.qkv, after every request’s
conv1d. The default returns an error.Source§fn prefill_gdn_full(
&self,
_state: &mut dyn LayerState,
_gdn_bufs: &GdnPrefillBuffers,
_ctx: &ForwardContext<'_>,
_stream: u64,
) -> Result<(), Error>
fn prefill_gdn_full( &self, _state: &mut dyn LayerState, _gdn_bufs: &GdnPrefillBuffers, _ctx: &ForwardContext<'_>, _stream: u64, ) -> Result<(), Error>
gdn_bufs.total_len tokens, reading packed QKV and gate/beta and writing output.
Default: nothing.Source§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<(), Error>
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<(), Error>
batched_meta. The default returns an
error; the Qwen3 attention layer overrides it.Source§fn prefill_gdn_full_batched(
&self,
_h_state_ptrs: DevicePtr,
_gdn_bufs: &GdnPrefillBuffers,
_batch_size: u32,
_chunk_len: u32,
_ctx: &ForwardContext<'_>,
_stream: u64,
) -> Result<(), Error>
fn prefill_gdn_full_batched( &self, _h_state_ptrs: DevicePtr, _gdn_bufs: &GdnPrefillBuffers, _batch_size: u32, _chunk_len: u32, _ctx: &ForwardContext<'_>, _stream: u64, ) -> Result<(), Error>
batch_size streams of equal length chunk_len in
one call. h_state_ptrs is a device array of one h-state pointer per stream
(TransformerModel::stage_h_state_ptrs), and the streams’ tokens lie back to back
in gdn_bufs. prefill_ssm_batched_layer calls it only when every stream has
the same length. The default returns an error.Source§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, Error>
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, Error>
cu_seqlens. Ok(true) if it ran; Ok(false) if not eligible, and the
caller then runs prefill_gdn_full per stream. Default Ok(false).Source§fn prefill_phase3(
&self,
_hidden: DevicePtr,
_residual: DevicePtr,
_num_tokens: usize,
_gdn_bufs: &GdnPrefillBuffers,
_token_offset: usize,
_ctx: &ForwardContext<'_>,
_stream: u64,
) -> Result<(), Error>
fn prefill_phase3( &self, _hidden: DevicePtr, _residual: DevicePtr, _num_tokens: usize, _gdn_bufs: &GdnPrefillBuffers, _token_offset: usize, _ctx: &ForwardContext<'_>, _stream: u64, ) -> Result<(), Error>
num_tokens tokens at token_offset,
gated RMSNorm of the GDN output with Z, the output projection, the residual add and
the FFN. Default: nothing.Source§impl LayerWeightSetup for Glm5NextLayer
impl LayerWeightSetup for Glm5NextLayer
Source§fn as_any_mut(&mut self) -> Option<&mut (dyn Any + 'static)>
fn as_any_mut(&mut self) -> Option<&mut (dyn Any + 'static)>
&mut dyn Any downcast hook for weight overlays applied after
construction, such as a LoRA install. Default None.Source§fn transpose_moe_for_prefill(
&mut self,
_gpu: &dyn GpuBackend,
_config: &ModelConfig,
) -> Result<(), Error>
fn transpose_moe_for_prefill( &mut self, _gpu: &dyn GpuBackend, _config: &ModelConfig, ) -> Result<(), Error>
factory/m2_setup.rs) calls it,
outside the hybrid and unified layouts, when free memory covers the full transpose
plus a safety margin. Default: nothing.Source§fn transpose_moe_gate_up_for_prefill(
&mut self,
_gpu: &dyn GpuBackend,
_config: &ModelConfig,
) -> Result<(), Error>
fn transpose_moe_gate_up_for_prefill( &mut self, _gpu: &dyn GpuBackend, _config: &ModelConfig, ) -> Result<(), Error>
transpose_moe_for_prefill, for gate and up only. Outside the
hybrid and unified layouts, the transpose pass calls it when the full transpose does
not fit but gate and up do. Default: nothing.Source§fn set_moe_down_transpose_scratch(
&mut self,
_scratch_packed: DevicePtr,
_scratch_scale: DevicePtr,
_packed_ptrs_t: DevicePtr,
_scale_ptrs_t: DevicePtr,
)
fn set_moe_down_transpose_scratch( &mut self, _scratch_packed: DevicePtr, _scratch_scale: DevicePtr, _packed_ptrs_t: DevicePtr, _scale_ptrs_t: DevicePtr, )
down_proj transpose
scratch that every MoE layer shares, after transpose_moe_gate_up_for_prefill.
Default: nothing.Source§fn transpose_moe_for_prefill_unified(
&mut self,
_gpu: &dyn GpuBackend,
_config: &ModelConfig,
) -> Result<(), Error>
fn transpose_moe_for_prefill_unified( &mut self, _gpu: &dyn GpuBackend, _config: &ModelConfig, ) -> Result<(), Error>
METRALE_UNIFIED_MOE_LAYOUT=1): build transposed
gate/up/down and free the untransposed copies. Decode reads the transposed copies
when MoeLayer::use_t_layout_for_decode holds. Default: nothing.Source§fn transpose_moe_for_prefill_hybrid(
&mut self,
_gpu: &dyn GpuBackend,
_config: &ModelConfig,
) -> Result<(), Error>
fn transpose_moe_for_prefill_hybrid( &mut self, _gpu: &dyn GpuBackend, _config: &ModelConfig, ) -> Result<(), Error>
METRALE_HYBRID_MOE_LAYOUT=1): build transposed
gate/up/down beside the originals and free nothing, so decode keeps the originals
and prefill uses the transposed copies. The transpose pass calls it only when twice
the full transpose fits. Default: nothing.Source§impl LayerWriteOnAccept for Glm5NextLayer
impl LayerWriteOnAccept for Glm5NextLayer
Source§fn gdn_woa_stash_seq_floats(&self) -> Option<usize>
fn gdn_woa_stash_seq_floats(&self) -> Option<usize>
None otherwise (the default). The model sizes the stash
from the largest answer and binds each GDN layer with Self::gdn_woa_bind.Source§fn gdn_woa_bind(&self, _flag: DevicePtr, _stash: DevicePtr, _seqs: usize)
fn gdn_woa_bind(&self, _flag: DevicePtr, _stash: DevicePtr, _seqs: usize)
seqs
sequences of Self::gdn_woa_stash_seq_floats f32 each). The model calls it
once, on the first write-on-accept request and before its graph decision, and
never moves the buffers afterwards (gdn_woa.rs).Source§fn gdn_carry_seq_floats(&self) -> Option<usize>
fn gdn_carry_seq_floats(&self) -> Option<usize>
None otherwise (the default).Source§fn gdn_carry_bind(&self, _binding: GdnCarryBinding)
fn gdn_carry_bind(&self, _binding: GdnCarryBinding)
Source§fn gdn_verify_run_batched(
&self,
_states: &[&mut (dyn LayerState + 'static)],
_kk: usize,
_gdn_wyn: bool,
_wy_tables: DevicePtr,
) -> Result<Option<bool>, Error>
fn gdn_verify_run_batched( &self, _states: &[&mut (dyn LayerState + 'static)], _kk: usize, _gdn_wyn: bool, _wy_tables: DevicePtr, ) -> Result<Option<bool>, Error>
states (adjacent sequences of
width kk, under the WY tables at wy_tables) through the batched conv and the
table-form WY, as this layer’s own batched verify decides it; None for a layer
without a batched GDN verify (the default). The circuit executor picks each run’s arm
with it.Source§fn gdn_fold_accepted(
&self,
_gpu: &dyn GpuBackend,
_h_table: DevicePtr,
_na_tab: DevicePtr,
_k_rows: usize,
_n: usize,
_stream: u64,
) -> Result<bool, Error>
fn gdn_fold_accepted( &self, _gpu: &dyn GpuBackend, _h_table: DevicePtr, _na_tab: DevicePtr, _k_rows: usize, _n: usize, _stream: u64, ) -> Result<bool, Error>
h_table is the layer’s WY pointer-table slice, na_tab a device u32[n]
of accepted row counts in batch order. Ok(false) when the layer did nothing, as
the default does.Source§impl TransformerLayer for Glm5NextLayer
impl TransformerLayer for Glm5NextLayer
Source§fn release_state(
&self,
state: &mut dyn LayerState,
gpu: &dyn GpuBackend,
) -> Result<()>
fn release_state( &self, state: &mut dyn LayerState, gpu: &dyn GpuBackend, ) -> Result<()>
2026-09-25: Frees a Glm5NextDsaState’s device buffers; the indexer cache is sized for
the configured maximum context (max_dsa_context), not the prompt. Any other state is
left alone, whatever the mixer: on the model path a KDA layer’s SsmLayerState holds SSM
pool addresses, which this layer does not own.
Source§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( &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 with the highway indexed by token: prompt token t uses highway slot
t, so every token’s streams are still there when the next layer reads them. The trait’s
default runs each token through decode, which uses slot 0 for all of them.
Errors when num_tokens exceeds ctx.buffers.max_batch_tokens(), the number of highway
slots in the buffer arena.
Source§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_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: K tokens of one sequence in one call, the speculative-verify body:
forward_k over highway slots 0..K, taking KDA state snapshots. The trait’s default
calls decode per token, and decode uses highway slot 0 for every token.
Source§fn alloc_state(&self, gpu: &dyn GpuBackend) -> Result<Box<dyn LayerState>>
fn alloc_state(&self, gpu: &dyn GpuBackend) -> Result<Box<dyn LayerState>>
trait_impl/meta.rs) calls it for every layer except a linear-attention layer that
uses the SSM pool, which gets pool addresses instead.Source§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 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<()>
hidden in place. Read moreSource§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<(), Error>
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<(), Error>
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.Source§fn decode_multi_seq<'a, 'b>(
&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<(), Error>where
'b: 'a,
fn decode_multi_seq<'a, 'b>(
&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<(), Error>where
'b: 'a,
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).Source§fn decode_verify_multi<'a, 'b>(
&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<(), Error>where
'b: 'a,
fn decode_verify_multi<'a, 'b>(
&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<(), Error>where
'b: 'a,
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.