Struct HybridModel 

pub struct HybridModel {
    pub cfg: ModelConfig,
    pub embd: EmbedHost,
    pub output_norm: GpuTensor,
    pub output: GpuTensor,
    pub layers: Vec<HybridLayer>,
    pub mtp: Option<MtpHead>,
    pub embd_gpu: OnceLock<CudaSlice<u8>>,
    pub gemma4_aux: Option<GemmaAux>,
    pub step35_aux: Option<Step35Aux>,
    pub prime_slabs: Mutex<HashMap<usize, Arc<Mutex<PrimeSlabs>>>>,
}

Fields

cfg: ModelConfig

embd: EmbedHost

output_norm: GpuTensor

output: GpuTensor

layers: Vec<HybridLayer>

mtp: Option<MtpHead>

embd_gpu: OnceLock<CudaSlice<u8>>

Lazily-uploaded DEVICE copy of the raw embed table (spec/graph hot loops gather rows on-device instead of host-dequant + htod). ~0.5GB; uploaded once on first use.

gemma4_aux: Option<GemmaAux>

step35_aux: Option<Step35Aux>

step35 (Step-3.7-Flash) model auxiliaries — Some iff cfg.step35.is_some().

prime_slabs: Mutex<HashMap<usize, Arc<Mutex<PrimeSlabs>>>>

PRIME ACTIVATION SLABS (piecewise-graph foundation, 2026-07-26): the layer loop’s seven trunk transients live in RESIDENT per-model buffers instead of per-call pool allocs — kills ~224 alloc/free API calls per prime AND freezes the Lt GEMM operand addresses (nvjet’s alignment-variant kernels become run-to-run stable once their pointers stop moving). Sized on first prime to the largest T seen. The map lock covers lookup/grow only; each device owns a separate slab lock so PP stages on distinct devices can drive their host-synchronized layer walks concurrently.

Implementations

impl HybridModel

pub fn load(e: &Engine, g: &GgufFile) -> Result<Self, Box<dyn Error>>

Load a hybrid (qwen35) model from GGUF. Thin byte-identical wrapper over load_from_source.

pub fn load_without_mtp( e: &Engine, g: &GgufFile, ) -> Result<Self, Box<dyn Error>>

Plain-generation loader. run-gen never calls the optional draft head, so avoid loading its weights and expert bank while preserving the model config and all trunk semantics.

pub fn load_from_source( e: &Engine, src: &dyn TensorSource, ) -> Result<Self, Box<dyn Error>>

Load a hybrid model from any TensorSource (GGUF or a safetensors HF checkpoint). The whole loop speaks ggml names; the source maps them (and, for safetensors, applies the SSM value transforms via the owned-buffer seam). The forward graph is untouched.

pub fn load_from_source_without_mtp( e: &Engine, src: &dyn TensorSource, ) -> Result<Self, Box<dyn Error>>

Source-backed twin of load_without_mtp, used by the safetensors/repack run-gen path.

pub fn ensure_embed_resident(&self, e: &Engine) -> Result<(), Box<dyn Error>>

Force the device embed table resident, FALLIBLY (F5 right-size ladder, 2026-08-05). The lazy embd_gpu.get_or_init(.. expect ..) sites panic the GPU worker on OOM; on a VRAM-tight rig a right-sized spec session that “fits” can leave too little for this ~hundreds-of-MB upload and die on its first prefill (observed: research/specpool-20260804/server-ladder-miss.log). The server calls this after each ladder landing so the biggest lazy transient surfaces as a catchable Err (shrink further / fall back) instead of a panic. No-op when the host-gather door (MEMRA_EMBED_DEV=0) is open or the table is already resident.

pub fn embed( &self, e: &Engine, tokens: &[u32], ) -> Result<CudaSlice<f32>, Box<dyn Error>>

impl HybridModel

pub fn forward( &self, e: &Engine, tokens: &[u32], ) -> Result<Vec<f32>, Box<dyn Error>>

Prefill forward over tokens; returns logits [T, n_vocab] (host f32).

pub fn forward_last( &self, e: &Engine, tokens: &[u32], ) -> Result<Vec<f32>, Box<dyn Error>>

Prefill that returns ONLY the last token’s logits — the common case (greedy/sample needs just the final position to start decode). Runs the trunk over all T, then the lm_head (output.weight, the largest matrix — 248320 rows) on the LAST hidden row ONLY, not all T. On a 512-token prompt this turns a [512,248320] GEMM into [1,248320] — the dominant prefill cost (nsys: ~99ms when done for all T). Bit-identical last-row logits to forward()[last].

