Struct TransformerModel 

pub struct TransformerModel { /* private fields */ }

2026-09-25: One loaded model: its layers, weights, kernel handles, KV and SSM pools and CUDA-graph caches. trait_impl implements Model for it.

Implementations

impl TransformerModel

pub fn new( config: ModelConfig, embed_tokens: DenseWeight, final_norm: DenseWeight, lm_head_weight: DenseWeight, lm_head_nvfp4: Option<QuantizedWeight>, lm_head_fp8: Option<Fp8DenseWeight>, mtp_lm_head_nvfp4: Option<QuantizedWeight>, lm_head_nvfp4_t: Option<(QuantizedWeight, u32)>, layers: Vec<Box<dyn TransformerLayer>>, buffers: BufferArena, kv_cache: PagedKvCache, mtp_weights: Vec<MtpWeights>, gpu: Box<dyn GpuBackend>, max_seq_len: usize, max_batch_size: usize, mtp_quant: MtpQuantization, use_speculative: bool, prefix_cache: Box<dyn PrefixCache>, mtp_vocab_size: u32, comm: Option<Arc<dyn CommBackend>>, self_speculative: bool, num_drafts: usize, vision_encoder: Option<VisionTower>, ssm_cache_slots: usize, ssm_checkpoint_interval: usize, ) -> Result<Self>

impl TransformerModel

pub fn decode_logits_fp32(&self) -> bool

2026-09-25: Whether the single-token decode lm_head writes FP32 logits to logits_fp32_buf (use_fp32_logits, which TransformerModel::new sets to false). Other lm_head paths always write BF16.

pub fn decode_logits_ptr(&self) -> DevicePtr

2026-09-25: The buffer the single-token decode lm_head writes: logits_fp32_buf when use_fp32_logits, else the shared BF16 logits buffer. Read it with the dtype decode_logits_fp32 reports.

impl TransformerModel

pub fn gpu_backend(&self) -> &dyn GpuBackend

2026-09-25: Borrow the GPU backend. The factory uses it to build the proposers installed after construction and to read allocation reports.

pub fn config_ref(&self) -> &ModelConfig

2026-09-25: Borrow the model config, for building the proposers that are installed after construction.

pub fn set_dflash_proposer(&mut self, proposer: Arc<dyn DraftProposer>)

2026-09-25: Install a proposer built after construction (the DeepSeek-V4 or GLM-5.3 MTP head, or a DFlash drafter), replacing any MTP proposer TransformerModel::new built. It also shrinks the MTP prefill capture buffer to the rows this proposer can use.

It does not check flag combinations; clap rejects --dflash with --speculative.

pub fn set_ngram_embedding(&mut self, ngram: NgramEmbedding)

2026-09-25: Install the fused n-gram input embedding (only the LongCat loader builds one). Once set, embed_ctx embeds through it and the bare-token embed refuses to run.

pub fn has_ngram_embedding(&self) -> bool

2026-09-25: True when this model fuses n-gram lookups into its input embedding.

impl TransformerModel

pub fn adapter_id_for_slot(&self, slot: i32) -> u64

2026-09-25: Adapter id for a request’s adapter_slot (< 0 resolves to the active slot): a hash of the slot’s name and generation, or 0 when no LoRA pool is resident or the slot is empty or out of range.

pub fn acquire_adapter_slot(&self, slot: i32) -> i32

2026-09-25: Take a ref on the pool slot a sequence uses (< 0 resolves to the active slot, as in Self::adapter_id_for_slot). Returns the resolved index; the caller stores it and releases that index at terminal free, so a rotate that changes active in between cannot move the release to another counter. Returns -1 (nothing acquired) when no LoRA pool is resident or the slot is out of range.

pub fn release_adapter_slot(&self, resolved: i32)

2026-09-25: Release a ref taken by Self::acquire_adapter_slot, by the resolved index it returned. -1 and no pool are no-ops.

pub fn set_lora_weights(&mut self, lora: Option<LoraWeights>) -> Result<()>

impl TransformerModel

pub fn rotate_lora_to(&mut self, name: &str) -> Result<()>

2026-09-25: Make the resident adapter name the active one: re-install its per-layer pairs with install_lora_layers, then destroy_lora_decode_graphs. Call it only at a scheduler-quiescent point.

Refuses, changing nothing, when no adapter is loaded, name is not resident, rotation is not armed (lora_rotatable), or the current active slot has in-flight sequences. If the install walk fails, the active slot has already switched.

impl TransformerModel

pub fn swap_lora_slot_from_peer( &mut self, peer_addr: &str, adapter_id: &str, adapter_name: &str, slot: usize, peft: PeftAdapterConfig, ) -> Result<()>

2026-09-25: Stage the adapter adapter_name (held by the peer at peer_addr as adapter_id) into pool slot in place and make it that slot’s resident adapter. Call it only at a scheduler-quiescent point.

Refuses when no pool is loaded, rotation is not armed (lora_rotatable), slot >= max_loras, or the slot has in-flight sequences. Otherwise the slot region is zeroed, the tensors land over RDMA, the slot’s pairs are rebuilt with the new r/scale, its tables are refreshed and its generation bumped. If the slot is active, its pairs are re-installed and the LoRA graph caches destroyed.

