Struct NllbGpuModel 

pub struct NllbGpuModel { /* private fields */ }

2026-09-25: The served NLLB encoder-decoder model. Fields are immutable after construction, behind a Mutex, or atomic. The decode scratch (dec) and the LoRA gate (lora_active) are shared by every sequence, so callers must not forward two sequences at once.

Implementations

impl NllbGpuModel

pub fn new( config: &ModelConfig, store: &WeightStore, gpu: Box<dyn GpuBackend>, lang: NllbLang, max_seq_len: usize, max_batch: usize, lora_dir: Option<&Path>, ) -> Result<Self>

2026-09-25: Build from the standard --model weight store and GPU backend. lang carries the language ids the server’s tokenizer resolved. max_seq_len sets the decoder KV depth, clamped to DEFAULT_CACHE_ROWS..=2048. Fails unless model.shared.weight is bf16.

Trait Implementations

impl ModelAdapters for NllbGpuModel

fn set_active_lora(&mut self, _name: &str) -> Result<()>

2026-09-25: Make the resident LoRA adapter name the active one. Takes &mut self, so it runs only while nothing else holds the model. Default: an error.

fn adapter_id_for(&self, _slot: i32) -> u64

2026-09-25: The prefix-cache adapter id for a LoRA slot selector: >= 0 names a resident slot, -1 the active adapter. Default 0, the base model’s id.

fn acquire_adapter_slot(&self, _slot: i32) -> i32

2026-09-25: Take a ref on the LoRA slot a sequence will use (called at prefill), resolving -1 to the active adapter as Self::adapter_id_for does. Returns the resolved slot, which the caller must pass to Self::release_adapter_slot, or -1 when no ref was taken. Default -1.

fn release_adapter_slot(&self, _resolved: i32)

2026-09-25: Release a ref taken by Self::acquire_adapter_slot, by the slot it returned. Default: no-op.

fn swap_lora_from_disk( &mut self, _dir: &Path, _name: &str, _slot: usize, ) -> Result<()>

2026-09-25: Load the LoRA adapter at dir into pool slot slot. Takes &mut self, so it runs only while nothing else holds the model. Default: an error.

fn promote_lora_from_peer( &mut self, _peer_addr: &str, _adapter_id: &str, _name: &str, _peft: PeftAdapterConfig, ) -> Result<(usize, Option<String>)>

2026-09-25: Fetch the adapter name (staged on peer_addr as adapter_id) from the peer into a pool slot and make it active; returns (slot, evicted adapter name). peft supplies the rank, alpha and scaling. Default: an error.

fn promote_lora_from_disk( &mut self, _adapter_dir: &Path, _name: &str, ) -> Result<(usize, Option<String>)>

2026-09-25: Load the adapter name from adapter_dir into a pool slot and make it active; returns (slot, evicted adapter name). Default: an error.

impl ModelCircuit for NllbGpuModel

fn set_forward(&self, sel: &ForwardSelect) -> Result<()>

2026-09-28: Make decode run sel from the next step. Captured decode graphs are dropped. Default: only Legacy is accepted.

fn forward_disclosure(&self) -> ForwardDisclosure

2026-09-28: The forward decode runs now. Default: legacy.

impl ModelDeviceFeed for NllbGpuModel

fn supports_device_token_feed(&self) -> bool

2026-09-25: Whether the model can run a decode step whose input ids live on the device (decode_batch_fed) and answer it with a masked argmax written into the feed cells (argmax_batch_to_feed). Default false.

fn decode_batch_fed( &self, _sources: &[FeedSource], _seqs: &mut [&mut SequenceState], _stream: u64, ) -> Result<DevicePtr>

2026-09-25: The decode step over seqs with each row’s input id read on the device from sources: the forward, KV writes and block allocation of decode_batch, but it does not push the token or advance seq_len, because the router does not know a fed row’s id at launch; the scheduler core does that once the previous step is committed. Position seq.seq_len is written. Default: an error.

fn argmax_batch_to_feed( &self, _logits: DevicePtr, _masks: &[RowMask], _dst: *mut u32, _event: u64, _stream: u64, ) -> Result<()>

