Struct QsaIndexer 

pub struct QsaIndexer { /* private fields */ }

Implementations

impl QsaIndexer

pub fn prefill_select( &self, st: &mut QsaSeqState, normed: DevicePtr, q_roped: DevicePtr, attn_ctx: DevicePtr, k_pool: DevicePtr, v_pool: DevicePtr, seq_block_table: &[u32], seq_start: usize, num_tokens: usize, nq: u32, block_size: u32, inv_sqrt_d: f32, scratch: DevicePtr, gpu: &dyn GpuBackend, stream: u64, ) -> Result<()>

2026-09-25: Per-query prefill selection for a prefill chunk. Chunk rows whose global position (seq_start + row) is at or past the inert bound get their attention context rows (pre-gate, pre-o_proj) overwritten with attention over their selected set, read from the paged KV cache, which must already hold every prior chunk plus this one. Rows below the bound keep the dense output, which is identical there. prefill_ingest must have run for this chunk (prefill_inner.rs calls it before attention). METRALE_QSA_NO_PREFILL_SELECT=1 makes this a no-op.

impl QsaIndexer

pub fn snapshot_aux( &self, st: &QsaSeqState, gpu: &dyn GpuBackend, stream: u64, ) -> Result<Vec<u8>>

2026-09-25: Aux blob: [ingested u64][pooled u64][raw_keys bf16 bytes]. Block keys are not serialised: restore_aux re-pools them from the raw keys with one qsa_block_pool launch.

pub fn restore_aux( &self, st: &mut QsaSeqState, blob: &[u8], gpu: &dyn GpuBackend, stream: u64, ) -> Result<()>

2026-09-25: Restore the blob from Self::snapshot_aux on a prefix-cache hit: upload the raw keys, reset the counters, re-pool the block keys.

impl QsaIndexer

pub fn new( qk_proj_w: DevicePtr, q_norm_w: DevicePtr, k_norm_w: DevicePtr, n_heads: usize, hd: usize, ratio: usize, budget: usize, rot: usize, theta: f32, eps: f32, hidden: usize, nkv_attn: usize, hd_attn: usize, gpu: &dyn GpuBackend, ) -> Result<Self>

pub fn new_seq_state(&self, gpu: &dyn GpuBackend) -> Result<QsaSeqState>

2026-09-25: One sequence’s indexer carry: counters + raw/pooled key buffers. The launch scratch stays in the layer and is shared by every sequence.

pub fn release_seq_state( &self, st: &mut QsaSeqState, gpu: &dyn GpuBackend, ) -> Result<()>

2026-09-25: Release one sequence’s indexer carry.

QsaSeqState holds bare DevicePtrs, so dropping the struct frees nothing; the max_tokens * hd * 2 raw-key and max_tokens/ratio * hd * 2 block-key buffers are freed only here.

Idempotent: each pointer is nulled once its free is attempted, so a second call (or a release after a partial failure) cannot double-free. A failure on the first buffer still attempts the second; the first error is returned.

pub fn inert_bound(&self) -> usize

2026-09-25: The largest visible prefix whose selection is all-visible: with at most budget + ratio - 1 visible tokens there are at most budget / ratio complete blocks, so every block is selected.

pub fn prefill_ingest( &self, st: &mut QsaSeqState, hidden: DevicePtr, num_tokens: usize, seq_start: usize, gpu: &dyn GpuBackend, stream: u64, ) -> Result<()>

2026-09-25: Ingest num_tokens prefill tokens starting at seq_start: project qk, store the raw keys, pool newly complete blocks. seq_start == 0 resets the counters. Errors when seq_start is not the ingested count (a prefix-cache skip) or the total exceeds METRALE_QSA_MAX_TOKENS.

pub fn decode_select( &self, st: &mut QsaSeqState, normed: DevicePtr, pos: usize, k_pool: DevicePtr, v_pool: DevicePtr, block_table_dev: DevicePtr, block_size: u32, gpu: &dyn GpuBackend, stream: u64, ) -> Result<Option<QsaSelection>>

2026-09-25: Decode-step ingest + selection for the token at pos (0-based; pos + 1 visible). None inside the inert bound (dense is exact). Errors when pos is not the ingested count or reaches max_tokens.