Struct NllbGpuModel 

pub struct NllbGpuModel { /* private fields */ }

The served NLLB encoder-decoder model. Send + Sync: every field is either immutable after construction or behind a Mutex; DevicePtr is a Copy device handle and the GPU-side buffers are driven only through &self.

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>

Build from the standard --model weight store + GPU backend. lang carries the tokenizer-resolved source/target language ids (resolved server-side, where the tokenizer lives). max_seq_len caps the decoder KV depth.

Trait Implementations

impl Model for NllbGpuModel

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

Batched decode: one forward_one per sequence into CONTIGUOUS logit rows 0..n (batch position i ↔ seqs[i]), the scheduler’s row contract. Each sequence’s own per-slot KV is looked up by slot_idx, so batch order is irrelevant. Sequences are processed serially on the default stream (shared decode scratch); the returned base pointer is [n, vocab].

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

Run prefill: process all prompt tokens through the model.

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

Process chunk_len tokens starting at chunk_start in the prompt. is_last_chunk runs final norm + LM head; intermediate chunks return DevicePtr::NULL. KV blocks alloc incrementally; SSM state carries across chunks; attention uses FA on chunk 0, paged decode after.

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

Run one decode step: process a single new token.

fn vocab_size(&self) -> usize

Vocab size (for sampler allocation).

fn supports_beam(&self) -> bool

True when this model implements run-to-completion beam search (Self::generate_beam_batch). Default false — only encoder-decoder translation models (NLLB) override it.

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

Run beam search to completion for each request, returning each one’s winning hypothesis token ids (EOS-terminated). Called from the prefill path for num_beams > 1 requests, bypassing the token-by-token decode loop. Default: unsupported.

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

Bind the GPU context to the current thread. Must be called from any thread other than the one that created the model.

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

Allocate a new SequenceState with SSM states.

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

Free all GPU resources associated with a sequence.

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

Move a sequence’s SSM states to a different pool slot.

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

Disown a retired sequence’s SSM pool slot after compact_sequence migrated it to a surviving sequence.

fn cache_sequence(&self, _seq: &SequenceState)

Insert the full token sequence (prompt + generated) into the prefix cache. Call BEFORE free_sequence() (block indices must still be valid). Benefits multi-turn agentic sessions that resend full history.

fn num_free_blocks(&self) -> usize

Number of free KV cache blocks available for allocation.

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

Copy logits from device to host buffer (for CPU-side sampling).

fn logits_buffer_ptr(&self) -> DevicePtr

Base pointer of the on-device logits buffer ([k, vocab] BF16 after decode_verify_graphed). Lets the scheduler read logits for temp sampling even though graphs bake in argmax.

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

GPU argmax: 4-byte D2H copy vs 304KB BF16 D2H + CPU argmax.

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

GPU batched argmax over [N, vocab] BF16; returns N token IDs.

fn hidden_after_norm(&self) -> DevicePtr

Return the hidden state after final norm from the last decode step.

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

L2-resident multi-token verification: per-position argmax token IDs; each token advances KV/SSM state. All tokens go through each layer before moving on so weights stay in L2.

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

Checkpoint SSM states before speculative verification.

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

Rollback SSM states after partial acceptance.

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

Speculative decoding via the model’s internal MTP proposer; falls back to regular decode when no proposer is wired up.

fn has_proposer(&self) -> bool

Check if speculative decoding is available (MTP or self-speculative).

fn has_self_speculative(&self) -> bool

Check if self-speculative decoding is enabled.

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

Eager decode skipping SSM layers. Used by self-speculative drafting. Returns logits pointer for argmax. Advances seq_len by 1.

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

CUDA-graphed K=2 verify: 2 tokens, capture-then-replay. Returns [verified_0, verified_1] argmax IDs. SSM intermediates saved for partial rollback via rollback_ssm_states.

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

CUDA-graphed K=3 verify (1 verified + 2 drafts). Returns 3 argmax IDs. SSM intermediates [0] and [1] are saved for partial rollback.

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

CUDA-graphed K=4 verify (1 verified + 3 drafts). Returns 4 argmax IDs. SSM intermediates [0..3] saved for partial rollback.

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