2026-09-25: Argmax over masks.len() rows of logits, applying each row’s two-id mask (argmax_feed.cu): writes the feed cells, copies the ids to the pinned dst with an async D2H on stream, then records event. It never synchronises; dst is valid once event completes. Default: an error.

fn reserve_decode_block(&self, _seq: &mut SequenceState) -> Result<usize>

2026-09-25: Allocate the KV block for the next decode position (seq.seq_len) if the block table lacks it; returns the blocks added (0 or 1). Err when the pool is empty. Default: an error.

fn release_decode_blocks( &self, _seq: &mut SequenceState, _blocks: usize, ) -> Result<()>

2026-09-25: Return the last blocks entries of the block table, added by reserve_decode_block for a discarded step, to the pool. Default: an error.

fn fed_step_settled(&self)

2026-09-25: A fed step has completed on the device, so the CUDA graph it replayed may be evicted. Default: no-op.

fn event_query(&self, _event: u64) -> Result<bool>

2026-09-25: Whether the work recorded against event has completed; never blocks. Default Ok(true).

fn event_synchronize(&self, _event: u64) -> Result<()>

2026-09-25: Block the host until the work recorded against event has completed. Default: Ok(()).

fn destroy_event(&self, _event: u64) -> Result<()>

fn copy_d2h_async( &self, _src: DevicePtr, _dst: &mut [u8], _stream: u64, ) -> Result<()>

2026-09-25: Async device-to-host copy on stream; dst must stay valid and unread until stream is synchronised. Default: an error.

fn alloc_host_pinned(&self, _bytes: usize) -> Result<*mut u8>

2026-09-25: Zeroed page-locked host memory for the router’s readback ring. Default: an error.

fn free_host_pinned(&self, _ptr: *mut u8, _bytes: usize) -> Result<()>

impl ModelDraft for NllbGpuModel

fn generate_speculative( &self, _tokens: &[u32], _params: &SamplingParams, _num_drafts: usize, ) -> Result<GenerateResult>

2026-09-25: Generate for one prompt with the model’s own draft proposer.

fn has_proposer(&self) -> bool

2026-09-25: Whether a draft proposer or self-speculative decoding is available.

fn has_self_speculative(&self) -> bool

fn decode_draft( &self, _token: u32, _seq: &mut SequenceState, _stream: u64, ) -> Result<DevicePtr>

2026-09-25: The draft decode step for self-speculative decoding; returns logits.

fn save_hidden_for_mtp(&self, _token_idx: usize, _stream: u64) -> Result<()>

2026-09-25: Copy the hidden state at token_idx into the MTP input buffer, which run_mtp_propose and run_mtp_propose_multi read. TransformerModel copies the hidden state before the final norm, because the MTP head applies its own norm.

fn run_mtp_propose( &self, _token: u32, _position: usize, _seq: &mut SequenceState, _stream: u64, ) -> Result<Option<u32>>

2026-09-25: One MTP draft from the saved hidden state; None without a proposer.

fn run_mtp_propose_multi( &self, _token: u32, _position: usize, _num_drafts: usize, _seq: &mut SequenceState, _stream: u64, _grammar_bitmask: Option<&[i32]>, ) -> Result<Vec<u32>>

2026-09-25: num_drafts MTP drafts from the hidden state saved by save_hidden_for_mtp.

fn trim_proposer_state( &self, _seq: &mut SequenceState, _num_accepted: usize, _stream: u64, ) -> Result<()>

2026-09-25: Tell the proposer how many drafts the last verify accepted, so it can trim its per-sequence state.

fn mtp_slot_draft_capacity(&self, _slot_idx: usize) -> usize

2026-09-25: The most drafts a speculative step may send to the sequence in SSM pool slot slot_idx without overflowing that slot’s per-token intermediate state. The scheduler uses the minimum over the active slots. Default usize::MAX (no SSM verify pools).

fn dflash_gamma(&self) -> Option<usize>

2026-09-25: The installed DFlash drafter’s block size γ, read from the drafter; the serve layer derives num_drafts = γ - 1 from it. Default None.

fn stash_verify_hidden_rows(&self, rows: &[usize], stream: u64) -> Result<()>

2026-09-25: Copy hidden row rows[i] of the last batched verify into stash slot i, before a propose overwrites the shared hidden buffer. Default: an error.

