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_keysresident (below), a row of this plane is read exactly once: by the pool-key build of the pool it belongs to. Every row underindex_pools_ready * poolis 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), andCacheSnapshotdoes 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 ofRrows withRa multiple ofpool(which keeps each pool contiguous modR) replaces 12 GiB with 60 MiB, EXACTLY: same rows, same kernel, different addresses, zero numeric cost, gated bygpu_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 isindex_head_dimf32 perpooltokens = 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: usizePools [0, index_pools_ready) of index_pool_keys hold FINAL keys.
index_pool: usizeRESIDENT 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
impl LatentKvLayer
Sourcepub fn truncate_index_pool_keys(&mut self, pool: usize)
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
impl LatentKvLayer
Sourcepub fn snapshot_plane(
&self,
e: &impl KvDev,
) -> Result<LatentPlaneSnapshot, Box<dyn Error>>
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
poolwas 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).
Sourcepub fn snapshot_tail(
&self,
e: &impl KvDev,
) -> Result<LatentTailCapture, Box<dyn Error>>
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.
Sourcepub fn snapshot_plane_at(
&self,
e: &impl KvDev,
cap: LatentTailCapture,
) -> Result<LatentPlaneSnapshot, Box<dyn Error>>
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).
Sourcepub fn validate_restore(
&self,
snap: &LatentPlaneSnapshot,
max_ctx: usize,
) -> Result<(), String>
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).
Sourcepub fn restore_plane(
&mut self,
e: &impl KvDev,
snap: &LatentPlaneSnapshot,
max_ctx: usize,
) -> Result<(), Box<dyn Error>>
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.