Save the post-norm hidden state at token_idx (0 or 1) to a dedicated MTP input buffer. Must precede run_mtp_propose — MTP overwrites shared buffers including norm_output.

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

Run the MTP proposer for one draft token off the saved hidden state. None when no proposer is wired.

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

Run the MTP proposer to generate multiple draft tokens.

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

Trim the MTP proposer’s KV cache after verification.

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

Release the device memory this model owns, in reverse construction order.

fn poll_innerq(&self)

Poll TQ+ InnerQ calibration for this model. Called once per prefill chunk. Default: a no-op, which is every model without a driver — the scheduler used to reach a process-wide OnceLock for this, which meant the driver could outlive the model whose device symbols it writes.

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>

Process N decode tokens + an M-token prefill chunk in one pass through the same weight loads. Returns decode logits [N, vocab] and prefill logits [1, vocab] (when is_last). Default: serial decode + prefill.

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

Process N concurrent prefill chunks in one forward pass (same weight load amortised across N streams). The default implementation falls back to a per-stream loop calling prefill_chunk — implementors that support kernel-level batched prefill should override this.

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

Like prefill_batch_chunk, but each finishing stream’s first-token logits land in row row_base + stream_idx of the shared logits arena instead of row stream_idx.

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

Generalised mixed forward: M decode tokens + N concurrent prefill chunks fused into one forward pass. Default: delegates to decode_batch + prefill_batch_chunk serially. Models that implement true mixed batching should override.

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

Normalize SSM h_state norms to prevent catastrophic state explosion during long chunked prefill. Called between chunks by the scheduler. Default: no-op (models without SSM layers don’t need normalization).

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

Per-layer chunked prefill: SSM layers use three phases (proj → single-launch GDN → post) so the recurrence sees the full sequence in one launch; attention layers use standard chunked prefill. Returns last-token logits. Default: single-chunk prefill (no SSM).

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

Runtime LoRA adapter rotation: select the resident adapter named name as active (re-points the delta pool pointers). MUST be called at a scheduler quiescent point (no in-flight decode). Graph-safety is via the eager-on-rotate gate. Default: unsupported (non-LoRA or non-rotatable).

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

Task #24: stable adapter_id (KV/prefix-cache identity) for a per-request pool-slot selector. slot follows SequenceState.adapter_slot: >= 0 picks that resident slot, -1 defers to the installed active adapter. The default (no LoRA) returns the base sentinel 0, keeping the prefix cache byte-identical to the pre-LoRA path.

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

Task #25: acquire a per-slot ref when a sequence begins using its adapter (at prefill), resolving -1 -> active like Self::adapter_id_for. Returns the RESOLVED pool index the ref was taken on (store it, release EXACTLY that index at terminal free — immune to a rotate changing active). Default (no LoRA) returns -1 “nothing acquired” so the release guard skips and the base path is byte-identical.

fn release_adapter_slot(&self, _resolved: i32)

Task #25: release a per-slot ref acquired by Self::acquire_adapter_slot, by the RESOLVED index it returned. -1 is a no-op. Default: no-op.

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

Runtime LoRA adapter dynamic-load: load the adapter at dir INTO pool slot and make it resident there (pool-size-1 per-request weight change). MUST be called at a scheduler quiescent point; needs rotation armed. Default: unsupported (non-LoRA or non-rotatable).

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

Task #27 (demand-driven promotion): RDMA-promote the adapter name (staged on peer_addr at adapter_id) from the peer into a cache pool slot and make it active, returning (slot, evicted_name). Runs at a scheduler quiescent point. peft supplies the r/alpha/scaling the peer manifest does not carry. Default: unsupported (non-LoRA / non-cuda).

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

Demand-driven DISK promotion (no RDMA/peer): load the adapter name from adapter_dir into a cache pool slot (LRU victim) and make it active, returning (slot, evicted_name). Local-disk sibling of Self::promote_lora_from_peer; the swap re-parses the dir’s adapter_config.json, so no peft arg. Runs at a scheduler quiescent point; needs rotation armed. Default: unsupported.

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

Dims for the --high-speed-swap orchestrator (installed thread-local after bind_gpu_to_thread). None for legacy/non-attention models.

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