fn stash_verify_catchup_rows(&self, slot_rows: &[(usize, usize)]) -> Result<()>

2026-09-25: Copy batched-verify rows into drafter catch-up stash slots, slot_rows[j] = (slot, row). Default: Ok(()).

fn run_mtp_catchup_batched( &self, tokens: &[Vec<u32>], first_slot: &[usize], first_pos: &[usize], seqs: &mut [&mut SequenceState], ) -> Result<usize>

2026-09-25: Append one drafter row per accepted draft before the next propose (see Self::stash_verify_catchup_rows); returns the rows written. Default: 0.

fn save_hidden_for_mtp_from_stash(&self, idx: usize, stream: u64) -> Result<()>

2026-09-25: Self::save_hidden_for_mtp from stash slot idx (written by Self::stash_verify_hidden_rows) instead of the live verify rows, which a propose may have overwritten. Default: an error.

fn run_mtp_propose_batched( &self, tokens: &[u32], positions: &[usize], stash_idx: &[usize], num_drafts: usize, seqs: &mut [&mut SequenceState], stream: u64, out_conf: Option<&mut Vec<Vec<f32>>>, ) -> Result<Option<Vec<Vec<u32>>>>

2026-09-25: num_drafts MTP drafts for each of tokens.len() sequences in one drafter forward per draft position. stash_idx[i] is the verify-stash slot holding sequence i’s hidden, and positions[i] its propose position, as for Self::run_mtp_propose_multi.

fn mtp_propose_batch_max(&self) -> usize

2026-09-25: The widest batch Self::run_mtp_propose_batched takes in one drafter forward. Default 1 (per sequence only).

fn save_hidden_for_catchup(&self, _token_idx: usize, _pos: usize) -> Result<()>

2026-09-25: Store a serially decoded token’s final hidden at pos in the drafter catch-up ring. Default: Ok(()).

fn save_dflash_hidden_for_propose( &self, _token_idx: usize, _stream: u64, ) -> Result<()>

2026-09-25: Capture row token_idx from every DFlash capture layer for the next propose. Default: Ok(()).

fn dflash_accept_append(&self, _seq: &mut SequenceState) -> Result<()>

2026-09-25: Append row 1 of the DFlash hidden scratch to the proposer context. Default: Ok(()).

fn dflash_eagle_accept_append(&self, _seq: &mut SequenceState) -> Result<()>

2026-09-25: Append scratch rows 0 and 1 to the DFlash context at positions seq_len - 2 and seq_len - 1, row 1 last, before the propose. Default: Ok(()).

fn dflash_eagle_kgamma_append( &self, _seq: &mut SequenceState, _num_accepted: usize, _base_pos: usize, ) -> Result<()>

2026-09-25: Append scratch rows 0..=num_accepted to the DFlash context at positions base_pos..=base_pos + num_accepted, row num_accepted last. Default: Ok(()).

fn dflash_serial_ctx_append(&self, _seq: &mut SequenceState) -> Result<()>

2026-09-25: Append the just-decoded token’s captured hidden (scratch row 0) to the DFlash context at position seq_len - 1, for a sequence decoded serially while speculation is suspended; otherwise that capture is overwritten and the context has a hole. It sets skip_next_decode_append so a later propose does not append it again. No-op without a DFlash proposer state. Default: Ok(()).

fn commit_ctx( &self, _seq: &mut SequenceState, _num_committed: usize, _base_pos: usize, _scratch_row: usize, ) -> Result<()>

2026-09-25: Append num_committed DFlash scratch rows, from scratch_row, at the end of the context, stamped with positions base_pos..base_pos + num_committed, after sliding the window if it would overflow. base_pos is a position, not a context row: the two differ after a slide. Default: Ok(()).

fn dflash_capture_band(&self) -> usize

2026-09-25: Rows per sequence band in the DFlash hidden scratch: sequence i of a batched verify captures into band i, so its commit_ctx scratch_row is i * dflash_capture_band(). Default 0 (no DFlash drafter).

fn read_deferred_draft_token(&self) -> Result<u32>

2026-09-25: The draft id the proposer left on the device. Default 0.

impl ModelEp for NllbGpuModel

fn ep_worker_step(&self, _slots: &mut [Option<SequenceState>]) -> Result<bool>

