Skip to main content

DflashKv

Struct DflashKv 

Source
pub struct DflashKv {
    pub k: Vec<CudaSlice<f32>>,
    pub v: Vec<CudaSlice<f32>>,
    pub len: usize,
    pub cap: usize,
    /* private fields */
}
Expand description

Draft KV cache (round-cost fix, 2026-07-13): per-layer normed+roped ctx K and raw ctx V, append-only in committed order. Block K/V land TRANSIENTLY at [len..len+b] each round (never committed — the reference crops them identically). Kills the per-round full-ctx projection recompute (first light was O(ctx)/round -> 7 tok/s).

Fields§

§k: Vec<CudaSlice<f32>>§v: Vec<CudaSlice<f32>>§len: usize§cap: usize

Implementations§

Source§

impl DflashKv

Source

pub fn new( e: &Engine, cfg: &DflashCfg, cap: usize, ) -> Result<Self, Box<dyn Error>>

Source

pub fn new_cold_at( e: &Engine, cfg: &DflashCfg, cap: usize, pos: usize, ) -> Result<Self, Box<dyn Error>>

A COLD DRAFTER at a restored trunk boundary (lane/spec-exclusions-20260902, the MEMRA_SPEC_WARM=1 arm): a fresh KV whose logical length is already pos (the restored prefix the trunk cache holds) but which owns NO ctx rows below it — floor == len == pos. The restored prefix’s tap features do not exist (the trunk planes hold K/V latents, not the tapped residual rows), and re-running the trunk to recover them is the prime the restore exists to skip; so instead the drafter starts with an empty context at the right absolute position and fills it from the suffix prime’s taps and every committed round from there, exactly as a cold session over a shorter prompt would. Row addressing (rope positions, kv.len == cache.pos at every round boundary) is identical to a tail import; only the attention floor differs.

The window_rows below the floor are zero-filled so a later export_tail (which starts at the floor anyway) can never publish uninitialised bytes even under a future geometry mistake — finite zeros are the same belt from_tail wears.

REFUSES under MEMRA_DFLASH2_SDPA_CLIP=0: the legacy full-scan kernel has no floor arm (d2_windowed_attn), and a cold drafter that scored 2048 zero rows at e^0 each would draft silently worse; the caller serves the plain hit by name instead.

Source

pub fn floor(&self) -> usize

The first ctx row this KV owns (doc on the field): 0 unless the KV was born as a cold drafter or imported from a short (floor-bearing) tail.

Source§

impl DflashKv

Source

pub fn export_tail(&self, e: &Engine, upto: usize) -> Option<DflashKvTail>

Copy out the readable tail ending at upto (see DflashKvTail). None when there is nothing to publish or an allocation fails — publication is always optional.

upto IS NOT self.len, and conflating them was the bug the first exactness-gate run caught: publication happens at the scheduler’s drain sweep, by which time the session has committed generated rows, so len had run 35 rows past the capture boundary and every restore was refused with draft KV len 30364 != prompt 30329. The trunk planes are copied at the capture pos for the same reason; the tail must agree with them.

Source

pub fn from_tail( e: &Engine, cfg: &DflashCfg, cap: usize, tail: &DflashKvTail, ) -> Option<Self>

Rebuild a draft KV from a published tail: a fresh allocation at cap, the tail copied back to the SAME absolute rows it came from, and len restored so the next round addresses positions exactly as a cold-primed session would.

Rows below tail.base are ZEROED, not left uninitialised. The clipped SDPA never reads below the block’s window floor, but the legacy full-scan kernel (MEMRA_DFLASH2_SDPA_CLIP=0, the rollback seam) scans EVERY row into the score and the output, relying on masked rows contributing exactly zero — an identity that holds only for finite data (0.0 * NaN = NaN, and an uninit K row can produce a NaN score that poisons the softmax sum). Zeros keep that identity on both kernel arms, so a clip rollback on a restore-armed box stays byte-exact instead of decoding silent garbage (review round 3). Rows above tail.len stay uninit — equally unwritten and unread in the cold path, so restored matches cold there.

This function still REFUSES rather than trusts the window math — if the tail does not cover the window, the caller gets None and must cold-prime.

FLOOR-BEARING TAILS (lane/spec-exclusions-20260902): a tail whose exporter had a context floor covers only [floor, len) and the import inherits floor = tail.base (the first row it actually owns), so the round attention never reads the zero rows below it. That needs the clipped kernel (d2_windowed_attn), so a floor-bearing tail under MEMRA_DFLASH2_SDPA_CLIP=0 refuses here by name. Full tails keep floor = 0 and the pre-lane program exactly.

Source

pub fn cfg_window_rows(&self) -> usize

Rows a restore must carry (see DsparkSpecSession::draft_tail_rows). Stored here because DflashKv owns the row geometry; the value comes from the drafter cfg.

Source

pub fn row_bytes(&self) -> usize

Bytes per row per layer (n_kv * head_dim * 4), the unit both the export and the import address rows in.

Source

pub fn n_layer(&self) -> usize

Number of draft layers, i.e. how many per-layer planes an export produces.

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.