Self::alloc_sequence told what this request can actually reach (prompt_len + max_tokens). Proposer state that scales with context is sized to THAT instead of --max-seq-len; see DraftProposer::alloc_state_for. Defaults to the unsized form.

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

FP32 logits flag (host buffer needs vocab*4 bytes, reinterpret &[f32]). True only for Gemma-4 dense single-token decode lm_head; default false.

fn has_ssm_layers(&self) -> bool

True when this model has recurrent SSM / Mamba layers whose h_state + conv_state are advanced in-place every decoded token.

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

Verify DRAFT capacity of the MTP state pools for a sequence occupying SSM pool slot slot_idx — the deepest num_drafts a speculative step may dispatch to it without overflowing its slot’s per-token H-intermediate allocation (tiered since 2026-08-16; SSOT ssm_reserve::verify_slot_h_intermediates). The scheduler clamps every spec step’s draft count to the MINIMUM capacity across the active slots. Default usize::MAX: no SSM verify pools to constrain (pure-attention models, spec off).

fn decode_rollback_ring_slots(&self) -> usize

Number of decode-rollback SSM snapshot slots reserved per active sequence (Phase-C). The scheduler’s per-sequence snapshot ring is sized from this. 0 (the default) means the model keeps no decode-rollback snapshots — appropriate for pure-attention models and for SSM models when the snapshot pool has no capacity reserved. SSM models with a populated pool override to ROLLBACK_RESTEER_CAP + 1.

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

Save seq’s live SSM h_state + conv_state (all SSM layers) into the decode-rollback snapshot slot ring_slot.

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

Restore seq’s SSM h_state + conv_state (all SSM layers) from the decode-rollback snapshot slot ring_slot previously written by Self::save_decode_ssm_snapshot.

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

The installed DFlash drafter’s block size γ, when one is installed. The serve layer derives num_drafts = γ - 1 from THIS (the head is the SSOT — it resolved the drafter config’s trained block size), never from a CLI default that may not match the checkpoint.

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

#155 iter3: during decode, save a block-aligned Marconi SSM snapshot at checkpoint-interval boundaries so the NEXT turn’s warm prefix-cache hit restores from decode-produced state near the conversation’s end — instead of replaying decode-produced tokens through the prefill kernel (the warm-hit drift ratchet, issue #155). Called from the scheduler after each decode step’s live SSM state is canonical (post-commit on the MTP path). Default no-op (non-hybrid models / caching disabled).

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

Whether Self::decode_verify_batched can run for ks.len() sequences at ks[i] verify rows each (one more than that sequence’s draft count; the K-vs-batch ladder passes 2..=4, and D-Cut makes the vector RAGGED — uniform is just the special case).

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

Batched K-row verify: ks.len() sequences × ks[i] rows in ONE eager forward (flat seq-major rows, tokens.len() == Σ ks). Weight matrices are read once for all Σ ks rows. Sequence i occupies rows [off_i, off_i + ks[i]) where off_i = Σ_{t<i} ks[t], holding [last_verified, d0, .., d_{ks[i]-2}]. Returns the Σ ks argmax IDs in the same flat order. On success each sequence’s tokens/seq_len advance by its own ks[i] (rewind is the caller’s verdict arithmetic, same as the per-seq path). On Err NO sequence state has been advanced.

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

Copy raw-hidden rows rows[i] of the just-run batched verify forward into stash slot i (verify_hidden_stash), BEFORE any propose clobbers the shared hidden_states buffer. Companion of Self::decode_verify_batched.

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

Stashed-row variant of Self::save_hidden_for_mtp: copy stash slot idx (written by Self::stash_verify_hidden_rows) into the MTP input buffer. Used by the batched-verify verdict path, whose propose calls have already overwritten the live verify rows.

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

Batched cross-sequence MTP propose for the batched K=4 verify path: num_drafts drafts for each of tokens.len() sequences, reading every drafter weight once per draft position instead of once per sequence. stash_idx[i] names the verify-stash slot holding sequence i’s accepted-position hidden (written by Self::stash_verify_hidden_rows); positions[i] is the propose position (post-rewind seq_len), matching the per-seq Self::run_mtp_propose_multi contract. Grammarless sequences only.