2026-09-25: EP worker step, on a rank above 0: receive a (seq_id, cmd) from rank 0 and run the command in slot seq_id. Returns Ok(false) when the worker should shut down.

fn is_ep(&self) -> bool

2026-09-25: Whether the model runs the multi-rank command protocol. The scheduler then turns off its single-GPU fused paths, the mixed forward among them. Default false.

fn ep_broadcast_cmd(&self, _cmd: u32) -> Result<()>

2026-09-25: Send a command word to every worker rank. Default: Ok(()).

fn ep_broadcast_cmd_for_seq(&self, _seq_id: u32, _cmd: u32) -> Result<()>

2026-09-25: Send a (seq_id, cmd) pair to every worker rank, as the first broadcast of a command; its follow-up broadcasts use Self::ep_broadcast_cmd. When Self::ep_protocol_v2 is false, seq_id is not sent. Default: Ok(()).

fn ep_protocol_v2(&self) -> bool

2026-09-25: Whether the EP protocol carries a seq_id with each command (METRALE_EP_PROTOCOL=v2). Default false.

fn ep_broadcast_tokens(&self, _tokens: &[u32]) -> Result<Vec<u32>>

2026-09-25: Broadcast a token array to every worker rank in one broadcast. Default: an empty Vec.

impl ModelForward for NllbGpuModel

fn decode_batch( &self, tokens: &[u32], seqs: &mut [&mut SequenceState], _stream: u64, ) -> Result<DevicePtr>

2026-09-25: Batched decode: one forward_one per sequence, serially on the default stream, into contiguous logit rows 0..n (batch position i ↔ seqs[i]); returns the base of [n, vocab]. Each sequence’s KV is found by its slot_idx. Fails when n is 0, differs from tokens.len(), or exceeds max_batch.

fn prefill( &self, tokens: &[u32], seq: &mut SequenceState, _stream: u64, ) -> Result<DevicePtr>

2026-09-25: Prefill the whole prompt and return the logits of its last position.

fn prefill_chunk( &self, tokens: &[u32], seq: &mut SequenceState, chunk_start: usize, chunk_len: usize, is_last_chunk: bool, _stream: u64, ) -> Result<DevicePtr>

2026-09-25: Prefill chunk_len prompt tokens from chunk_start. Only a chunk with is_last_chunk returns logits (of the last prompt position); any other chunk returns DevicePtr::NULL.

fn decode( &self, token: u32, seq: &mut SequenceState, _stream: u64, ) -> Result<DevicePtr>

2026-09-25: One decode step for one sequence; returns the logits for the next token.

fn supports_beam(&self) -> bool

2026-09-25: Whether this model implements Self::generate_beam_batch. Default false; NllbGpuModel returns true.

fn generate_beam_batch(&self, reqs: &[BeamReq]) -> Result<Vec<Vec<u32>>>

2026-09-25: Run beam search to completion for each request and return each winning hypothesis. The scheduler calls it at prefill for num_beams > 1, instead of the token-by-token decode loop. Default: an error.

fn mixed_forward( &self, decode_tokens: &[u32], decode_seqs: &mut [&mut SequenceState], prefill_tokens: &[u32], prefill_seq: &mut SequenceState, prefill_chunk_start: usize, prefill_chunk_len: usize, prefill_is_last: bool, stream: u64, ) -> Result<MixedForwardResult>

2026-09-25: One decode step for each of decode_seqs plus one prefill chunk. Default: decode_batch (skipped when there are no decode tokens), then prefill_chunk.

