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 gemma4_batch_on() -> bool

Kill-switch seam for the gemma4 dense-31B batched decode arm. DEFAULT ON since the 2026-08-16 owner flip (“if the performance are so strong in favor… we serve the correctness and best performance”): the arm’s exactness battery is green at B=4/8, the served identity gate is byte-exact vs eager at c1/c4, and the served aggregate read 55→257 tok/s c16 on the NVFP4mix artifact at 450W (SERVED-AGGREGATE.md). MEMRA_GEMMA4_BATCH=0 forces the eager per-session path (the rollback); 1 is the old opt-in spelling, still accepted. Any OTHER value REFUSES LOUD at first use — a mis-typed kill switch must not silently pick a serving path.

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 gemma_spec_session_new( &self, e: &Engine, d: &mut GemmaDraft, prompt: &[u32], max_ctx: usize, ) -> Result<GemmaSpecSession, Box<dyn Error>>

Open a burst-scoped gemma spec session: prime the prompt, park the first predicted token as pending. Mirrors generate_spec_gemma’s entry verbatim (trim-adapt learn point 1, the PRIME_MIN_T split, the post-norm h convention).

pub fn gemma_spec_session_burst( &self, e: &Engine, d: &mut GemmaDraft, sess: &mut GemmaSpecSession, target: usize, k: usize, eos: &[u32], ) -> Result<(Vec<u32>, usize, usize), Box<dyn Error>>

One serve burst: run complete spec rounds until >= target NEW tokens have been emitted this burst (overshoot committed and returned) or EOS lands. Returns (tokens emitted this burst in order, drafted, accepted). The round body is the EAGER arm of generate_spec_gemma, kept behaviorally identical under default env (adapt/floor/pmin/in-round-cut logic verbatim) — gemma-spec-session-gate enforces byte-equality of the emitted stream against the one-shot at every burst width.

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).