fn mtp_propose_batch_max(&self) -> usize

Widest batch Self::run_mtp_propose_batched can carry in ONE drafter forward per draft position. 1 = per-sequence only. Schedulers chunk their propose groups by this — never by a constant.

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

DFlash K=γ graphed verify (γ+1 tokens). Specialization of the K=2/3/4 pattern for arbitrary K. Default impl falls back to eager decode_verify. Models can override for CUDA-graph speedup keyed by (slot_idx, K).

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

DFlash γ-token verification: 1 verified + γ drafts → per-position argmax. Variable-length γ (vs fixed K=2/3/4) because it’s a drafter config field. CUDA-graph capture keyed by (slot_idx, tokens.len()). Default routes to decode_verify_graphed_kgamma.

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

DFlash fused decode+verify: one M=(1+k) forward replacing separate M=1 decode + M=k verify on the DFlash path.

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

ATLAS_MTP_CATCHUP: ring-capture a serially decoded token’s final hidden at pos for the drafter catch-up feed. Default no-op.

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

Capture hidden_states[token_idx] from every DFlash capture layer into dflash_hidden_save. Called after gamma verify Phase 3 D2H sync (bonus position known). No-op when DFlash is disabled.

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

Append the accepted draft’s hidden state (row 1 of dflash_hidden_save) into the proposer context. Base primitive for both legacy and Eagle paths. Default no-op for models without a DFlash drafter.

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

EAGLE-fix (K=2 accept): append row 0 @ N then row 1 @ N+1 BEFORE propose so forward_block conditions on row 1 (the hidden that generated bonus). Default no-op for models without a DFlash drafter.

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

EAGLE-fix (K=gamma): append rows 0..=num_accepted at positions base_pos..=base_pos+num_accepted. Row num_accepted is appended LAST -> freshest ctx slot = the hidden that generated the bonus (EAGLE). Default no-op for models without a DFlash drafter.

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

Ctx-holes fix (serial decode): append the just-decoded token’s captured per-layer hidden (dflash_hidden_save row 0, filled by try_dflash_capture inside the decode layer loop) into the seq’s DFlash ctx accumulator, stamped at its true position (seq.seq_len - 1, matching propose.rs’s decode-append convention).

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

Unified DFlash ctx commit (ATLAS_DFLASH_UNIFIED_CTX=1). Copies num_committed scratch rows (dflash_hidden_save rows scratch_row..scratch_row+num_committed) into ctx_hidden_acc at the CURRENT TAIL (ctx_len), stamping RoPE positions base_pos..base_pos+num_committed, folding the watermark slide in first. base_pos is the RoPE position, NOT the acc row index (they diverge after a watermark slide — DDD §4.1 landmine). scratch_row is 0 on every single-sequence path; batched decode (n>1) captures ALL batch rows, so seq i commits from scratch row i. The single structural replacement for the ~5 fragmented appends. Default no-op for models without a DFlash drafter.

fn dflash_capture_band(&self) -> usize

Rows per per-sequence capture BAND in the DFlash hidden scratch (γ+1). Sequence i of a batched K=γ verify captures into band i, so its commit_ctx scratch_row is i * dflash_capture_band(). Returning the model’s own stride keeps the capture and the commit from ever disagreeing. 0 when there is no DFlash drafter.

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

Read the draft token ID stored on GPU by the last run_mtp_propose_multi call (which used embed_from_argmax to write the draft embedding and token ID directly on GPU). Returns 0 if no proposer is available.

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

Encode images through the vision encoder and store embeddings for the next prefill.

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

Batched vision encode across N requests’ images in ONE forward_batched call (block GEMM weights read once over Σpatches). per_request[i] is request i’s images. Returns one (patch_row_offset, grid_index_offset, num_images, patch_row_count) per request, in request order, locating its slice of the shared packed buf_out. Default: no-op (text models).

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

Set the co-dispatched batched-ViT slice base for the NEXT prefill_chunk (row offset into buf_out, grid index offset, image count owned). Pass (0,0,0) to reset to the legacy single-request behaviour. Default: no-op.

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

