spark_model/

traits.rs

// SPDX-License-Identifier: AGPL-3.0-only

//! Model trait (SDD: single trait, multiple implementations possible).
//!
//! The Model trait defines the interface for running inference. Business
//! logic (scheduler, engine) programs against this trait, not concrete types.

use spark_runtime::gpu::DevicePtr;

use crate::layer::LayerState;
use crate::speculative::ProposerState;

/// Result of a mixed forward pass (decode + prefill in one pass).
pub struct MixedForwardResult {
    /// Logits for decode sequences: [N, vocab_size] BF16.
    /// NULL if no decode sequences.
    pub decode_logits: DevicePtr,
    /// Logits for the prefill sequence's last token: [1, vocab_size] BF16.
    /// NULL if `is_last_chunk` was false (intermediate chunk, no logits).
    pub prefill_logits: DevicePtr,
}

/// Per-stream input slice for batched prefill.
///
/// One of these per concurrent prefilling stream — `prefill_batch_chunk` and
/// `mixed_forward_batch` accept a `&mut [PrefillSlice<'_>]` and process all
/// streams' chunks in a single forward pass. See Q12 in
/// `/workspace/atlas-internal/qwen-refactor/notes.md` for the bug this
/// fixes (concurrent prefills serialized through `prefilling.first_mut()`
/// in the scheduler, causing 5× asymmetric TTFT).
pub struct PrefillSlice<'a> {
    /// Full prompt tokens for this stream.
    pub prompt_tokens: &'a [u32],
    /// Per-stream sequence state (KV blocks, SSM slot, etc.).
    pub seq: &'a mut SequenceState,
    /// Token offset into `prompt_tokens` where this chunk starts.
    pub chunk_start: usize,
    /// Number of tokens in this chunk.
    pub chunk_len: usize,
    /// Whether this is the final chunk for this stream (controls whether
    /// the model emits last-token logits for sampling).
    pub is_last_chunk: bool,
}

/// Result of a fully-batched mixed forward pass: M decode tokens + N prefill
/// chunks in one pass.
pub struct MixedBatchResult {
    /// Logits for decode lanes: [M, vocab] BF16. NULL if no decode lanes.
    pub decode_logits: DevicePtr,
    /// Logits per prefill stream — one DevicePtr per stream in the input
    /// slice, in the same order. Each entry is `[1, vocab]` BF16 when that
    /// stream's chunk was `is_last_chunk`, or NULL otherwise.
    pub prefill_logits: Vec<DevicePtr>,
}

/// Per-sequence paged attention metadata for chunked prefill.
///
/// Positions and slots remain chunk-local, but the paged block table and
/// running sequence length can persist across chunks so we only upload the
/// changed tail instead of rebuilding the full page metadata every time.
pub struct ChunkedPrefillPageMetadata {
    /// Device buffer holding the sequence block table as raw 32-bit entries.
    pub block_table: DevicePtr,
    /// Device buffer holding the running paged-prefill sequence length.
    pub seq_len: DevicePtr,
    /// Total block-table entries allocated for this prompt.
    pub block_capacity: usize,
    /// Number of block-table entries already uploaded to `block_table`.
    pub uploaded_blocks: usize,
}

