Struct PleLayer 

pub struct PleLayer { /* private fields */ }

Implementations

impl PleLayer

pub fn snapshot_aux( &self, st: &PleSeqState, gpu: &dyn GpuBackend, stream: u64, ) -> Result<Vec<u8>>

Marconi aux blob: [hist_len u32][history u32s][conv f32 bytes]. The whole per-sequence carry — a prefix hit restoring KV+SSM without this would run the n-gram hash on the PREVIOUS request’s history.

pub fn restore_aux( &self, st: &mut PleSeqState, blob: &[u8], gpu: &dyn GpuBackend, stream: u64, ) -> Result<()>

Restore the blob from Self::snapshot_aux on a prefix-cache hit.

impl PleLayer

pub fn prestage( &self, st: &mut PleSeqState, tokens: &[u32], gpu: &dyn GpuBackend, stream: u64, ) -> Result<()>

Hoisted per-step HOST work for decode under CUDA graphs: the n-gram hash, the NVMe fault-in and the slot upload into the stable slots_dev buffer. All three are capture-illegal (the upload reads pageable memory, which invalidates a recording graph with status 901), so the scheduler calls this BEFORE graph replay/capture — the same phasing decode_a already gives the token_ids upload. forward then consumes prestaged_va and enqueues only stable-buffer kernels.

History advances HERE; the prestaged forward must not advance it again.

pub fn release_seq_state( &self, st: &mut PleSeqState, gpu: &dyn GpuBackend, ) -> Result<()>

Release one sequence’s PLE carry.

Same shape of defect as the QSA indexer carry: conv is a bare DevicePtr, so dropping PleSeqState frees nothing. Individually small (~147 KB) and below the 32 MB allocation-trace threshold, which is exactly why it stayed invisible — but it is one per SSM layer (36 on qwen4_exp) per sequence, and it never comes back.

Idempotent: conv is nulled once freed.

impl PleLayer

pub fn new( dims: PleIdDims, head_dim: usize, hidden: usize, hc_mult: usize, k_size: usize, dilation: usize, eps: f32, weights: PleWeights, table: NgramTable, max_tokens: usize, gpu: &dyn GpuBackend, ) -> Result<Self>

pub fn new_seq_state(&self, gpu: &dyn GpuBackend) -> Result<PleSeqState>

Allocate one sequence’s PLE carry (conv buffer + empty history). reset runs on first use (fresh), so contents start undefined.

pub fn rearm(&self, st: &mut PleSeqState)

Restore the prestaged state after a failed CUDA-graph capture attempt (the eager replay re-runs forward, which consumed prestaged_va).

pub fn forward_row( &self, st: &mut PleSeqState, highway_row: DevicePtr, ids: &[u32], ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

Inject into highway [T, hc_mult*hidden] FP32, in place.

fresh starts a new sequence (prefill from position 0). One highway ROW with an explicit id — the multi-seq decode entry (ctx.host_token_ids holds the whole batch; the caller slices).

pub fn forward( &self, st: &mut PleSeqState, highway: DevicePtr, num_tokens: usize, fresh: bool, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>