EP worker step: receive a (seq_id, cmd) preamble from rank 0 and execute the command in the addressed slot.

fn is_ep(&self) -> bool

Check whether expert parallelism (EP) is enabled (multi-GPU MoE).

fn decode_logits_fp32(&self) -> bool

True when single-token decode lm_head writes FP32 logits to a dedicated FP32 scratch buffer (rather than the shared BF16 logits buffer). Callers that consume those logits must read from Self::decode_logits_ptr using 4 bytes/element. Defaults false; only Gemma-4 dense overrides today (gated by ATLAS_GEMMA4_FP32_LMHEAD=1).

fn decode_logits_ptr(&self) -> DevicePtr

Buffer pointer the single-token decode lm_head last wrote to. The returned dtype is FP32 when Self::decode_logits_fp32 is true, BF16 otherwise. The default impl returns the shared BF16 logits buffer used by every existing model. Override on models that route the lm_head output through an FP32 scratch (Gemma-4 + softcap).

fn is_mla(&self) -> bool

Multi-head Latent Attention guard. When true, chunked prefill MUST run as a single chunk — Atlas has no paged-MLA prefill kernel and multi-chunk MLA silently corrupts attention output (see Mistral-Small-4 2026-05-01 sweep: 8K collapses to “The\nThe…”).

fn hc_mult(&self) -> usize

mHC hyper-connection stream count (0 = no highway). Non-zero means the batched GDN decode paths are UNWIRED for this model (they carry their own residual, which the highway replaces — see qwen3_ssm::hc::refuse_batched_under_hc); the scheduler must clamp concurrency to 1 until the batched highway lands (Avarok #753 item B).

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

Tokens per paged-KV block, or None when the model has no paged KV. The scheduler uses this to land a prefill chunk boundary exactly on the block boundary a warm turn will match at (see spark_runtime::ssm_tail_boundary).

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

EP broadcast: send a command (u32) to all worker ranks.

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

EP broadcast: send a (seq_id, cmd) pair to all worker ranks.

fn ep_protocol_v2(&self) -> bool

Returns true if this model’s EP comm path is using the v2 protocol (slot-aware seq_id preamble). Default false — pre-PR behaviour.

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

EP bulk broadcast: send an array of u32 tokens to all worker ranks. Uses a single NCCL broadcast instead of per-token broadcasts.

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

Launch SSM state checkpoint D2D copies on a secondary CUDA stream.

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

Launch SSM state rollback + checkpoint on the secondary stream.

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

Wait for all work on the secondary stream to complete.

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

Item #2 (STree-style in-place verify commit): commit the surviving prefix of a verify pass directly onto the canonical h_state / conv_state. Full accept (num_accepted == k) is a no-op (the kernel’s final state is already live); partial accept is a single index-select of h_state_intermediates[num_accepted-1]. No-op default for backends without the dual-buffer SSM state. Runs on secondary_stream; pair with sync_secondary.

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

Save KV blocks + SSM state to writer. Does NOT free resources.

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

Restore KV blocks + SSM state from reader into an allocated sequence.

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

Whether tokens contains a vision pad token for this model — i.e. the KV at those positions came from image/video EMBEDDINGS that a plain token re-prefill cannot reproduce. Decode-time preemption uses this to exclude vision sequences from the requeue-with-re-prefill path (the spill path, which saves KV verbatim, stays eligible). Default false: pure-text models are always re-prefillable.

fn num_total_blocks(&self) -> usize

Total KV blocks in the paged cache (denominator for occupancy gauges). Default 0 for backends without a paged cache.

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

Reclaim up to num_blocks blocks from the prefix cache, returning how many actually became free.

fn default_stream(&self) -> u64

Return the default CUDA stream handle.

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

Create a new CUDA stream (for overlapping prefill with decode).

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

Create a CUDA event (for inter-stream synchronization).

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

Record an event on a stream (marks a point in the stream’s work).

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

Make a stream wait for an event (GPU-side sync, CPU does not block).

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

Block the host until all work submitted to stream has completed. Used by mixed_forward_batch to retire the decode pass (which runs on the default stream) before the batched prefill reuses the shared arena buffers on another stream (#110). Default no-op for non-CUDA mocks.