/// Sequence state tracked across decode steps.
pub struct SequenceState {
    /// Token IDs generated so far (including prompt).
    pub tokens: Vec<u32>,
    /// Block table for paged KV cache (indices into PagedKvCache).
    pub block_table: Vec<u32>,
    /// Current sequence length (prompt + generated).
    pub seq_len: usize,
    /// Per-layer state (EmptyLayerState for attention, SsmLayerState for SSM).
    pub layer_states: Vec<Box<dyn LayerState>>,
    /// Per-sequence state for speculative decoding proposer (None if no proposer).
    pub proposer_state: Option<Box<dyn ProposerState>>,
    /// SSM state pool slot index. Used for CUDA graph stability — all sequences
    /// at the same slot_idx use the same fixed GPU addresses. Derived from
    /// `ssm_slot` at claim time (the guard is the authority on release
    /// responsibility; this index is the authority on pool-offset math).
    pub slot_idx: usize,
    /// RAII guard that returns `slot_idx` to the SSM pool's free list on drop.
    /// Guarantees the slot is released on EVERY sequence-exit path — including
    /// abort/cancel, decode error, swap-out failure, and panic/unwind — not
    /// just the explicit `free_sequence`/`compact_sequence` sites, which
    /// `take()`/`migrate()` the guard so the release happens EXACTLY once.
    /// `None` for models without an SSM pool (e.g. the unit-test mock).
    pub(crate) ssm_slot: Option<crate::model::ssm_pool::SlotGuard>,
    /// Marconi: token position up to which SSM state is valid from a snapshot.
    /// Set on chunk 0's prefix cache lookup, read by subsequent chunks to skip
    /// computation for tokens already covered by the snapshot + KV cache.
    pub marconi_skip_to: usize,
    /// Marconi exact-hit: snapshot slot when the *entire* prompt matched a
    /// leaf snapshot (`matched == total`). On this path the last prompt
    /// token is re-run for logits, which double-advances the SSM recurrent
    /// state; `finalize_last` uses this to re-restore the pristine state@N
    /// and emit the first token's logits from the snapshot's stashed hidden
    /// instead. `None` for all other paths.
    pub marconi_exact_snap: Option<usize>,
    /// Session hash for SSM snapshot isolation. Set by the scheduler before
    /// prefill. The model uses this to tag saved snapshots and verify ownership
    /// before restoring. 0 = no session tracking (legacy behavior).
    pub session_hash: u64,
    /// Ownership stamp for the SINGLE-SLOT whole-prompt hidden capture
    /// (`mtp_prefill_hidden`). Written by `try_mtp_prefill_capture` when THIS
    /// sequence's chunk 0 (re)starts the capture, with the model's monotonic
    /// capture generation. `ensure_drafter_context` prefills the drafter only
    /// while the stamp still matches the model's current generation — at
    /// C>=2 interleaved prefills restart the shared capture, and without this
    /// check a sequence's first propose could pair ITS tokens with ANOTHER
    /// sequence's captured hiddens (poisoned drafter KV; blind is strictly
    /// better than poisoned). 0 = never owned a capture.
    pub mtp_capture_gen: u64,
    /// Ownership ticket for the shared hidden-row interval
    /// (`mtp_store_range`), drawn at `alloc_sequence` from the same atomic
    /// that issues capture generations.
    ///
    /// Distinct from `mtp_capture_gen` because that one is assigned ONLY under
    /// `chunk_start == 0`, and a warm turn never starts at 0 — so it is `0` for
    /// the entire life of exactly the sequences the carry path serves, and
    /// would make every warm sequence look like the same owner. This is drawn
    /// unconditionally at admission. `0` = drawn outside `alloc_sequence` (the
    /// mock and test fakes), and never matches anything.
    pub mtp_store_gen: u64,
    /// Per-adapter prefix-cache namespace (adapter-correct KV). Folded into the
    /// prefix hash so two adapters that share a token prefix never reuse each
    /// other's blocks. `0` = base / no adapter (a strict no-op in the fold, so
    /// behavior is byte-identical until a LoRA path stamps a non-zero id).
    pub adapter_id: u64,
    /// Persistent paged metadata for chunked prefill, allocated lazily on the
    /// first chunk that needs paged attention.
    pub chunked_prefill_meta: Option<ChunkedPrefillPageMetadata>,
    /// Number of prompt tokens MATCHED by the chunk-0 prefix-cache lookup
    /// (block-aligned). Load-bearing for block-ref accounting: `free_sequence`
    /// release, `cache_sequence` double-bump avoidance, and the chunked-prefill
    /// KV write floor all key off it. NOT what gets reported to clients —
    /// a match can be found and then discarded (see `reused_prefix_tokens`).
    pub cached_prefix_tokens: usize,
    /// Number of prompt tokens whose KV this request actually READ from the
    /// prefix cache instead of recomputing. Stamped after the skip decision;
    /// read by the scheduler to populate
    /// `usage.prompt_tokens_details.cached_tokens`.
    ///
    /// Atlas #919: this used to be `cached_prefix_tokens`, which reports the
    /// lookup result — so a request that matched 48 tokens and then recomputed
    /// all of them (no SSM snapshot / exact-leaf bypass / declined Marconi
    /// restore) advertised `cached_tokens: 48` next to a full-prefill log line.
    pub reused_prefix_tokens: usize,
    /// Number of `block_table` entries that came FROM the prefix cache on this
    /// sequence's lookup (`matched_blocks.len()`). The cache already holds its
    /// own "+1" KV ref on each of those blocks, and eviction returns exactly ONE
    /// ref per radix node — so re-bumping them in `cache_sequence` would add a
    /// ref nothing can ever release, permanently pinning the whole reused prefix
    /// on every warm turn until the pool wedges. 0 when there was no cache hit.
    pub cached_prefix_blocks: usize,
    /// The matched prefix token IDs (`tokens[..cached_prefix_tokens]`) stashed
    /// at prefix-lookup time. `free_sequence` releases the prefix cache's radix
    /// refs over these when `tokens` is too short to cover the prefix — i.e. a
    /// prefill that matched a prefix (bumping radix refs) then FAILED to
    /// allocate its suffix, so `tokens` was never populated. Without this the
    /// `release(&tokens)` on the failure path is a no-op and the matched radix
    /// nodes stay pinned at ref≥2 forever → the pool progressively wedges. Empty
    /// on the common path (no match / success releases over the full `tokens`).
    pub prefix_ref_tokens: Vec<u32>,
    /// Whether the chunk-0 prefix-cache lookup already ran for this sequence.
    ///
    /// The lookup is NOT idempotent: it bumps radix refs, `inc_ref`s each
    /// matched KV block and PUSHES it onto `block_table`. It also runs BEFORE
    /// `ensure_blocks_through_prefill`, so a chunk-0 prefill that fails to
    /// allocate its suffix (KV exhausted) leaves all of that applied. The
    /// preempt-and-retry in `run_standard_chunk_loop` re-enters `prefill_chunk`
    /// for the SAME chunk, which would run the lookup a second time — appending
    /// the matched blocks to `block_table` again (so `block_table[i]` no longer
    /// maps to logical block `i`) and taking a second radix ref that the single
    /// `release` in `free_sequence` can never balance. This flag makes the
    /// re-entry a no-op that replays chunk 0's original decision.
    pub prefix_lookup_applied: bool,
    /// The `skip` half of the chunk-0 lookup's return value, replayed verbatim
    /// when `prefix_lookup_applied` short-circuits a retry.
    pub prefix_lookup_skip: bool,
    /// Token count of an SSM anchor near this prompt's end that the pool
    /// holds for it: the tail-split checkpoint saved during THIS prefill, or
    /// the checkpoint a warm prefill restored from. Cleared at chunk 0; read
    /// by `finalize_last` to decide whether the exact prefill-end leaf earns
    /// a pool slot (`prefill_b::exact_leaf`).
    pub tail_checkpoint_tokens: Option<usize>,
    /// Contiguous prefix length (in tokens, from position 0) whose paged KV is
    /// guaranteed fully written for THIS sequence — either reused from a valid
    /// prefix-cache match or written by a real prefill pass this turn. Updated
    /// per chunk in `prefill_b_proc_range`. The prefix-cache insert path caps
    /// the cached complete-block count to `kv_valid_tokens / block_size` so a
    /// block whose K/V was never written (e.g. the `proc_count==1` decode
    /// shortcut skips an entire trailing chunk) is NEVER inserted with stale V.
    /// Without this cap, stale (donor/zeroed) V in trailing complete blocks
    /// gets cached and read by the next turn's full-attention layers, making
    /// cache-ON decode nondeterministic at temperature 0 (see fix/in-think-
    /// tool-call-leak prefix-cache stale-V diagnosis).
    pub kv_valid_tokens: usize,
    /// #155 iter3: block index (`seq_len / block_size`) of the most recent
    /// decode-time Marconi checkpoint. Dedups re-saving the same boundary
    /// across consecutive decode steps. 0 until the first decode checkpoint.
    pub last_decode_ckpt_block: usize,
    /// Original prompt token count, set at the first prefill and never
    /// mutated by decode. Used by `cache_sequence` to split seq.tokens into
    /// prompt (already inserted + ref-bumped by prefill) vs generated
    /// (needs a fresh bump so `release` in `free_sequence` leaves the
    /// cache's baseline ref intact). 0 before the first prefill.
    pub prompt_len: usize,
    /// Disk-block-ID list for `--high-speed-swap` (Phase 6.1.c).
    /// Each entry is a stable disk-side identifier that outlives HBM block
    /// recycling. `disk_block_ids` grows monotonically with the sequence
    /// and represents its **full historical block list**. IDs are
    /// layer-agnostic — the same ID indexes a slot in every layer's
    /// on-disk file. Empty when `--high-speed-swap` is disabled.
    ///
    /// **Sliding-window invariant** (Phase 6.3): in HSS mode `block_table`
    /// is the suffix `disk_block_ids[hss_window_start()..]`, so
    /// `disk_block_ids.len() == hss_window_start() + block_table.len()`.
    /// Both vectors are grown together by the alloc helper; the offload
    /// helper only fills layer K/V data (no length growth). When
    /// `block_table.len() == cap` and a new logical block is needed, the
    /// alloc helper drops `block_table[0]` (frees the physical HBM block
    /// back to the pool) but keeps `disk_block_ids[0]` — the evicted
    /// block's data lives on at that disk_id for streaming reads.
    pub disk_block_ids: Vec<u32>,
    /// Per-attention-layer offload progress tracker for `--high-speed-swap`
    /// (Phase 6.1.d critical fix). `disk_last_offloaded_per_layer[L]` is
    /// the number of `disk_block_ids` entries this attention layer has
    /// successfully offloaded to its on-disk file. Each layer maintains
    /// its own counter because each layer writes its own K/V independently;
    /// without per-layer tracking, only the first layer to encounter a new
    /// block would offload, leaving subsequent layers' on-disk slots
    /// uninitialised. Length equals the model's attention layer count;
    /// empty when HSS is disabled.
    pub disk_last_offloaded_per_layer: Vec<u32>,
    /// Legacy /v1/completions echo+logprobs: Some(k) = during prefill,
    /// project every prompt position's hidden state and record the actual
    /// next token's logprob plus top-k alternatives. Set by the scheduler
    /// before prefill; None = zero-cost (the collection helper
    /// early-returns). Requests with this set bypass the prefix cache so
    /// every position has a live hidden row.
    pub collect_prompt_logprobs: Option<u8>,
    /// Accumulated across prefill chunks: one entry per prompt position
    /// i in [0, prompt_len-1) scoring tokens[i+1]. The final prompt
    /// position (whose target is the first GENERATED token) is excluded.
    pub prompt_logprobs: Vec<PromptTokenLogprob>,
    /// M2 per-request LoRA routing: the adapter POOL SLOT this sequence's
    /// requests select (NOT `slot_idx`, which is the KV/SSM pool slot). `-1`
    /// (the default for every existing path) means "defer to the installed
    /// active adapter" — so an unset request is byte-identical to today. Set
    /// once from `InferenceRequest::adapter_slot()` at prefill; read by
    /// `decode_batch` to build the per-step device `seq_slot[N]` buffer the
    /// batched bgmv routes on.
    pub adapter_slot: i32,
    /// Task #25 (slot ref_count): the RESOLVED LoRA pool slot this sequence holds
    /// a ref on (`-1` = none / not acquired — the default and every non-LoRA
    /// path). Set at the prefill acquire (and re-acquire on swap-in resume) to
    /// the index `Model::acquire_adapter_slot` returned; the terminal free
    /// releases EXACTLY this index (not a re-resolved `adapter_slot`, which would
    /// mis-decrement if `active` rotated between prefill and finish) and zeroes
    /// it back to `-1` so release fires exactly once per acquire. Stored resolved
    /// (not raw) so it also guards the non-scheduler alloc paths (which never
    /// acquire) from an underflow.
    pub acquired_adapter_slot: i32,
    /// NLLB / M2M-100 per-request translation source-language token id (the
    /// encoder-input prefix). `0` = use the deployment default (`--src-lang`).
    /// Unused by every other model type.
    pub src_lang_id: u32,
    /// NLLB / M2M-100 per-request target-language token id (`forced_bos`).
    /// `0` = use the deployment default (`--tgt-lang`). Unused by other models.
    pub tgt_lang_id: u32,
    /// NLLB beam search: number of beams for this request (`1` = greedy,
    /// disables the beam path). Unused by every other model type.
    pub num_beams: u32,
    /// NLLB beam search: length penalty applied to hypothesis scores
    /// (`1.0` = neutral). Unused by other models.
    pub length_penalty: f32,
    /// NLLB beam search: stop as soon as `num_beams` finished hypotheses
    /// exist (`false` = exhaust `max_new`). Unused by other models.
    pub early_stopping: bool,
}

