Skip to main content

LatentKvLayer

Struct LatentKvLayer 

Source
pub struct LatentKvLayer {
    pub rows: CudaSlice<f32>,
    pub width: usize,
    pub len: usize,
    pub len_d: CudaSlice<i32>,
    pub index_rows: Option<CudaSlice<f32>>,
    pub index_width: usize,
    pub index_ring_rows: Option<usize>,
    pub index_pool_keys: Option<CudaSlice<f32>>,
    pub index_pools_ready: usize,
    pub index_pool: usize,
}
Expand description

Per-MLA-layer latent KV plane (DESIGN.md §3.2). ONE row per token, width elements wide, where width == StatePlan::LatentKvCache { width } == kv_lora_rank + rope_head_dim: row = [ rmsnorm(c_kv) : kv_rank | rope(k_pe) : d_rope ] There is NO V plane — V is the FIRST kv_rank elements of the SAME row, and every query head streams that one row (MQA). NoPE models (glm5_next, rope_head_dim 0) have width == kv_rank and no k_pe tail.

f32, UNQUANTIZED, deliberately: increment 4 is the correctness arm and its gate is maxdiff against the memra_engine::mla f32 oracle, whose c_kv is f32. DESIGN.md §3.2’s eventual q8_0 latent row (576 = 18 blocks, V view boundary 512 = 16 blocks, both on a 32-element boundary) is a later increment; quantizing here would fork the plane from the oracle it is gated against. The % 32 == 0 KVQUANT constraint therefore does NOT apply to this plane.

Fields§

§rows: CudaSlice<f32>

[max_ctx * width] f32, row-major by token.

§width: usize§len: usize§len_d: CudaSlice<i32>

Device mirror of len, kept in lock-step exactly like KvLayer::len_d.

§index_rows: Option<CudaSlice<f32>>

DSA k-pool indexer state, [max_ctx * index_width] f32 row-major by token: row = [ k_norm(wk(x)) : index_head_dim | index_kpool_compress_gate(x) : index_head_dim ] None when the layer declares index_width == 0 (no k-pool indexer). The reference’s past_key_values.update_indexer carries the same two channels; its third (a per-token validity flag) is DELIBERATELY absent — this cache is single-sequence and unpadded, so every row below len is valid and pooling starts at token 0, which is exactly the scope memra_reference::kpool_allowed_tokens documents for itself. A batched/padded arm needs that channel back and its own gate.

len above is authoritative for BOTH planes: they are appended in the same call and must never carry independent lengths.

MEMORY — the TAIL RING, and it SHIPPED (index_ring_rows, MEMRA_DSA_INDEX_RING). Flat, this plane is 2 * index_head_dim = 256 f32 = 1 KiB per token per layer, i.e. 12 GiB (12.88 GB) over glm5_next’s 12 MLA layers at 1M — larger than the latent plane’s share of the same budget is comfortable with. Two reductions were considered:

  • f16/bf16 rows — DECLINED. The rows feed the pool-key softmax, whose output feeds the ReLU score, whose ties the selection order depends on. Halving the mantissa moves scores, and moved scores move which pools win a tie — the one thing the gates forbid. It would need its own selection-parity gate at serving scale before it could ship, and it buys 6 GiB where the option below buys 11.94.
  • A tail ring — the real answer, and the one implemented. With index_pool_keys resident (below), a row of this plane is read exactly once: by the pool-key build of the pool it belongs to. Every row under index_pools_ready * pool is therefore PROVABLY DEAD — this cache has exactly ONE in-call reader of the plane (Engine::mla_kpool_pool_keys) and ONE writer (Engine::mla_index_append), and CacheSnapshot does not carry latent planes at all, so nothing else can observe a lapped row. (snapshot_plane, lane/glm5-prefix-latent, is a second reader BETWEEN calls, and it reads only the LIVE tail window [index_pools_ready * pool, len) — the liveness argument is unchanged.) The plane only has to hold the incomplete tail plus whatever slice of the current call is in flight, so a ring of R rows with R a multiple of pool (which keeps each pool contiguous mod R) replaces 12 GiB with 60 MiB, EXACTLY: same rows, same kernel, different addresses, zero numeric cost, gated by gpu_kpool_tail_ring_wraps_and_matches_the_flat_plane. Net before: 12 GiB here + 1.5 GiB of pool keys. Net after: 1.56 GiB, an 8.7x cut, because a pool key is index_head_dim f32 per pool tokens = 32 f32/token against 256.