pub fn prime_cache( &self, e: &Engine, tokens: &[u32], cache: &mut Cache, queued_after: usize, ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn Error>>

BATCHED PROMPT PRIME (the measured #1 e2e gap, e2e-image-1): forward_last’s batched prefill body EXTENDED to leave a DECODE-READY cache behind — vs the tokenwise prime’s ~102/38 tok/s (9B/27B) decode_step loop, this runs the whole prompt at prefill throughput. (a) full-attn layers append their T post-RoPE K/V rows into cache.kv[il] via the SAME per-row quantize kernel as the decode append (bit-identical cache bytes per row); (b) linear layers run STATEFULLY from the cache’s current recurrent state (zero at a fresh prime): carried-ring conv (ssm_conv1d_tm_state) + ONE gdn_scan(state_in, state_out) whose internal sequential t-loop equals T chained T=1 steps — but with the NORMAL prefill matmul dispatch (GEMM at m>=16), NOT the decode-exact MMVQ the spec verify uses (prime is a prefill-regime pass; the run-gen prefill==decode argmax gate is the accuracy authority, exactly as for forward_last); (c) cache.pos/KV len/len_d advance by T. Returns (last-row logits host, h_seed = last-row PRE-output_norm hidden [n_embd], hiddens = the full pre-output_norm hidden stack [T, n_embd] — generate_spec’s prompt_h). FRESH-PROMPT ONLY (cache.pos == 0): the fa_prefill tiles attend within tokens alone. forward_last itself stays untouched (kernel-check / run-gen gate on it).

queued_after (lane/tick-seg, 2026-08-07): the number of prompt tokens of the SAME REQUEST that the caller will prime in LATER calls — 0 when this call is the whole request (every single-shot caller). Serve splits a long prompt across SEVERAL prime_cache calls (one per scheduler tick, plus the prefix-cache LCP split), and the request’s absolute end position seq_end = cache.pos + t + queued_after steers step35’s SWA prefill arm — computing it per CALL made the arm a function of the tick budget (budgets 512/256/64 DIFFER 1.813e0 vs monolithic, greedy diverging at step 6; dark lanes default to 256 AND cap by live SLO headroom, so identical judge requests primed differently under load — research/tick-seg-20260807, receipt in research/step35-chunkfix-20260807 §9). The parameter is what prime_cache structurally lacked: it cannot know from tokens and cache alone whether more of the request is coming. A SESSION CONTINUATION (a NEW user turn primed onto a live cache) is a NEW request — its arithmetic is keyed to its own extent, so those callers pass 0; only a caller that SPLITS one request across calls passes the remainder.

pub fn prime_cache_overlaid( &self, e: &Engine, tokens: &[u32], cache: &mut Cache, queued_after: usize, overlay: Option<&EmbedOverlay>, ) -> Result<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn Error>>

prime_cache with a vision embedding overlay (lane/vision): image merger outputs replace the <|image_pad|> token embeddings at prompt-relative positions before the trunk walk — the mixed-embedding prime. Text-only callers use prime_cache (overlay None, byte-identical path). v1 scope: the serial chunk walk only — PP prime arms and gemma4 refuse loudly (the vision serving box is single-GPU).

pub fn prime_slabs_get( &self, e: &Engine, t: usize, n_embd: usize, n_ff_max: usize, ) -> Result<Arc<Mutex<PrimeSlabs>>, Box<dyn Error>>

(cache.pos + i). Returns (last-row logits, h_seed, this chunk’s hidden stack [T, n_embd]). See HybridModel::prime_slabs — the eager prime’s resident trunk transients. PER-DEVICE since lane/pp-leverb (2026-08-08): the map is keyed by the allocating engine’s CUDA ordinal — under the prime stage split each stage’s range walks through its OWN slabs on its own device (a dev0 slab dereferenced by a dev1 kernel would be a peer read per GEMM operand, the exact class Lever B removes). Single-device rigs see one entry, byte-identical behavior.

pub fn prime_chunk_captured( &self, e: &Engine, x_in: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize, cache: &mut Cache, len_d: &CudaSlice<i32>, logits_out: &mut CudaSlice<f32>, h_seed_out: &mut CudaSlice<f32>, ) -> Result<(), Box<dyn Error>>

CAPTURE-SAFE prime trunk (task #14 increment 1, ARCHITECTURE-H100.md design v2): prime_chunk’s layer stack with every capture hazard hoisted — x is the PRE-EMBEDDED input in a STABLE graph-input buffer (embed + its 8MB htod stay eager, one launch), pos_d is a baked device param (fresh prime = 0..T, constant per bucket), logits/h_seed/hidden stay DEVICE-resident (no dtoh), and cache.pos is NOT advanced (host state — the replay wrapper owns it). Body mirrors prime_chunk (the prime-graph-gate pins them together). KNOWN smoke-scope gap: append host-len bookkeeping still runs on the host per call — the real replay path moves the write slot to the len_d device counter (increment 3). GRAPH-OUTPUT CONTRACT: results are COPIED into caller-provided stable buffers (logits_out [n_vocab], h_seed_out [n_embd]) — every internal allocation drops INSIDE the capture region (alloc+free node pairs). Retaining an in-capture allocation across end_capture makes instantiate throw INVALID_VALUE (smoke finding 2026-07-26), and under AUTO_FREE_ON_LAUNCH its address wouldn’t survive a launch anyway — the decode GraphSession’s pre-allocated-output pattern is the law here.

pub fn prime_cache_batch( &self, e: &Engine, prompts: &[&[u32]], caches: &mut [&mut Cache], ) -> Result<Vec<(Vec<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn Error>>

Cross-request BATCHED fresh prime (task #13, design in ARCHITECTURE-H100.md): the trunk’s token-parallel ops (embed, norms, adds, ffn, projection GROUPS) run once on the CONCATENATION of B sequences — GEMMs at m = sum_T, the continuous-batching win the serving-lane bench attributed the remaining vLLM gap to. The stateful mixer CORES (QK-norm/RoPE/FA/append, conv/GDN scans) run per sequence on split projection buffers (D2D row copies; the mixers’ own out-projections stay per-seq this increment). CONTINUATION primes (increment (b), 2026-07-30): cache.pos > 0 seqs batch the projections/FFN/lm_head exactly like fresh; the mixer cores take the per-seq CONTINUATION arms (Full: core_inner with carried pos_d + fa_prefill_view over the quantized past; Linear: the stateful pad_view twin — the same state carry the chunked single-seq prime rides). The fresh-only favl/gdn-vl fast paths stay byte-identical (gated on !carried). gemma4 models have no continuation prime (v0 monolithic fresh) — carried gemma4 batches return Err (caller falls back to single-chunk serving). NUMERIC CONFIG: a concat GEMM tiles K differently than per-seq GEMMs — same class as every prefill GEMM change; prime_batch_gate arbitrates (argmax + stream battery).

pub fn full_attn( &self, e: &Engine, fa: &FullAttnLayer, h: &CudaSlice<f32>, pos_d: &CudaSlice<i32>, t: usize, il: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Full-attention mixer with QK-norm, partial RoPE, sigmoid output gate (qwen35 :257-336).

il = layer index: step35 needs it (per-layer n_head / rope width / window / gate) and routes to its own mixer. Every other arch ignores it (uniform geometry).

pub fn linear_attn( &self, e: &Engine, la: &LinearAttnLayer, h: &CudaSlice<f32>, t: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Linear-attention (Gated DeltaNet) mixer (qwen35 :338-470).

impl HybridModel

pub fn moe_ffn_il( &self, e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize, il: u16, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

MoE FFN (EDGE-1). z: [T, n_embd] (already post-attention-normed). Returns moe_out [T, n_embd]. Node-for-node vs llama.cpp build_moe_ffn + qwen35moe::build_layer_ffn.

il is the trunk layer index — the residency-cache key prefix (a gate-expert of layer 3 is a different 860160-byte block than the same expert of layer 7).

Routing: host softmax+sort (default) OR the fused router kernel (MEMRA_FUSED_ROUTER). Dispatch: stage-every-token into 3 scratch slots (default) OR the SLRU residency cache (MEMRA_MOE_CACHE). The cache-HIT weight path is bit-identical to stage-every-token (§B.3). Convenience wrapper used by the hybrid trunk/MTP loops: pulls dims + max-block from self.

pub fn moe_ffn_il_prefill( &self, e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, t: usize, il: u16, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Prefill twin: Step35 promotes expert-grouped dispatch by default while decode/spec callers keep moe_ffn_il and therefore retain their existing dispatch class.

pub fn moe_ffn_il_zq8( &self, e: &Engine, m: &MoeWeights, z: &CudaSlice<f32>, zq8: Option<&(CudaSlice<i8>, CudaSlice<f32>)>, t: usize, il: u16, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Decode-path twin with a PRE-QUANTIZED z (from add_rms_norm_zq8): threads (zq, zd) into the t=1 dev arm so the per-layer standalone quantize_q8_1 launch folds away. Identical bytes (the fused kernel reproduces quantize_q8_1 exactly); every other path ignores the pair.

pub fn stage1_h2d_per_token(&self) -> u64

Stage-1 (no-cache) per-DECODE-TOKEN H2D bytes: every routed block re-staged every layer every token = sum over MoE layers of n_used * (gate+up+down expert_stride). The §D.4 PCIe baseline.

pub fn save_cpu_expert_residency_profile( &self, e: &Engine, path: &Path, ) -> Result<(), Box<dyn Error>>

Persist the frozen residency set so a later process can restage it directly and skip the profiling warmup. Plain text: a versioned header binding slot geometry, then one layer proj ex triple per line. A mismatched or stale profile is rejected at load (header check) or degrades to fewer restaged blocks (per-id checks); either way the post-freeze argmax gate still validates the serving assignment.

pub fn restore_cpu_expert_residency_profile( &self, e: &Engine, path: &Path, ) -> Result<bool, Box<dyn Error>>

Restage a saved freeze profile and freeze immediately, skipping the profiling warmup. Returns false (leaving the cache untouched for a normal warmup) when the profile is missing or its header does not match this model’s slot geometry.

pub fn freeze_cpu_expert_residency( &self, e: &Engine, ) -> Result<(), Box<dyn Error>>

Freeze the heterogeneous CPU/GPU split after the caller’s discarded profile warmup.

pub fn ffn_act( e: &Engine, cfg: &ModelConfig, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, act: &mut CudaSlice<f32>, n: usize, ) -> Result<(), Box<dyn Error>>

FFN activation dispatch: swigluoai (clamped, alpha/limit) when cfg.m3 says so, else the standard SiLU*up. One seam so every FFN site (dense, routed expert, shared expert) follows the model’s activation exactly.

NO-il FORM: cannot apply step35’s PER-LAYER SwiGLU clamp. Only call it from a site whose layer provably has no live limit (dense-FFN layers, MTP blocks) — ffn_act_lim is the form for anything that can land on a clamped layer.

pub fn start_moe_prefetch_predictor( &self, e: &Engine, cfg: &ModelConfig, ) -> Result<(), Box<dyn Error>>

Start the prediction-guided prefetch worker (MEMRA_MOE_PREFETCH=depth). Call after residency freeze: the worker filters against a static snapshot of the frozen HBM set. Builds a fully-owned per-layer table (host router copies via one-time DtoH, bias, active mask, prebuilt projection descriptors) so no model reference escapes.

pub fn moe_route_sigmoid_host_public( logits: &[f32], t: usize, n_expert: usize, n_used: usize, bias: Option<&[f32]>, sf: f32, route_norm: bool, active: Option<&[bool]>, ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn Error>>

Sigmoid-routing oracle shared by the prefetch predictor and kernel-check: identical selection math to the rollback runtime, applied to host-computed logits.

impl HybridModel

pub fn gemma4_decode_step_dc( &self, e: &Engine, token_d: &CudaSlice<u32>, pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>, embd_qt: i32, embd_rb: usize, cache: &mut Cache, n_vocab: usize, cap_bucket_max: Option<(usize, usize)>, ) -> Result<CudaSlice<u32>, Box<dyn Error>>

gemma4 DEVICE-COUNTER decode step (graph arc): token id + rope pos + KV lengths live in device counters; ZERO varying host kernel args. cap_bucket_max = Some(bucket) for graph capture (host mirrors untouched, n_splits from bucket, full-buffer KV views) / None for the eager-dc gate path (host mirrors advanced, live geometry — bit-identical target = gemma4_decode_step_h’s token stream). V1 scope: t_kv <= sliding_window (no window views in-graph; the driver gates).

pub fn gemma4_decode_step_dc_into( &self, e: &Engine, token_d: &CudaSlice<u32>, pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>, embd_qt: i32, embd_rb: usize, cache: &mut Cache, n_vocab: usize, cap_bucket_max: Option<(usize, usize)>, tok_out: &mut CudaSlice<u32>, ) -> Result<(), Box<dyn Error>>

CAPTURE body: argmax lands in the PERSISTENT tok_out (same buffer = same address on every replay; pass token_d itself for the self-feeding graph loop).

pub fn g4_dc_slots(&self, e: &Engine) -> Result<G4DcSlots, Box<dyn Error>>

Persistent transient slots for the ALLOC-FREE captured dc step (the graph door): every buffer the step produces per token lives here, allocated ONCE pre-capture, so the captured graph carries zero cuMemAllocAsync/Free nodes (the 226us/launch tax, osrt 2026-07-23). Sized for the model’s max per-layer shapes. Build the slot set (call OUTSIDE any capture).

pub fn gemma4_decode_step_dc_slotted( &self, e: &Engine, token_d: &CudaSlice<u32>, pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>, embd_qt: i32, embd_rb: usize, cache: &mut Cache, n_vocab: usize, cap_bucket_max: Option<(usize, usize)>, sl: &mut G4DcSlots, tok_out: &mut CudaSlice<u32>, ring: Option<(&mut CudaSlice<u32>, usize)>, ) -> Result<(), Box<dyn Error>>

ALLOC-FREE dc step (capture body): kernel-for-kernel mirror of gemma4_decode_step_dc_into at t=1 with every transient slot-fed. Dense gemma4 only (12B/31B; uniform q4_0 trunk guarantees the fused2/3 arms).

pub fn gemma4_generate_graph( &self, e: &Engine, prompt_pos: usize, first_token: u32, cache: &mut Cache, max_new: usize, eos: &[u32], on_token: impl FnMut(u32) -> bool, ) -> Result<(Vec<u32>, StopReason), Box<dyn Error>>

gemma4 GRAPH-REPLAY greedy loop: per (swa-key, global-key) fa bucket, capture ONE full dc step (self-feeding: argmax writes token_d in-graph) and replay it — one graph launch per token, one 4B dtoh. V1 scope: whole generation under the sliding window (no window views in-graph); caller gates and falls back to the dc-eager loop.

impl HybridModel

pub fn is_gemma4_e4b(&self) -> bool

pub fn gemma4_e4b_decode_step_t_am_dev( &self, e: &Engine, tok_d: &CudaSlice<u32>, t: usize, pos0: usize, cache: &mut Cache, ) -> Result<(CudaSlice<u32>, CudaSlice<f32>), Box<dyn Error>>

E4B batched VERIFY (device tokens, the spec round’s t=K+1 step): t rows through the e4b trunk (per-row causal attention; own-KV layers append t rows host-len, KV-shared layers ride their targets), per-row device argmax + the POST-output_norm hidden stack (the drafter’s h convention). Advances cache.pos/kvl.len by t — the spec round rolls back rejected rows (shared layers have no KvLayer, so the plain rewind covers exactly the layers that appended).

pub fn gemma4_e4b_decode_step_dcg( &self, e: &Engine, token_d: &mut CudaSlice<u32>, pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>, embd_qt: i32, embd_rb: usize, cache: &mut Cache, n_vocab: usize, bucket: usize, ) -> Result<(), Box<dyn Error>>

E4B GRAPH-CAPTURABLE dc step: same trunk as the dc step but token_d is updated IN PLACE (self-feeding replay) and every launch arg is a device counter — pos from pos_d (inc’d in-stream), KV slots from len_d (advanced in-stream), attention from fa_decode_dc at bucket. Host mirrors (cache.pos / kvl.len) advance in the caller’s replay loop. UNDER-WINDOW regime only (the caller gates pos + budget < window).

pub fn gemma4_e4b_decode_step_dc( &self, e: &Engine, token_d: &CudaSlice<u32>, pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>, embd_qt: i32, embd_rb: usize, cache: &mut Cache, n_vocab: usize, ) -> Result<CudaSlice<u32>, Box<dyn Error>>

E4B DEVICE-COUNTER decode step (the dc serving arm): token id rides token_d, the greedy argmax lands in the returned device buffer — 4B/token host traffic. The layer stack is gemma4_e4b_trunk_core, i.e. the SAME functions the eager chain runs (stream identity by construction, not by twin-kernel parity). Host KV mirrors advance like the 26B dc-eager arm (window views are host math); len_d stays synced by the caller’s entry sync + the appends here don’t read it. Graph capture is NOT wired (no cap_bucket_max) — the E4B graph arc comes after the perf gates.

impl HybridModel

pub fn decode_step( &self, e: &Engine, token: u32, cache: &mut Cache, ) -> Result<Vec<f32>, Box<dyn Error>>

One decode step for token at cache.pos; returns logits [n_vocab] (host f32). Advances cache.

pub fn decode_step_aux( &self, e: &Engine, token: u32, cache: &mut Cache, aux_layers: &[usize], ) -> Result<(Vec<f32>, Vec<CudaSlice<f32>>), Box<dyn Error>>

EAGLE3 aux-hidden capture (EAGLE-PLAN N1): one decode step that ALSO returns the trunk residual-stream x taken AFTER each of the blocks in aux_layers (the EAGLE3 encoder feeds these 3 layer hiddens through fc). Returns (logits[n_vocab] host, aux: Vec<[n_embd] dev>), one device buffer per requested aux layer, in aux_layers order. The captured tensor is the residual x produced by that block (x2 at the loop tail), cloned before the next block overwrites it — cheap (one clone_dtod of [n_embd] per aux layer). T=1 decode regime.

pub fn decode_step_hy3_layer0_stages( &self, e: &Engine, token: u32, cache: &mut Cache, ) -> Result<(Vec<f32>, Hy3Layer0Stages), Box<dyn Error>>

Diagnostic-only Hy3 layer-0 trace through the real eager T=1 serving path. Besides the final block residual, this captures the attention output before its residual add, the after-attention residual, and the dense-MLP output before the final residual add.

pub fn decode_step_h( &self, e: &Engine, token: u32, cache: &mut Cache, ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn Error>>

Like decode_step, but ALSO returns the trunk’s hidden state x taken BEFORE the final output_norm (MTP-PLAN §A: this is h_seed for the NextN head). Device buffer [n_embd].

pub fn decode_step_h_ppn_deferred( &self, e: &Engine, token: u32, cache: &mut Cache, ) -> Result<PendingLogits, Box<dyn Error>>

M2 increment 3 (DEFERRED READBACK — the pipelining seed): the ppN step WITHOUT the terminal logits D2H. Returns PendingLogits (device logits + completion event + the runtime’s dedicated readback stream); the caller keeps 2+ tokens in flight by enqueueing step t+1 BEFORE waiting step t (with MEMRA_PP_OVERLAP=1 the double-buffered boundary slots actually alternate, so stage 0 of t+1 runs under stage 1..N-1 of t; the slot ev_tx/ev_rx chain keeps each token’s math fully event-ordered either way — enqueueing deeper than 2 is CORRECT, the slots simply serialize device-side).

EXACTNESS CONTRACT: per-token logits are BIT-IDENTICAL to the serial arm — same kernels, same per-token event order; only the host-side wait moves (scheduling change, never math). The pipelined replay arm of ppn-gate proves it per step.

NOT produced here (both are trunk COPIES — no math feeding the logits changes): h_seed and the MEMRA_DUMP_HN diagnostic tap. The serving loop decides their deferred form when it adopts this API.

The caller advances the token stream, so cache.pos advances at ENQUEUE (host state; device work is event-ordered regardless).

pub fn decode_step_lockstep( &self, e: &Engine, tokens: &[u32], caches: &mut [Cache], ) -> Result<Vec<Vec<f32>>, Box<dyn Error>>

LOCKSTEP MULTI-STREAM decode (lane-3 M1): m independent streams advance one token each through a single per-layer walk. Per-stream math is identical to decode_step_h (same fusion chain, same mixer and FFN calls against that stream’s own Cache), so each stream’s token sequence is bit-identical to its single-stream run. The lockstep order puts the m streams’ layer-il MoE calls adjacent in time, so one stream’s expert-cache fill serves its siblings within the step — the measured cross-stream io amortization (1.12x/1.32x/1.66x at m=2/4/8) lands without batching attention or the CPU ABI.

pub fn decode_step_dc( &self, e: &Engine, token_d: &CudaSlice<u32>, pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>, embd_qt: i32, embd_row_bytes: usize, cache: &mut Cache, n_vocab: usize, ) -> Result<CudaSlice<u32>, Box<dyn Error>>

DEVICE-COUNTER decode step (CUDA-GRAPH-PLAN Phase 2). A clone of decode_step_h that removes the two per-step VARYING host kernel-args by reading them from device counters:

1. the KV-append write slot -> per-layer kvl.len_d (device i32[1])
2. the fa_decode t_kv bound -> the same kvl.len_d after inc_seqlen plus it keeps the token id + rope pos DEVICE-RESIDENT (embed_gather_device, device rope pos, argmax_token_device). NO graph capture yet — runs the kernels eagerly through the counter path. Must be BIT-IDENTICAL to decode_step_h’s token stream (the gate).

Args: token_d = resident device token id [1] (this step’s input token); pos_d = resident device rope pos i32[1] (== cache.pos at entry; INCREMENTED in-path); embd_gpu = resident embed table; (qt,row_bytes) from EmbedHost::qt_and_row_bytes. Returns the NEXT token id device buffer. cache.pos and each kvl.len/kvl.len_d are advanced to match decode_step_h.

pub fn decode_step_dc_cap( &self, e: &Engine, token_d: &mut CudaSlice<u32>, pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>, embd_qt: i32, embd_row_bytes: usize, cache: &mut Cache, n_vocab: usize, bucket_max: usize, ) -> Result<(), Box<dyn Error>>

CAPTURE body for CUDA-graph replay (CUDA-GRAPH-PLAN Phase 3). One full decode step enqueued entirely on e.stream() with ZERO host sync and ZERO per-step varying host kernel-args:

• embed reads the PERSISTENT device token_d (last step’s argmax), writes scratch x.
• full-attn layers size n_splits from bucket_max (fixed for this capture); the kernel reads the ACTUAL t_kv from the device counter kvl.len_d. KV append + device-counter inc happen in-graph. The host kvl.len/cache.pos are NOT advanced here (the driver advances the host mirrors once per replay; only the DEVICE counters advance inside the graph).
• linear-attn layers use the persistent-state variant (copy-back, stable pointers).
• lm_head -> parallel 2-pass argmax (argmax_partial_f32+argmax_final_f32) writes the next id into the PERSISTENT token_d.
• inc_seqlen(pos_d) advances the rope-pos device counter in-graph. Captured ONCE per bucket_max; replayed for every t_kv in that bucket. Bit-identical to eager when bucket_max reproduces eager’s n_splits for the replayed t_kv (the bucket-key contract).

pub fn decode_step_dc_cap_masked( &self, e: &Engine, token_d: &mut CudaSlice<u32>, pos_d: &mut CudaSlice<i32>, embd_gpu: &CudaSlice<u8>, embd_qt: i32, embd_row_bytes: usize, cache: &mut Cache, n_vocab: usize, bucket_max: usize, mask: Option<(&CudaSlice<u32>, usize)>, ) -> Result<(), Box<dyn Error>>

decode_step_dc_cap + GRAMMAR MASK (constrained decoding): with mask = Some((buf, words)), mask_logits_f32 bans the packed bitset’s unset ids IN the captured graph — a stable-pointer read between lm_head and the in-graph argmax (the KV-pointer pattern: contents change per step, address is baked). None is bit-for-bit the unmasked capture.

pub fn generate_graph( &self, e: &Engine, gs: &mut GraphDecodeState, prompt: &[u32], max_new: usize, ) -> Result<Vec<u32>, Box<dyn Error>>

CUDA-GRAPH decode driver (CUDA-GRAPH-PLAN Phase 3). Primes the prompt EAGERLY (device-counter decode_step_dc, advancing host + device counters together), then generates max_new tokens by CUDA-graph REPLAY: per step it picks the t_kv bucket key, captures a graph on first sight of that key (re-using the SAME persistent counters/cache so replays continue the sequence), and replays. The argmax-written next token stays device-resident in gs.token_d; we read back only the [1] u32 after each launch (the gate compares it; a real server can defer this). Returns the generated token ids. Greedy. Bit-identical to eager decode_step (the gate).

CAPTURE STATE HYGIENE: capture_graph runs the step body 3x (2 warmup + 1 capture), each of which mutates the device KV/conv/ssm/counter state. We SNAPSHOT the cache + device counters + token id before capturing and RESTORE them after, so the 3 throwaway runs leave zero residue and replay resumes from the true pre-capture state.

pub fn graph_session_new( &self, e: &Engine, prompt: &[u32], max_new: usize, ) -> Result<(GraphSession, u32), Box<dyn Error>>

Step-wise CUDA-graph decode session (ARCHITECTURE-H100.md graph-serving lane, 2026-07-26): generate_graph’s prime+capture lifted into a long-lived session so a SERVING scheduler can replay ONE step per tick instead of blocking a whole generation. Serving policy (measured): graphs win only at B=1 (214 solo vs 425 aggregate batched-eager at B=4) — this is the single-interactive-session path. Capture discipline is generate_graph’s verbatim: event tracking must be OFF for every buffer the graph references (new() toggles it), capture at bucket_max = pos + max_new + 1, fa geometry retuned per step (fa_apply, FP lockstep with eager).

pub fn graph_session_from_cache( &self, e: &Engine, cache: Cache, first_token: u32, max_new: usize, ) -> Result<(GraphSession, u32), Box<dyn Error>>

GraphSession over an ALREADY-PRIMED cache (round 35): keeps the chunked-prefill TTFT. graph_session_new’s token-wise re-prime made solo long-prompt promotion a net ~3x END-TO-END LOSS (measured live: 871-tok prompt + 400 gen = 6.4s vs ~2.2s eager). Device counters sync from host state; capture recipe unchanged. Requires event tracking OFF (engine default; MEMRA_EVT=1 callers must not use this — the primed cache’s buffers would carry events, illegal inside capture).

pub fn graph_session_from_cache_masked( &self, e: &Engine, cache: Cache, first_token: u32, max_new: usize, mask_init: Option<&[u32]>, ) -> Result<(GraphSession, u32), Box<dyn Error>>

graph_session_from_cache + GRAMMAR MASK (constrained decoding, 2026-08-03): mask_init = Some(packed bitset) allocates the session’s stable mask buffer (tracking is OFF here — capture-legal), seeds it with the FIRST step’s mask, and captures mask_logits_f32 into the graphed step. The caller re-uploads contents per step via GraphSession::upload_mask — same stable-pointer discipline as the KV len_d counters. None = the unmasked session, byte-identical.

pub fn graph_session_recapture_pub( &self, e: &Engine, sess: &mut GraphSession, ) -> Result<(), Box<dyn Error>>

Measurement door for graph_session_recapture (graph-allocfree-probe): the capture path timed WITHOUT the prompt prime. Same call the live step() makes at a kernel-class crossing.

pub fn generate( &self, e: &Engine, prompt: &[u32], max_new: usize, ) -> Result<Vec<u32>, Box<dyn Error>>

Greedy generation: prime with prompt tokens (decode them in sequence to build state), then generate max_new tokens. Returns the generated token ids. (Back-compat: greedy, no EOS/stop — used by the decode==prefill validation gate. New code uses generate_with.)

pub fn generate_with<F: FnMut(u32) -> bool>( &self, e: &Engine, prompt: &[u32], params: &GenParams, sampler: &mut Sampler, on_token: F, ) -> Result<GenOutput, Box<dyn Error>>

The reusable serving generation API (BASE-3). Primes the prompt, then samples up to params.max_new tokens, stopping on EOS, any stop-token, or the context-length guard. Calls on_token(id) after each emitted token (for streaming; return false to stop early). Returns GenOutput { tokens, stop_reason }. Does NOT detokenize — the caller (which owns the tokenizer) handles text + stop-STRING matching on the detokenized tail.

pub fn linear_attn_decode( &self, e: &Engine, la: &LinearAttnLayer, h: &CudaSlice<f32>, cache: &mut Cache, il: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Linear-attention decode: conv with ring-buffer state, GDN scan carrying SSM state.

pub fn linear_attn_decode_pre( &self, e: &Engine, la: &LinearAttnLayer, h: &CudaSlice<f32>, hq: &CudaSlice<i8>, hd: &CudaSlice<f32>, cache: &mut Cache, il: usize, persistent: bool, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

PRE-QUANTIZED-INPUT variant (DECODE attn-input NORM-FUSION lever): the caller passes the post-attn-norm activation ALREADY q8_1-quantized (hq,hd) (produced by rms_norm_q8_1, fusing the attn_norm + the mixer’s internal quantize_q8_1). Skips the internal quantize. Caller GUARANTEES the projections are q8_1-fast. persistent selects the capture-safe state plumbing. BIT-IDENTICAL to linear_attn_decode(h) when (hq,hd)==quantize_q8_1(rms_norm(x)*w).

impl HybridModel

pub fn decode_batch_cap() -> usize

Batched-decode width cap. 8 = the exactness-tier default (see the assert below); MEMRA_DECODE_BATCH_CAP overrides for tier-probe measurement, clamped to 32.

pub fn decode_batch_exact16_ok(&self) -> bool

EXACT-16 TIER admission (increment 3a, 2026-08-01, 5090 receipts research/batched-tick-inc3-20260801): true iff EVERY matmul the batched decode step runs has a per-(token,row) bit-exact kernel class at m=9..16 under the verify_exact scope — i.e. the batched-mmvq b16 family (32-thread warp reduce, the exact m=1 mmvq program per column) or the e4m3 grid.y=m mmvq catch-all. Q8_0 qualifies only with the split-plane mirror (rp4, MEMRA_Q8RP): its b16 kernel exists only as the _rp twin. Float matmuls (cuBLASLt, n-dependent reductions) and MoE FFNs disqualify the model. Measured attribution for WHY the naked m=16 tier is not exact: the m>=16 arms (MMQ int8-MMA mul_mat_q — MEMRA_PP_Q8MMQ default-on — and qmatvec_gemm, both block-scale f32) and the m=9..15 dp4a tail (128-thread two-level reduce) all break per-row bit-identity vs isolated decode (gate2 step-0 bit-diffs, maxdiff ~1.3-2.3e-1).

pub fn b1_fast_on() -> bool

Opt-in/A-B seam for the eager B=1 fusion program. MEMRA_SERVE_B1FAST=1 sends an eligible solo tick through that program; unset/other values keep B=1 on the generic batched body, the same numeric class used at B>=2.

EXACTNESS, stated precisely (measured on-box 2026-08-05, sm_120 q9 NVFP4-MTP): the fast path is BIT-IDENTICAL TO decode_step_h — decode-batch-gate’s STRICT gate1 (--mode strict) PASSes with it ON and FAILs with it OFF at maxdiff 1.591e-1. It is deliberately NOT bit-identical to the batched body: the two carry a decode-config FP-composition gap (same class gate1’s config mode measures). That gap became correctness-visible under live load: Step35, Q35-MoE, and finally dense Q27 all produced load-history-dependent token streams, including early EOS, when a request crossed between the two programs. The generic body is therefore the correctness default; the eager program remains available only for fixed-solo A/Bs. Historical token-stream/performance receipts: research/servepath-p2-20260805 (greedy 150 ids + seeded-sampled identical to the run-gen oracle AND cross-arm, so the gap is sub-token here as designed).

Read fresh (an AtomicU8 memo, not a OnceLock): decode-batch-gate flips this seam BETWEEN gates in-process — gate1 needs the fast path ON to prove bit-identity, gate2 needs it pinned OFF to keep testing the batched body. A latch-once read would bake whichever gate ran first, so the gate could never test both sides. The memo caches the parse but set_b1_fast invalidates it.

pub fn set_b1_fast(on: bool)

Test/gate seam: force the B=1 fast path on or off for the rest of the process, overriding the env. Used by decode-batch-gate to exercise the opt-in eager arm and pin gate2’s default reference arm.

pub fn b1_fast_arch_eligible(&self) -> bool

Whether this architecture may switch a live serving row onto the eager B=1 fusion class. Qwen35-MoE must stay on the batched trunk at every width: its eager and batched hybrid/MoE walks are each deterministic, but crossing B=1 -> B>=2 changes greedy token ids and can introduce an early EOS (Q35 sellgate, 2026-08-12).

pub fn decode_step_batch( &self, e: &Engine, tokens: &[u32], caches: &mut [&mut Cache], ) -> Result<Vec<Vec<f32>>, Box<dyn Error>>

One batched greedy-decode step over B independent sequences. tokens[b] is sequence b’s input token; caches[b] its private cache (position, quantized KV, GDN/conv state). Returns the B logits rows (host, [n_vocab] each). Each cache’s pos/len advance exactly as decode_step_h would.

pub fn decode_step_batch_sampled( &self, e: &Engine, tokens: &[u32], caches: &mut [&mut Cache], samp: &[Option<DevSamp>], ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn Error>>

decode_step_batch + DEVICE-SIDE SAMPLING for eligible rows (the batched-tick lever, 2026-08-01): the host sampler’s temp-path is O(n_vocab) with a full-vocab exp per row (measured 1.36 ms/row at the 9B’s 248320 vocab = 10.9 ms/tick at B=8 — the single largest component of the serving tick). Here each requested row samples ON DEVICE between the lm_head matmul and the logits D2H: temp <= 0 (greedy): the 2-pass device argmax — bit-identical to host argmax (argmax-gate contract, same kernels as the dc serving path). temp > 0: gumbel_perturb(seed, ctr, temp) + the same argmax = ONE categorical draw from softmax(logits/temp) — the sampled-spec Philox machinery. Deterministic per (seed, ctr) and INDEPENDENT of batch composition (the isolation contract; decode-batch-gate gate3). NOTE: the draw stream differs from the host sampler’s SplitMix64 (distribution-equal, seed-deterministic, NOT byte-equal to the old host draws) — greedy rows are unchanged bit-exact. samp[bi] = Some((temp, seed, ctr)) requests a device sample for row bi; the full logits rows are still returned (worker keeps last_logits semantics + fallback rows).

pub fn decode_step_batch_sampled_lean( &self, e: &Engine, tokens: &[u32], caches: &mut [&mut Cache], samp: &[Option<DevSamp>], lean: bool, ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn Error>>

decode_step_batch_sampled + LEAN LOGITS (increment 2 component 3, 2026-08-01): with lean, device-sampled rows SKIP the [n_vocab] logits D2H (9.4%/32.5% of the pre-/post-inc2 tick profile) — their returned row is EMPTY. The audit-mapped consumers: (a) the next tick’s host sample — never fires, device_next carries the token; (b) the graph-promotion argmax — reads only prefill logits (generated empty); (c) the KV-reuse pool park at retire — the REAL consumer, served by a per-cache device park: the row is dtod-copied into cache.last_logits_dev (device bandwidth) and D2H’d ONCE at retire by the worker. Rows without a device sample keep a per-row D2H. lean=false is bit-for-bit the previous method (gates + non-serving callers).

pub fn decode_step_batch_sampled_lean_masked( &self, e: &Engine, tokens: &[u32], caches: &mut [&mut Cache], samp: &[Option<DevSamp>], masks: &[Option<(&CudaSlice<u32>, usize)>], lean: bool, ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn Error>>

decode_step_batch_sampled_lean + GRAMMAR MASKS (constrained decoding, 2026-08-03): masks[bi] = Some((packed_bitset, words)) bans every unset-bit vocab id on row bi (mask_logits_f32, -FLT_MAX) BETWEEN the lm_head matmul and the device sampler, so a constrained row rides the SAME device-sample/lean-logits tick as everyone else — no full-row D2H, no host O(n_vocab) sample. Contract: a masked row must also request a device sample. The row’s PRISTINE logits are preserved for their consumers before the in-place ban: lean rows park the unmasked row into cache.last_logits_dev (the retire-time reuse-pool park stays unmasked — continuations resume grammar-free, the v1 host-path contract), non-lean rows D2H the unmasked row. masks = &[] is bit-for-bit the unmasked method.

pub fn decode_step_batch_sampled_lean_masked_scheduled( &self, e: &Engine, tokens: &[u32], caches: &mut [&mut Cache], samp: &[Option<DevSamp>], masks: &[Option<(&CudaSlice<u32>, usize)>], lean: bool, dual_wave_mid: usize, ) -> Result<(Vec<Vec<f32>>, Vec<Option<u32>>), Box<dyn Error>>

Worker-scheduled twin of Self::decode_step_batch_sampled_lean_masked. The worker supplies the balanced dual-wave boundary it used when forming this tick. Direct engine callers keep the automatic midpoint above; the explicit seam makes scheduler chunking and engine execution one checked contract instead of two coincident width calculations.

pub fn step35_batch_on() -> bool

Rollback seam for the step35 batched decode arm (lane/step35-batched-decode, 2026-08-08). Default ON; MEMRA_STEP35_BATCH=0 caps serving at B=1 and makes the batched bodies return Err. Since lane/cx-b1fix, PP-N also refuses the eager B=1 numeric class, so the seam disables PP-N Step35 decode rather than serving unstable bytes. Also the b2geo35 gate’s CANARY seam — the live assertions must fail under it.

pub fn moesd_target_forward( &self, e: &Engine, tokens: &[u32], batch: usize, gamma: usize, caches: &mut [&mut Cache], ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Standalone MoESD target forward. This entrypoint is not used by serving: it widens the existing Step-3.7 batched layer walk to B*gamma rows while preserving one causal KV chain per session. It returns device logits and performs no sampling or logits D2H, matching the target-model term T_T measured by the paper.

impl HybridModel

pub fn generate_spec_dflash( &self, e: &Engine, draft: &DflashDraft, prompt: &[u32], max_new: usize, eos: &[u32], ) -> Result<Vec<u32>, Box<dyn Error>>

impl HybridModel

pub fn generate_spec_eagle( &self, e: &Engine, draft: &Eagle3Draft, prompt: &[u32], max_new: usize, k: usize, ) -> Result<(Vec<u32>, usize, usize), Box<dyn Error>>

Greedy EAGLE3 speculative decode (EAGLE-PLAN N6). Token-identical to generate(prompt,n) but drafts K tokens with the separate EAGLE3 draft, then verifies them in ONE batched target forward. Verify/accept/snapshot/rollback are REUSED from the MTP path (decode_step_t, cache.snapshot/rollback). Returns (tokens, total_drafted, total_accepted).

impl HybridModel

pub fn gemma4_draft_step( &self, e: &Engine, d: &GemmaDraft, token: u32, h: &CudaSlice<f32>, pos: usize, cache: &Cache, ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn Error>>

One drafter step: (token, h[2816 device]) at absolute position pos over the FROZEN main cache. Returns (draft logits host [n_vocab], h_next [2816 device]).

pub fn gemma4_draft_step_greedy( &self, e: &Engine, d: &GemmaDraft, token: u32, h: &CudaSlice<f32>, pos: usize, cache: &Cache, ) -> Result<(u32, CudaSlice<f32>), Box<dyn Error>>

Greedy draft step: like gemma4_draft_step but the token argmax stays on device — host sees 4 bytes (no 1MB logits dtoh per draft). Returns (token, h_next).

impl HybridModel

pub fn generate_spec_gemma( &self, e: &Engine, d: &mut GemmaDraft, prompt: &[u32], max_new: usize, k: usize, eos: &[u32], ) -> Result<Vec<u32>, Box<dyn Error>>

gemma4 MTP greedy spec loop: prime the prompt, then rounds of (chained K-token draft over the frozen main cache) + (ONE batched verify) + longest-prefix accept + KV rollback. Returns generated tokens; prints acceptance stats.

impl HybridModel

pub fn gemma4_generate_plain_graph( &self, e: &Engine, cache: &mut Cache, last: u32, max_new: usize, eos: &[u32], ) -> Result<Vec<u32>, Box<dyn Error>>

PLAIN-DECODE CUDA-GRAPH loop (gemma4, greedy): one captured verify-trunk step (t=1, device tokens/pos/lens) replayed per token — the launch-gap eraser the decode decomposition demanded (2026-07-23: ~2.3ms/token idle at 128 launches). Self-feeding: argmax -> tok_d -> next embed; counters advance in-graph via spec_rollback_stream(base=1, acc=0). Tokens land in a device ring; ONE host sync per drain window. Captures are keyed on the (rung, window-side, f512-side) regime (the round-graph hint law); regime-crossing stretches run the same body eagerly. Caller guarantees: gemma4, greedy, shared_kv_layers == 0, prompt already primed (cache.pos = prompt len, host kvl.len mirrors set).

impl HybridModel

pub fn decode_step_t( &self, e: &Engine, tokens: &[u32], pos0: usize, cache: &mut Cache, ) -> Result<Vec<f32>, Box<dyn Error>>

Batched target verify forward over tokens at positions pos0..pos0+T (§D.3, T=K+1). Returns ALL T logit columns (host f32, [T*n_vocab]); appends T cols to every full-attn KV and advances every linear-attn recur state by T steps (the recur steps are SEQUENTIAL T=1). Advances cache.pos by T.

pub fn decode_step_t_h( &self, e: &Engine, tokens: &[u32], pos0: usize, cache: &mut Cache, ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn Error>>

Like decode_step_t but ALSO returns the LAST column’s pre-output_norm hidden (h_seed for the next draft round). This lets partial-accept replay run as ONE batched T=(n_acc+1) forward (single weight read) instead of n_acc+1 separate T=1 decode_steps (n_acc+1 weight reads). At batch=1 decode is bandwidth-bound, so batching the replay is THE MTP profitability lever.

pub fn decode_step_t_h_emb( &self, e: &Engine, tokens: &[u32], pos0: usize, cache: &mut Cache, embd_dev: Option<(&CudaSlice<u8>, i32, usize)>, ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn Error>>

Like decode_step_t_h with an optional RESIDENT embed table (spec hot loop): device gather instead of host dequant + [T, n_embd] f32 htod. Bit-identical rows.

pub fn decode_step_t_h_emb_dev( &self, e: &Engine, tokens: &[u32], pos0: usize, cache: &mut Cache, embd_dev: Option<(&CudaSlice<u8>, i32, usize)>, ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn Error>>

DEVICE-LOGITS verify forward (spec device-argmax lever): identical kernel chain to decode_step_t_h_emb but returns the [T, n_vocab] logits ON DEVICE — the accept walk argmaxes each column on-device and reads back ONE [T] u32 instead of dtoh’ing the full T x n_vocab f32 block (~1-4 MB + T host argmaxes, every round). Kernel dispatch is UNCHANGED (same decode-exact kernels); only the post-logits transfer moves.

pub fn decode_step_t_aux2( &self, e: &Engine, tokens: &[u32], pos0: usize, cache: &mut Cache, aux_layers: &[usize], pred_col: Option<usize>, ) -> Result<(Vec<f32>, Vec<CudaSlice<f32>>, Option<Vec<CudaSlice<f32>>>), Box<dyn Error>>

EAGLE3 aux-capturing verify forward over tokens (T) — mirrors decode_step_t_h exactly (same KV append, same causal verify, same recur advance) but ALSO clones the aux residual- stream hiddens (blocks in aux_layers) for TWO columns: the LAST column (always) and the optional pred_col (the EAGLE seed = bonus’s predecessor). Returns (all_T_logits host, last_col_aux, pred_col_aux?). Used by the EAGLE3 orchestrator’s commit.

pub fn plain_session_kv_bytes_per_token(&self) -> usize

Context-linear bytes for a plain serving session’s trunk cache.

pub fn plain_session_kv_shape(&self) -> (usize, usize, usize)

(logical bytes/token, ring-capped bytes/token, ring row cap) for exact admission.

pub fn spec_session_kv_bytes_per_token(&self) -> usize

Context-linear bytes for a speculative serving session: trunk cache plus persistent MTP scratch. With no MTP head this equals the plain coefficient.

pub fn spec_session_kv_shape(&self) -> (usize, usize, usize)

Spec twin of HybridModel::plain_session_kv_shape; Step35’s persistent MTP scratch is capped by the same SWA ring rows as the trunk.

pub fn new_session( &self, e: &Engine, max_ctx: usize, ) -> Result<SpecSession, Box<dyn Error>>

Greedy MTP speculative decode (§B). Token-identical to generate(prompt, max_new) but uses the NextN head to draft K tokens then verifies them in one batched target forward. Returns (generated tokens, total_drafted, total_accepted) so the caller can report acceptance rate. k = draft length per round.

GRAPH DRAFT (stage 2 of graph-grade spec): when the model is all-Dense and the MTP head is Dense (no MoE host readbacks), the fixed-shape T=1 MTP forward is CUDA-graph-captured ONCE and replayed per draft step — the ~40 eager launches per drafted token collapse into one graph dispatch; only the 4-byte token id (and 4-byte p-min confidence) round-trip per step. Event tracking is disabled for the whole call (generate_graph pattern) so every buffer the captured graph references is event-free; the spec loop is strictly single-stream. MEMRA_SPEC_NOGRAPH=1 forces the eager draft chain. SAMPLED mode (MEMRA_SPEC_TEMP>0) has its OWN capture (gumbel-perturbed in-graph argmax, device Philox event counter, persistent q retention) — graph-vs-eager sampled streams are bit-identical for the same (seed, prompt, K, temp); see the sampled-graph setup in generate_spec_inner2. Multi-turn session: trunk cache + MTP draft scratch persist across generate calls, so turn N+1 primes ONLY its new suffix (the 124k-conversation daily pattern — re-priming a 32k history costs ~54s; a suffix prime costs seconds). APPEND-ONLY by construction: the hybrid linear-attn states are in-place (no position index), so a session can extend but never rewind — committed is the exact token list whose state the caches hold (includes any overshoot tokens past max_new; the caller renders from committed, not its own echo).

pub fn optipipe_compare_session_state( &self, e: &Engine, reference: &SpecSession, candidate: &SpecSession, ) -> Result<OptiForkStateIdentity, Box<dyn Error>>

Forced-gate exact state comparison. This intentionally reads the real live prefixes from their owning PP devices: matching emitted ids alone would miss a stale len_d, recurrent snapshot, or draft-KV row that only corrupts the following round.

pub fn spec_rewind_to_checkpoint( &self, e: &Engine, sess: &mut SpecSession, ) -> Result<Option<usize>, Box<dyn Error>>

SESSION-AFFINITY REWIND (lane/session-affinity, 2026-08-05): roll sess back to its retained prompt-end checkpoint, so a request whose prompt matches committed[..rewind_pos()] exactly can resume there and prime only its own delta.

EXACTNESS. After this returns, the session is byte-for-byte the state it was in AT that boundary: full-attn KV truncated to it (append-only, position-addressed), GDN conv/ssm restored from the device copy taken there, draft scratch length reset, committed truncated, last_h = the boundary’s predecessor anchor. That is precisely the state a fresh prime of committed[..pos] would have produced, so the following suffix prime and every burst after it are identical to a cold run of the same token stream — the committed-tokens-authoritative contract.

next_pred and pending_tok are CLEARED: both describe generation past the boundary, which the rewind discards. The caller therefore must supply a non-empty suffix (a rewound session cannot serve an empty-suffix continuation burst — there is nothing to continue). The persistent draft graph survives: it bakes only session-stable pointers (the scratch KV, the resident embedding), none of which the rewind moves.

The checkpoint is CONSUMED (turn_ckpt taken): its snapshot buffers are freed here, and this turn’s own prime installs a fresh one at the new prompt end. Returns the position rewound to, or None when the session holds no checkpoint (caller: full re-prime).

pub fn spec_grow_and_rewind_to_checkpoint( &self, e: &Engine, sess: &mut SpecSession, target_cap: usize, ) -> Result<Option<usize>, Box<dyn Error>>

Grow a parked speculative session to target_cap and rewind it to its retained turn checkpoint without re-priming the checkpoint prefix.

The trunk cache is restored exactly like a plain grown cache: append-only full-attention KV rows come from the parked cache, while recurrent state comes from the checkpoint’s owned snapshot. The MTP scratch is also context-linear and its rows below the checkpoint remain authoritative, so they are copied into a fresh larger scratch before its length is truncated. Pointer-baking draft graphs are dropped and recaptured on the next burst.

All fallible work completes before sess is mutated. A failed allocation or copy leaves the parked session intact, allowing the caller one reclaim-and-retry attempt.

pub fn spec_flush_pending( &self, e: &Engine, sess: &mut SpecSession, ) -> Result<(), Box<dyn Error>>

Commit a carried pending bonus (see SpecSession::pending_tok): one T=1 trunk pass (its logits’ argmax becomes next_pred) + the draft-KV fill at the carried anchor — byte-identical to the pre-carry session tail. Required before a non-empty-suffix prime, a sampled turn, or parking a session for pool reuse. No-op without a pending.

pub fn spec_pipe_available(&self, e: &Engine) -> bool

Reduced-matrix admission for increment 1. This deliberately does not change the PP-2 serving policy: the worker calls it only after MEMRA_SPEC_PIPE=1 and an explicit spec session already exist.

pub fn generate_spec_session_pair( &self, e: &Engine, sess_a: &mut SpecSession, max_new_a: usize, k_a: usize, sess_b: &mut SpecSession, max_new_b: usize, k_b: usize, ) -> Result<((Vec<u32>, usize, usize), (Vec<u32>, usize, usize)), Box<dyn Error>>

Two warm greedy continuation bursts over one PP-2 interval coordinator. The two existing generate_spec_inner2 call stacks own all per-session round locals; only phase issue order changes. No callback is accepted in increment 1 — the worker publishes each completed burst.

pub fn generate_spec_session( &self, e: &Engine, sess: &mut SpecSession, suffix: &[u32], max_new: usize, k: usize, ) -> Result<(Vec<u32>, usize, usize), Box<dyn Error>>

One spec-decode turn on a live session. suffix = the NEW tokens only (turn N+1’s user message rendered through the chat template continuation). Returns (new tokens emitted, drafted, accepted); session.committed grows by suffix + emitted.

pub fn generate_spec_session_sampled( &self, e: &Engine, sess: &mut SpecSession, suffix: &[u32], max_new: usize, k: usize, sampling: Option<SpecSampling>, on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>, ) -> Result<(Vec<u32>, usize, usize), Box<dyn Error>>

Serve-path sampled spec: routes the burst through the rejection-sampling verify with per-SESSION Philox continuity (sess.sctr/uctr). None = env-driven (CLI) or greedy. Filters (top-k/p/min-p) apply SYMMETRICALLY to draft q and verify p — distribution-exact for the filtered target (feat/filtered-spec).

on_commit (sse-cadence, 2026-08-05): called with each newly-emitted slice of the output — once right after the prime’s first token, then once per round commit — so a streaming caller can flush text at round cadence instead of once per burst. The slices are disjoint, in order, and concatenate to exactly the returned token vec. Emission- timing only: token bytes, session state, and exactness are untouched.

The returned bool is a CONTINUE-VERDICT (admission yield, 2026-08-06): false ends the burst at the current round boundary, exactly as if max_new had been reached — the caller’s scheduler regains control without waiting the burst out. Burst size is content-neutral (spec-levers battery), so an early exit moves WHEN the burst returns, never what tokens say. The slice may be EMPTY (a poll-only boundary — round-stream drains and the defensive tail flush can land with nothing new committed).

pub fn generate_spec_session_sampled_prime_split( &self, e: &Engine, sess: &mut SpecSession, suffix: &[u32], max_new: usize, k: usize, sampling: Option<SpecSampling>, prime_split: Option<usize>, on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>, ) -> Result<(Vec<u32>, usize, usize), Box<dyn Error>>

Serve-only cold-prime segmentation twin. prime_split is the same stable boundary the plain worker would honor before entering its sub-floor tokenwise tail; warm continuations pass None and stay on the existing zero-prime path.

pub fn generate_spec_session_constrained( &self, e: &Engine, sess: &mut SpecSession, suffix: &[u32], max_new: usize, k: usize, sampling: Option<SpecSampling>, constraint: Option<&mut dyn SpecConstraint>, on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>, ) -> Result<(Vec<u32>, usize, usize), Box<dyn Error>>

generate_spec_session_sampled + GRAMMAR (constrained decoding, 2026-08-03): the hook truncates acceptance at the first grammar-illegal token AFTER the exactness verify (grammar is an extra rejection rule, ordering like the batched-verify twins) and replaces an illegal bonus with the MASKED argmax of the target’s own verify column — token-identical to constrained plain greedy decode. GREEDY only (the worker routes sampled constrained to plain decode). Acceptance under tight grammars may drop (drafter is unconstrained); that is measured, not hidden.

pub fn generate_spec_session_constrained_prime_split( &self, e: &Engine, sess: &mut SpecSession, suffix: &[u32], max_new: usize, k: usize, sampling: Option<SpecSampling>, constraint: Option<&mut dyn SpecConstraint>, prime_split: Option<usize>, on_commit: Option<&mut dyn FnMut(&[u32]) -> bool>, ) -> Result<(Vec<u32>, usize, usize), Box<dyn Error>>

pub fn generate_spec( &self, e: &Engine, prompt: &[u32], max_new: usize, k: usize, ) -> Result<(Vec<u32>, usize, usize), Box<dyn Error>>

pub fn extract_dspark_anchors( &self, e: &Engine, tokens: &[u32], anchor_positions: &[usize], gamma: usize, top_k: usize, chunk: usize, temperature: f32, ) -> Result<Vec<DsparkAnchorRecord>, Box<dyn Error>>

Anchor-bounded DSpark target extraction. The trunk sees the exact generated token tape; only requested hidden rows and target-logit rows cross PCIe. An anchor token at p pairs with the pre-output-norm h[p-1] carrier, exactly as the existing replay/NextN path does.

pub fn replay_acceptance( &self, e: &Engine, tokens: &[u32], k: usize, stride: usize, chunk: usize, hdump: Option<&mut File>, ) -> Result<(Vec<(usize, Vec<u32>, Vec<u32>)>, Vec<u32>), Box<dyn Error>>

TEACHER-FORCED REPLAY ACCEPTANCE (hqmtp MTP-heal protocol): walk a FIXED token sequence and, at sampled positions, compare the MTP head’s K-token draft chain against the trunk’s own teacher-forced greedy predictions. Nothing is generated — the context is the corpus text itself, so (a) degenerate self-generated loops cannot inflate acceptance and (b) two arms (bf16 ceiling vs NVFP4) score on IDENTICAL contexts, isolating the quant-induced head/hidden-state mismatch from text drift.

Per eval position p (context = tokens[0..=p], predecessor pairing as in spec decode): draft_j = chain token j from (tokens[p], h_{p-1}), then its own drafts — the exact eager spec-decode chain (same mtp_head_forward_dev, same rope positions). target_j = teacher-forced greedy pick for position p+1+j (argmax of the trunk logits at forced context tokens[0..p+j]). For j==0 this equals live spec acceptance; for j>=1 live verify would condition on the drafts, here it conditions on the corpus — deterministic and arm-comparable by design.

Returns (rows, bg): one (p, drafts[k], targets[k]) row per eval position (ascending p), plus the full teacher-forced greedy track bg (bg[i] = greedy pick for position i, i>=1) so harnesses can cross-check runs (e.g. different chunk sizes must give identical bg).

hdump: when Some, every position’s pre-output_norm trunk hidden (the exact rows the draft-KV fill pairs from) streams to the file as little-endian f32 [t_total, n_embd] — the head-distillation extraction (hqmtp): the ENGINE is the source of truth for trunk hiddens (HF torch reproductions of the hybrid trunk measured only ~0.5 greedy agreement vs this path — not usable as a training-data source).

impl HybridModel

pub fn prime_graph_new( &self, e: &Engine, bucket: usize, ) -> Result<PrimeGraph, Box<dyn Error>>

Capture the fresh-prime graph for bucket tokens (13-15ms measured). Manual staged capture — capture_graph_retained’s keeper path trips on the prime (smoke finding 4).

pub fn prime_graph_run( &self, e: &Engine, pg: &mut PrimeGraph, tokens: &[u32], session: &mut Cache, ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn Error>>

Replay the graph for tokens (len <= bucket) and copy the outputs into session (a FRESH cache: pos == 0). Returns host logits (the prefill_tick contract).