impl SequenceState {
    /// A detached, host-only sequence state: no GPU resources, no SSM
    /// slot, no layer states, every counter zeroed. The single source
    /// for the "empty sequence" field defaults — construction sites
    /// that own real resources build on top of it instead of repeating
    /// the full literal (NLLB's `alloc_sequence`, the engine-test
    /// mock), so a new field gets ONE default site. Also the only way
    /// for other crates to construct a `SequenceState` at all (e.g.
    /// the scheduler's lifecycle unit tests): `ssm_slot` is
    /// crate-private by design.
    pub fn host_only(slot_idx: usize) -> Self {
        SequenceState {
            tokens: Vec::new(),
            block_table: Vec::new(),
            seq_len: 0,
            layer_states: Vec::new(),
            proposer_state: None,
            slot_idx,
            ssm_slot: None,
            marconi_skip_to: 0,
            marconi_exact_snap: None,
            session_hash: 0,
            mtp_capture_gen: 0,
            // Not from `alloc_sequence`, so it owns no hidden rows.
            mtp_store_gen: 0,
            adapter_id: 0,
            chunked_prefill_meta: None,
            cached_prefix_tokens: 0,
            reused_prefix_tokens: 0,
            cached_prefix_blocks: 0,
            prefix_ref_tokens: Vec::new(),
            prefix_lookup_applied: false,
            tail_checkpoint_tokens: None,
            prefix_lookup_skip: false,
            kv_valid_tokens: 0,
            last_decode_ckpt_block: 0,
            prompt_len: 0,
            disk_block_ids: Vec::new(),
            disk_last_offloaded_per_layer: Vec::new(),
            collect_prompt_logprobs: None,
            prompt_logprobs: Vec::new(),
            // -1 = defer to the installed active adapter (see field docs).
            adapter_slot: -1,
            // -1 = no LoRA slot ref held until prefill acquires (Task #25).
            acquired_adapter_slot: -1,
            src_lang_id: 0,
            tgt_lang_id: 0,
            num_beams: 1,
            length_penalty: 1.0,
            early_stopping: false,
        }
    }