R DOES NOT BOUND THE PER-CALL t, and getting that wrong is what shipped a regression (lane/glm53-ring-sizing, 2026-08-28). The first cut sized R against the largest t a CHUNKED prefill could hand one call; glm5_next primes MONOLITHICALLY, so its t is the whole prompt and the guard refused every prompt past R: 4630 usable tokens inside a configured 8192. mla_kpool_indices now DRAINS the ring inside the call, appending what fits and building the pool keys that frees, so R is a working-set choice (INDEX_RING_WORKING_ROWS) with a floor of one pool and nothing else.

The indexer’s pool is the one input the state plan does NOT carry, the same gap that makes index_pool_keys a lazy allocation below. So pool is not used to size the ring: the allocator books index_ring_rows PHYSICAL rows and the engine rounds that DOWN to a multiple of pool on first use, so the effective ring is always >= rows - pool + 1.

§index_width: usize§index_ring_rows: Option<usize>

PHYSICAL rows of index_rows when it is a tail ring; None when the plane is flat (max_ctx rows, absolute row addressing). The EFFECTIVE ring is this rounded down to a multiple of the indexer’s pool, computed by the engine — see the field doc above.

§index_pool_keys: Option<CudaSlice<f32>>

RESIDENT DSA pool-key plane, [max_ctx / pool * index_head_dim] f32 row-major by pool.

LAYOUT: index_pool_keys[p * d + c] is channel c of the collapsed key of pool p, i.e. of cache rows [p * pool, (p + 1) * pool). d is index_width / 2 (the indexer’s head dim); pool comes from the layer’s MlaIndexerGeom and is NOT in the state plan, which is why this buffer is allocated on FIRST USE by the engine rather than by the allocator below.

INVALIDATION RULE, and it is the whole point: a pool’s key is a function of exactly its own pool rows of index_rows plus the layer’s constant kpool_ape. index_rows is APPEND-ONLY — a row is written once, when its token is appended, and never rewritten — so once a pool’s LAST row lands the key is FINAL and is never recomputed. index_pools_ready is how many leading pools hold such final keys; each call builds only [index_pools_ready, len / pool) and then advances it. The incomplete tail is NOT a pool and has no key: rows [len / pool * pool, len) reach the query through the selection kernel’s always_tail append, recomputed every call.

The rule therefore has exactly ONE trigger: if len ever DECREASES (a rewind that overwrites already-pooled rows), index_pools_ready must be clamped to len / pool by the same code that shortens len. Use truncate_index_pool_keys for that. Today len is written in two places (HybridModel::mla_attn_cached, and restore_plane on a FRESH layer) and only ever grows — Cache::rollback does not touch the latent planes at all — so no caller needs the clamp yet; mla_kpool_indices asserts the invariant on every call so a future rewind that forgets it fails loudly instead of selecting against stale keys, and snapshot_plane/validate_restore assert it at both prefix-cache seams.

§index_pools_ready: usize

Pools [0, index_pools_ready) of index_pool_keys hold FINAL keys.

§index_pool: usize

RESIDENT copy of the indexer’s pool (tokens per k-pool), the one geometry input the state plan does NOT carry (the same gap that makes index_pool_keys a lazy allocation). 0 until the engine’s first indexer call writes it (mla_attn_cached, which refuses a nonzero value that disagrees with the loaded geometry rather than overwriting it). The latent-plane snapshot/restore path (lane/glm5-prefix-latent, 2026-08-30) reads it to address the tail ring and size the restored key plane; it refuses to capture a plane whose pool is still unknown.

Implementations§

Source§

impl LatentKvLayer

Source

pub fn truncate_index_pool_keys(&mut self, pool: usize)

Shorten the resident pool-key plane to what len still justifies. Call from any path that REDUCES len; pools at or above len / pool may have been built over rows the rewind is about to overwrite, so their keys are no longer final.

Source§

impl LatentKvLayer

Source

pub fn snapshot_plane( &self, e: &impl KvDev, ) -> Result<LatentPlaneSnapshot, Box<dyn Error>>

