Struct TransformerModel 

pub struct TransformerModel { /* private fields */ }

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>, 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<VisionEncoder>, ssm_cache_slots: usize, ssm_checkpoint_interval: usize, ) -> Result<Self>

impl TransformerModel

pub fn decode_logits_fp32(&self) -> bool

True when single-token decode lm_head writes FP32 logits to logits_fp32_buf. Callers that consume those logits (sampler) MUST read with the matching dtype. Prefill / batched-decode lm_head still produce BF16, so this only applies to the lm_head (single-token) return value.

pub fn decode_logits_ptr(&self) -> DevicePtr

Buffer pointer the single-token decode lm_head last wrote to. FP32 scratch when use_fp32_logits, otherwise the shared BF16 logits buffer. Callers that previously hard-coded self.buffers.logits() after self.lm_head(...) must use this so the sampler reads the correct buffer dtype (the BF16 buffer is stale/empty in the FP32 path because lm_head writes elsewhere). Pair with logits_ptr_is_fp32 / decode_logits_fp32 for dtype-aware reads.

impl TransformerModel

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

Borrow the GPU backend for post-construction wiring (e.g. installing a DFlash proposer that needs to allocate paged KV caches against the same GPU the target uses).

pub fn config_ref(&self) -> &ModelConfig

Borrow the model config for post-construction wiring (e.g. building the DeepSeek-V4 MTP proposer, which needs hidden_size / kv_lora_rank / qk_rope_head_dim to size its private MLA KV cache).

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

Install a DFlash drafter as the active proposer, replacing whatever MTP proposer (if any) TransformerModel::new built. The target’s hidden-state capture buffer is already allocated when the config’s dflash_capture_layers is non-empty (factory.rs populates it before construction), so this method only swaps the proposer slot.

Mutually exclusive with --speculative MTP at the CLI level (clap conflicts_with); this method does not enforce that — the caller is expected to have validated the flag combination already.

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

Install the fused n-gram input embedding (LongCat family). Once set, every embedding site routes through it instead of the plain embed_tokens gather.

pub fn has_ngram_embedding(&self) -> bool

True when this model fuses n-gram lookups into its input embedding.

impl TransformerModel

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

Install a startup-static LoRA adapter (post-construction, mirroring Self::set_dflash_proposer). Walks the model layers by GLOBAL index — LoraWeights.layers is indexed the same way — and copies each adapted layer’s K/V/O (+ optional gate/up/down) pairs into the Qwen3AttentionLayer (which routes FFN pairs into its dense FFN component). M0: layers only STORE the adapter; base output is unchanged until the M1 compute insertions read it. Task #24: stable adapter_id for a per-request pool-slot selector. Returns the base sentinel 0 when no LoRA pool is resident (byte-identical base), else the NAME-derived id of the resolved slot (-1 -> active). Resolved here at prefill time because LoraWeights.active can rotate between HTTP request resolution and prefill.

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

Task #25: acquire a per-slot ref for a sequence beginning to use its adapter (called at prefill, resolving -1 -> active exactly like Self::adapter_id_for_slot). Returns the RESOLVED pool index the ref was taken on — the caller stores it and releases EXACTLY that index at terminal free, so an intervening rotate changing active cannot make release hit a different counter. -1 (“nothing acquired”) when no LoRA pool is resident or the slot is out of range → byte-identical no-op base.

pub fn release_adapter_slot(&self, resolved: i32)

Task #25: release a ref acquired by Self::acquire_adapter_slot, by the RESOLVED index it returned. -1 and no-pool are no-ops (base path).

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

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

Runtime adapter rotation (eager-on-rotate). Selects the resident adapter named name as ACTIVE: re-points every layer’s LoraPair (a/b DevicePtr + rank/scale) to that slot’s sub-region, then clears the decode-graph caches defensively (empty under forced eager). MUST be called at a scheduler QUIESCENT point (no in-flight decode reading the old slot). Graph-safety rests on lora_rotatable forcing eager decode — this method never re-captures a graph.

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

RDMA-swap the adapter named adapter_name (staged on $ATLAS_LORA_PEER at adapter_id) INTO pool slot, in place, then make it that slot’s resident adapter. Byte-identical to a disk pack (the loader does the same F16/F32→BF16 convert + B row-repack). MUST be called at a scheduler QUIESCENT point (no in-flight decode reading slot). Re-zeroes the slot sub-region first (a reused slot may hold the prior adapter’s bytes), then rebuilds the slot’s LoraLayerWeights with the NEW adapter’s r/scale — re-installing if the swapped slot is currently active. Requires rotation armed (ATLAS_LORA_ROTATE/$ATLAS_LORA_PEER) so decode is eager.

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