fn prefill_batch_chunk( &self, streams: &mut [PrefillSlice<'_>], stream: u64, ) -> Result<Vec<DevicePtr>>

2026-09-25: Prefill one chunk for each of streams. Returns one entry per stream, in order: the last-token logits when that stream’s chunk is its last, else DevicePtr::NULL. Default: prefill_chunk per stream.

fn prefill_batch_chunk_rows( &self, streams: &mut [PrefillSlice<'_>], stream: u64, _row_base: usize, ) -> Result<Vec<DevicePtr>>

2026-09-25: Self::prefill_batch_chunk, with finishing stream i’s logits written to row row_base + i of the shared logits arena instead of row i.

fn mixed_forward_batch( &self, decode_tokens: &[u32], decode_seqs: &mut [&mut SequenceState], prefill_streams: &mut [PrefillSlice<'_>], stream: u64, ) -> Result<MixedBatchResult>

2026-09-25: One decode step for each of decode_seqs plus one prefill chunk for each of prefill_streams. Default: decode_batch, a synchronise of the default stream, then prefill_batch_chunk_rows with the prefill rows placed after the decode rows.

fn normalize_ssm_states(&self, _seq: &SequenceState, _stream: u64) -> Result<()>

2026-09-25: Normalise the sequence’s SSM h_state. The scheduler calls it between prefill chunks. Default: no-op.

fn prefill_twophase( &self, tokens: &[u32], seq: &mut SequenceState, _chunk_size: usize, stream: u64, ) -> Result<DevicePtr>

2026-09-25: Prefill the whole prompt in one call, with chunk_size as the chunk length when an implementation splits it. Returns the last position’s logits. Default: one prefill_chunk over all of tokens.

fn is_mla(&self) -> bool

2026-09-25: Whether chunked prefill must run as one chunk. TransformerModel answers requires_single_chunk_prefill: MLA models (kv_lora_rank > 0) except glm5_next and kimi_k3. Default false.

fn hc_mult(&self) -> usize

2026-09-25: mHC hyper-connection stream count (config.hc_mult in TransformerModel); 0 means no highway. Default 0.

fn kv_block_size(&self) -> Option<usize>

2026-09-25: Tokens per paged-KV block, or None without a paged KV cache. The scheduler uses it to end a prefill chunk at metrale_gpu_runtime::ssm_tail_boundary. Default None.

fn prefill_tail_split(&self, _tokens: &[u32]) -> Option<usize>

2026-09-27: The token at which the prefill of tokens is split so that an SSM snapshot lands at prefill_plan::tail_split_point, or None when this model does not split it. The scheduler ends a non-last chunk there (prefill_plan::plan_chunk_len). Default None.

impl ModelLifecycle for NllbGpuModel

fn bind_gpu_to_thread(&self) -> Result<()>

2026-09-25: Bind the GPU context to the current thread. The scheduler thread and the EP worker call it before using the model.

fn alloc_sequence(&self) -> Result<SequenceState>

2026-09-25: A new sequence with its layer states.

fn free_sequence(&self, seq: &mut SequenceState) -> Result<()>

2026-09-25: Release the sequence’s KV blocks, prefix-cache refs, SSM pool slot and LoRA slot ref.

fn compact_sequence( &self, seq: &mut SequenceState, new_slot: usize, ) -> Result<()>

2026-09-25: Move the sequence’s SSM state (h_state and conv_state of every SSM layer) to pool slot new_slot.

fn detach_slot_for_reuse(&self, seq: &mut SequenceState)

2026-09-25: Give up a retired sequence’s SSM pool slot after compact_sequence moved a surviving sequence into it: sets slot_idx to usize::MAX and empties the slot guard, so freeing or dropping this sequence does not release the slot. Call it right after that compact_sequence, before any step that can drop the sequence.

fn cache_sequence(&self, _seq: &SequenceState)

2026-09-25: Insert the sequence’s tokens (prompt and generated) into the prefix cache. Call before Self::free_sequence, while its blocks are still valid.

fn num_free_blocks(&self) -> usize

2026-09-25: Free KV blocks. Default 0.

fn kv_block_bytes(&self) -> usize

2026-10-01: Bytes one block index of the paged cache holds across every layer, K and V, so num_total_blocks() * kv_block_bytes() is the pool’s size. 0 for a model without a paged cache. No default: a model that has a pool must say what its blocks cost.

fn teardown(&mut self) -> Result<()>

2026-09-25: Release the device memory this model owns, newest first. The scheduler calls it once at shutdown (core/finish.rs). Default: Ok(()), for models that own no pooled device memory; a model that owns pools and keeps the default leaks them.

fn poll_innerq(&self)

2026-09-25: Poll this model’s TQ+ InnerQ calibration driver; the scheduler calls it once per prefill step. Default: no-op.

fn high_speed_swap_dims(&self) -> Option<ModelDims>

2026-09-25: Model dimensions for the --high-speed-swap orchestrator; None for a model that does not support it (the default).

fn alloc_sequence_for(&self, budget_tokens: usize) -> Result<SequenceState>

2026-09-25: Self::alloc_sequence for a request that can reach at most budget_tokens (prompt_len + max_tokens), so proposer state that scales with context can be sized to that (DraftProposer::alloc_state_for). Default: Self::alloc_sequence.

fn save_sequence_state( &self, _seq: &SequenceState, _writer: &mut dyn Write, ) -> Result<()>

2026-09-25: Write the sequence’s KV blocks and SSM state to writer, in a format the model owns, without freeing anything. Default: an error.

fn restore_sequence_state( &self, _seq: &mut SequenceState, _num_blocks: usize, _reader: &mut dyn Read, ) -> Result<()>

2026-09-25: Read state written by Self::save_sequence_state into an allocated sequence, allocating num_blocks KV blocks. Default: an error.

fn num_total_blocks(&self) -> usize

2026-09-25: Total KV blocks in the paged cache. Default 0.

fn reclaim_prefix_blocks(&self, _num_blocks: usize) -> usize

2026-09-25: Evict up to num_blocks blocks from the prefix cache; returns how many became free. Swap-in and preemption resume check num_free_blocks() before they allocate, so they call this to reach blocks the cache holds. 0 means nothing was evictable, and the caller stops asking. Default 0.

fn reclaim_prefix_blocks_for( &self, prompt: &[u32], adapter_slot: i32, blocks_needed: usize, ) -> PrefixReclaim

2026-09-30: Reclaim prefix-cache blocks until num_free_blocks() covers what a prefill of prompt under adapter_slot allocates: blocks_needed (the caller’s count for the whole prompt, headroom included) less the blocks of the cached prefix that prefill will reuse. An empty prompt means the prefill reuses nothing (prompt logprobs). Admission and preemption resume call it before they allocate. Default: no prefix is reused, and Self::reclaim_prefix_blocks runs until blocks_needed is free or it returns 0.

impl ModelLogits for NllbGpuModel

fn vocab_size(&self) -> usize

fn copy_logits_to_host( &self, logits_ptr: DevicePtr, dst: &mut [u8], ) -> Result<()>

2026-09-25: Copy one row of logits to the host; dst is vocab_size * 2 bytes for BF16 logits.

fn logits_buffer_ptr(&self) -> DevicePtr

2026-09-25: The device logits buffer, where a verify pass leaves its [k, vocab] rows.

fn argmax_on_device(&self, logits_ptr: DevicePtr, _stream: u64) -> Result<u32>

2026-09-25: Argmax of one logits row, computed on the device.

fn argmax_batch( &self, logits_ptr: DevicePtr, n: usize, _stream: u64, ) -> Result<Vec<u32>>

2026-09-25: Argmax of each of n logits rows on the device.

fn hidden_after_norm(&self) -> DevicePtr

2026-09-25: The post-final-norm hidden state of the last decode step.

fn logits_ptr_is_fp32(&self, _logits_ptr: DevicePtr) -> bool

2026-09-25: Whether logits_ptr holds FP32 logits (vocab * 4 bytes per row). Default false; TransformerModel returns false too, because its use_fp32_logits is fixed to false in TransformerModel::new.

fn decode_logits_fp32(&self) -> bool

2026-09-25: Whether single-token decode writes FP32 logits, read from Self::decode_logits_ptr at 4 bytes per element. Default false; TransformerModel returns its use_fp32_logits, which is always false.

fn decode_logits_ptr(&self) -> DevicePtr

2026-09-25: The buffer single-token decode last wrote its logits to: FP32 when Self::decode_logits_fp32 is true, else BF16. The default panics, so a model must override it; TransformerModel does.

impl ModelSsmState for NllbGpuModel

fn checkpoint_ssm_states(&self, _seq: &mut SequenceState) -> Result<()>

fn rollback_ssm_states(&self, _seq: &mut SequenceState, _n: usize) -> Result<()>

2026-09-25: Restore the SSM state to where num_accepted tokens of the last verify were accepted.

fn has_ssm_layers(&self) -> bool

2026-09-25: Whether the model has recurrent SSM layers, whose state is advanced in place each token. Lowering seq_len rewinds paged KV but not that state, so a decode rollback must also restore it from a snapshot. Default false.

fn decode_rollback_unsupported(&self) -> bool

2026-09-26: Whether some layer keeps per-sequence state that lowering the KV cursor does not rewind (LayerCapabilities::decode_rollback_unsupported); the scheduler then declines a decode rollback (rollback_to_boundary). Default false.

fn decode_rollback_ring_slots(&self) -> usize

2026-09-25: Decode-rollback SSM snapshot slots per active sequence; the scheduler sizes each sequence’s snapshot ring from it. 0 (the default) means no decode-rollback snapshots. TransformerModel returns the depth ssm_reserve::decode_rollback_ring_slots chose.

fn save_decode_ssm_snapshot( &self, _seq: &SequenceState, _ring_slot: usize, ) -> Result<()>

2026-09-25: Save seq’s SSM h_state and conv_state (every SSM layer) into its decode-rollback ring slot ring_slot, in 0..decode_rollback_ring_slots(). Default: Ok(()).

fn restore_decode_ssm_snapshot( &self, _seq: &SequenceState, _ring_slot: usize, ) -> Result<()>

2026-09-25: Restore seq’s SSM state from ring slot ring_slot, written earlier by Self::save_decode_ssm_snapshot. Default: Ok(()).

fn decode_marconi_checkpoint(&self, _seq: &mut SequenceState)

2026-09-25: During decode, save a block-aligned SSM snapshot at a checkpoint boundary, so the next turn’s prefix-cache hit can restore from decode-produced state. The scheduler calls it after a step’s SSM state is final. Default: no-op.

fn start_checkpoint_async(&self, seq: &mut SequenceState) -> Result<()>

2026-09-25: Start the SSM checkpoint copies on the secondary stream; call Self::sync_secondary before the next verify. Default: the synchronous Self::checkpoint_ssm_states.

fn start_rollback_and_checkpoint_async( &self, seq: &mut SequenceState, num_accepted: usize, ) -> Result<()>

2026-09-25: Roll the SSM state back to num_accepted and checkpoint it again, on the secondary stream. Default: the synchronous rollback, then checkpoint.

fn sync_secondary(&self) -> Result<()>

2026-09-25: Wait for the secondary stream’s work. Default: Ok(()).

fn commit_accepted_prefix( &self, _seq: &mut SequenceState, _num_accepted: usize, _k: usize, ) -> Result<()>

2026-09-25: Commit the accepted prefix of a verify onto the live SSM h_state and conv_state, on the secondary stream (pair with Self::sync_secondary). Default: Ok(()).

fn ssm_snapshot_occupancy(&self) -> Option<(u32, u32)>

2026-09-25: SSM snapshot-pool occupancy (used, total), or None without a snapshot pool (the default; TransformerModel also returns None when the pool has no slots).

impl ModelStreams for NllbGpuModel

fn default_stream(&self) -> u64

2026-09-25: The default stream. Default 0.

fn create_stream(&self) -> Result<u64>

2026-09-25: A new stream. Default Ok(0).

fn create_event(&self) -> Result<u64>

2026-09-25: A new event. Default Ok(0).

fn record_event(&self, _event: u64, _stream: u64) -> Result<()>

2026-09-25: Record event on stream. Default: Ok(()).

fn stream_wait_event(&self, _stream: u64, _event: u64) -> Result<()>

2026-09-25: Make stream wait for event on the device; the host does not block. Default: Ok(()).

fn synchronize(&self, _stream: u64) -> Result<()>

2026-09-25: Block the host until stream’s work has completed. Default: Ok(()).

impl ModelVerify for NllbGpuModel

fn decode_verify( &self, _t: &[u32], _s: &mut SequenceState, _st: u64, ) -> Result<Vec<u32>>

2026-09-25: Run tokens through the model as one verify pass and return the argmax id at each position; every token advances the KV and SSM state.

fn decode_verify_graphed( &self, _t: &[u32; 2], _s: &mut SequenceState, _st: u64, ) -> Result<[u32; 2]>

2026-09-25: K=2 verify ([last, draft]); returns the argmax id at each position.

fn decode_verify_graphed_k3( &self, _t: &[u32; 3], _s: &mut SequenceState, _st: u64, ) -> Result<[u32; 3]>

2026-09-25: K=3 verify (1 verified token and 2 drafts); returns 3 argmax ids.

fn decode_verify_graphed_k4( &self, _t: &[u32; 4], _s: &mut SequenceState, _st: u64, ) -> Result<[u32; 4]>

2026-09-25: K=4 verify (1 verified token and 3 drafts); returns 4 argmax ids.

fn can_batch_verify(&self, _ks: &[usize]) -> bool

2026-09-25: Whether Self::decode_verify_batched can run ks.len() sequences with ks[i] verify rows each (the sequence’s draft count plus one; the rows may differ per sequence). Default false, and there is no default batched implementation: a loop of per-sequence verifies would leave only the last sequence’s rows in the shared logits buffer.

fn decode_verify_batched( &self, tokens: &[u32], ks: &[usize], seqs: &mut [&mut SequenceState], stream: u64, opts: VerifyBatchedOpts, ) -> Result<Vec<u32>>

2026-09-25: Verify ks.len() sequences in one forward. Rows are flat and sequence-major (tokens.len() == Σ ks): sequence i holds [last_verified, d0, ..] in rows off_i..off_i + ks[i], with off_i = Σ_{t<i} ks[t]. Returns the Σ ks argmax ids in the same order. On success each sequence advances by its own ks[i], and the caller rewinds by its verdict.

fn decode_verify_graphed_kgamma( &self, tokens: &[u32], seq: &mut SequenceState, stream: u64, ) -> Result<Vec<u32>>

2026-09-25: DFlash verify of tokens (1 verified token and γ drafts). Default: Self::decode_verify.

fn decode_verify_dflash( &self, tokens: &[u32], seq: &mut SequenceState, stream: u64, ) -> Result<Vec<u32>>

2026-09-25: DFlash verify. Default: Self::decode_verify_graphed_kgamma.

fn decode_and_verify_fused( &self, tokens: &[u32], seq: &mut SequenceState, stream: u64, ) -> Result<Vec<u32>>

2026-09-25: DFlash decode and verify in one forward: tokens[0] is the accepted token and tokens[1..] the draft block. Default: Self::decode_verify_graphed_kgamma.

fn gdn_fold_accepted( &self, _slots: &[usize], _accepted_rows: &[u32], _k_rows: usize, ) -> Result<bool>

2026-09-25: After a batched verify run with write_on_accept, commit the GDN h state of each slot slots[i] at accepted_rows[i] verify rows (anchor included), in the verify’s batch order. Ok(true) means the h states were committed here, and commit_accepted_prefix then restores only conv state for those slots. Ok(false) means nothing was folded. Default: Ok(false).

impl ModelVision for NllbGpuModel

fn prepare_vision_embed(&self, _images: &[VisionItem]) -> Result<()>

2026-09-25: Encode images with the vision encoder and keep the embeddings for the next prefill. Call it before prefill_chunk for a prompt with image pad tokens. Default: Ok(()).

fn ep_sync_vision_embeds(&self, _tokens: &[u32]) -> Result<()>

2026-09-25: Give every rank the same merged vision rows for this prompt. Otherwise only rank 0 has the encoder’s rows, the other ranks keep the pad-token embedding, and the per-layer all-reduce mixes the two. The head calls it after the prompt broadcast and the worker at the matching point, so the collective pairs. Default: Ok(()).

fn prepare_vision_embed_batched( &self, _per_request: &[Vec<VisionItem>], ) -> Result<Vec<(usize, usize, usize, usize)>>

2026-09-25: Encode the images of several requests in one encoder pass. Returns, per request in order, (patch_row_offset, grid_index_offset, num_images, patch_row_count) in the shared output buffer. Default: an empty Vec.

fn set_vision_slice_base( &self, _row_base: usize, _grid_base: usize, _owned_images: usize, )

2026-09-25: Where the next prefill_chunk reads its vision rows in the batched encoder output: row offset, grid index offset and image count; (0, 0, 0) resets. Default: no-op.

fn tokens_contain_vision_pad(&self, _tokens: &[u32]) -> bool

2026-09-25: Whether tokens holds a vision pad token, whose KV came from image embeddings that re-prefilling the tokens cannot reproduce. Decode preemption then does not pick the sequence for requeue-and-re-prefill unless it can be spilled. Default false.

impl Model for NllbGpuModel