Deep-copy this layer’s latent-plane state OUT of a live session cache. Stream-ordered on the implementor’s worker stream, like every other prefix-capture copy. Errors instead of capturing anything a restore could not make whole:

  • len == 0 (the caller records an unexecuted layer as absent instead),
  • an indexer plane whose pool was never resolved,
  • index_pools_ready != len / pool — a capture off a drained call boundary would publish keys that are behind or ahead of their rows (the finality invariant).
Source

pub fn snapshot_tail( &self, e: &impl KvDev, ) -> Result<LatentTailCapture, Box<dyn Error>>

EAGER half of the deferred boundary capture (doc on LatentTailCapture): copy out only what generation will destroy — the incomplete tail-ring rows — plus the boundary metadata the publish-time slice validates against. Same preconditions as snapshot_plane (len > 0, resolved pool, the pools-ready drain invariant); the big planes are NOT copied here.

Source

pub fn snapshot_plane_at( &self, e: &impl KvDev, cap: LatentTailCapture, ) -> Result<LatentPlaneSnapshot, Box<dyn Error>>

DEFERRED half of the boundary capture: complete a LatentPlaneSnapshot at the captured boundary by slicing the append-only planes (rows [0..cap.len), FINAL pool keys [0..cap.index_pools_ready * d)) from the LIVE layer and moving the eagerly captured tail in. Every disagreement between the capture and the live layer refuses — a publication is an optimization and must never publish planes it cannot prove are the boundary’s (the append-only-below-boundary invariant is what makes the slice legal: the glm5 verify rollback truncates to the accepted length, never below the prime boundary, and pool keys are final the instant their last row lands).

Source

pub fn validate_restore( &self, snap: &LatentPlaneSnapshot, max_ctx: usize, ) -> Result<(), String>

Device-independent half of the restore preflight: every shape/identity/bounds check, no copies, so the caller can validate EVERY layer before the first byte moves (a malformed entry must never leave a half-restored cache for a fallback to consume).

Source

pub fn restore_plane( &mut self, e: &impl KvDev, snap: &LatentPlaneSnapshot, max_ctx: usize, ) -> Result<(), Box<dyn Error>>

Deep-copy a snapshot INTO this freshly allocated layer: latent rows at [0..len), len + device mirror, and (for indexer-bearing layers) the resident key plane sized to the SESSION’s capacity — exactly the capacity_tokens / pool * d sizing mla_kpool_indices books, so the next call keeps it resident instead of reallocating (a reallocation resets index_pools_ready and, under the ring, the rows to rebuild the keys from are gone) — plus index_pools_ready and the live tail rows at their physical ring (or flat) addresses. Validation runs first; a shape error moves no bytes.

Auto Trait Implementations§

Blanket Implementations§

Source§

impl<T> Any for T
where T: 'static + ?Sized,

Source§

fn type_id(&self) -> TypeId

Gets the TypeId of self. Read more
Source§

impl<T> Borrow<T> for T
where T: ?Sized,

Source§

fn borrow(&self) -> &T

Immutably borrows from an owned value. Read more
Source§

impl<T> BorrowMut<T> for T
where T: ?Sized,

Source§

fn borrow_mut(&mut self) -> &mut T

Mutably borrows from an owned value. Read more
Source§

impl<T> From<T> for T

Source§

fn from(t: T) -> T

Returns the argument unchanged.

Source§

impl<T, U> Into<U> for T
where U: From<T>,

Source§

fn into(self) -> U

Calls U::from(self).

That is, this conversion is whatever the implementation of From<T> for U chooses to do.

Source§

impl<T> Same for T

Source§

type Output = T

Should always be Self
Source§

impl<T, U> TryFrom<U> for T
where U: Into<T>,

Source§

type Error = !

The type returned in the event of a conversion error.
Source§

fn try_from(value: U) -> Result<T, !>

Performs the conversion.
Source§

impl<T, U> TryInto<U> for T
where U: TryFrom<T>,

Source§

type Error = <U as TryFrom<T>>::Error

The type returned in the event of a conversion error.
Source§

fn try_into(self) -> Result<U, <U as TryFrom<T>>::Error>

Performs the conversion.