Task #27 (demand-driven promotion): promote the adapter adapter_name (staged on peer_addr at adapter_id) from the peer into a CACHE-region pool slot and make it ACTIVE, returning (slot, evicted_name). Runs on the scheduler thread at a QUIESCENT point (the only place per-slot ref_count is authoritative). Victim policy (pure select_victim_slot): a never-filled placeholder first, else the LRU idle (ref_count == 0) cache slot, else POOL_FULL (retryable — a busy slot is NEVER evicted). The underlying Self::swap_lora_slot_from_peer re-checks ref_count>0 and bails as a backstop, and bumps the slot generation so #24 KV stays correct. Making the promoted slot active mirrors the rotate/load control plane so the delta actually applies under batch-1 (the per-slot bgmv route tables are still dormant — compute reads the installed active adapter).

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

Demand-driven DISK promotion (no RDMA/peer): load the adapter at adapter_dir (named name) into a CACHE-region pool slot (LRU victim) and make it ACTIVE, returning (slot, evicted_name). Local-disk analog of Self::promote_lora_slot_from_peer — same victim policy (pure select_victim_slot: never-filled placeholder first, else the LRU idle (ref_count == 0) cache slot, else POOL_FULL, retryable — a busy slot is NEVER evicted) and the same make-active control plane, but the inner swap reads the adapter from disk instead of the peer. Self::swap_lora_slot_from_disk re-parses the dir’s adapter_config.json (so no peft arg), re-checks ref_count>0 as a backstop, and bumps the slot generation so #24 KV stays correct. Runs on the scheduler thread at a QUIESCENT point; requires rotation armed (ATLAS_LORA_ROTATE=1) — the inner swap enforces it.

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

Disk-swap the adapter at adapter_dir INTO pool slot, in place, then make it that slot’s resident adapter (re-installing onto the layer structs if the slot is currently active). The local-disk analog of Self::swap_lora_slot_from_peer — same audit + pack + re-point, no RDMA. This is the pool-size-1 dynamic-load path: load a DIFFERENT adapter into the single slot at runtime (per-request weight change). MUST be called at a scheduler QUIESCENT point (no in-flight decode reading slot) and needs rotation armed (ATLAS_LORA_ROTATE=1/$ATLAS_LORA_PEER) so decode is eager and no captured graph replays the swapped slot’s stale pointers.

impl TransformerModel

Release every pool this model owns, newest first.

Construction order is buffers → kv cache → ssm pools → derived, so release runs the reverse. Teardown is used rather than a hand-rolled sequence because it attempts every resource even after one fails: a half-torn-down GPU is worse than a reported error.

NOT released here: the weights. build_model takes store: &WeightStore and the layers only copy pointers out of it, so this model does not own them — the host that retained the store releases it after this returns.

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

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

fn poll_innerq(&self)

Poll this model’s own InnerQ driver. A miss is logged, never fatal — it is a diagnostic lever, not part of serving.

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

Q12 Phase 4b override. The concrete dispatcher routes ineligible batches to its sequential path before state mutation. Errors from an admitted kernel batch must propagate: retrying sequentially can reapply prefix-cache and KV state.

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

Mixed-step variant: shift the finishing streams’ logits rows clear of the decode lanes. See the trait docs for the aliasing this prevents.

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

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

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 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 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 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 decode( &self, token: u32, seq: &mut SequenceState, _stream: u64, ) -> Result<DevicePtr>

Run one decode step: process a single new token.

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

Run batched decode: process one token per sequence.

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

Vocab size (for sampler allocation).

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, 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 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 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 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 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_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 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, tokens: &[u32], seq: &mut SequenceState, stream: 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, num_accepted: usize, ) -> Result<()>

Rollback SSM states after partial acceptance.

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 generate_speculative( &self, prompt_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 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 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 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 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 free_sequence(&self, seq: &mut SequenceState) -> Result<()>

Free all GPU resources associated with a sequence.

fn decode_verify_graphed( &self, tokens: &[u32; 2], seq: &mut SequenceState, _stream: 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, tokens: &[u32; 3], seq: &mut SequenceState, _stream: 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, tokens: &[u32; 4], seq: &mut SequenceState, _stream: 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 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_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_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 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_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 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_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 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 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 trim_proposer_state( &self, seq: &mut SequenceState, num_accepted: usize, _stream: u64, ) -> Result<()>

Trim the MTP proposer’s KV cache after verification.

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

Number of free KV cache blocks available for allocation.

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

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

impl Send for TransformerModel

impl Sync for TransformerModel