pub fn promote_lora_slot_from_peer( &mut self, peer_addr: &str, adapter_id: &str, adapter_name: &str, peft: PeftAdapterConfig, ) -> Result<(usize, Option<String>)>

2026-09-25: Promote the adapter adapter_name from the peer into a cache-region pool slot and make it active; returns (slot, evicted_name). Call it only at a scheduler-quiescent point.

The victim (select_victim_slot) is a never-filled placeholder first, else the least recently used idle (ref_count == 0) cache slot. With every cache slot busy it fails with POOL_FULL; a busy slot is never evicted. Self::swap_lora_slot_from_peer re-checks the ref count and bumps the slot generation.

pub fn promote_lora_slot_from_disk( &mut self, adapter_dir: &Path, name: &str, ) -> Result<(usize, Option<String>)>

2026-09-25: Self::promote_lora_slot_from_peer with the adapter loaded from adapter_dir (named name) instead of a peer: same victim choice, same make-active step, returns (slot, evicted_name). Self::swap_lora_slot_from_disk reads the directory’s adapter_config.json, refuses when rotation is not armed, re-checks the ref count and bumps the slot generation. Call it only at a scheduler-quiescent point.

pub fn swap_lora_slot_from_disk( &mut self, adapter_dir: &Path, name: &str, slot: usize, ) -> Result<()>

2026-09-25: Load the adapter at adapter_dir into pool slot in place and make it that slot’s resident adapter, using the directory’s own adapter_config.json. If the slot is active, its pairs are re-installed and the LoRA graph caches destroyed. Call it only at a scheduler-quiescent point.

Refuses when rotation is not armed (lora_rotatable) or the slot has in-flight sequences. pack_store_into_slot also refuses adapters with router/expert deltas or token-overlay tensors.

impl TransformerModel

pub fn install_declared_lm_head_w8a8(&mut self) -> Result<()>

2026-09-28: Install the W8A8 head when the declared-precision policy wants the head run W8A8 (WeightQuantPolicy::fp8_decode_act("lm_head") == Some(Fp8)) and the model loaded the FP8 head (lm_head_fp8, the declared head format under declared).

pub fn set_lm_head_w8a8(&mut self, ctx: W8a8Ctx) -> Result<()>

2026-09-28: Run the FP8 head W8A8 (the checkpoint’s declared precision) from now on. Needs lm_head_fp8, and kernels and scratch that serve one row.

impl TransformerModel

pub fn install_declared_lm_head_w4a16_rows(&mut self) -> Result<()>

2026-10-02: Install the W4A16 row-tile head when the declared-precision policy declares an NVFP4 head with 16-bit activations, the model loaded the NVFP4 head, the target ships the kernels, the head fits their shape contract, and the vocab is not sharded.

impl TransformerModel

pub fn set_lm_head_q6k(&mut self, blocks: DevicePtr) -> Result<()>

2026-09-25: Install the raw Q6_K head. blocks is the store’s lm_head.weight pointer, [vocab, hidden] in Q6_K super-blocks. Fails when hidden_size is not a multiple of 256.

impl TransformerModel

pub fn adopt_weight_store(&mut self, store: WeightStore)

2026-09-25: Hand the model the ledger of its own weights, for teardown.

Trait Implementations

impl Drop for TransformerModel

fn drop(&mut self)

Executes the destructor for this type.

impl ModelAdapters for TransformerModel

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, 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 TransformerModel

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 TransformerModel

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 TransformerModel

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 generate_speculative( &self, prompt_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 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 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 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 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 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 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_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 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_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 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_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 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 read_deferred_draft_token(&self) -> Result<u32>

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

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.

impl ModelEp for TransformerModel

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 TransformerModel

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

2026-09-25: prefill_batch_chunk_dispatch sends a batch that is ineligible, or whose cache plan is not admitted, to its per-stream loop before any sequence is mutated. An error from an admitted kernel batch is returned: a per-stream retry could apply prefix-cache and KV state twice.

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

2026-09-25: Mixed-step variant: finishing streams write their logits from row row_base, clear of the decode lanes. The trait docs describe the aliasing this prevents.

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 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 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 decode_batch( &self, tokens: &[u32], seqs: &mut [&mut SequenceState], stream: u64, ) -> Result<DevicePtr>

2026-09-25: One decode step for each sequence; returns logits [seqs.len(), vocab].

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

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

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

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

impl ModelLifecycle for TransformerModel

fn poll_innerq(&self)

2026-09-25: Poll this model’s InnerQ calibration driver, if it has one. A maybe_finalize error is logged as a warning and not returned.

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

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

fn num_total_blocks(&self) -> usize

2026-09-25: Total KV blocks in the paged cache. 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 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 TransformerModel

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

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

fn rollback_ssm_states( &self, seq: &mut SequenceState, num_accepted: 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 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).

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

impl ModelStreams for TransformerModel

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 TransformerModel

fn decode_verify( &self, tokens: &[u32], seq: &mut SequenceState, stream: 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, tokens: &[u32; 2], seq: &mut SequenceState, _stream: 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, tokens: &[u32; 3], seq: &mut SequenceState, _stream: 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, tokens: &[u32; 4], seq: &mut SequenceState, _stream: 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_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).

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.

impl ModelVision for TransformerModel

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

impl Send for TransformerModel

impl Sync for TransformerModel