    /// SSM-pool slot index for this sequence, if it has GDN/SSM (linear-attn)
    /// layers. Used by the scheduler to order the decode batch by slot so the
    /// batched-recurrent SSM + CUDA-graph contiguity invariant holds
    /// (position i ↔ pool_base + i*stride). `None` for pure-attention models.
    #[inline]
    pub fn ssm_slot_idx(&self) -> Option<usize> {
        self.ssm_slot.as_ref().and_then(|g| g.idx())
    }

    /// Phase 6.3 sliding-window helper: the absolute logical block index
    /// of `block_table[0]`. Returns 0 when `--high-speed-swap` is off
    /// (`disk_block_ids` is empty then; `block_table` is the full history).
    /// Derived rather than stored — the invariant
    /// `disk_block_ids.len() == hss_window_start() + block_table.len()`
    /// is maintained by the alloc helper and asserted by the offload
    /// helper, so no separate field is needed.
    #[inline]
    pub fn hss_window_start(&self) -> usize {
        self.disk_block_ids
            .len()
            .saturating_sub(self.block_table.len())
    }

    /// Map an absolute logical block index → physical HBM block id.
    /// Returns `None` when the block has been evicted to disk-only
    /// (the caller should route attention through the HSS orchestrator's
    /// `attend_layer_on_stream` for that position). With HSS off,
    /// `hss_window_start()` is 0 and this is a direct lookup.
    #[inline]
    pub fn physical_block_for(&self, abs_block_idx: usize) -> Option<u32> {
        let ws = self.hss_window_start();
        if abs_block_idx < ws {
            return None;
        }
        self.block_table.get(abs_block_idx - ws).copied()
    }
}

/// Model trait for forward pass execution.
///
/// Implementations: `TransformerModel` (all architectures).
///
/// # Safety
///
/// `Send + Sync` is required by `Box<dyn Model>` usage patterns.
/// `Sync` safety: the model is exclusively accessed from the scheduler
/// thread. The `unsafe impl Sync` on `TransformerModel` documents this
/// single-thread invariant — do NOT share `&dyn Model` across threads.
mod logprobs;
mod model;
pub use logprobs::*;
pub use model::{BeamReq, EpCommandFailed, Model, padded_batch_n};