Skip to main content

Engine

Struct Engine 

Source
pub struct Engine {
    pub gpu: Gpu,
    pub copy_stream: Arc<CudaStream>,
    /* private fields */
}
Expand description

Engine device context: CUDA context, stream, loaded kernel modules, cuBLASLt (via runtime::Gpu).

Fields§

§gpu: Gpu§copy_stream: Arc<CudaStream>

EDGE-1 §C.2: dedicated H2D copy stream for async prefetch (event-synced to the compute stream).

Implementations§

Source§

impl Engine

Source

pub fn mmq_supports(&self, w: &GpuTensor) -> bool

True if w should take a vendored MMQ GEMM under the current env policy (see mmq_w4a8_enabled): NVFP4 needs in_f % 64 == 0, Q4_K/Q5_K need in_f % 256 == 0.

Source

pub fn qmatvec_mmq( &self, w: &GpuTensor, x: &CudaSlice<f32>, m: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Unified vendored-MMQ dispatch: routes to the NVFP4 or Q4_K/Q5_K launcher by qtype. Caller MUST have checked mmq_supports(w). x is the RAW f32 activation.

Source

pub fn qmatvec_mmq_iq4xs_raw( &self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Bare IQ4_XS dense MMQ launch (no macro-scale) — also the kernel_check gate entry.

Source

pub fn qmatvec_mmq_q45k_raw( &self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, qtype: i32, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Bare Q4_K/Q5_K MMQ launch (no macro-scale) — also the kernel_check accuracy-gate entry. Conventional xy-tiling only (the vendored stream-K arm — MEMRA_MMQ_STREAMK — was removed 2026-07-08: 1.11x per-GEMM but its k-split f32 reorder flipped the model argmax gate; rig5090.jsonl 2026-07-03 has the record).

Source

pub fn qmatvec_mmq_q8_0_raw( &self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Bare Q8_0 int8-MMA MMQ launch (no macro-scale) — the kernel_check accuracy-gate entry and the qmatvec_mmq dispatch body. Conventional xy-tiling only (no stream-K / fixup scratch).

Source

pub fn accprobe_act_bytes(&self, in_f: usize, m: usize) -> usize

Accumulator-instrument bytes for a pre-quantized block_q8_1_mmq activation buffer (cu/mmq_q8_0_f32acc.cu). The caller synthesizes that buffer itself — see accprobe_gemm.

Source

pub fn accprobe_gemm( &self, w_q8_0: &CudaSlice<u8>, act_q: &CudaSlice<u8>, m: usize, in_f: usize, out_f: usize, f32acc: bool, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Run one arm of the Q1 accumulator instrument. f32acc=false is the Q8_0 MMQ floor’s GEMM verbatim (s32 accumulate); f32acc=true is the byte-identical kernel with the f8f6f4 f32 accumulate. act_q is a PRE-QUANTIZED block_q8_1_mmq buffer of at least accprobe_act_bytes(in_f, m) bytes — keeping the quantizer out of the timed region is the point, so this wrapper does not build it. Research instrument: the output is not a numeric claim.

Source

pub fn mmq_act_begin(&self)

Open a quantize-once sharing window for the NEXT activation (quantize-once seam): sibling Q4_0 MMQ matmuls on the SAME input (q/k/v; gate/up) quantize its D4 scratch once. Safe by construction: a hit requires the same window epoch AND the same (ptr, m, in_f) — the caller opens a window while it holds the shared input alive, so its address can neither change nor be recycled inside the window. Paths that never call this never hit the cache.

Source

pub fn qmatvec_mmq_q4_0_raw( &self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, rp: bool, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Bare Q4_0 int8-MMA MMQ launch (no macro-scale) — the kernel_check accuracy-gate entry and the qmatvec_mmq dispatch body. rp selects the weight layout (MEMRA_Q4RP split-plane vs raw ggml 18B blocks) — pure address remap, bit-identical output.

Source

pub fn qmatvec_mmq_nvfp4( &self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, scale: f32, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Run the vendored NVFP4 MMQ prefill GEMM from raw weight bytes + f32 activation. y[m, out_f] = x[m, in_f] @ W^T. The per-tensor NVFP4 macro-scale is FOLDED into the MMQ write-back epilogue (was a separate scale_inplace launch + full y round-trip per matmul). Same elementwise multiply -> bit-identical to the two-launch form. x is the RAW f32 activation (the launcher quantizes it to block_fp4_mmq internally).

Source

pub fn qmatvec_mmq_nvfp4_raw( &self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Bare MMQ launch (no macro-scale) — for the kernel_check accuracy gate.

Source

pub fn qmatvec_mmq_nvfp4_raw_v1( &self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Bare MMQ launch on the PRE-PORT activation quantizer (per-sub-block UE4M3 scale only, no per-token row amax). The numeric oracle for the two-level quantizer: kernel-check runs both and reports the accuracy delta, so the port’s value is measured rather than asserted.

Source

pub fn qmatvec_mmq_nvfp4_raw_res( &self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, residual_k: i32, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Bare MMQ launch with an explicit residual-channel count — for the kernel-check k sweep.

Source

pub fn qmatvec_mmq_nvfp4_w4a8( &self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, scale: f32, rp: bool, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

STAGE 2 W4A8 MMQ NVFP4: same tile as the W4A4 path, but weight FP4 is LUT-dequantized to int8 at tile-load and the activation stays q8_1 int8 — the accuracy-safe rung. Macro-scale folded into the write-back epilogue (bit-identical to a post-matmul scale_inplace). rp selects the weight layout (A6 split-plane vs GGUF blocks) — bit-identical output.

Source

pub fn qmatvec_mmq_nvfp4_w4a8_raw( &self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Bare W4A8 MMQ launch (no macro-scale, GGUF layout) — for the kernel_check accuracy gate.

Source

pub fn qmatvec_mmq_nvfp4_w4a8_raw_rp( &self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Bare W4A8 MMQ launch on an A6 split-plane repacked weight — the rp-loader bit-identity gate compares this against qmatvec_mmq_nvfp4_w4a8_raw on the same weight.

Source

pub fn qmatvec_mmq_fp8_blk( &self, w_e4m3: &CudaSlice<u8>, blk_scales: &CudaSlice<f32>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

PER-BLOCK FP8 MMQ prefill GEMM (cu/mmq_fp8_blk.cu). w_e4m3 is the raw checkpoint e4m3 plane [out_f x in_f] and blk_scales the device f32 grid [ceil(out_f/128) x ceil(in_f/128)] — no re-quantization of either.

Source

pub fn qmatvec_mmq_fp8_blk_scaled( &self, w_e4m3: &CudaSlice<u8>, blk_scales: &CudaSlice<f32>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, scale: f32, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Source

pub fn fp8_blk_nan_count( &self, w_e4m3: &CudaSlice<u8>, ) -> Result<u32, Box<dyn Error>>

Count e4m3 NaN codes (magnitude 0x7F) in a device e4m3 plane. 0 is the precondition for routing that tensor through qmatvec_mmq_fp8_blk (hardware decodes them to NaN, the host/ARM B’ reference to 0.0).

Source

pub fn mmq_iq_quantize_act( &self, x: &CudaSlice<f32>, in_f: usize, n_tokens: usize, ) -> Result<CudaSlice<u8>, Box<dyn Error>>

Quantize token-major f32 activation [n_tokens, in_f] to the block_q8_1_mmq (D4) scratch the IQ expert-MMA kernel consumes. Returns the scratch buffer (one per proj input per layer).

Source

pub fn mmq_iq_fused_act_quant( &self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, in_f: usize, n_tokens: usize, act_kind: i32, ) -> Result<CudaSlice<u8>, Box<dyn Error>>

Fused act-epilogue (research lever #3): silu/gelu(gate)up + D4 quantize in one launch — replaces moe_pairs_{silu,gelu}_mul + mmq_iq_quantize_act without materializing the f32 act buffer (saves one full write + one full read pass over [n_pairs x n_ff]). Scratch bytes are BYTE-IDENTICAL to the two-pass path (kernel-check iq fused act+quant gates it). act_kind: 0 = silumul (qwen35moe), 1 = gelu_tanh*mul (gemma4).

Source

pub fn mmq_iq_experts( &self, table: &CudaSlice<u64>, proj: i32, n_expert: usize, ex_ids: &CudaSlice<i32>, ex_off: &CudaSlice<i32>, ex_pairs: &CudaSlice<i32>, pair_tok: &CudaSlice<i32>, act_scratch: &CudaSlice<u8>, in_f: usize, out_f: usize, n_active: usize, n_pairs: usize, n_tokens: usize, qtype: i32, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Expert-segmented IQ3_S/IQ4_XS int8-MMA MMQ (the m16n8k16.s8 analog of moe_pairs_matvec_q8_dec). Same CSR inputs (table/ex_ids/ex_off/ex_pairs/pair_tok) + a pre-quantized q8_1_mmq activation scratch (from mmq_iq_quantize_act over n_tokens). y = [n_pairs, out_f] pair-major.

Source

pub fn moe_f16g_act( &self, x: &CudaSlice<f32>, pair_tok: Option<&CudaSlice<i32>>, in_f: usize, n_pairs: usize, ) -> Result<(CudaSlice<u8>, CudaSlice<f32>), Box<dyn Error>>

Gather+convert the activation to f16 pair-major [n_pairs, in_f] for the grouped GEMM, normalized per row by its amax (raw f16 overflows on gemma’s activation spikes — round 46 NaN find). Returns (act_f16, row_scales) — the scales fold back into the GEMM output. pair_tok = None when the input is already pair-major.

Source

pub fn moe_f16_grouped( &self, table: &CudaSlice<u64>, proj: i32, n_expert: usize, ex_ids: &CudaSlice<i32>, ex_off_host: &[i32], ex_off_dev: &CudaSlice<i32>, act_f16: &CudaSlice<u8>, act_scale: &CudaSlice<f32>, in_f: usize, out_f: usize, n_active: usize, n_pairs: usize, qtype: i32, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

One projection through the grouped f16 lane: dequant the active experts’ rows to an f16 workspace, then ONE grouped GEMM over the CSR groups (variable m per expert). y = f32 [n_pairs, out_f] pair-major — same layout as mmq_iq_experts. MEMRA_MOE_F16G=1: cublasGemmGroupedBatchedEx (+ h2f pass + per-projection sync — the grouped API runs on internal streams unordered with ours, round-47 ledger). MEMRA_MOE_F16G=2: single-kernel grouped GEMM on the engine stream (round 49) — the row scale folds into the kernel epilogue; no f16 C, no h2f, NO sync (ordered by construction). f16-MIRROR numeric class either way (argmax/spec gated, not byte-identity). Errors on unsupported qtype (caller keeps the MMQ arm as fallback).

Source

pub fn moe_f16g_gemm_sk_raw( &self, w_f16: &CudaSlice<u8>, act_f16: &CudaSlice<u8>, row_scale: &CudaSlice<f32>, ex_off_host: &[i32], ex_off_dev: &CudaSlice<i32>, in_f: usize, out_f: usize, n_pairs: usize, shape_sel: i32, cross: i32, tail: i32, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Raw sk grouped-GEMM entry for kernel-check (“f16g-sk” section): explicit shape/cross instead of the env policy. shape_sel < 0 = the round-49 grid-scan rollback arm; else the round-51 problem-visitor split at cross (1 forces all-128, i32::MAX all-32). tail: 1 = the deep tail (32x64x64 3-stage, lane/sk-tail-form) on sub-cross groups, 0 = the round-51 2-stage 32x64x32 tail. w_f16 = [n_active][out_f][in_f] f16 bytes, act_f16 = [n_pairs][in_f] f16 bytes.

Source

pub fn moe_kq_gemm_sk_raw( &self, table: &CudaSlice<u64>, proj: i32, n_expert: usize, ex_ids: &CudaSlice<i32>, act_f16: &CudaSlice<u8>, row_scale: &CudaSlice<f32>, ex_off_host: &[i32], ex_off_dev: &CudaSlice<i32>, in_f: usize, out_f: usize, n_pairs: usize, qtype: i32, row_bytes: usize, cross: i32, tail: i32, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Raw direct-from-quant sk grouped-GEMM entry for kernel-check (“f16g-kq-direct”): explicit cross/tail instead of the env policy. table = device u64 pointer table (proj-major, [n_proj][n_expert] — same contract as moe_f16_grouped), ex_ids = active-expert ids (device). Visitor forms only (the C side rejects anything else).

Source

pub fn moe_f16g_dequant_raw( &self, table: &CudaSlice<u64>, proj: i32, n_expert: usize, ex_ids: &CudaSlice<i32>, in_f: usize, out_f: usize, n_active: usize, qtype: i32, row_bytes: usize, ) -> Result<CudaSlice<u8>, Box<dyn Error>>

Raw dequant-workspace entry for kernel-check: dequant the active experts’ rows to a fresh f16 workspace via the same kernel moe_f16_grouped uses (the direct loaders’ bitwise reference).

Source§

impl Engine

Source

pub fn f16_scratch_swap(&self, new: Option<F16Scratch>) -> Option<F16Scratch>

Swap the resident f16 scratch (task #14 capture isolation). Returns the previous contents; pass them back to restore.

Source

pub fn try_f16_gemm( &self, w: &GpuTensor, x: &CudaSlice<f32>, m: usize, ) -> Result<Option<CudaSlice<f32>>, Box<dyn Error>>

FP16 prefill GEMM for a weight carrying the f16 mirror: y[m,out] = x[m,in] @ (fp16 W)^T, f32 accumulate. Returns None when the weight has no mirror (caller falls through to MMQ).

Source

pub fn qmatvec_gemm_f16_raw( &self, w16: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Bare FP16 GEMM launch on an fp16 mirror — also the kernel_check gate entry.

Source

pub fn f16_act( &self, x: &CudaSlice<f32>, nelem: usize, in_f: usize, ) -> Result<CudaSlice<u8>, Box<dyn Error>>

f32 -> fp16 activation convert into a fresh buffer (matmul_group: ONE convert feeds every mirror-carrying weight in the group; the standalone per-GEMM converts were ~250 launches/prime of gap-cluster fuel, nsys 2026-07-26).

Source

pub fn try_f16_gemm_pre_into( &self, w: &GpuTensor, xh: &CudaSlice<u8>, m: usize, y: &mut CudaSlice<f32>, ) -> Result<bool, Box<dyn Error>>

_into twin of try_f16_gemm_pre (piecewise-slab plumbing): the GEMM writes into a caller-provided buffer (a resident slab view) instead of a fresh allocation — the FFI has always taken the y pointer; only the wrapper allocated. Returns Ok(false) when the weight has no mirror (caller falls back and copies).

Source

pub fn try_f16_gemm_pre_into_off( &self, w: &GpuTensor, xh: &CudaSlice<u8>, m: usize, y: &mut CudaSlice<f32>, off_elems: usize, ) -> Result<bool, Box<dyn Error>>

_into at a ROW OFFSET (task #16): the batched prime’s per-seq out-GEMMs write straight into the concat mixed trunk at offs[s] — removing the per-seq gather copy. off_elems must keep the pointer’s alignment class (n_embd rows do).

Source

pub fn try_f16_gemm_pre( &self, w: &GpuTensor, xh: &CudaSlice<u8>, m: usize, ) -> Result<Option<CudaSlice<f32>>, Box<dyn Error>>

FP16 GEMM on a pre-converted activation — the matmul_group arm. Same contract as try_f16_gemm minus the convert.

Source

pub fn build_q8_f16_raw( &self, bytes: &CudaSlice<u8>, in_f: usize, out_f: usize, ) -> Result<CudaSlice<u8>, Box<dyn Error>>

Raw fp16 mirror build from GGUF Q8_0 device bytes (gates/benches; also the loader’s worker via build_q8_f16).

Source

pub fn build_q8_f16(&self, t: &mut GpuTensor) -> Result<(), Box<dyn Error>>

Load-time fp16 mirror pass for one tensor (hybrid.rs calls this under MEMRA_PP_F16=1, next to build_q8_rp4). No-op unless 2D Q8_0 with integral rows and budget headroom.

Source

pub fn build_q4_f16_raw( &self, bytes: &CudaSlice<u8>, in_f: usize, out_f: usize, ) -> Result<CudaSlice<u8>, Box<dyn Error>>

Q4_0 twin of build_q8_f16_raw (18B blocks, campaign A 2026-07-31).

Source

pub fn build_q5k_f16_raw( &self, bytes: &CudaSlice<u8>, in_f: usize, out_f: usize, ) -> Result<CudaSlice<u8>, Box<dyn Error>>

Q5_K twin (176B superblocks, round 49b). Also the kernel_check gate entry for the Q5_K f16-mirror class.

Source

pub fn build_q4k_f16_raw( &self, bytes: &CudaSlice<u8>, in_f: usize, out_f: usize, ) -> Result<CudaSlice<u8>, Box<dyn Error>>

Q4_K twin of build_q6k_f16_raw (144B superblocks, round 49). Also the kernel_check gate entry for the Q4_K f16-mirror class.

Source

pub fn build_q6k_f16_raw( &self, bytes: &CudaSlice<u8>, in_f: usize, out_f: usize, ) -> Result<CudaSlice<u8>, Box<dyn Error>>

Source§

impl Engine

Source

pub fn try_fp8_gemm( &self, w: &GpuTensor, x: &CudaSlice<f32>, m: usize, ) -> Result<Option<CudaSlice<f32>>, Box<dyn Error>>

FP8 prefill GEMM for a weight carrying the fp8 operand: y[m,out] = x[m,in] @ (e4m3 W)^T with the per-batch act scale and per-tensor weight_scale folded in-GEMM. Returns None when the env is off or the weight has no fp8 operand (caller falls through to the Q8_0 path).

Source§

impl Engine

Source

pub fn try_fp8_blk_mmq( &self, w: &GpuTensor, x: &CudaSlice<f32>, m: usize, ) -> Result<Option<CudaSlice<f32>>, Box<dyn Error>>

PER-BLOCK FP8 MMQ prefill GEMM for a weight carrying a block-128 fp8 operand: y[m,out] = x[m,in] @ (e4m3 W)^T with each [128x128] weight block scaled by its own f32. Returns None when the env is off, the weight has no block-128 fp8 operand, the shape is unsupported, or the NaN precondition fails (caller falls through to the Q8_0 floor).

Source§

impl Engine

Source

pub fn fp8_blk_q8_0_bytes(out_f: usize, in_f: usize) -> usize

Q8_0 slab byte count for an [out_f, in_f] block-128 FP8 weight.

Source

pub fn fp8_blk_dequant_q8_0( &self, f8: &[u8], grid: &[f32], out_f: usize, in_f: usize, ) -> Result<CudaSlice<u8>, Box<dyn Error>>

ARM B’ load-time pass: upload the raw e4m3 codes + the block-128 scale grid, dequant on the GPU, and return the Q8_0 slab (byte-identical to the host re-encode). f8 is the checkpoint’s row-major [out_f x in_f] codes; grid is the row-major [ceil(out_f/128) x ceil(in_f/128)] f32 scale grid (F8BlockGrid order, verbatim).

Source

pub fn fp8_blk_dequant_q8_0_dev( &self, f8: &CudaSlice<u8>, grid: &CudaSlice<f32>, out_f: usize, in_f: usize, ) -> Result<CudaSlice<u8>, Box<dyn Error>>

DEVICE-RESIDENT twin of fp8_blk_dequant_q8_0 (lane/fp8-blk128-decode): identical kernel, identical output bytes, but the e4m3 codes and the scale grid are ALREADY on the device and there is no trailing synchronize.

Both differences matter to its caller (try_e4m3_blk_prefill, per prefill call rather than once per load): the host arm’s two htods would re-upload a weight that is already resident, and its synchronize would stall the CUDA owner thread on every prefill projection. Stream ordering is sufficient without it — the dequant and the Q8_0 GEMM that consumes dst are issued to the SAME stream, so the GEMM cannot observe a partially written slab.

Source§

impl Engine

Source

pub fn new(ordinal: usize) -> Result<Self, Box<dyn Error>>

Source

pub fn ctx(&self) -> &Arc<CudaContext>

Source

pub fn pool_cached_bytes(&self) -> usize

Bytes the async pool holds MAPPED but NOT LIVE (reserved - used), i.e. freed blocks parked in the pool because Engine::new pins RELEASE_THRESHOLD to u64::MAX above.

Why this is a public engine surface: mem_get_info’s free DOES NOT SEE these bytes — they are mapped to this process, so free counts them as gone, yet the very next alloc_u8 is satisfied from them without touching free at all. Any admission or budget decision that reads free alone therefore under-counts real headroom by exactly this amount. Effective allocatable headroom is free + pool_cached_bytes().

MEASURED SIZE (c=64 serve burst, 9B NVFP4 + draft, 24GB card, 2026-08-06): 34-89 MB during the burst — SMALL. The admission gate adds it because a term that can only ever under-count headroom does not belong in a gate that queues real work, but the honest reading of this number is that pool caching is NOT where a long-running server’s VRAM hides on this path: reserved ~= used throughout, so the memory the driver reports as gone is genuinely LIVE (see pool_reserved_used for the diagnostic pair).

Returns 0 if the pool cannot be queried (never a false-positive headroom claim).

Source

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

Raw async-pool occupancy: (RESERVED_MEM_CURRENT, USED_MEM_CURRENT) in bytes. Reserved is what the pool has mapped from the driver; used is what is live inside it. Exposed for admission/VRAM diagnostics — the pair distinguishes “memory is parked in the pool and free cannot see it” (reserved >> used) from “memory is genuinely held live by some owner” (reserved ~= used), which are opposite bugs with opposite fixes. (0, 0) if the pool cannot be queried.

Source

pub fn stream(&self) -> Arc<CudaStream>

Ambient stream (by value since M1-PP2 increment 2): the thread’s pp2 stage stream when a stage scope is active, else the main compute stream — see Gpu::stream.

Source

pub fn gkv_on() -> bool

FP8-GLOBALS switch (MEMRA_GEMMA_GKV, default ON): gemma global (hd512) layers keep their KV in e4m3 — the dequant-latency arc (HANDOVER). Windowed layers stay q8_0/q5_1.

Source

pub fn wkv_on() -> bool

FP8-WINDOWED switch (MEMRA_GEMMA_WKV — measured 2026-07-12 in a validity-gated window: 1.7k 174.1-174.4 vs 168.6-169.4 default (+3%), 4.9k 158.7-160.4; vs llama same-window 159.5-160.2 / 140.6 = 1.09x / 1.13x): gemma windowed (hd256 SWA) layers hold e4m3 KV and ride the format-aware v4 lane from the kf8vf8 module. SERVING-MODE DEFAULT (2026-07-12, the 31B spec unlock): fp8-windowed KV GUTS the MTP drafter’s acceptance — its single swa attention reads the windowed cache and e4m3 noise flips its argmaxes (31B short accept .758 -> 1.000 with q8/q5, spec 88 -> 122.7 vs llama-mtp 112; depth .59 -> .78; 26B depth .57 -> .89). So the default keys on serving intent: SPEC serving (MEMRA_DRAFT set) -> OFF, plain -> ON (its depth-plain +3% stands). Explicit MEMRA_GEMMA_WKV always wins. GKV (globals) stays ON for both — no acceptance cost measured.

Source

pub fn kv_fp8_on() -> bool

QWEN FP8-KV switch (MEMRA_KV_FP8 explicit; else the per-model KV_FP8_FORCE door set at model load; else OFF). Non-gemma full-attn layers hold e4m3 K/V via the kf8vf8 module. Per-model verdict 2026-07-12: 9B +0.7-4% scaling with depth, 27B flat, 35B −2% (fp8 format-gates its v3 dp4a lane) — so the 9B class defaults ON (adopted 2026-07-28 with the deferred acceptance battery), others stay OFF.

Source

pub fn scatter_trim_logits( &self, src: &CudaSlice<f32>, d2t: &CudaSlice<u32>, dst: &mut CudaSlice<f32>, d_vocab: usize, n_vocab: usize, ) -> Result<(), Box<dyn Error>>

Scatter trimmed draft logits into full-vocab space: dst = -inf everywhere, then dst[d2t[i]] = src[i]. Two launches (fill, scatter) — no grid-wide sync needed.

Source

pub fn filter_stats( &self, x: &CudaSlice<f32>, row_stride: usize, rows: &CudaSlice<i32>, out_th: &mut CudaSlice<f32>, out_z: &mut CudaSlice<f32>, out_max: &mut CudaSlice<f32>, n: usize, nrow: usize, temp: f32, top_k: i32, top_p: f32, min_p: f32, ) -> Result<(), Box<dyn Error>>

Per-row filtered-softmax stats: out[r] = (threshold_e, renorm_mass_e, row_max) for the filter (top_k, top_p, min_p) at temp. Rows index into x with row_stride f32s.

Source

pub fn softmax_gather_filtered( &self, x: &CudaSlice<f32>, row_stride: usize, ids: &CudaSlice<u32>, rows: &CudaSlice<i32>, th: &CudaSlice<f32>, z: &CudaSlice<f32>, out: &mut CudaSlice<f32>, n: usize, npair: usize, temp: f32, ) -> Result<(), Box<dyn Error>>

out[pair] = filtered-softmax prob of ids[pair] in row rows[pair] (th/z per PAIR).

Source

pub fn residual_sample_filtered( &self, p: &CudaSlice<f32>, q: Option<&CudaSlice<f32>>, n: usize, temp: f32, seed: u64, stream_pos: u32, p_stats: (f32, f32, f32), q_stats: (f32, f32, f32), out_tok: &mut CudaSlice<u32>, ) -> Result<(), Box<dyn Error>>

Filtered residual sample: token ~ norm(max(0, fp - fq)) with fp/fq the filtered softmaxes.

Source

pub fn gumbel_perturb_filtered( &self, x: &CudaSlice<f32>, y: &mut CudaSlice<f32>, n: usize, seed: u64, stream_pos: u32, temp: f32, row_max: f32, th: f32, ) -> Result<(), Box<dyn Error>>

Gumbel-max draw from the FILTERED distribution (masked perturb; argmax after).

Source

pub fn penalize_logits( &self, x: &mut CudaSlice<f32>, hist: &CudaSlice<u32>, n_hist: usize, rep: f32, freq: f32, present: f32, n: usize, ) -> Result<(), Box<dyn Error>>

Keskar penalties applied IN PLACE to a logits buffer: history token ids get rep-divided/multiplied + freq*count + presence subtracted. Symmetric p/q usage keeps filtered rejection sampling exact for the penalized target.

Source

pub fn penalize_logits_rows( &self, x: &mut CudaSlice<f32>, hist: &CudaSlice<u32>, n_hist: usize, rep: f32, freq: f32, present: f32, n: usize, nrow: usize, ) -> Result<(), Box<dyn Error>>

Rows variant: penalize nrow contiguous rows of length n in one launch.

Source

pub fn wpf_level() -> u32

WEIGHT PREFETCH (SOTA item 3, 2026-07-13, DEFAULT ON): during a bandwidth-idle window (the fa launch reads KV, not weights) prefetch the NEXT matvec’s decode-plane bytes into L2 so it reads L2-warm. Value-free scheduling op — same class as prefetch_l2 (numerics untouched by construction). Wired only where it measured positive: the E4B dc attn arm (+0.65%). 26B (flat — MoE ffn dominates), 31B (−0.2% — decode at the DRAM wall) and the ffn gate/up cascade (−1% — 29MB/layer floods the fill path) all probed and NOT wired. MEMRA_WPF=0 rollback seam.

Source

pub fn set_verify_exact(&self, on: bool)

PDL launch arm (SOTA item 2, 2026-07-13, DEFAULT ON): the six MEMRA_PDL_ENTRY glue kernels launch through cuLaunchKernelEx with PROGRAMMATIC_STREAM_SERIALIZATION — the grid launches while the predecessor drains (~120ns/kernel back, pdl_probe), the kernels’ entry grid-dep sync restores read order (SASS-audited: ACQBULK precedes every LDG in all six). Valid windows: E4B +1.0-1.2% (128 AND 384-tok gens); 26B/31B/qwen flat no-harm. Battery: kernel-check GREEN, run-gen tokens IDENTICAL x3 gemma, spec 64/64 E4B K=1/4/8 + 26B/31B K=4 + qwen PASS. Works eager AND under capture (capture encodes native programmatic edges — the post-capture edge-REWRITE arm died: engine graphs hold cuMemAllocAsync alloc nodes, edge edits on those return CUDA_ERROR_NOT_SUPPORTED). MEMRA_PDL=0 rollback seam. See the verify_exact field. Scoped by the dflash round around its t=16 verify.

Source

pub fn qkv_append_on() -> bool

m=1 norm+rope+append fold seam (2026-07-23): MEMRA_QKV_APPEND=0 reverts to the fused-norm-rope + standalone-append pair (the exact-oracle bisect arm).

Source

pub fn pdl_wb_on() -> bool

PDL wave-B1a seam: the four dense-glue kernels (rms_norm_f32, add_rms_norm_f32, add_scale_rms_norm_q8_1, quantize_q8_1). MEMRA_PDL_WB=0 reverts alone.

Source

pub fn pdl_mmvq_on() -> bool

PDL wave-A seam: the mmvq matvec PDL launches only (the six glue kernels keep their own MEMRA_PDL master seam). MEMRA_PDL_MMVQ=0 reverts wave-A alone — the per-model no-harm bisect knob.

Source

pub fn pdl_on() -> bool

Source

pub fn prefetch_weight_l2(&self, w: &GpuTensor) -> Result<(), Box<dyn Error>>

L2-prefetch a quant weight’s DECODE plane (the rp4 split-plane mirror when present — that is what the m<=8 dispatch reads — else the raw block bytes). No-op on float arms.

Source

pub fn gather_row_bf16( &self, table: &CudaSlice<u8>, tok: &CudaSlice<u32>, idx: usize, dst: &mut CudaSlice<f32>, ncols: usize, ) -> Result<(), Box<dyn Error>>

DSpark markov chain ops (dflash lane): gather one bf16 row of a [V, rank] table by the DEVICE token id at tok[idx] into f32.

Source

pub fn add_row_inplace( &self, logits: &mut CudaSlice<f32>, bias: &CudaSlice<f32>, n: usize, row_off: usize, ) -> Result<(), Box<dyn Error>>

logits[row_off .. row_off+n] += bias[0..n] (in place, one row).

Source

pub fn prefetch_l2( &self, p: &CudaSlice<u8>, n: usize, ) -> Result<(), Box<dyn Error>>

L2 prefetch of a device byte range (latency-hiding arc; value-free scheduling op).

Source

pub fn router_gemv( &self, w: &CudaSlice<f32>, x: &CudaSlice<f32>, n_embd: usize, n_experts: usize, t: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

MoE router GEMV (MEMRA_ROUTER_KERNEL): deterministic warp-per-(expert,token) f32 dot. Different FP order than the cuBLAS path it replaces — battery-gated numeric config.

Source

pub fn router_gemv_form( &self, w: &CudaSlice<f32>, x: &CudaSlice<f32>, n_embd: usize, n_experts: usize, t: usize, w8: bool, batch: bool, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Form-explicit router GEMV launch (kernel-check bit-identity gate + crossover bench force both forms; batch requires w8).

Source

pub fn rows_permute( &self, src: &CudaSlice<f32>, idx: &CudaSlice<i32>, nrows: usize, ncols: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

f32 row permute: dst[idx[i], :] = src[i, :] (grouped-GEMM CSR -> pair-id reorder).

Source

pub fn sigmoid_dot_rows( &self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, n_embd: usize, t: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

shexp gate fused dot: g[tok] = sigmoid(dot(x[tok,:], w)) — replaces the per-layer cuBLASLt m=1 GEMM + separate sigmoid launch on the qwen35moe decode path (the splitKreduce x40/step dig, 2026-07-31). One fold order for every t, so the t=1 decode chain and the small-t spec-verify chain match per row by construction.

Source

pub fn spec_rollback_stream( &self, len_ptrs: &CudaSlice<u64>, pos_start: &CudaSlice<i32>, acc: &CudaSlice<u32>, base: usize, n_rows: usize, ) -> Result<(), Box<dyn Error>>

ROUND-STREAM stream rollback: all counters <- pos_start + base + n_acc.

Source

pub fn plain_tok_ring( &self, vam: &CudaSlice<u32>, pos_start: &CudaSlice<i32>, base: usize, ring: &mut CudaSlice<u32>, ) -> Result<(), Box<dyn Error>>

PLAIN-DECODE GRAPH ring store: ring[(pos_start - base) % cap] = vam[0].

Source

pub fn spec_ring_commit( &self, vtok: &CudaSlice<u32>, acc: &CudaSlice<u32>, brk: &CudaSlice<u32>, ring: &mut CudaSlice<u32>, pend: &mut CudaSlice<u32>, ) -> Result<(), Box<dyn Error>>

ROUND-STREAM stage (c) 4 epilogue: ring commit + tiny counter copies.

Source

pub fn i32_copy_add( &self, src: &CudaSlice<i32>, dst: &mut CudaSlice<i32>, delta: i32, ) -> Result<(), Box<dyn Error>>

Source

pub fn u32_copy( &self, src: &CudaSlice<u32>, dst: &mut CudaSlice<u32>, ) -> Result<(), Box<dyn Error>>

Source

pub fn spec_adapt_k( &self, acc: &CudaSlice<u32>, brk: &mut CudaSlice<u32>, floor: usize, cap: usize, ) -> Result<(), Box<dyn Error>>

ROUND-GRAPH adaptive depth: brk[0] <- clamp(acc[0] + 1, floor, cap) — the host adaptive policy as a captured device op (policy-identical: the accept walk depth caps acceptance exactly like drafting fewer tokens).

Source

pub fn spec_accept_greedy_dc( &self, preds: &CudaSlice<u32>, vtok: &CudaSlice<u32>, last_pred: &CudaSlice<u32>, brk: &CudaSlice<u32>, out: &mut CudaSlice<u32>, ) -> Result<(), Box<dyn Error>>

ROUND-STREAM stage (c) 3: accept walk fully device-driven (brk + assembled vtok).

Source

pub fn pos_iota( &self, pos0: &CudaSlice<i32>, out: &mut CudaSlice<i32>, t: usize, ) -> Result<(), Box<dyn Error>>

ROUND-STREAM stage (c) 2: verify-chain device-pos entries.

Source

pub fn append_kv_quantized_rows_dc( &self, k_rows: &CudaSlice<f32>, v_rows: &CudaSlice<f32>, kc: &mut CudaSlice<u8>, vc: &mut CudaSlice<u8>, t0_dev: &CudaSlice<i32>, t: usize, kv_dim_k: usize, kv_dim_v: usize, k_tok_bytes: usize, v_tok_bytes: usize, g: bool, ) -> Result<(), Box<dyn Error>>

Source

pub fn append_kv_quantized_row_dc_inc( &self, k_row: &CudaSlice<f32>, v_row: &CudaSlice<f32>, kc: &mut CudaSlice<u8>, vc: &mut CudaSlice<u8>, t0_dev: &mut CudaSlice<i32>, kv_dim_k: usize, kv_dim_v: usize, k_tok_bytes: usize, v_tok_bytes: usize, g: bool, ) -> Result<(), Box<dyn Error>>

t=1 dc append with a FUSED len_d increment (wave 5c) — one launch replaces append_rows_dc + inc_seqlen. Single block (read-before-inc ordering).

Source

pub fn pack_tok_p( &self, tok: &CudaSlice<u32>, p: &CudaSlice<f32>, out: &mut CudaSlice<u32>, slot: usize, ) -> Result<(), Box<dyn Error>>

ROUND-STREAM: draft-chain pack + in-graph d2t remap (see kernels.cu headers).

Source

pub fn tok_map_u32( &self, tok: &mut CudaSlice<u32>, map: &CudaSlice<u32>, ) -> Result<(), Box<dyn Error>>

Source

pub fn spec_assemble_verify( &self, tokp: &CudaSlice<u32>, pend: &CudaSlice<u32>, d2t: Option<&CudaSlice<u32>>, vtok: &mut CudaSlice<u32>, brk: &mut CudaSlice<u32>, p_min: f32, k: usize, pmin0: bool, ) -> Result<(), Box<dyn Error>>

ROUND-STREAM stage (c) 1: device verify-token assembly + p-min break derivation.

Source

pub fn ssm_conv_ring_rebuild_dc( &self, qkv_tm: &CudaSlice<f32>, ring_old: &CudaSlice<f32>, conv_state: &mut CudaSlice<f32>, conv_dim: usize, acc: &CudaSlice<u32>, base: usize, t_v: usize, d_conv: usize, ) -> Result<(), Box<dyn Error>>

ROUND-STREAM stage (b) 3b: recur-restore twins with device-j (see hybrid.cu headers).

Source

pub fn gdn_scan_s128_dc( &self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>, g: &CudaSlice<f32>, beta: &CudaSlice<f32>, state_in: &CudaSlice<f32>, state_out: &mut CudaSlice<f32>, o: &mut CudaSlice<f32>, n_head: usize, acc: &CudaSlice<u32>, base: usize, t_v: usize, scale: f32, ) -> Result<(), Box<dyn Error>>

Source

pub fn spec_rollback_kv( &self, len_ptrs: &CudaSlice<u64>, saved: &CudaSlice<i32>, acc: &CudaSlice<u32>, base: usize, n_layer: usize, ) -> Result<(), Box<dyn Error>>

ROUND-STREAM stage (b) 3a: device per-layer KV-len rollback (see spec_rollback_kv).

Source

pub fn spec_seed_gather( &self, vx: &CudaSlice<f32>, fill_prev: &CudaSlice<f32>, acc: &CudaSlice<u32>, h_seed: &mut CudaSlice<f32>, base: usize, n_embd: usize, ) -> Result<(), Box<dyn Error>>

ROUND-STREAM stage (b): device next-round seed gather (see spec_seed_gather header). Caller D2Ds h_seed into fill_prev after (both slots carry the same value in every arm).

Source

pub fn spec_accept_greedy( &self, preds: &CudaSlice<u32>, draft: &CudaSlice<u32>, last_pred: u32, base: usize, k_round: usize, out: &mut CudaSlice<u32>, ) -> Result<(), Box<dyn Error>>

ROUND-STREAM stage (a): device greedy accept walk (see spec_accept_greedy header).

Source

pub fn gumbel_perturb( &self, x: &CudaSlice<f32>, y: &mut CudaSlice<f32>, n: usize, seed: u64, stream_pos: u32, temp: f32, ) -> Result<(), Box<dyn Error>>

y = x/temp + Gumbel(Philox(seed, stream_pos)) over n logits (then run device argmax on y = one categorical sample at temperature temp). temp<=0: y = x (pure copy).

Source

pub fn mask_logits_col( &self, logits: &mut CudaSlice<f32>, mask: &CudaSlice<u32>, col: usize, n: usize, mask_words: usize, ) -> Result<(), Box<dyn Error>>

GRAMMAR TOKEN MASK (constrained decoding, lane/constrained-full): ban every vocab id whose bit is unset in the packed llguidance bitset, IN PLACE on row col of a stacked [B, n_vocab] logits buffer. mask = the SimpleVob u32 words H2D’d verbatim (~n_vocab/8 bytes/step — trivial on PCIe); ids >= 32*mask_words (padded lm_head tail) are banned too, the device twin of constrained::apply_mask. Banned value -FLT_MAX == the argmax/gumbel kernels’ init sentinel, so a fully-banned tail can never win and ordering matches the host -inf mask bit-for-bit for every finite logit.

Source

pub fn gumbel_perturb_col( &self, x: &CudaSlice<f32>, col: usize, y: &mut CudaSlice<f32>, n: usize, seed: u64, stream_pos: u32, temp: f32, ) -> Result<(), Box<dyn Error>>

Column-col twin of gumbel_perturb over stacked logits [B, n_vocab] (the batched serving tick’s device sampler): y = x[col]/temp + gumbel(seed, stream_pos, lane). SAME kernel/Philox mapping as gumbel_perturb — bit-identical perturbation for the same (seed, stream_pos, temp) regardless of which batch column the row sits in (the lane index is the in-row position; col only moves the input pointer). That pointer-invariance IS the serving isolation contract for sampled rows.

Source

pub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn Error>>

In-graph sampling-event counter bump (spec_sample.cu kernel 5): ctr[0] += 1. The sampled graph-draft chain replays with FIXED kernel args, so the Philox event counter must be DEVICE data — the host seeds it once per round; every replay bumps it before the perturb reads it (counter is data, not state — graph-replay-safe).

Source

pub fn gumbel_perturb_ctr( &self, x: &CudaSlice<f32>, y: &mut CudaSlice<f32>, n: usize, seed: u64, ctr: &CudaSlice<u32>, temp: f32, ) -> Result<(), Box<dyn Error>>

Graph-capturable gumbel_perturb: the sampling-event counter comes from DEVICE memory (ctr[0]) instead of a host scalar. Identical math to gumbel_perturb at stream_pos == ctr[0] (same Philox call, same lane mapping) — the eager and graph sampled chains produce bit-identical perturbations for the same (seed, counter, temp).

Source

pub fn softmax_gather( &self, x: &CudaSlice<f32>, row_stride: usize, ids: &CudaSlice<u32>, rows: &CudaSlice<i32>, out: &mut CudaSlice<f32>, n: usize, npair: usize, temp: f32, ) -> Result<(), Box<dyn Error>>

out[pair] = softmax_temp(x[rows[pair]])[ids[pair]] for npair (row, id) pairs; rows index into x with row_stride f32s per row. temp<=0: out = 1.0 iff id is the row argmax (smallest-index tie-break — matches the argmax-gate contract).

Source

pub fn residual_sample( &self, p: &CudaSlice<f32>, q: Option<&CudaSlice<f32>>, n: usize, temp: f32, seed: u64, stream_pos: u32, out_tok: &mut CudaSlice<u32>, ) -> Result<(), Box<dyn Error>>

Sample token from norm(max(0, softmax_temp(p) - softmax_temp(q))) (q = None -> plain categorical from softmax_temp(p)). Row stats (max, sumexp at temp) must be precomputed (softmax_gather’s pass-1 values; see spec.rs caller). Deterministic fixed-order CDF walk.

Source

pub fn with_moe_cache<R>( &self, max_block_bytes: usize, f: impl FnOnce(&mut MoeSlotCache, &Engine) -> Result<R, Box<dyn Error>>, ) -> Result<R, Box<dyn Error>>

Access the shared MoE residency cache (EDGE-1 §B), building it on first use under MEMRA_MOE_CACHE. The closure runs while the lock is held — keep it to lookup/admit/issue, not the GEMM. max_block_bytes sizes the slots (largest of gate/up/down). Returns the closure’s result. If MEMRA_MOE_CACHE is unset this is never called (the caller checks the env first).

Source

pub fn freeze_moe_cache(&self)

Freeze the already-built MoE residency set. This never constructs a cache: callers use it only after a real prefill has populated the machine-specific CPU/GPU working set.

Source

pub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>>

The current residency set as (layer, proj, ex) triples, or None if no cache was built. Never constructs a cache.

Source

pub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool

A frozen heterogeneous CPU/GPU expert split cannot use Hy3’s ordinary batched prefill efficiently: T>=PRIME_MIN_T bypasses the CPU backend and transiently rereads every missing expert through the GPU spill path. Replay the short prompt through decode after freezing, while leaving the profiling warmup’s established batched behavior untouched. (pub: run-gen’s #46 batched-prime gate skips itself when generation will take the tokenwise arm anyway.)

Source

pub fn moe_cache_enabled() -> bool

True if the MoE residency cache is enabled (MEMRA_MOE_CACHE set).

Source

pub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)>

Snapshot the MoE cache counters (hits, misses, staged_bytes, n_slots) for the §D.4 PCIe gate. Returns None if the cache was never built (disabled or no MoE forward ran).

Source

pub fn cpu_expert_stats( &self, ) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)>

Experimental CPU expert backend counters: completed layer calls, experts served, and the sum of backend wall nanoseconds. The timer includes explicit disk->RAM fills on cache misses; callers compare a before/after snapshot around a decode window.

Source

pub fn cpu_expert_predictor_stats(&self) -> (u64, u64)

Caller-blocked nanoseconds at CPU expert joins. Compare before/after snapshots to measure the backend tail that resident-GPU expert work did not hide.

Source

pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64>

Source

pub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)>

CPU-routed expert selections grouped by how many of their three projections were already resident in HBM. This makes otherwise-stranded partial residency visible to tuning runs.

Source

pub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)>

Positioned-read proof-backend counters: (reads, bytes, read_errors, short_reads, mmap_fallbacks, buffer_waits, ring_full).

Source

pub fn moe_cache_reset_counters(&self)

Reset the MoE cache perf counters (to separate warmup from steady-state windows).

Source

pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn Error>>

Source

pub fn htod_bytes_padded( &self, v: &[u8], pad: usize, ) -> Result<CudaSlice<u8>, Box<dyn Error>>

htod_bytes with a mapped (uninit) tail pad: the wide-load expert dots read up to 6B past the final q4_0 block through their aligned window — the bytes never reach a result (funnelshift discards them) but must be mapped memory.

Source

pub fn copy_into( &self, dst: &mut CudaSlice<f32>, off: usize, src: &CudaSlice<f32>, len: usize, ) -> Result<(), Box<dyn Error>>

Device-to-device copy of src into dst[off..off+len] (f32). For in-place KV append.

Source

pub fn copy_u8_into( &self, dst: &mut CudaSlice<u8>, off: usize, src: &CudaSlice<u8>, len: usize, ) -> Result<(), Box<dyn Error>>

View a sub-range of a device buffer (for attending over [0..len) of a KV cache). u8 twin of copy_into (D2D byte-range copy at an offset).

Source

pub fn htod_u8_into( &self, dst: &mut CudaSlice<u8>, off: usize, src: &[u8], ) -> Result<(), Box<dyn Error>>

H2D write of src into dst[off..off+src.len()] (u8). In-place row updates for the adaptive trim head: no realloc, so captured graphs keep their baked addresses.

Source

pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> CudaView<'a, f32>

Source

pub fn view_u8_range<'a>( &self, b: &'a CudaSlice<u8>, start: usize, end: usize, ) -> CudaView<'a, u8>

View the first len BYTES of a u8 device buffer (quantized KV cache: [0..t_kv*tok_bytes)). Byte-range view (gemma4 R6 window offset into the quantized KV stream).

Source

pub fn view_u8<'a>(&self, b: &'a CudaSlice<u8>, len: usize) -> CudaView<'a, u8>

Source

pub fn append_kv_quantized( &self, k_row: &CudaSlice<f32>, v_row: &CudaSlice<f32>, kc: &mut CudaSlice<u8>, vc: &mut CudaSlice<u8>, t: usize, kv_dim_k: usize, kv_dim_v: usize, k_tok_bytes: usize, v_tok_bytes: usize, g: bool, ) -> Result<(), Box<dyn Error>>

Append-quantize ONE token’s post-RoPE K (q8_0) and V (q5_1) into the resident byte caches at token index t (KVQUANT-PLAN §C). One CTA (one warp) per 32-element block; the kernel writes the f16 scale(s) + packed quants for K and V. k_row/v_row are f32 [kv_dim_k]/[kv_dim_v].

Source

pub fn append_kv_quantized_dc( &self, k_row: &CudaSlice<f32>, v_row: &CudaSlice<f32>, kc: &mut CudaSlice<u8>, vc: &mut CudaSlice<u8>, t_dev: &CudaSlice<i32>, kv_dim_k: usize, kv_dim_v: usize, k_tok_bytes: usize, v_tok_bytes: usize, g: bool, ) -> Result<(), Box<dyn Error>>

Device-counter variant of append_kv_quantized (CUDA-GRAPH-PLAN Phase 2): the write slot t is read from t_dev[0] (a resident device i32[1]) instead of a host int arg, so the launch args are FIXED across decode steps (graph-capturable). Identical quant math.

Source

pub fn append_kv_quantized_rows( &self, k_rows: &CudaSlice<f32>, v_rows: &CudaSlice<f32>, kc: &mut CudaSlice<u8>, vc: &mut CudaSlice<u8>, t0: usize, t: usize, kv_dim_k: usize, kv_dim_v: usize, k_tok_bytes: usize, v_tok_bytes: usize, g: bool, ) -> Result<(), Box<dyn Error>>

Append-quantize T token rows in one shot (BATCHED PROMPT PRIME). k_rows/v_rows are token-major [T, kv_dim] post-RoPE f32; rows land at cache slots t0..t0+T. Default = the batched _rows kernel: one (nblk, T) launch whose per-(block,token) warp program is the per-token append kernel verbatim -> every written row is BIT-IDENTICAL to T sequential append_kv_quantized_view calls (kernel_check pins the bytes). MEMRA_PRIME_APPEND_LOOP=1 forces the T-launch per-row loop (the A/B seam that measured the launch overhead).

Source

pub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn Error>>

Increment a device i32[1] counter in place (p[0] += 1) via the resident inc_i32 kernel. Used to advance the device-resident seqlen/pos counters inside the decode-dc path (and, later, inside a captured graph) without a host round-trip.

Source

pub fn append_kv_quantized_view( &self, k_row: &CudaView<'_, f32>, v_row: &CudaView<'_, f32>, kc: &mut CudaSlice<u8>, vc: &mut CudaSlice<u8>, t: usize, kv_dim_k: usize, kv_dim_v: usize, k_tok_bytes: usize, v_tok_bytes: usize, g: bool, ) -> Result<(), Box<dyn Error>>

Like append_kv_quantized but k_row/v_row are CudaViews (one token’s row sliced out of a token-major [T, kv_dim] activation buffer — the MTP verify path appends T tokens).

Source

pub fn copy_view_into( &self, dst: &mut CudaSlice<f32>, off: usize, src: &CudaView<'_, f32>, len: usize, ) -> Result<(), Box<dyn Error>>

Device-to-device copy of a CudaView src into dst[off..off+len] (f32). Like copy_into but the source is a sub-view (e.g. one column of a token-major activation buffer).

Source

pub fn clone_dtod( &self, src: &CudaSlice<f32>, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Real device-to-device COPY of src into a freshly allocated buffer (NOT an Arc clone). Used for cache snapshots (MTP-PLAN §D.4): CudaSlice::clone() only bumps a refcount and would alias the live buffer; this allocs new device memory and memcpy_dtod’s the contents.

Source

pub fn dtod_copy_view( &self, src: &CudaView<'_, f32>, dst: &mut CudaSlice<f32>, ) -> Result<(), Box<dyn Error>>

D2D row extraction: copy a view (e.g. one row of a [B, n] batch buffer) into dst. Stream-ordered, async — decode_batch’s per-sequence row plumbing.

Source

pub fn dtod_copy_view_i8( &self, src: &CudaView<'_, i8>, dst: &mut CudaSlice<i8>, ) -> Result<(), Box<dyn Error>>

D2D i8 twin of dtod_copy_view (q8_1 activation rows).

Source

pub fn dtod_copy_into( &self, src: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, offset: usize, ) -> Result<(), Box<dyn Error>>

D2D row placement: copy src into dst[offset .. offset+src.len()].

Source

pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn Error>>

Uninitialized i8 device buffer (decode_batch q8_1 row scratch).

Source

pub fn qmatvec( &self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, qtype: i32, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Resident-quantized linear (Stage-A: f32 dequant-in-kernel). y[m,out]=x[m,in]@W[out,in]^T.

Source

pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn Error>>

Allocate a reusable u8 GPU scratch buffer (for staged expert weights).

Source

pub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn Error>>

Uninitialized u8 scratch — skips alloc_zeros’ memset. ONLY for staging buffers whose read range is fully overwritten by a stage_expert H2D before any kernel reads it (LAUNCH-STRUCTURE STAGE 2: the per-layer MoE scratch trio was 3 dead ~1MB memsets per layer per decode token).

Source

pub fn memset_zeros_view( &self, dst: &mut CudaViewMut<'_, f32>, ) -> Result<(), Box<dyn Error>>

Zero a SUB-RANGE of an f32 buffer (CudaViewMut) — the row-sized memset the moe_out memset-elision uses for tokens that fall off the gdec fast path (LAUNCH-STRUCTURE STAGE 2).

Source

pub fn stage_expert( &self, host_bytes: &[u8], scratch: &mut CudaSlice<u8>, off: usize, ) -> Result<(), Box<dyn Error>>

EDGE-1 staging: copy host_bytes (a sub-slice of a HostExps buffer) into scratch at byte offset off (async H2D on the default stream). Length is host_bytes.len(). The qmatvec_view that reads scratch[off..] is enqueued on the SAME stream after this, so ordering is guaranteed without an explicit sync (Stage-1; Stage-2 prefetch on a 2nd stream would require an event).

Source

pub fn moe_router_topk( &self, logits: &CudaSlice<f32>, t: usize, n_expert: usize, n_used: usize, ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn Error>>

EDGE-1 §A: fused MoE router. logits is the router output [t, n_expert] (device, f32, the gate_inp @ z result). Returns (sel_idx [t, n_used] i32, sel_w [t, n_used] f32): the top-k expert ids (DESC by prob, ascending-index tiebreak) and renormalized weights. Replaces the host dtoh + softmax-256 + stable DESC top-8 sort + renorm (hybrid_forward.rs ~281-298). One CTA per token row, 256 threads (one per expert).

Source

pub fn moe_router_topk_scaled( &self, logits: &CudaSlice<f32>, t: usize, n_expert: usize, n_used: usize, ex_scale: &CudaSlice<f32>, ) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn Error>>

gemma4 twin: per-expert output scale folded into the topk renorm write (replaces the separate moe_w_exscale launch; value chain identical: (w/ws) * s[sel]).

Source

pub fn moe_router_topk_host( &self, logits: &CudaSlice<f32>, t: usize, n_expert: usize, n_used: usize, ) -> Result<(Vec<u32>, Vec<f32>), Box<dyn Error>>

LAUNCH-STRUCTURE STAGE 1 (2026-07-05): fused router + SINGLE-SYNC host readback. The old MEMRA_FUSED_ROUTER path lost 2% at t=1 because it paid TWO full stream syncs (dtoh_i32 then dtoh, each = clone_dtoh + synchronize) + two alloc_zeros memsets per MoE layer, where the host route pays ONE sync on the 1KB logits dtoh. This variant: uninit outputs (kernel fully overwrites), both DtoH copies issued ASYNC into a persistent PINNED host staging buffer (flags=0 — cacheable, NOT cudarc’s WRITECOMBINED default, so the host-side reads of sel/w stay cached), then ONE synchronize. Numerics identical to moe_router_topk (same kernel).

Source

pub fn stage_expert_async( &self, host_bytes: &[u8], scratch: &mut CudaSlice<u8>, off: usize, ) -> Result<CudaEvent, Box<dyn Error>>

EDGE-1 §C.2: async H2D of host_bytes into scratch[off..] on the COPY stream, returning a recorded event the compute stream can wait on before the dependent GEMM. Used for in-token expert prefetch (pipeline by one). host_bytes should be pinned for a true DMA (§C.1).

Source

pub fn compute_wait(&self, ev: &CudaEvent) -> Result<(), Box<dyn Error>>

Make the compute stream wait for an async copy event (the consumer side of stage_expert_async).

Source

pub fn qmatvec_view( &self, w: &CudaSlice<u8>, range: Range<usize>, x: &CudaView<'_, f32>, m: usize, in_f: usize, out_f: usize, qtype: i32, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

qmatvec over a byte sub-range of a (resident/scratch) CudaSlice holding ONE expert matrix. x is a CudaView (a sliced row of z, or a sliced activation). Reuses the validated qmatvec_f32 dequant path (NOT a fast path — the correctness gate). The CudaView base+offset pointer is honored by the launch arg.

Source

pub fn moe_gate_up_silu8_q8( &self, gp: WPtr8, up: WPtr8, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, in_f: usize, n_ff: usize, n_used: usize, qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

STAGE-2 GROUPED DECODE (2026-07-04): one MoE layer’s gate+up+SiLU for all n_used routed experts of ONE token in ONE launch (replaces 8x qmatvec(gate) + 8x qmatvec(up) + 8x silu_mul = 24 launches). gp/up are the 8 expert weight-block device pointers (SLRU cache slots — fixed-address, stable for the launch). Returns act [n_used, n_ff]. BIT-IDENTICAL to the sequential chain: each dot reproduces qmatvec_f32’s exact 256-thread reduction; the SiLU epilogue is silu_mul_f32’s exact expression (see kernel header). dp4a q8 twins (MoE expert dp4a arc, 2026-07-06): same contract as the _f32 versions but consume a PRE-QUANTIZED q8_1 activation. FP-order differs from _f32 (int dot + warp tree) — the argmax/stream-identity battery arbitrates; MEMRA_MOE_Q8=0 restores f32.

Source

pub fn moe_down8_fma_q8( &self, dp: WPtr8, w: F32x8, aq2: &CudaSlice<i8>, ad2: &CudaSlice<f32>, dst: &mut CudaViewMut<'_, f32>, in_f: usize, out_f: usize, n_used: usize, qt: i32, rb: usize, ) -> Result<(), Box<dyn Error>>

Source

pub fn qmatvec_expert_q8( &self, w: &CudaSlice<u8>, range: Range<usize>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, qtype: i32, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

q8 sequential expert matvec (staged path twin of qmatvec_view for IQ3_S/IQ4_XS).

Source

pub fn moe_gate_up_silu8( &self, gp: WPtr8, up: WPtr8, x: &CudaView<'_, f32>, in_f: usize, n_ff: usize, n_used: usize, qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Source

pub fn moe_down8_fma_into( &self, dp: WPtr8, w: F32x8, act: &CudaSlice<f32>, dst: &mut CudaViewMut<'_, f32>, in_f: usize, out_f: usize, n_used: usize, qt: i32, rb: usize, ) -> Result<(), Box<dyn Error>>

STAGE-2 GROUPED DECODE: one MoE layer’s down-proj + weighted accumulation for all n_used routed experts in ONE launch (replaces 8x qmatvec(down) + 8x axpy = 16 launches), writing the token’s moe_out row DIRECTLY (dst is the zeroed row; the in-kernel slot-ordered __fmaf_rn chain starting at 0.0f reproduces the sequential axpy_f32 accumulation into the zeroed row bit-for-bit — the A2 byte-identity scheme at m=1).

Source

pub fn moe_pairs_matvec_q8( &self, table: &CudaSlice<u64>, proj: i32, pair_tok: &CudaSlice<i32>, pair_ex: &CudaSlice<i32>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, in_f: usize, out_f: usize, n_expert: usize, n_pairs: usize, qtype: i32, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of moe_gate_up_silu8 for FULLY-RESIDENT layers. The expert ids come from the router kernel’s DEVICE sel output (no DtoH) and the weight pointers from the per-layer device table [3, n_expert] of slot base addresses. BIT-IDENTICAL math (same grid/block/reduction; only the pointer/id source differs). dp4a q8 twin of the _dev pair (resident-experts arc).

GEOMETRY VARIANTS (multirow/occupancy arc 2026-07-05): all outputs are BIT-IDENTICAL to the base one-warp-per-(row,slot) kernel (same expert_dot_g g-order + warp tree per row; down’s FMA chain stays slot-ordered serial). Seams: MEMRA_MOE_DEVQ8_GU = 0(base) | 1 | 2 | 4 -> _r{1,2,4} multirow twin (RPW rows/warp) | s2 (gate/up warp split) | s2z (s2 + WPB rows packed per block) | gs4 (gate/up x low/high-group 4-warp split, nsb==64 only) | u64 (nsb==64 unrolled ILP twin, geometry unchanged) MEMRA_MOE_DEVQ8_WPB = warps per block for _r twins / z-rows for s2z (default 4) MEMRA_MOE_DEVQ8_DOWN = auto(default: w8h2 when in_f==512 & n_used<=8 — measured +3.8% decode on 35B/G7e) | 0 (base one-warp serial-slot) | 1 | 2 | 4 -> _w8r{1,2,4} slot-parallel twin | h2 (half-warp dual-row, nsb==16 only) | w8h2 (h2 x slot-parallel) MoE PREFILL pair-batch matvec: one launch covers all (token,expert) pairs for one proj.

Source

pub fn moe_pairs_matvec_q8_em( &self, table: &CudaSlice<u64>, proj: i32, ex_ids: &CudaSlice<i32>, ex_off: &CudaSlice<i32>, ex_pairs: &CudaSlice<i32>, pair_tok: &CudaSlice<i32>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, in_f: usize, out_f: usize, n_expert: usize, n_active: usize, n_pairs: usize, qtype: i32, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Expert-major pair matvec (weight-reuse across each expert’s token group).

Source

pub fn moe_pairs_matvec_q8_dec( &self, table: &CudaSlice<u64>, proj: i32, ex_ids: &CudaSlice<i32>, ex_off: &CudaSlice<i32>, ex_pairs: &CudaSlice<i32>, pair_tok: &CudaSlice<i32>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, in_f: usize, out_f: usize, n_expert: usize, n_active: usize, n_pairs: usize, qtype: i32, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Source

pub fn moe_pairs_gelu_mul( &self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, n: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Source

pub fn moe_pairs_silu_mul( &self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, n: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Source

pub fn moe_pairs_scatter( &self, y_down: &CudaSlice<f32>, pair_w: &CudaSlice<f32>, tok_pair_off: &CudaSlice<i32>, tok_pair_ids: &CudaSlice<i32>, moe_out: &mut CudaSlice<f32>, t: usize, n_embd: usize, ) -> Result<(), Box<dyn Error>>

Source

pub fn moe_gate_up_gelu8_dev_q8( &self, table: &CudaSlice<u64>, sel: &CudaView<'_, i32>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, in_f: usize, n_ff: usize, n_used: usize, n_expert: usize, qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

gemma4 GELU twin of moe_gate_up_silu8_dev_q8 (base geometry — slot-packed j8/j8r2 twins probed 2026-08-01 g26 decode dig: bit-identical rows, -2.5%/-2.9% whole-model decode x3 interleaved -> refuted and killed; research/g26-decode-20260801/receipts.md).

Source

pub fn moe_gate_up_gelu8_dev_q8_rows( &self, table: &CudaSlice<u64>, sel: &CudaSlice<i32>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, t: usize, in_f: usize, n_ff: usize, n_used: usize, n_expert: usize, qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

gemma4 GELU rows twin (verify): one launch over (n_ff, n_used, t).

Source

pub fn moe_gate_up_gelu8_dev_q8_csr( &self, table: &CudaSlice<u64>, sel: &CudaSlice<i32>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, n_pairs: usize, in_f: usize, n_ff: usize, n_used: usize, n_expert: usize, qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

gemma4 GELU CSR twin (verify dedup: owner block serves every pair of its expert).

Source

pub fn moe_down8_fma_dev_q8_rows_g( &self, table: &CudaSlice<u64>, sel: &CudaSlice<i32>, w: &CudaSlice<f32>, aq2: &CudaSlice<i8>, ad2: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, t: usize, in_f: usize, out_f: usize, n_used: usize, n_expert: usize, qt: i32, rb: usize, ) -> Result<(), Box<dyn Error>>

gemma4 generic down rows twin (verify): one launch over (out_f, 1, t).

Source

pub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn Error>>

rp_q4 microprobe (2026-07-10 verify-trunk lever): b4 GGUF-block layout vs the Q4_0 split-plane twin on the wq-class shape. Returns (blk_us, rp_us) after asserting bitwise identity. Bench-only surface (rp_q4_probe bin); no production dispatch reads this.

Source

pub fn build_q4_rp4(&self, t: &mut GpuTensor) -> Result<(), Box<dyn Error>>

Build the Q4_0 split-plane decode mirror for a 2D Quant tensor (device-side permutation, q4_0_split_rp_build). Raw bytes stay resident (prefill/gemm/Stage-A); the m<=8 decode dispatch prefers the mirror (_rp twins). No-op unless (Q4_0, 2D, mirror absent). VRAM cost == the tensor’s weight size. MEMRA_Q4RP=0 disables at the call sites.

Source

pub fn build_q8_rp4(&self, t: &mut GpuTensor) -> Result<(), Box<dyn Error>>

Q8_0 twin of build_q4_rp4 (H100 coalescing fix, 2026-07-26 ncu: GGUF 34B-stride weight loads hold Max Bandwidth at 41-46%; the split mirror makes them aligned 16B ldcs). Raw bytes stay resident (prefill GEMM/MMQ/fused m=1 launches read GGUF layout); the mmvq/batched decode arms prefer the mirror via rp4. Bit-identical outputs.

Source

pub fn build_q8_rp4_raw( &self, bytes: &CudaSlice<u8>, in_f: usize, out_f: usize, ) -> Result<CudaSlice<u8>, Box<dyn Error>>

Raw rp-mirror build for gates/benches: split GGUF Q8_0 bytes into the qplane+dplane mirror without a GpuTensor (same kernel the loader path above uses).

Source

pub fn build_q4k_rp4(&self, t: &mut GpuTensor) -> Result<(), Box<dyn Error>>

K-quant twins of build_q8_rp4 (H100 K-quant coalescing fix, 2026-08-01 ncu on the q27 Q4_K_M decode: q4_K mmvq DRAM 41-54% with 65% excessive sectors, q6_K 40% with 78% — the 144B/210B superblock strides land every 4B weight load off-sector). The mirror re-packs each tensor into planes (q4_K: qs ++ 16B meta; q6_K: ql ++ qh ++ scales ++ d — same total bytes) so every quant fetch is an aligned 16B ldcs. Raw bytes stay resident (prefill GEMM/dequant/Stage-A read GGUF layout); the mmvq/batched decode arms prefer the mirror via rp4. Bit-identical outputs.

Source

pub fn build_q6k_rp4(&self, t: &mut GpuTensor) -> Result<(), Box<dyn Error>>

Source

pub fn build_kq_rp4_raw( &self, bytes: &CudaSlice<u8>, in_f: usize, out_f: usize, qtype: i32, ) -> Result<CudaSlice<u8>, Box<dyn Error>>

Raw K-quant rp-mirror build for gates/benches (same kernels the loader path uses).

Source

pub fn kqrp_enabled() -> bool

MEMRA_KQRP seam: the K-quant (q4_K/q6_K) split-plane decode mirrors at model load. Default follows the Q8RP convention — ON on the Hopper lane (80GB pays the mirror VRAM), OFF elsewhere (a 24GB card cannot hold model + mirror + KV for the big trunks).

Source

pub fn build_q4_rp_swap( &self, t: &mut GpuTensor, ) -> Result<bool, Box<dyn Error>>

IN-PLACE split-plane swap (the 31B dense arc): build the split layout and REPLACE the GGUF bytes (zero extra steady-state VRAM — the transient peak is one tensor’s size). The tensor’s rp flag then routes every consumer (mmvq/batched _rp twins, the qmatvec_gemm_q4_0_rp prefill kernel). Callers gate on the fast path being active — the Stage-A f32 oracle (MEMRA_FAST=0) reads GGUF layout and must never see a swap.

Source

pub fn q4rp_enabled() -> bool

MEMRA_Q4RP seam (default ON): the Q4_0 split-plane decode mirror at model load.

Source

pub fn copy_rows_strided( &self, src: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, row_elems: usize, n_rows: usize, src_stride: usize, src_off: usize, ) -> Result<(), Box<dyn Error>>

gemma4-E4B: dense [t][row_elems] gather of layer il’s rows from the strided prologue buffer ([t][n_layer][n_epl]; off = iln_epl, stride = n_layern_epl).

Source

pub fn u32_set_k( &self, dst: &mut CudaSlice<u32>, v: u32, idx: usize, ) -> Result<(), Box<dyn Error>>

Async device u32 store (value rides the kernel ARG — no host-memory transfer/sync).

Source

pub fn i32_add_k( &self, d: &mut CudaSlice<i32>, v: i32, ) -> Result<(), Box<dyn Error>>

counter += v (device-slot append advance; the +1 twin is inc_seqlen).

Source

pub fn i32_iota_from( &self, ctr: &CudaSlice<i32>, dst: &mut CudaSlice<i32>, n: usize, ) -> Result<(), Box<dyn Error>>

pos rows from a device counter: dst[i] = ctr[0] + i (verify-stream rope positions).

Source

pub fn u32_map_k( &self, buf: &mut CudaSlice<u32>, map: &CudaSlice<u32>, idx: usize, ) -> Result<(), Box<dyn Error>>

In-place trim-id translate: buf[idx] = map[buf[idx]] (FR-Spec d2t, async single-slot).

Source

pub fn u32_pack2( &self, a: &CudaSlice<u32>, off_a: usize, n1: usize, b_in: &CudaSlice<u32>, n2: usize, out: &mut CudaSlice<u32>, ) -> Result<(), Box<dyn Error>>

Pack a[off..off+n1] ++ b[0..n2] into one buffer (single dtoh follows).

Source

pub fn moe_w_exscale( &self, w: &mut CudaSlice<f32>, sel: &CudaSlice<i32>, s: &CudaSlice<f32>, n: usize, ) -> Result<(), Box<dyn Error>>

gemma4 R3 device fold: w[i] *= s[sel[i]] over the router’s [n] (sel, w) pair.

Source

pub fn moe_w_scale_by_expert( &self, w: &mut CudaSlice<f32>, sel: &CudaSlice<i32>, macros: &CudaSlice<f32>, n_expert: usize, n: usize, ) -> Result<(), Box<dyn Error>>

Down-projection macro fold: w[i] = macros[2n_expert + sel[i]] on the device router weights (one launch per MoE layer, only for macro-carrying artifacts — see MoeWeights).

Source

pub fn moe_gate_up_silu8_dev_q8( &self, table: &CudaSlice<u64>, sel: &CudaView<'_, i32>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, in_f: usize, n_ff: usize, n_used: usize, n_expert: usize, qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize, macros: &CudaSlice<f32>, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Source

pub fn moe_down8_fma_dev_q8( &self, table: &CudaSlice<u64>, sel: &CudaView<'_, i32>, w: &CudaView<'_, f32>, aq2: &CudaSlice<i8>, ad2: &CudaSlice<f32>, dst: &mut CudaViewMut<'_, f32>, in_f: usize, out_f: usize, n_used: usize, n_expert: usize, qt: i32, rb: usize, ) -> Result<(), Box<dyn Error>>

Source

pub fn moe_gate_up_silu8_dev_q8_rows( &self, table: &CudaSlice<u64>, sel: &CudaSlice<i32>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, t: usize, in_f: usize, n_ff: usize, n_used: usize, n_expert: usize, qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize, macros: &CudaSlice<f32>, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

SMALL-M VERIFY rows twin (MEMRA_SPEC_M2, lane/spec-m2): ONE launch covers all t tokens of the spec verify’s MoE dev gate/up (grid.z = token) — the _v geometry per token, with tok-offset sel/aq/ad/act pointers matching the serial loop’s slices. BIT-IDENTICAL per token (see the kernel header). aq/ad are the BATCHED z-quantize ([t, in_f] rows — quantize_q8_1’s per-32-block program is row-independent, so batched rows == the serial loop’s per-token quantize_q8_1_view bytes). Returns act [t, n_used, n_ff].

Source

pub fn moe_down8_fma_dev_q8_rows( &self, table: &CudaSlice<u64>, sel: &CudaSlice<i32>, w: &CudaSlice<f32>, aq2: &CudaSlice<i8>, ad2: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, t: usize, in_f: usize, out_f: usize, n_used: usize, n_expert: usize, qt: i32, rb: usize, ) -> Result<(), Box<dyn Error>>

SMALL-M VERIFY rows twin of the down proj: w8h2v geometry per token on a grid.z token axis. Caller gates the w8h2v shape contract (in_f == 512, n_used <= 8) — same gate as the AUTO dispatch in moe_down8_fma_dev_q8. aq2/ad2 = batched act quantize ([t*n_used, in_f] rows). dst rows are FULLY overwritten per token.

Source

pub fn moe_gate_up_silu8_dev_q8_csr( &self, table: &CudaSlice<u64>, sel: &CudaSlice<i32>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, n_pairs: usize, in_f: usize, n_ff: usize, n_used: usize, n_expert: usize, qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

CSR gate/up v3 (owner-scan dedup, no build kernel): qtypes {IQ4_XS, IQ3_S} (caller gates), grid.y = pair index; the first pair of each expert serves all its pairs. Bit-identical to moe_gate_up_silu8_dev_q8_v_rows (explicit-intrinsic accumulate).

Source

pub fn moe_down8_fma_dev_q8_variant( &self, variant: &str, table: &CudaSlice<u64>, sel: &CudaView<'_, i32>, w: &CudaView<'_, f32>, aq2: &CudaSlice<i8>, ad2: &CudaSlice<f32>, dst: &mut CudaViewMut<'_, f32>, in_f: usize, out_f: usize, n_used: usize, n_expert: usize, qt: i32, rb: usize, ) -> Result<(), Box<dyn Error>>

TEST SEAM (down8 lane 2026-07-08): launch a down dev_q8 variant BY NAME with its canonical geometry, bypassing the env-cached dispatch so moe-devq8-check can byte- compare variants in one process. Variants: “base”, “w8h2”, “w8h2r2”, “w8h2v”, “w8h2r2v”.

Source

pub fn moe_gate_up_silu8_dev_q8_variant( &self, variant: &str, table: &CudaSlice<u64>, sel: &CudaView<'_, i32>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, in_f: usize, n_ff: usize, n_used: usize, n_expert: usize, qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

TEST SEAM (down8 lane): gate_up twin of the above. Variants: “base”, “v”.

Source

pub fn moe_gate_up_silu8_dev( &self, table: &CudaSlice<u64>, sel: &CudaView<'_, i32>, x: &CudaView<'_, f32>, in_f: usize, n_ff: usize, n_used: usize, n_expert: usize, qt_g: i32, qt_u: i32, rb_g: usize, rb_u: usize, macros: &CudaSlice<f32>, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Source

pub fn moe_down8_fma_dev( &self, table: &CudaSlice<u64>, sel: &CudaView<'_, i32>, w: &CudaView<'_, f32>, act: &CudaSlice<f32>, dst: &mut CudaViewMut<'_, f32>, in_f: usize, out_f: usize, n_used: usize, n_expert: usize, qt: i32, rb: usize, ) -> Result<(), Box<dyn Error>>

LAUNCH-STRUCTURE STAGE 3: device-dispatch twin of moe_down8_fma_into — expert ids AND renormalized weights read from the router kernel’s device output. BIT-IDENTICAL chain.

Source

pub fn axpy_into( &self, src: &CudaSlice<f32>, alpha: f32, dst: &mut CudaViewMut<'_, f32>, n: usize, ) -> Result<(), Box<dyn Error>>

dst[i] += alpha * src[i], i in 0..n. dst is a CudaViewMut (a row of moe_out).

Source

pub fn add_scaled_rows( &self, src: &CudaSlice<f32>, scale: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, ncols: usize, nrows: usize, ) -> Result<(), Box<dyn Error>>

dst[rncols + c] += src[rncols + c] * scale[r]. Per-row scalar accumulate (shared expert).

Source

pub fn gather_rows( &self, src: &CudaSlice<f32>, idx: &CudaSlice<i32>, dst: &mut CudaSlice<f32>, ncols: usize, m_e: usize, ) -> Result<(), Box<dyn Error>>

Gather m_e rows from src[T, ncols] into dst[m_e, ncols] using index array idx[m_e].

Source

pub fn scatter_slot( &self, src: &CudaSlice<f32>, tok_idx: &CudaSlice<i32>, slot_idx: &CudaSlice<i32>, weight: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, wbuf: &mut CudaSlice<f32>, ncols: usize, n_used: usize, m_e: usize, ) -> Result<(), Box<dyn Error>>

Scatter expert outputs into per-token slots: dst[tok_idx[r], slot_idx[r], :] = src[r, :] * weight[r]. dst is [T, n_used, ncols], zero-initialized. Each (expert, token) pair maps to a unique slot. Scatter expert outputs into per-token slots (raw copy, no weight multiply). Weight stored into wbuf[tok*n_used + slot] for FMA in reduce step.

Source

pub fn reduce_slots( &self, slots: &CudaSlice<f32>, wbuf: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, ncols: usize, n_used: usize, t: usize, ) -> Result<(), Box<dyn Error>>

Reduce n_used slots per token: dst[t, col] = sum_s slots[t, s, col]. Reduce n_used slots per token: dst[t, col] = sum_s FMA(wbuf[t,s], slots[t,s,col], acc). Uses FMA for bit-identity with the sequential axpy path.

Source

pub fn quantize_q8_1_view( &self, x: &CudaView<'_, f32>, m: usize, in_f: usize, ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn Error>>

Stage-B: quantize activation [m,in] f32 -> q8_1 (int8 qs + per-block f32 scale). Quantize an activation [m, in_f] to q8_1 (int8 qs + per-32 f32 scale). Public so the forward can quantize a SHARED activation ONCE and feed it to several matmuls (gate+up share z; q/k/v and wqkv/gate/beta/alpha share h) — quantize_q8_1 was 13.5% of decode GPU time, ~half of it redundant re-quantization of the same row. quantize_q8_1 over a CudaView (a sliced z-row) — same kernel, offset-honoring arg.

Source

pub fn quantize_q8_1( &self, x: &CudaSlice<f32>, m: usize, in_f: usize, ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn Error>>

Source

pub fn quantize_fp4_act( &self, x: &CudaSlice<f32>, m: usize, in_f: usize, ) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn Error>>

Stage-C FP4: quantize activation [m,in] f32 -> e2m1 nibbles (aq4: u32 [m, in/8]) + per-16 UE4M3 scale (ad4: u8 [m, in/16]), the layout the mxf4nvf4 block-scale GEMM B-operand wants. in_f must be a multiple of 64 (one NVFP4 K-block). One thread per (token, 16-block).

Source

pub fn qmatvec_gemm_nvfp4_fp4( &self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, row_bytes: usize, scale: f32, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Stage-C FP4 GEMM (NVFP4 weights): native mxf4nvf4 block-scale tensor-core matmul. Feeds raw e2m1 weight nibbles + raw UE4M3 micro-scales directly to mma.sync.m16n8k64 (762 TFLOP/s peak, 3.5x int8). Activation x is quantized to FP4 e2m1 here. NVFP4 per-tensor macro-scale applied post (scale==1.0 -> no-op). bytes = raw NVFP4 weight rows. Used by the MEMRA_FP4 prefill path.

Source

pub fn qmatvec_gemm_nvfp4_fp4_raw( &self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Test entry (kernel_check): run the FP4 GEMM from raw bytes; NO macro-scale (caller compares bare).

Source

pub fn qmatvec_q8_0_fast( &self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Stage-B: Q8_0 weight x q8_1 activation int8 dp4a matmul. y[m,out]=x@W^T.

Source

pub fn qmatvec_q4_K_fast( &self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Stage-B: Q4_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.

Source

pub fn qmatvec_q6_K_fast( &self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Stage-B: Q6_K weight x q8_1 activation int8 dp4a (decode, symmetric).

Source

pub fn qmatvec_q5_K_fast( &self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Stage-B: Q5_K weight x q8_1 activation int8 dp4a (decode). Min-offset via q8_1 sum term.

Source

pub fn qmatvec_q3_K_fast( &self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Stage-B: Q3_K weight x q8_1 activation int8 dp4a (decode, symmetric).

Source

pub fn qmatvec_nvfp4_fast_rp( &self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

A6 split-plane twin of qmatvec_nvfp4_fast (weights repacked; used by the rp gates).

Source

pub fn qmatvec_nvfp4_fast( &self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Stage-B: NVFP4 weight x q8_1 activation int8 dp4a (decode, symmetric, codebook lookup).

Source

pub fn qmatvec_iq4_XS_fast( &self, w: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Stage-B (optional perf): IQ4_XS codebook int8 dp4a.

Source

pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn Error>>

Source

pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn Error>>

Source

pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn Error>>

i8 upload (moe-devq8-check: synthetic q8_1 activation bytes).

Source

pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn Error>>

Source

pub fn dtoh_view( &self, d: &CudaView<'_, f32>, ) -> Result<Vec<f32>, Box<dyn Error>>

View twin of dtoh (lean-logits component 3: D2H one row of a [B, n_vocab] stack).

Source

pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn Error>>

Source

pub fn dtoh_pair( &self, a: &CudaSlice<f32>, b: &CudaSlice<f32>, ) -> Result<(Vec<f32>, Vec<f32>), Box<dyn Error>>

Queue two f32 device-to-host copies on the compute stream, then establish one host boundary for both. Hy3’s CPU/GPU expert split needs the router logits and the MoE input; issuing them together avoids a second stream synchronization in every trunk layer.

Source

pub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn Error>>

Device-to-host copy of an i32 buffer (fused-router sel_idx readback).

Source

pub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn Error>>

Device-to-host copy of a u8 buffer (used to read back the quantized KV cache for validation).

Source

pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn Error>>

Source

pub fn prob_of_token_device( &self, logits: &CudaSlice<f32>, tok: &CudaSlice<u32>, n_vocab: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

GPU-resident greedy argmax (CUDA-GRAPH-PLAN Phase 1): logits[n_vocab] -> token id in a resident device u32 [1]. PARALLEL 2-pass (RANK1 LEVER): the old single-CTA scan (one 256-thread block on one SM over 248K logits) was memory-starved at ~426us/token. Now pass 1 fans NB=256 blocks across the SMs to saturate HBM, pass 2 reduces the NB partials. Bit-identical to host argmax (smallest index on tie). The whole point is NOT to dtoh logits — only a [1] u32 is read back (or kept resident for graph replay). Returns the device token buffer. Softmax probability of the (already-argmaxed) token tok under logits — the spec-decode p-min confidence signal. 2-pass like the parallel argmax; returns a device [1] f32.

Source

pub fn prob_of_token_device_col( &self, logits: &CudaSlice<f32>, tok_all: &CudaSlice<u32>, tok_idx: usize, p_out: &mut CudaSlice<f32>, p_idx: usize, n_vocab: usize, ) -> Result<(), Box<dyn Error>>

Like prob_of_token_device but writes into a PERSISTENT p_out buffer (stable pointer). Required for CUDA-graph capture of the draft chain: the captured prob kernels must write where the host reads the p-min confidence between replays. Same kernels, same math. Slot-addressed twin of prob_of_token_device_into: token read from tok_all[tok_idx] (a view at the slot), probability written to p_out[p_idx] — same two kernels, the pointers just land mid-buffer. Zero-sync (gemma confidence-adaptive draft depth).

Source

pub fn prob_of_token_device_into( &self, logits: &CudaSlice<f32>, tok: &CudaSlice<u32>, p_out: &mut CudaSlice<f32>, n_vocab: usize, ) -> Result<(), Box<dyn Error>>

Source

pub fn argmax_token_device( &self, logits: &CudaSlice<f32>, n_vocab: usize, ) -> Result<CudaSlice<u32>, Box<dyn Error>>

Source

pub fn argmax_token_device_into( &self, logits: &CudaSlice<f32>, tok: &mut CudaSlice<u32>, n_vocab: usize, ) -> Result<(), Box<dyn Error>>

Like argmax_token_device but writes into a PERSISTENT tok buffer (stable pointer) instead of allocating a fresh one. Required for CUDA-graph capture: the captured argmax must write the next token into the SAME device buffer the next replay’s embed_gather reads, so the buffer pointer is baked once and the token id never round-trips to host inside steady state. The pass-1 partials scratch (argmax_partials) is also a resident stable-pointer buffer so both captured passes bake fixed addresses.

Source

pub fn argmax_token_device_col( &self, logits: &CudaSlice<f32>, col: usize, n_vocab: usize, toks: &mut CudaSlice<u32>, out_idx: usize, ) -> Result<(), Box<dyn Error>>

Column-col device argmax over a stacked verify-logits buffer [t, n_vocab] (spec accept walk): toks[out_idx] = argmax(logits[col*n_vocab .. (col+1)*n_vocab]). SAME 2-pass kernels and tie-break contract as argmax_token_device_into (bit-identical to host argmax, argmax_gate-validated) — only the input pointer (a column view) and the output slot differ. Lets the accept walk read ONE [t] u32 instead of dtoh’ing the full [t, n_vocab] logits.

Source

pub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn Error>>

Read back a device u32 buffer (the spec accept walk’s [t] per-column argmax tokens).

Source

pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn Error>>

Source

pub fn htod_u32_into( &self, dst: &mut CudaSlice<u32>, src: &[u32], ) -> Result<(), Box<dyn Error>>

Allocate a zeroed device u32 buffer (persistent spec-loop prediction slots). H2D into an EXISTING u32 buffer (stable pointer — the per-step grammar-mask upload: contents change every step, the address must not, so a captured graph can read it).

Source

pub fn alloc_u32_zeroed( &self, n: usize, ) -> Result<CudaSlice<u32>, Box<dyn Error>>

Source

pub fn embed_gather_device_into( &self, embd: &CudaSlice<u8>, token_d: &CudaSlice<u32>, x_out: &mut CudaSlice<f32>, n_embd: usize, qtype: i32, row_bytes: usize, ) -> Result<(), Box<dyn Error>>

embed_gather into a PERSISTENT x_out buffer (stable pointer) for CUDA-graph capture (the embed output starts the per-step kernel chain and must be at a fixed address across replays).

Source

pub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn Error>>

Read a [1] i32 device counter (pos / seqlen) back to host. Tiny D2H + sync.

Source

pub fn i32_set_k( &self, dst: &mut CudaSlice<i32>, v: i32, ) -> Result<(), Box<dyn Error>>

Set a [1] i32 device counter IN PLACE (keeps the buffer pointer stable — required for the graph-resident pos/seqlen counters whose addresses are baked into captured graphs). Restores the counter value after the throwaway capture warmups corrupt it. ASYNC i32 single-slot store (value rides the kernel arg — no host-memory transfer/sync). The graph-arc device-len counters use this; set_i32_one below is the SYNCING pageable copy (fine at stream-idle boundaries, poison mid-round).

Source

pub fn set_i32_one( &self, d: &mut CudaSlice<i32>, v: i32, ) -> Result<(), Box<dyn Error>>

Source

pub fn set_u32_one( &self, d: &mut CudaSlice<u32>, v: u32, ) -> Result<(), Box<dyn Error>>

Set a [1] u32 device buffer IN PLACE (stable pointer) — for the resident token_d counter during priming / capture-state restore.

Source

pub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn Error>>

Read back a [1] u32 device buffer (the argmax token). One tiny D2H + sync.

Source

pub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn Error>>

Upload raw bytes to a resident device u8 buffer (e.g. the embed table for device gather).

Source

pub fn embed_gather_device( &self, embd: &CudaSlice<u8>, token_d: &CudaSlice<u32>, n_embd: usize, qtype: i32, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Embed-from-device (CUDA-GRAPH-PLAN Phase 1): gather+dequant the row for the token id in token_d[0] from the resident embed table embd -> x_out[n_embd]. Bit-identical to host EmbedHost::gather (same per-dtype deq). No host round-trip of the token id.

Source

pub fn embed_gather_device_t( &self, embd: &CudaSlice<u8>, tokens: &[u32], n_embd: usize, qtype: i32, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

T-token device embed gather (spec verify/replay): tokens uploaded as a tiny [T] u32 htod, rows dequanted on-device -> x[T, n_embd]. Replaces host per-row dequant + Tn_embd4B htod (nsys: 84% of spec API time was HtoD). Bit-identical rows (same per-dtype deq).

Source

pub fn embed_gather_device_tv( &self, embd: &CudaSlice<u8>, tok_v: &CudaView<'_, u32>, t: usize, n_embd: usize, qtype: i32, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

T-token embed gather from a DEVICE token buffer (round-stream stage c: the verify tokens are assembled on-device from the draft-chain pack slots; no host round trip). Same kernel as embed_gather_device_t — bit-identical rows. embed_gather over a token VIEW (spec round: tokens live in the round’s batch buffer).

Source

pub fn embed_gather_device_td( &self, embd: &CudaSlice<u8>, tok_d: &CudaSlice<u32>, t: usize, n_embd: usize, qtype: i32, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Source

pub fn uninit_q8_pair( &self, n: usize, ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn Error>>

Public f32 uninitialized scratch (see alloc_uninit). For decode/forward scratch a kernel fully overwrites. SAFETY: producing kernel must write every element before any read. Uninitialized q8_1 activation pair (int8 + per-32 scales) — the fa combine q8-emit consumers alloc through this (m=1 decode arms).

Source

pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn Error>>

Source

pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn Error>>

i8 uninitialized scratch (same contract as uninit).

Source

pub fn rms_norm3( &self, x: &CudaSlice<f32>, w0: &CudaSlice<f32>, w1: &CudaSlice<f32>, w2: &CudaSlice<f32>, d0: &mut CudaSlice<f32>, d1: &mut CudaSlice<f32>, d2: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(), Box<dyn Error>>

RMSNorm: x[ncols,nrows] row-major, weight[ncols] -> dst. One block/row, 256 threads. gemma4: 3 rms_norms of the SAME input in one launch (one reduction, three weights). Per-output bit-identical to three rms_norm calls (verbatim reduction/scale chain).

Source

pub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool

gemma4 fused q/k/v head norms (one launch, per-row rms_norm_f32-verbatim). True when the warp-per-row qkv norm would engage for (rows, ncols) — the emit lane piggybacks on the same conditions.

Source

pub fn rms_norm_qkv_w4b( &self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>, wq: &CudaSlice<f32>, wk: &CudaSlice<f32>, wv: &CudaSlice<f32>, dq: &mut CudaSlice<f32>, dk: &mut CudaSlice<f32>, dv: &mut CudaSlice<f32>, dvb: &mut CudaSlice<u8>, ncols: usize, rq: usize, rk: usize, eps: f32, vf16: bool, ) -> Result<(), Box<dyn Error>>

w4 norm with bf16 V EMIT (31B glue lane): the v segment also writes its normed rows as bf16 (the FA V operand — bit-identical to a post-hoc f32_to_bf16). Prefill-depth only.

Source

pub fn rms_norm_qkv( &self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>, wq: &CudaSlice<f32>, wk: &CudaSlice<f32>, wv: &CudaSlice<f32>, dq: &mut CudaSlice<f32>, dk: &mut CudaSlice<f32>, dv: &mut CudaSlice<f32>, ncols: usize, rq: usize, rk: usize, eps: f32, ) -> Result<(), Box<dyn Error>>

Source

pub fn rms_norm2x( &self, a: &CudaSlice<f32>, bb: &CudaSlice<f32>, wa: &CudaSlice<f32>, wb: &CudaSlice<f32>, da: &mut CudaSlice<f32>, db: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(), Box<dyn Error>>

gemma4 fused pair of rms_norms over two different inputs (same width).

Source

pub fn softcap( &self, y: &mut CudaSlice<f32>, cap: f32, n: usize, ) -> Result<(), Box<dyn Error>>

gemma4 R4: in-place final-logit softcap y = cap*tanh(y/cap).

Source

pub fn mask_ids_rows( &self, y: &mut CudaSlice<f32>, ids: &CudaSlice<i32>, n_ids: usize, n_vocab: usize, t: usize, ) -> Result<(), Box<dyn Error>>

gemma4 suppress-token mask: y[row][ids[j]] = -inf over t logits rows (fixed-arg launch — graph-capture safe; NOT monotonic like softcap, so it must run before any argmax).

Source

pub fn add_scale_rms_norm( &self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, c: f32, w: &CudaSlice<f32>, res: &mut CudaSlice<f32>, dst: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(), Box<dyn Error>>

gemma4: res = (a+b)*c AND dst = rms_norm(res, w) in one launch.

Source

pub fn add_scale_rms_norm_q8_1( &self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, c: f32, w: &CudaSlice<f32>, res: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn Error>>

gemma4: res = (a+b)*c AND the next layer’s attn_norm EMITTED q8_1 in one launch. Quantize epilogue bit-identical to quantize_q8_1 (the rms_norm_q8_1 form).

Source

pub fn add_scale_rms_norm_q8_1_into( &self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, c: f32, w: &CudaSlice<f32>, res: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, out_q: &mut CudaSlice<i8>, out_d: &mut CudaSlice<f32>, ) -> Result<(), Box<dyn Error>>

Slot-fed add_scale_rms_norm_q8_1 twin (alloc-free capture lane).

Source

pub fn rms_pre_add_scale_rms_norm_q8_1( &self, a: &CudaSlice<f32>, wa: &CudaSlice<f32>, b_in: &CudaSlice<f32>, c: f32, w: &CudaSlice<f32>, res: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn Error>>

E4B glue fusion: rms(a, wa) prologue + the add_scale_rms_norm_q8_1 program — one launch replaces the per-layer rms_norm_f32(y) + emit pair in the PLE tail.

Source

pub fn gelu_tanh_mul_q8_1( &self, gate: &CudaSlice<f32>, up: &CudaView<'_, f32>, act: &mut CudaSlice<f32>, ncols: usize, nrows: usize, ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn Error>>

GELU(tanh)*up with the activation emitted q8_1 alongside f32 (glue-fusion lane): the consumer matmul rides matmul_pre, killing its standalone quantize_q8_1 launch.

Source

pub fn gelu_tanh_mul_q8_1_into( &self, gate: &CudaSlice<f32>, up: &CudaView<'_, f32>, act: &mut CudaSlice<f32>, ncols: usize, nrows: usize, out_q: &mut CudaSlice<i8>, out_d: &mut CudaSlice<f32>, ) -> Result<(), Box<dyn Error>>

Slot-fed gelu_tanh_mul_q8_1 twin (alloc-free capture lane; incl. the PDL arm).

Source

pub fn add_rms_norm3_q8z( &self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, w0: &CudaSlice<f32>, w1: &CudaSlice<f32>, w2: &CudaSlice<f32>, res: &mut CudaSlice<f32>, out1: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, ) -> Result<((CudaSlice<i8>, CudaSlice<f32>), (CudaSlice<i8>, CudaSlice<f32>)), Box<dyn Error>>

gemma4: add + rms_norm3 with outputs 0/2 emitted q8_1 (zsh + moe_in) and 1 f32 (router).

Source

pub fn add_rms_norm3( &self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, w0: &CudaSlice<f32>, w1: &CudaSlice<f32>, w2: &CudaSlice<f32>, res: &mut CudaSlice<f32>, d0: &mut CudaSlice<f32>, d1: &mut CudaSlice<f32>, d2: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(), Box<dyn Error>>

gemma4: res = a+b AND the three rms_norms of res in one launch.

Source

pub fn add_scale( &self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, c: f32, dst: &mut CudaSlice<f32>, n: usize, ) -> Result<(), Box<dyn Error>>

dst = (a + b) * c (residual add + layer scale, one launch).

Source

pub fn rms_norm( &self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(), Box<dyn Error>>

Source

pub fn rms_norm_decode( &self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(), Box<dyn Error>>

RMS-norm with blockDim=1024 — BIT-IDENTICAL to the fused rms_norm_q8_1 and add_rms_norm_q8_1 kernels’ sum-of-squares reduction. The spec verify path MUST use this to match decode’s FP accumulation order: the standard rms_norm at blockDim=256 has a different per-thread stride (ncols/256 partials vs ncols/1024 partials) and therefore a different shfl-tree reduction that can shift scale = rsqrt(sum/n + eps) by ULPs, causing divergence through the GDN scan and argmax flips on the 9B text prompt. The underlying rms_norm_f32 kernel supports any blockDim (generic reduce with shared[32]).

Source

pub fn rms_norm_q8_1( &self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn Error>>

DECODE GLUE-FUSION LEVER: z = rms_norm(x)*w emitted DIRECTLY as q8_1 (no f32 z materialized, no standalone quantize_q8_1 launch). Returns (out_q [nrowsncols i8], out_d [nrowsnblk f32]) ready to feed matmul_pre. BIT-IDENTICAL to rms_norm + quantize_q8_1. ncols % 32 == 0.

Source

pub fn rms_norm_q8_1_into( &self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, q: &mut CudaSlice<i8>, d: &mut CudaSlice<f32>, ) -> Result<(), Box<dyn Error>>

Slot-fed rms_norm_q8_1 twin (alloc-free capture lane): identical launch (incl. the PDL arm), caller-owned outputs.

Source

pub fn quantize_q8_1_into( &self, x: &CudaSlice<f32>, m: usize, in_f: usize, q: &mut CudaSlice<i8>, d: &mut CudaSlice<f32>, ) -> Result<(), Box<dyn Error>>

Slot-fed quantize_q8_1 twin (alloc-free capture lane).

Source

pub fn add_rms_norm_q8_1( &self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, w: &CudaSlice<f32>, res: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn Error>>

DECODE GLUE-FUSION LEVER: res = a+b; z = rms_norm(res)*w with z emitted as q8_1. res is still written (the post-ffn residual add reads it). Fuses add_rms_norm + quantize_q8_1. Returns (out_q, out_d) for matmul_pre. BIT-IDENTICAL. ncols % 32 == 0.

Source

pub fn add_rms_norm( &self, a: &CudaSlice<f32>, b: &CudaSlice<f32>, w: &CudaSlice<f32>, res: &mut CudaSlice<f32>, dst: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(), Box<dyn Error>>

RANK3 LEVER (add+rmsnorm fuse): res = a + b; dst = rms_norm(res) * w in ONE launch. Fuses e.add(a,b,res) + e.rms_norm(res,w,dst), removing one launch + one HBM read of the residual per residual+norm pair. BIT-IDENTICAL to the two-kernel sequence (same IEEE add, same reduction).

Source

pub fn rms_pre_add_rms_norm( &self, a: &CudaSlice<f32>, wa: &CudaSlice<f32>, b: &CudaSlice<f32>, w: &CudaSlice<f32>, res: &mut CudaSlice<f32>, dst: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(), Box<dyn Error>>

E4B glue fusion: rms(a, wa) prologue + add_rms_norm — folds the post-attn norm into the tail entry (res = rms(a)*wa + b; dst = rms(res)*w).

Source

pub fn rms_pre_add_rms_norm_q8z( &self, a: &CudaSlice<f32>, wa: &CudaSlice<f32>, b: &CudaSlice<f32>, w: &CudaSlice<f32>, res: &mut CudaSlice<f32>, dst: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn Error>>

wave-2 fold: rms(a,wa) + add + ffn-norm with zsh EMITTED q8_1 (fused2 consumes it).

Source

pub fn build_q4_out_concat3( &self, w0: &GpuTensor, w1: &GpuTensor, w2: &GpuTensor, ) -> Result<Option<GpuTensor>, Box<dyn Error>>

wave-4b: OUT-dim concat of three Q4_0 tensors (same in_features; rows are independent blocks, so the concat is a D2D byte concat of the GGUF-layout planes). Returns None off-class (non-Q4_0, mismatched widths, or any tensor already rp-swapped in place).

Source

pub fn rms_norm_qkv_rope_cat( &self, qkv: &CudaSlice<f32>, wq: &CudaSlice<f32>, wk: &CudaSlice<f32>, wv: &CudaSlice<f32>, q: &mut CudaSlice<f32>, k: &mut CudaSlice<f32>, v: &mut CudaSlice<f32>, head_dim: usize, rq: usize, rk: usize, pos: &CudaSlice<i32>, nh_q: usize, nh_k: usize, base: f32, freq_scale: f32, ff: Option<&CudaSlice<f32>>, eps: f32, ) -> Result<(), Box<dyn Error>>

wave-4b: the qkv-cat twin — one contiguous [rq+2*rk, hd] input from the concat matvec.

Source

pub fn rms_norm_qkv_rope( &self, q0: &CudaSlice<f32>, k0: &CudaSlice<f32>, v0: &CudaSlice<f32>, wq: &CudaSlice<f32>, wk: &CudaSlice<f32>, wv: &CudaSlice<f32>, q: &mut CudaSlice<f32>, k: &mut CudaSlice<f32>, v: &mut CudaSlice<f32>, head_dim: usize, rq: usize, rk: usize, pos: &CudaSlice<i32>, nh_q: usize, nh_k: usize, base: f32, freq_scale: f32, ff: Option<&CudaSlice<f32>>, eps: f32, ) -> Result<(), Box<dyn Error>>

wave-3 fold: rms_norm_qkv + rope_neox2 in ONE launch (n_dims == head_dim; ff nullable).

Source

pub fn rms_norm_qkv_rope_append_dc( &self, q0: &CudaSlice<f32>, k0: &CudaSlice<f32>, v0: &CudaSlice<f32>, wq: &CudaSlice<f32>, wk: &CudaSlice<f32>, wv: &CudaSlice<f32>, q: &mut CudaSlice<f32>, k: &mut CudaSlice<f32>, v: &mut CudaSlice<f32>, head_dim: usize, rq: usize, rk: usize, pos: &CudaSlice<i32>, nh_q: usize, nh_k: usize, base: f32, freq_scale: f32, ff: Option<&CudaSlice<f32>>, eps: f32, kc: &mut CudaSlice<u8>, vc: &mut CudaSlice<u8>, t_dev: &CudaSlice<i32>, k_tok_bytes: usize, v_tok_bytes: usize, g: bool, ) -> Result<(), Box<dyn Error>>

FUSED norm+rope+APPEND (m=1 decode, 2026-07-23): one launch replaces the rms_norm_qkv_rope + append_kv_quantized_dc pair. Kernel lives in the flash fatbins (format-flavored quant tail) — g must mirror the append path’s flavor exactly.

Source

pub fn add_q8_1( &self, a: &CudaSlice<f32>, b: &CudaSlice<f32>, res: &mut CudaSlice<f32>, ncols: usize, nrows: usize, ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn Error>>

wave-2 fold: a + b with the sum emitted q8_1 alongside f32.

Source

pub fn rms_pre_add_q8_1( &self, a: &CudaSlice<f32>, wa: &CudaSlice<f32>, b: &CudaSlice<f32>, res: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn Error>>

E4B FFN-tail exit fusion (glue wave 5): resid = b + rms(a, wa) emitted f32 + q8_1 pair in ONE launch — replaces rms_norm(a,wa->sn) + add_q8_1(sn,b). Same rms_block() config as both parents (bit-identity: identical reduction + quad-walk quantize).

Source

pub fn l2_v2_on(ncols: usize) -> bool

L2 norm per row (head_dim), no weight. PREFILL l2 dispatch (round 27): the warp-per-row float4 v2 when the numeric-config seam allows (MEMRA_L2_V2, default ON, d_state==128 only); else the strided kernel.

Source

pub fn l2_norm_pp( &self, x: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, dst16: Option<&mut CudaSlice<u8>>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(), Box<dyn Error>>

Source

pub fn l2_norm( &self, x: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(), Box<dyn Error>>

Source

pub fn l2_norm_decode( &self, x: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(), Box<dyn Error>>

L2-norm with blockDim=32 (warp-tree reduction) — BIT-IDENTICAL to gdn_prep_decode_f32’s per-warp L2 norm. The verify path MUST use this to match decode’s FP accumulation order: l2_norm at blockDim=256 produces a different shfl-tree reduction of the 128-element squared-sum (pairwise tree vs serial-4-then-warp-tree), causing ULP differences that propagate through gdn_scan and flip argmax on marginal logits.

Source

pub fn rope_neox( &self, x: &mut CudaSlice<f32>, pos: &CudaSlice<i32>, head_dim: usize, n_dims: usize, n_heads: usize, n_tokens: usize, freq_base: f32, freq_scale: f32, ) -> Result<(), Box<dyn Error>>

RoPE NEOX in-place. x:[head_dim, n_heads, n_tokens], pos:[n_tokens].

Source

pub fn rope_neox_ff( &self, x: &mut CudaSlice<f32>, pos: &CudaSlice<i32>, head_dim: usize, n_dims: usize, n_heads: usize, n_tokens: usize, freq_base: f32, freq_scale: f32, ff: &CudaSlice<f32>, ) -> Result<(), Box<dyn Error>>

RoPE NEOX with per-dim freq factors (gemma4 global layers, rope_freqs.weight [n_dims/2]).

Source

pub fn rope_neox2( &self, q: &mut CudaSlice<f32>, k: &mut CudaSlice<f32>, pos: &CudaSlice<i32>, head_dim: usize, n_dims: usize, nh_q: usize, nh_k: usize, n_tokens: usize, freq_base: f32, freq_scale: f32, ff: Option<&CudaSlice<f32>>, ) -> Result<(), Box<dyn Error>>

gemma4: rope q and k in one launch (per-row chain = rope_neox / rope_neox_ff verbatim).

Source

pub fn gelu_tanh_mul( &self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, n: usize, ) -> Result<(), Box<dyn Error>>

gemma4 R1: dst = GELU_tanh(gate) * up.

Source

pub fn silu_mul( &self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, n: usize, ) -> Result<(), Box<dyn Error>>

Source

pub fn silu_mul_f16out( &self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, dst16: &mut CudaSlice<u8>, n: usize, ) -> Result<(), Box<dyn Error>>

f16out twin of silu_mul (task #17): the epilogue also emits the fp16 GEMM operand for the down projection — kills the standalone convert pass. Bit-identical class.

Source

pub fn silu_mul_scaled( &self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, gs: f32, us: f32, dst: &mut CudaSlice<f32>, n: usize, ) -> Result<(), Box<dyn Error>>

FFN SwiGLU epilogue fusion (RANK3 LEVER 2): dst = silu(gate*gs) * (up*us) in ONE launch, folding the per-tensor NVFP4 macro-scale (gs,us) that would otherwise be two separate scale_inplace launches on the gate/up matmul outputs. BIT-IDENTICAL to scale_inplace(gate,gs); scale_inplace(up,us); silu_mul(gate,up,dst) — identical float ops in identical order. For non-NVFP4 weights gs==us==1.0 -> identical to silu_mul. Net: -2 launches per dense FFN layer (the gate+up post-matmul scales).

Source

pub fn swigluoai_mul_scaled( &self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, gs: f32, us: f32, alpha: f32, limit: f32, dst: &mut CudaSlice<f32>, n: usize, ) -> Result<(), Box<dyn Error>>

swigluoai (MiniMax-M3 / GPT-OSS): clamped SwiGLU epilogue, math 1:1 vs llama.cpp ggml_cuda_op_swiglu_oai_single. dst = swish_alpha(clamp(gate*gs)) * (1 + clamp(up*us)). gs/us fold the NVFP4 macro-scales exactly like silu_mul_scaled.

Source

pub fn silu_mul_scaled_q8_1( &self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, gs: f32, us: f32, n: usize, ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn Error>>

RANK2 LEVER (q8_1 quant-fold): SwiGLU epilogue that EMITS the q8_1 quantization of act directly (aq int8 [n] + ad f32 [n/32]), so ffn_down’s standalone quantize_q8_1 launch is removed — the down-proj activation has one consumer, so the quant folds into the producer for free (no extra HBM read; no f32 act write). gs/us fold the gate/up NVFP4 macro-scales like silu_mul_scaled. BIT-IDENTICAL q8_1 to silu_mul_scaled(…) then quantize_q8_1(…). Only valid when ffn_down uses the q8_1 dp4a/mmvq path; the caller checks uses_q8_1_fast(ffn_down). n must be a multiple of 32 (n_ff always is).

Source

pub fn add( &self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, n: usize, ) -> Result<(), Box<dyn Error>>

Source

pub fn mul( &self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, n: usize, ) -> Result<(), Box<dyn Error>>

Source

pub fn matmul( &self, w: &GpuTensor, x: &CudaSlice<f32>, m: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Unified weight-tensor matmul: dispatches quant tensors to qmatvec (weights packed) and float tensors to cuBLASLt. y[m,out] = x[m,in] @ W[out,in]^T.

Source

pub fn uses_q8_1_fast(&self, w: &GpuTensor) -> bool

True if w would take the int8-dp4a fast path under MEMRA_FAST (so its activation can be pre-quantized once and shared across sibling matmuls via matmul_pre).

Source

pub fn matmul_pre( &self, w: &GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, x_fallback: &CudaSlice<f32>, m: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

matmul with a PRE-QUANTIZED q8_1 activation (aq,ad from quantize_q8_1). Skips the per-matmul re-quantize so sibling matmuls that share an input (gate+up share z; q/k/v + wqkv/gate/beta/alpha share h) quantize ONCE. Caller MUST have checked uses_q8_1_fast(w); falls back to plain matmul otherwise (Stage-A / Float / non-fast).

Source

pub fn matmul_decode_exact( &self, w: &GpuTensor, x: &CudaSlice<f32>, m: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

DECODE-EXACT matmul at any m: guarantees the SAME warp-per-row (MMVQ, 32-thread) FP accumulation order as the T=1 decode path for EVERY token row. The spec-decode verify MUST use this for linear-attn projections to be bit-identical to greedy decode. The dp4a kernel (128 threads, two-level reduction) used by matmul/matmul_pre at m>=5 has a different shfl-tree shape that produces ULP differences propagating through gdn_scan into argmax flips. The MMVQ kernel with grid.y=m already processes each row independently (same 32-thread warp reduce as m=1); this method just forces that path unconditionally.

Source

pub fn matmul_decode_exact_pre( &self, w: &GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

DECODE-EXACT matmul from a PRE-QUANTIZED q8_1 activation (batched-verify epilogue re-fuse, lane/vt-fixes fix 2, 2026-08-03): the EXACT matmul_decode_exact dispatch for q8_1-fast Quant tensors, with the caller’s (aq, ad) replacing the internal quantize_q8_1. quantize_q8_1 is deterministic (same input bytes -> same q8 bytes), so sharing one quantize across sibling matmuls of the same activation — or consuming the q8 emitted by a fused epilogue (rms_norm_q8_1 / add_rms_norm_q8_1 / silu_mul_scaled_q8_1 / gated_rmsnorm_q8_1, all kernel-check-pinned bit-identical to their unfused chains) — cannot change any dispatched kernel’s input bytes. Caller MUST guarantee uses_q8_1_fast(w) (the fused epilogues only exist on that path).

Source

pub fn matmul_decode_exact_dual_pre( &self, w0: &GpuTensor, w1: &GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize, ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn Error>>

DUAL gate+up batched matvec from a PRE-QUANTIZED activation, macro-scales DEFERRED (lane/vt-fixes fix 2): same eligibility as matmul_decode_exact_dual, but the caller’s (aq, ad) replaces the internal quantize and the NVFP4 per-tensor scales are RETURNED instead of applied via two scale_inplace launches — the fused SwiGLU epilogue (silu_mul_scaled_q8_1) folds them, exactly like the m=1 decode chain does. Deferring is value-exact: y[i]*s inline in the epilogue is the same IEEE multiply scale_inplace would store (f32 store/load round-trips are exact). None -> caller falls back to the per-tensor path.

Source

pub fn matmul_decode_exact_dual( &self, w0: &GpuTensor, w1: &GpuTensor, x: &CudaSlice<f32>, m: usize, ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn Error>>

DUAL gate+up BATCHED matvec at verify t=2..8 (lane/verify-economics, 2026-08-02): ONE launch computes both FFN projections of a verify batch — same activation, same shape, blockIdx.y selects the tensor. Per (tensor, token, row) the kernel body is the single batched program on the SAME layout (split-plane rp: b2 rp / b4 rpr2 / b8 rpr2; GGUF: b2 base / b4 r2 / b8 r2) -> BIT-IDENTICAL to the two single matmul_decode_exact launches (kernel-check gates bitwise on both layouts; run-spec K=1..8 arbitrates e2e). The one activation quantize replaces two IDENTICAL quantizes of the same x (same kernel, same input -> same q8_1 bytes), and the two independent weight streams in one grid restore the memory-level parallelism the two-launch form loses to tail drain + launch gap (m=1 dual_mr2 precedent: DRAM 40% -> 47-50% on the 27B pair). Some((y0, y1)) only when both tensors are NVFP4, the SAME layout (both rp or both GGUF, no rp4 mirror), identical (in_f, out_f, row_bytes), q8_1-fast, and m in 2..=4 (the b2/b4 tiers = verify T for K=1..3, the profitable-K window — the b8 dual measured FLAT vs the rpsc singles x3 interleaved, research/verify-economics-20260802, and was killed per doctrine). None -> caller runs the two singles. MEMRA_SPEC_DUAL_T=0 rollback.

Source

pub fn qmatvec_batched_dual_raw( &self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, row_bytes: usize, rp: bool, ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn Error>>

Launch body of the dual batched twins from raw NVFP4 weight bytes + a pre-quantized q8_1 activation (kernel-check’s bit-equivalence entry; matmul_decode_exact_dual’s core). mcols tier = batched_mcols(m); macro-scale NOT applied. rp selects the split-plane twins (both buffers must be the repacked layout).

Source

pub fn matmul_pre_dual_noscale( &self, w0: &GpuTensor, w1: &GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize, ) -> Result<Option<((CudaSlice<f32>, f32), (CudaSlice<f32>, f32))>, Box<dyn Error>>

Like matmul_pre but RETURNS THE RAW (un-macro-scaled) matmul output together with the per-tensor NVFP4 scale, instead of applying scale_inplace internally. Used by the fused SwiGLU epilogue (RANK3 LEVER 2) so the gate/up scales fold into one silu_mul_scaled launch. Some((y_raw, scale)) only on the m==1 decode fast path (mmvq / dp4a) where the scale is a separate post-launch op we can defer; returns None for every other path (prefill GEMM, FP4 GEMM, Stage-A, Float) so the caller falls back to the scaled matmul_pre + silu_mul. DUAL gate+up NVFP4 matvec (mm-fusion): ONE launch computes both projections (same activation, same shape) — grid.y selects the tensor. Bit-identical per element to two mr2 launches at m=1. Returns (gate_raw, up_raw) un-scaled (caller folds the two macro scales into the SwiGLU epilogue, same as the matmul_pre_noscale contract). None unless both tensors are NVFP4 q8_1-fast with identical (in_f, out_f, row_bytes) and m==1.

Source

pub fn matmul_q8_fused2( &self, w0: &GpuTensor, w1: &GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn Error>>

FUSED Q8_0 m=1 matvec PAIR with UNEQUAL out_f (trunk launch-fusion, 2026-07-05). Folds two same-input q8_0 projections (35B trunk: wqkv+wqkv_gate 8192/4096, gate_shexp+up_shexp 512/512) into ONE launch via a block-offset split (blocks [0,nb0) -> w0, rest -> w1) — the dual-mr2 recipe with the same-out_f restriction lifted. Per (tensor,row) the kernel body is qmatvec_q8_0_mmvq VERBATIM -> BIT-IDENTICAL to two separate m=1 launches. Returns None when ineligible (not both Q8_0 / in_f mismatch / MEMRA_MMVQ off / MEMRA_Q8_DUAL=0) — caller falls back to the per-tensor path.

Source

pub fn matmul_q8_fused2_x( &self, w0: &GpuTensor, w1: &GpuTensor, x: &CudaSlice<f32>, ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn Error>>

f32-activation entry for the fused2 pair: quantizes x to q8_1 ONCE then runs the fused launch — replaces two matmul(w, x, 1) calls that would each re-quantize the same x (35B shared-expert gate+up per MoE layer per token). Same bits: quantize_q8_1 is deterministic, the fused body is the MMVQ kernel verbatim. None when ineligible (the callers’ m==1-under-MEMRA_FAST dispatch would take MMVQ; anything else falls back).

Source

pub fn qmatvec_q8_fused2_raw( &self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, x: &CudaSlice<f32>, in_f: usize, out0: usize, out1: usize, row_bytes: usize, ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn Error>>

Test entry for the kernel_check gate: launch the fused2 kernel from raw weight bytes, quantizing the f32 activation internally (mirrors qmatvec_mmvq_raw; no env gating).

Source

pub fn matmul_q4_fused3( &self, w0: &GpuTensor, w1: &GpuTensor, w2: &GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn Error>>

FUSED Q8_0 m=1 matvec TRIPLE (wq+wk+wv on the 35B full-attn layers: out_f 8192/512/512). Same block-offset recipe as matmul_q8_fused2 with three ranges. BIT-IDENTICAL per (tensor,row) to three separate m=1 MMVQ launches. FUSED Q4_0 m=1 TRIPLE (gemma q/k/v — same quantized input; per (tensor,row) chain identical to the mr2 kernel). Returns None unless all three are Q4_0 with equal in_f.

Source

pub fn matmul_q4_fused3_into( &self, w0: &GpuTensor, w1: &GpuTensor, w2: &GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, y0: &mut CudaSlice<f32>, y1: &mut CudaSlice<f32>, y2: &mut CudaSlice<f32>, ) -> Result<bool, Box<dyn Error>>

Slot-fed fused3 twin (alloc-free capture lane): identical launch, caller-owned outputs. Returns Ok(false) when the fused path is unavailable (caller falls back).

Source

pub fn matmul_q4_fused2( &self, w0: &GpuTensor, w1: &GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn Error>>

FUSED Q4_0 m=1 PAIR (gemma shared gate+up).

Source

pub fn matmul_q4_fused2_into( &self, w0: &GpuTensor, w1: &GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, y0: &mut CudaSlice<f32>, y1: &mut CudaSlice<f32>, ) -> Result<bool, Box<dyn Error>>

Slot-fed fused2 twin (alloc-free capture lane): identical launch, caller-owned outputs.

Source

pub fn matmul_q4_fused2_batched( &self, w0: &GpuTensor, w1: &GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize, ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn Error>>

BATCHED fused2 (2026-07-13, megakernel-microcosm probe): gate+up b-tier matvecs in ONE segmented-grid launch — the up segment fills SMs as the gate segment drains (the per-launch tail waves behind the 6x-falsified b-tier plateau). Bit-identical per row to two mr2_rp launches. rp layout required; m in 2..=8 (b16 has no twin).

Source

pub fn matmul_q4_fused3_batched( &self, w0: &GpuTensor, w1: &GpuTensor, w2: &GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize, ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn Error>>

BATCHED fused3 (see matmul_q4_fused2_batched): three-segment single launch for the verify qkv triple. Same-in_f q4_0 rp tensors, m in 2..=8. Bit-identical per row.

Source

pub fn matmul_q8_fused3( &self, w0: &GpuTensor, w1: &GpuTensor, w2: &GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn Error>>

Source

pub fn qmatvec_q8_fused3_raw( &self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, b2: &CudaSlice<u8>, x: &CudaSlice<f32>, in_f: usize, out0: usize, out1: usize, out2: usize, row_bytes: usize, ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn Error>>

Test entry for the kernel_check gate: fused3 from raw weight bytes (internal q8_1 quant).

Source

pub fn matmul_q8_fused2_t( &self, w0: &GpuTensor, w1: &GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize, ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn Error>>

BATCHED twin of matmul_q8_fused2 for the verify t=2-4 tier (MEMRA_SPEC_FUSED_T call sites, lane/close35b): ONE launch computes both same-input Q8_0 projections for m tokens. Per (tensor,token,row) the kernel body is q8_0_mmvq_batched VERBATIM with the identical row mapping (Q8_0’s batched_variant is always “base”) -> BIT-IDENTICAL to the two per-tensor _b2/_b4 launches matmul_decode_exact dispatches at m=2-4, with the caller’s single shared q8_1 activation replacing two per-call re-quantizes (quantize_q8_1 is deterministic -> same bytes). None when ineligible (m outside 2..=4 / not both Q8_0 / in_f mismatch / MEMRA_MMVQ=0 / MEMRA_Q8_DUAL=0 / MEMRA_NO_BATCHED set — the last keeps dispatch parity: without batched kernels decode-exact runs grid.y=m MMVQ, and the fused twin must not introduce a batched program the reference path would not run).

Source

pub fn qmatvec_q8_fused2_t_raw( &self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out0: usize, out1: usize, row_bytes: usize, ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn Error>>

Test entry for the kernel_check gate: fused2 batched from raw weight bytes (internal q8_1 quant of the [m, in_f] activation), no env gating.

Source

pub fn matmul_q8_fused3_t( &self, w0: &GpuTensor, w1: &GpuTensor, w2: &GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize, ) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn Error>>

BATCHED twin of matmul_q8_fused3 (wq+wk+wv at verify t=2-4). Same contract as matmul_q8_fused2_t with three ranges.

Source

pub fn qmatvec_q8_fused3_t_raw( &self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, b2: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out0: usize, out1: usize, out2: usize, row_bytes: usize, ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn Error>>

Test entry for the kernel_check gate: fused3 batched from raw weight bytes.

Source

pub fn q8_ffn_fuse2_on(&self) -> bool

Rollback seam for the Q8_0 dense-FFN gate+up fusion arm in matmul_pre_dual_noscale (lane/q27-deepdive, 2026-08-05). Default ON; MEMRA_Q8_FFN_FUSE2=0 restores the two-matmul_pre_noscale pair. Read once — the dispatch must not vary within a run.

Source

pub fn e4m3_dual_on(&self) -> bool

Rollback seam for the F8-E4M3 launch-fusion arm (lane/fp8-decode-v1, 2026-08-05). Default ON; MEMRA_E4M3_DUAL=0 restores the per-tensor m=1/batched launches.

Source

pub fn qmatvec_e4m3_blk_mmvq( &self, bytes: &CudaSlice<u8>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, scales: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, row_bytes: usize, scale_cols: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

BLOCK-128 e4m3 MMVQ launcher (qmatvec_e4m3_blk_mmvq, lane/fp8-blk128-decode 2026-08-05). The per-block-dequant twin of qmatvec_mmvq’s QT_F8_E4M3 arm: same grid/block decomposition (warp per output row, ROWS_PER_BLOCK warps per block, grid.y = m), same q8_1 activation, but the weight scale is a resident [rows, cols] f32 grid read per k128 block inside the kernel instead of one scalar folded at the write. It cannot share qmatvec_mmvq’s body because that launcher’s arg list is fixed at (bytes, aq, ad, y, in_f, out_f, m, row_bytes [, scale]).

mr and rp have no analogue here (no split-plane e4m3 layout exists), so there is exactly one kernel and no name table — a shape this cannot serve must be refused at LOAD, not here.

Source

pub fn qmatvec_e4m3_blk_mmvq_into( &self, bytes: &CudaSlice<u8>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, scales: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, row_bytes: usize, scale_cols: usize, y: &mut CudaSlice<f32>, ) -> Result<(), Box<dyn Error>>

Slot-fed twin of qmatvec_e4m3_blk_mmvq (caller-owned output; the alloc-free capture lane).

Source

pub fn qmatvec_e4m3_blk_mmvq_raw( &self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, scales: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, row_bytes: usize, scale_cols: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Test entry for the kernel_check exactness gate: the block-128 e4m3 MMVQ from raw bytes with an internal q8_1 quantize (mirrors qmatvec_mmvq_raw).

Source

pub fn qmatvec_e4m3_fused2_raw( &self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, x: &CudaSlice<f32>, in_f: usize, out0: usize, out1: usize, row_bytes: usize, ws0: f32, ws1: f32, ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn Error>>

Test entries for the kernel_check bit-parity gate: fused e4m3 launches from raw weight bytes with internal q8_1 quantize, no env gating (mirrors qmatvec_q8_fused*_raw).

Source

pub fn qmatvec_e4m3_fused3_raw( &self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, b2: &CudaSlice<u8>, x: &CudaSlice<f32>, in_f: usize, out0: usize, out1: usize, out2: usize, row_bytes: usize, ws0: f32, ws1: f32, ws2: f32, ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn Error>>

Source

pub fn qmatvec_e4m3_fused2_t_raw( &self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out0: usize, out1: usize, row_bytes: usize, ws0: f32, ws1: f32, ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn Error>>

Source

pub fn qmatvec_e4m3_fused3_t_raw( &self, b0: &CudaSlice<u8>, b1: &CudaSlice<u8>, b2: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out0: usize, out1: usize, out2: usize, row_bytes: usize, ws0: f32, ws1: f32, ws2: f32, ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn Error>>

Source

pub fn matmul_pre_noscale( &self, w: &GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize, ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn Error>>

Source

pub fn mmvq_supports(&self, qtype: i32) -> bool

True if qtype has a warp-per-row MMVQ decode kernel AND MEMRA_MMVQ is set. Only the 4 daily-hot dtypes (Q8_0, Q4_K, Q6_K, NVFP4) — others keep the _dp4a matvec (oracle/fallback).

Source

pub fn qmatvec_mmvq( &self, bytes: &CudaSlice<u8>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, qtype: i32, row_bytes: usize, scale: f32, rp: bool, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

PERF-3 warp-per-row MMVQ launcher (decode m=1 hot path). block=(32,ROWS_PER_BLOCK,1): one warp owns one output row, warp-only _shfl reduction (no smem barrier). Bit-equivalent to qmatvec*_dp4a up to f32 reduction order. Pre-quantized q8_1 activation (aq,ad). NVFP4 per-tensor macro-scale applied post (scale==1.0 for other dtypes -> no-op).

Source

pub fn qmatvec_mmvq_into( &self, bytes: &CudaSlice<u8>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, qtype: i32, row_bytes: usize, scale: f32, rp: bool, y: &mut CudaSlice<f32>, ) -> Result<(), Box<dyn Error>>

Slot-fed MMVQ twin (alloc-free capture lane): full policy body, caller-owned output.

Source

pub fn qmatvec_mmvq_raw( &self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, qtype: i32, row_bytes: usize, rp: bool, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Test entry for the kernel_check bit-equivalence gate: run the warp-per-row MMVQ directly from raw weight bytes (quantize the f32 activation x to q8_1 internally). NVFP4 per-tensor macro-scale is NOT applied (caller compares bare, like qmatvec_*_fast). Mirrors qmatvec_gemm_raw.

Source

pub fn batched_supports(&self, qtype: i32) -> bool

True if qtype has a batched weight-resident (_b2/_b4) matvec kernel. These mirror the _mmvq kernels but iterate the m token columns INSIDE one warp/row, so the weight bytes leave HBM/L2 once for m tokens (vs grid.y=m re-reading m times). The 5 daily-hot dtypes have them.

Source

pub fn iq_fast_enabled() -> bool

IQ4_XS trunk fast seam: MEMRA_IQ_FAST=0 reverts non-expert IQ4_XS matmuls to the Stage-A f32 oracle path. Default ON since 2026-08-02 (research/kat-anomaly-20260802/): the old opt-in default left every IQ4_XS-trunk artifact (KAT-Coder IQ4_XS: attn_qkv/attn_gate/ ssm_out/shexp, ~0.52GB re-read per decode tick) on the oracle kernel — decode 106.7 -> 193.4 tok/s (x5 interleaved), pp512 228 -> 697, same bytes, via qmatvec_iq4_XS_dp4a. The supported artifacts carry IQ4_XS only in EXPERT banks (their own dispatch, not this seam), so this admission is dispatch-unchanged for every non-IQ4_XS-trunk model.

Source

pub fn b8_enabled() -> bool

b8 tier seam: MEMRA_B8=0 keeps m=5..8 on the per-m grid.y=m path (m=2..4 batched dispatch unaffected). Default ON — the K=4..7 spec-verify weight-read-once fix.

Source

pub fn batched_mcols(m: usize) -> usize

Compile-time column batch for a runtime m: 2 -> b2, 3..4 -> b4, 5..8 -> b8.

Source

pub fn sm_count(&self) -> i32

BATCHED weight-tile-resident matvec from a PRE-QUANTIZED q8_1 activation (the m=2-8 verify/MTP win). One warp walks the weight row ONCE, dp4a vs all m activation columns -> weight HBM/L2 traffic 1x for m tokens (vs grid.y=m re-reading it m times). mcols ∈ {2,4,8} is the compile-time batch; m must be <= mcols (the c >= m columns are masked in-kernel). y is [m, out_f] token-major. NVFP4 per-tensor macro-scale applied post (scale==1.0 for other dtypes -> no-op). BIT-IDENTICAL per (token,row) to qmatvec_*_mmvq.

NVFP4 VARIANT DISPATCH: the batched NVFP4 kernel measured memory-LATENCY bound on the real 27B verify (ncu –set full, 12 steady launches: long_scoreboard 18-30 stalls/issue vs <=1.7 for every other reason, DRAM only 41-51% active, lg_throttle 0.7, L1 hit 94% — ONE 6-LDG weight wavefront in flight per warp is the binding constraint, NOT bandwidth and NOT the column-unroll break). Two exactness-free fixes, chosen PER SHAPE from the DRAM-cold 8-copy msweep on all six 27B shapes (2026-07-03): pf = next-g weight-prefetch double-buffer (48 regs, occupancy intact) — wins everywhere it applies for b4 (-3..-14%), never loses; r2 = two rows/warp (67 regs -> 7 resident blocks/SM) — the bigger win (-8.5..-30%) but wave-quantization-sensitive: with the grid halved to ceil(out_f/8) blocks, a fractional straggler wave (waves in ~1.05-1.5) costs a full extra latency round on a latency-bound kernel (27B ffn_down 640 blocks / 574 resident = 1.11 waves: +17%), while <=1 wave (9B ffn_down 0.89: -30%) or >=2 waves (tail amortized; qkv 2.2: -8.5%, ffn_gate 3.8: -12.5%) win. For b2, r2 wins on DEEP k-loops (in_f>=6144: -8..-19%) where the 2-col body starves weight MLP hardest; pf measured negative. b4: r2 when waves(out_f) <= 1 (and grid fills >=half the SMs) or >= 2, else pf. b2: in_f>=6144 -> r2, else base. MEMRA_MMVQ_BV=base|pf|r2|pfr2 forces one variant everywhere (A/B + rollback seam). All variants BIT-IDENTICAL per (token,row): same dp4a order, scales, adg factor, reduce — only load issue time and the row->warp mapping change (kernel-check gates all of them). rp = the weight buffer is the A6 SPLIT-PLANE repacked layout (NVFP4 only): the same wave-aware auto rule applies, mapped onto the _rp twins (rp/rpr2/rpr2w8 mirror pf/r2/r2w8 — regs 44/67/64 land in the same residency classes). The variant the batched dispatch will pick for this (shape, m, mcols, layout) — exposed so gates can distinguish bit-identical variants (bit-bad==0 required) from the k-split family (deterministic but k-reduce-order-shifted: rel<1e-3 + run-to-run bit-identity required). Device SM count (cached) — grid-fill policy input.

Source

pub fn batched_variant( &self, _m: usize, in_f: usize, out_f: usize, qtype: i32, row_bytes: usize, mcols: usize, rp: bool, ) -> &'static str

Source

pub fn qmatvec_mmvq_batched( &self, bytes: &CudaSlice<u8>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, qtype: i32, row_bytes: usize, mcols: usize, scale: f32, rp: bool, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Source

pub fn qmatvec_batched_raw( &self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, qtype: i32, row_bytes: usize, mcols: usize, rp: bool, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

BATCHED weight-tile-resident matvec from raw weight bytes (quantizes the f32 activation x to q8_1 internally; macro-scale NOT applied — caller compares bare, like qmatvec_*_fast). For the kernel_check bit-equivalence gate. mcols ∈ {2,4,8}. Works for Q8_0/Q4_K/Q5_K/Q6_K/NVFP4.

Source

pub fn qmatvec_nvfp4_batched_raw( &self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, row_bytes: usize, mcols: usize, rp: bool, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Back-compat NVFP4-only batched raw launcher (used by older gates). Delegates to the generic one.

Source

pub fn rms_norm_f16out( &self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, dst16: &mut CudaSlice<u8>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(), Box<dyn Error>>

rms_norm + fused fp16 twin (task #14): f32 output verbatim rms_norm + the fp16 copy the f16-mirror GEMM group would otherwise produce with a standalone convert launch. BIT-IDENTICAL end-to-end (same reduction, same __float2half values).

Source

pub fn add_rms_norm_f16out( &self, a: &CudaSlice<f32>, b: &CudaSlice<f32>, w: &CudaSlice<f32>, res: &mut CudaSlice<f32>, dst: &mut CudaSlice<f32>, dst16: &mut CudaSlice<u8>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(), Box<dyn Error>>

add+norm(+f16out) fusion for the prefill trunk (round 28; add_rms_norm precedent — bit-identical to add_f32 -> rms_norm_f16out). block_dim matches rms_norm_f16out’s.

Source

pub fn matmul_group_xh( &self, ws: &[&GpuTensor], x: &CudaSlice<f32>, xh: &CudaSlice<u8>, m: usize, ) -> Result<Vec<CudaSlice<f32>>, Box<dyn Error>>

matmul_group with a PRE-EMITTED fp16 activation (task #14: the producer norm fused the convert). Mirror-less members fall back to matmul on the f32 activation.

Source

pub fn gdn_pad_mask( &self, beta: &mut CudaSlice<f32>, g_log: &mut CudaSlice<f32>, len_d: &CudaSlice<i32>, h: usize, t: usize, ) -> Result<(), Box<dyn Error>>

task #14 pad-proofing: zero beta/g_log at rows >= len_d[0] (pads become identity GDN steps). Layouts [T, H].

Source

pub fn row_gather_dev( &self, src: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, len_d: &CudaSlice<i32>, ncols: usize, ) -> Result<(), Box<dyn Error>>

task #14 pad-proofing: dst[ncols] = src row (len_d[0]-1) — device-indexed last-row gather for the padded prime graph’s h_seed/hlast.

Source

pub fn matmul_group( &self, ws: &[&GpuTensor], x: &CudaSlice<f32>, m: usize, ) -> Result<Vec<CudaSlice<f32>>, Box<dyn Error>>

Grouped matmul: several weights consuming ONE activation (hybrid layers: the GDN 4-tuple wqkv/gate/beta/alpha, attention q/k/v, ffn gate/up). Semantics identical to calling matmul per weight; the f16-mirror arm converts the activation ONCE for the whole group instead of once per GEMM (the standalone converts were ~250 launches/prime of small-kernel gap fuel — nsys 2026-07-26). Any member without a mirror (or with a different in_f) falls back to its own matmul — behavior unchanged.

Source

pub fn matmul_group_multi( &self, ws: &[&GpuTensor], xs: &[&CudaSlice<f32>], ms: &[usize], ) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn Error>>

Cross-request grouped matmul (task #13): run ONE projection group over the CONCATENATION of several sequences’ activations (m = sum of per-seq rows — the GEMM-batch win vLLM gets from continuous batching), then split each output back into per-seq buffers. Zero view plumbing: gather/scatter are stream-ordered D2D copies (~us at prime sizes). NUMERIC CONFIG NOTE: a GEMM at m=sum tiles K differently than per-seq GEMMs — argmax-gated like every prefill GEMM change.

Source

pub fn gemm_supports(&self, w: &GpuTensor) -> bool

True if w’s qtype has a batched tensor-core GEMM kernel (the prefill T>1 root fix). Only the 4 daily-hot dtypes: Q8_0, Q4_K, Q6_K, NVFP4. NVFP4 needs in_f % 64 == 0. DEFAULT-ON (2026-06-28): measured pp512 9B-NVFP4 = 1413 tok/s WITH this GEMM vs 298 with the dp4a fallback (4.7x) AND MORE accurate (prefill logit maxdiff 0.159 vs dp4a 0.55, both argmax MATCH). The int8 tensor-core GEMM is unconditional (its historical MEMRA_GEMM opt-in gate shipped with Phase 0 — mma + smem swizzle + cp.async — and was removed). Prefill-only (m>=GEMM_M_THRESHOLD); m=1 decode keeps dp4a/MMVQ (this returns true but matmul only calls it at m>=threshold). Portable CUDA targets always use the correctness fallback; on sm_120a, MEMRA_NO_GEMM forces that same dp4a fallback (the bit-reference).

Source

pub fn qmatvec_gemm( &self, w: &GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Batched tensor-core int8 GEMM with a PRE-QUANTIZED q8_1 activation (aq,ad). The prefill (T>1) root fix: decode each weight 32-block to int8 in shared memory ONCE per (row-tile, K-step) and reuse it across all BN tokens via mma.sync.m16n8k32.s8 — amortizing the weight read/decode N-fold (vs the dp4a matvec’s per-token re-read). s32 accumulate is exact vs dp4a; only the final f32 block-scale rounding differs. Caller MUST have checked gemm_supports(w). y[m,out] token-major. NVFP4 per-tensor macro-scale applied post.

Source

pub fn qmatvec_gemm_raw( &self, bytes: &CudaSlice<u8>, x: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, qtype: i32, row_bytes: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Test entry: run the GEMM directly from raw weight bytes + qtype (no GpuTensor). Quantizes the f32 activation x to q8_1 internally then launches the tensor-core GEMM. NVFP4 per-tensor macro-scale is NOT applied here (caller passes it separately, like the dp4a path). Used by kernel_check for the bit-equivalence gate vs qmatvec_*_dp4a.

Source

pub fn qmatvec_gemm_q8_0_wgmma_raw( &self, rp4: &CudaSlice<u8>, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize, in_f: usize, out_f: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

H100 warpgroup GEMM raw entry (task 8): launch qmatvec_gemm_q8_0_wgmma on an rp4 split-plane mirror + pre-quantized (aq, ad) activation planes. One warpgroup (128 thr) owns a 64x64 C tile; grid (out_f/64, ceil(m/64)). out_f % 64 == 0 REQUIRED (row loads and dplane scale reads are unguarded); the token edge is guarded in-kernel. Standalone harness verdict (tools/bench_q8_gemm_wgmma.cu, 4096x4096x512): rel 1.6e-05 vs CPU ref, 179us vs the portable mma kernel’s 688us (3.84x, unpipelined).

Source

pub fn scale_inplace( &self, y: &mut CudaSlice<f32>, s: f32, n: usize, ) -> Result<(), Box<dyn Error>>

y[i] *= s. NVFP4 per-tensor macro-scale broadcast over the whole output.

Source

pub fn bf16_to_f32( &self, data: &CudaView<'_, u8>, n: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

MEMRA_FULL_PREC dequant-on-use: expand a bf16-resident weight (GpuTensor::FloatBf16, raw bf16 bytes) to a transient f32 scratch of n elements, which then feeds the existing f32 cuBLASLt GEMV. The scratch is freed when the caller drops it, so peak VRAM = resident bf16 weights + ONE (largest) weight’s f32 expansion + activations. SLOW IS FINE (research mode).

Source

pub fn linear_decode_exact( &self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, m_tokens: usize, in_f: usize, out_f: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

On-device linear: y[m,out] = x[m,in] @ W[out,in]^T, weights row-major [out,in] (ggml). cuBLASLt col-major mapping (see memra_runtime::Gpu::linear_f32 for the derivation). DECODE-EXACT float linear: per-column m=1 cuBLASLt calls. cuBLASLt’s reduction split is n-dependent (lt_ndep probe: m=1 vs m=2 col0 differs every bit), so spec-verify batches must not batch float matmuls the T=1 decode chain runs at m=1. Used by the small-t MoE router/shexp sites and matmul_decode_exact’s Float arm.

Source

pub fn linear( &self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, m_tokens: usize, in_f: usize, out_f: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Source

pub fn sdpa_naive( &self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>, o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool, ) -> Result<(), Box<dyn Error>>

Naive SDPA. Q:[head_dim,n_head,T], K/V:[head_dim,n_head_kv,T_kv] -> O:[head_dim,n_head,T].

Source

pub fn sdpa_naive_w( &self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>, o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool, window: usize, ) -> Result<(), Box<dyn Error>>

Windowed sdpa_naive twin (gemma4 R6): masks keys older than q_pos-(window-1).

Source

pub fn sdpa_naive_view( &self, q: &CudaSlice<f32>, k: &CudaView<'_, f32>, v: &CudaView<'_, f32>, o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool, ) -> Result<(), Box<dyn Error>>

SDPA where K/V are CudaViews into a resident KV cache (decode hot path, no host round-trip).

Source

pub fn fa_dequant_kv_view_f32( &self, k: &CudaView<'_, u8>, v: &CudaView<'_, u8>, kf: &mut CudaSlice<f32>, vf: &mut CudaSlice<f32>, kv_dim_k: usize, kv_dim_v: usize, t_kv: usize, k_tok_bytes: usize, v_tok_bytes: usize, g: bool, ) -> Result<(), Box<dyn Error>>

Correctness fallback for quantized resident K/V views. Dequantizes K and V once into f32 workspaces, then calls sdpa_naive. This is an explicit API: the optimized prefill view dispatch remains unchanged, so callers can use it as a reference or compatibility path. Dequant a quantized KV view into caller-owned f32 buffers (one grid-stride launch). g picks the kf8vf8-module stamp for e4m3 caches (same flag contract as fa_decode/ fa_prefill_view). Used by the E4B shared-KV prefill arms (2026-07-31) to feed the f32 fa_prefill_w / fa_prefill_hd512 twins from the target layer’s quantized rows.

Source

pub fn sdpa_naive_quantized_view( &self, q: &CudaSlice<f32>, k: &CudaView<'_, u8>, v: &CudaView<'_, u8>, o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool, k_tok_bytes: usize, v_tok_bytes: usize, ) -> Result<(), Box<dyn Error>>

Source

pub fn fa_prefill( &self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>, o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool, ) -> Result<(), Box<dyn Error>>

Hand-written FlashAttention prefill (sm_120, FA-2 online softmax on validated mma.sync, head_dim 256 or 128 (template-stamped twins), GQA, causal). Replaces sdpa_naive for T>1. Q/K/V/O [head_dim, n_head(_kv), T].

Source

pub fn fa_prefill_w( &self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>, o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool, window: usize, ) -> Result<(), Box<dyn Error>>

Windowed FA prefill (gemma4 SWA layers past the sliding window, hd256): fa_prefill’s exact dispatch (pp default, MEMRA_FA_FLOOR seam) with the sliding-window mask + tile skip in-kernel. Replaces the O(T*T_kv) scalar sdpa_naive_w on the prime path.

Source

pub fn fa_prefill_w_pre( &self, qb: &CudaSlice<u8>, kb: &CudaSlice<u8>, vb: &CudaSlice<u8>, o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool, window: usize, v_f16: bool, ) -> Result<(), Box<dyn Error>>

Windowed FA prefill with PRE-CONVERTED bf16 operands (producer-emitted; 31B glue lane). Launches the P1 stamp directly — callers guarantee qb/kb/vb hold the exact bf16 of q/k/v.

Source

pub fn fa_prefill_w_arm( &self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>, o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool, window: usize, f32_stage: bool, floor: bool, ) -> Result<(), Box<dyn Error>>

Windowed FA prefill with the stage arm FORCED — the kernel_check bit-identity entry.

Source

pub fn fa_prefill_hd512( &self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>, o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool, ) -> Result<(), Box<dyn Error>>

hd512 FA prefill (gemma4 GLOBAL layers): BLOCK_Q=32 x 2 warps, Q staged in smem, grid.z = 2 O-halves (each CTA computes the full 512-dim scores, accumulates half the V dims). Replaces the scalar sdpa_naive on the prime path’s globals.

Source

pub fn fa_prefill_hd512_pre( &self, qb: &CudaSlice<u8>, kb: &CudaSlice<u8>, vb: &CudaSlice<u8>, o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool, v_f16: bool, ) -> Result<(), Box<dyn Error>>

hd512 single-pass FA with PRE-CONVERTED bf16 operands (producer-emitted).

Source

pub fn fa_prefill_hd512_arm( &self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>, o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool, f32_stage: bool, sp: bool, f16pv: bool, ) -> Result<(), Box<dyn Error>>

hd512 FA prefill with the stage/sp arms FORCED — the kernel_check gate entry (fa_prefill_hd512 picks the arms from MEMRA_FA512_STAGE / MEMRA_FA512_SP).

Source

pub fn rope_neox2_bf16e( &self, q: &mut CudaSlice<f32>, k: &mut CudaSlice<f32>, qb: &mut CudaSlice<u8>, kb: &mut CudaSlice<u8>, pos: &CudaSlice<i32>, head_dim: usize, n_dims: usize, nh_q: usize, nh_k: usize, n_tokens: usize, base: f32, freq_scale: f32, ff: Option<&CudaSlice<f32>>, ) -> Result<(), Box<dyn Error>>

rope_neox2 with bf16 EMIT (31B glue lane): identical rope math/stores plus the post-rope values written as bf16 — the FA q/k operands come from this launch (bit-identical to the separate f32_to_bf16 the FA entries would run).

Source

pub fn f32_to_bf16( &self, x: &CudaSlice<f32>, n: usize, ) -> Result<CudaSlice<u8>, Box<dyn Error>>

Flat f32 -> bf16 conversion into a fresh scratch buffer (2 bytes/elem). n % 4 == 0 (float4 in, 4x bf16 out). Feeds the bf16-staged hd512 FA prefill.

Source

pub fn f32_to_f16( &self, x: &CudaSlice<f32>, n: usize, ) -> Result<CudaSlice<u8>, Box<dyn Error>>

Source

pub fn bf16_to_f16( &self, xb: &CudaSlice<u8>, n: usize, ) -> Result<CudaSlice<u8>, Box<dyn Error>>

bf16 bytes -> f16 bytes, n elements (the f16-P/V door’s V re-encode on the emit lane).

Source

pub fn bf16_to_f16_into( &self, xb: &CudaSlice<u8>, n: usize, y: &mut CudaSlice<u8>, ) -> Result<(), Box<dyn Error>>

Same conversion into a caller-owned (pooled) buffer; y.len() >= n*2.

Source

pub fn fa_prefill_vl8( &self, seqs: &[FaSeqVl], head_dim: usize, n_head: usize, n_head_kv: usize, scale: f32, ) -> Result<(), Box<dyn Error>>

task #18 (attn side): varlen FA — bf16 K/V mirrors (2 launches) + ONE fa_prefill_bf16kv launch for every fresh sequence. Same per-block math as the per-seq path (bit-gateable). Caller guarantees: fresh causal (T_kv == T), head_dim in {256, 128}, bf16kv lane on.

Source

pub fn attn_pre_vl8( &self, seqs: &[AttnPreVl], wq: &CudaSlice<f32>, wk: &CudaSlice<f32>, head_dim: usize, rope_dims: usize, n_head: usize, n_head_kv: usize, eps: f32, freq_base: f32, freq_scale: f32, kv_dim_k: usize, kv_dim_v: usize, k_tok_bytes: usize, v_tok_bytes: usize, ) -> Result<(), Box<dyn Error>>

task #18 (attn pre-FA): varlen split + QK-norm + RoPE + KV-append — FOUR launches for every fresh sequence (was 6 x B, plus the q/k/v split copies which the view inputs remove entirely). Fresh-only (append at t0=0, RoPE pos = token index).

Source

pub fn fa_prefill_view( &self, q: &CudaSlice<f32>, k: &CudaView<'_, u8>, v: &CudaView<'_, u8>, o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool, k_tok_bytes: usize, v_tok_bytes: usize, g: bool, ) -> Result<(), Box<dyn Error>>

FA prefill where K/V are QUANTIZED CudaViews into the resident byte KV cache (the T=K verify path, MTP-PLAN §D.3). Uses fa_prefill_q (inline-dequant during stage-to-smem). The view’s base+offset pointer is honored; the kernel reads [0..t_kv*tok_bytes). Q is the T fresh query rows; t = T, t_kv = cache len. k_tok_bytes/v_tok_bytes are the per-token byte strides.

Source

pub fn fa_prefill_view_ws( &self, q: &CudaSlice<f32>, k: &CudaView<'_, u8>, v: &CudaView<'_, u8>, o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize, t: usize, t_kv: usize, scale: f32, causal: bool, k_tok_bytes: usize, v_tok_bytes: usize, g: bool, ) -> Result<(), Box<dyn Error>>

ARC B (2026-07-05): dequant-once chunk-prime FA. Same contract as fa_prefill_view, but instead of every (q-block, head) CTA re-dequanting the whole quantized KV stream inline (T/64 x n_head redundant at chunk prime — 30.5% of the 32k prime wall), dequant the full [t_kv, kv_dim] K and V ONCE into a resident bf16 workspace (fa_dequant_kv_ws_bf16), then run fa_prefill_qw (the bf16-workspace twin) over it. EXACT: the workspace holds the same _float2bfloat16(dq*_elem(…)) values fa_prefill_q stages to smem, and the twin’s MMA/ softmax/PV code is byte-identical -> bit-identical O (kernel_check pins bitdiff=0). The workspace allocation is REUSED across layers/chunks (grown to the largest shape); contents are rewritten per call. MEMRA_PRIME_DEQW=0 falls back to fa_prefill_view (callers gate).

Source

pub fn fa_decode( &self, q: &CudaSlice<f32>, k: &CudaView<'_, u8>, v: &CudaView<'_, u8>, o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize, t_kv: usize, scale: f32, k_tok_bytes: usize, v_tok_bytes: usize, ) -> Result<(), Box<dyn Error>>

FA decode (T=1 split-K) over the resident QUANTIZED KV cache (q8_0 K / q5_1 V) as u8 views. Replaces sdpa_naive_view for decode; inline-dequants per element. k_tok_bytes/v_tok_bytes are the per-token byte strides (differ: q8_0=34nblk, q5_1=24nblk per token).

Source

pub fn fa_decode_kvmod( &self, q: &CudaSlice<f32>, k: &CudaView<'_, u8>, v: &CudaView<'_, u8>, o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize, t_kv: usize, scale: f32, k_tok_bytes: usize, v_tok_bytes: usize, g: bool, ) -> Result<(), Box<dyn Error>>

Source

pub fn fa_decode_batch_seqs_v4( &self, q: &CudaSlice<f32>, kv_ptrs: &CudaView<'_, u64>, pos_seq: &CudaSlice<i32>, o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize, b_n: usize, t_kv_max: usize, scale: f32, split_keys: usize, k_tok_bytes: usize, v_tok_bytes: usize, ) -> Result<(), Box<dyn Error>>

BATCHED-TICK increment 2: ONE fa_decode launch covering ALL B sequences of the batched decode step (blockIdx.z = sequence). Per-seq K/V cache bases ride a device pointer table (kv_ptrs, [2B] interleaved k0,v0,…); per-seq key bounds ride the tick’s position table (pos_seq, T_kv = pos+1). v4-lane only: the CALLER (decode_batch) gates every row through fa_seqs_eligible AND one fa_split_keys rung (split_keys), so each sequence’s split partition, key walk and combine order reproduce its per-seq eager v4 program exactly (kernel-check pins seqs-vs-loop bit identity; decode-batch-gate strict pins the whole tick vs decode_step_h). q is the stacked [B, n_head, head_dim] tick buffer read in place (no per-seq q copies); o is written [B, n_head, head_dim] in place (no per-seq a copies).

Source

pub fn append_kv_quantized_seqs( &self, k_rows: &CudaSlice<f32>, v_rows: &CudaSlice<f32>, kv_ptrs: &CudaView<'_, u64>, pos_seq: &CudaSlice<i32>, b_n: usize, kv_dim_k: usize, kv_dim_v: usize, k_tok_bytes: usize, v_tok_bytes: usize, ) -> Result<(), Box<dyn Error>>

BATCHED-TICK increment 2: z-batched decode KV append — one launch appends this step’s B rows, each into ITS OWN sequence cache at slot pos_seq[z], through the same [2B] interleaved pointer table the seqs FA reads. Each (block, z) warp executes the per-token appender’s exact warp program on row z of the stacked [B, kv_dim] k/v — written cache bytes are BIT-IDENTICAL to the B per-seq calls it replaces (kernel-check pins the bytes). Default flash module only (callers exclude fp8-KV).

Source

pub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool

True iff the MULTI-ROW verify FA (fa_decode_rows) is usable for a verify batch whose FIRST row attends base_len + 1 keys: every row must take the SAME kernel eager decode would (the vec path) — mirrors fa_decode’s gate exactly (MEMRA_NO_FA_VEC + FA_VEC_MIN_TKV + head_dim), evaluated at the MINIMUM row bound so no row could have picked scalar. MEMRA_FA_ROWS_OFF=1 is the A/B + fallback seam (per-row loop).

Source

pub fn fa_decode_rows( &self, q: &CudaSlice<f32>, k: &CudaView<'_, u8>, v: &CudaView<'_, u8>, o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize, base_len: usize, t: usize, scale: f32, k_tok_bytes: usize, v_tok_bytes: usize, base_dev: Option<(&CudaSlice<i32>, i32)>, kv_shared: bool, g: bool, q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>, ) -> Result<(), Box<dyn Error>>

MULTI-ROW verify FA: run fa_decode_vec_q’s EXACT per-row program for T causal query rows (row r attends keys [0..base_len+r+1)) in ONE kernel launch with grid.z = row, plus ONE row-batched combine. Replaces the T separate (fa_decode + combine) launches of the spec verify — same per-row split partition (n_splits_r = ceil(t_kv_r/split_keys), the fa_split_keys formula), same key-walk order, same reduce shapes => bit-identical outputs per row (kernel-check pins rows-vs-loop byte identity; run-spec is the end gate). Caller must have checked fa_rows_eligible(base_len, head_dim). q is the verify’s token-major [T, n_head, head_dim] stack; o is written [T, n_head, head_dim].

Source

pub fn fa_decode_rows_w( &self, q: &CudaSlice<f32>, k: &CudaView<'_, u8>, v: &CudaView<'_, u8>, o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize, base_dev: &CudaSlice<i32>, base_plus: i32, t: usize, scale: f32, window: usize, k_tok_bytes: usize, v_tok_bytes: usize, q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>, ) -> Result<(), Box<dyn Error>>

WINDOWED verify rows (gemma R6 deep-ctx): every row attends exactly window keys — bit-identical per row to the T=1 decode’s fa_decode over the window VIEW. Caller gates base_len + 1 >= window (no under-window rows) and head_dim == 256 (v4 stamp).

Source

pub fn fa_decode_rows_dc( &self, q: &CudaSlice<f32>, k: &CudaView<'_, u8>, v: &CudaView<'_, u8>, o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize, base_dev: &CudaSlice<i32>, t_kv_upper: usize, t: usize, scale: f32, k_tok_bytes: usize, v_tok_bytes: usize, base_plus: i32, g: bool, ) -> Result<(), Box<dyn Error>>

ROUND-STREAM stage (c): fa rows with the causal base from a device counter. Two lanes: v3 (qwen stream, fa_v3_active) and v4 (gemma hd256 burst — rows_v4_dc, g-module aware); t_kv_upper sizes splits/partials — the same one-sp-for-all-rows approximation class the host rows path already uses (battery-arbitrated); actual per-row bounds derive in-kernel from the counter (+ base_plus, v4 lane only — v3’s kernel has no plus arg).

Source

pub fn fa_decode_dc( &self, q: &CudaSlice<f32>, k: &CudaView<'_, u8>, v: &CudaView<'_, u8>, o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize, t_kv_dev: &CudaSlice<i32>, bucket_max: usize, scale: f32, k_tok_bytes: usize, v_tok_bytes: usize, g: bool, ) -> Result<(), Box<dyn Error>>

Device-counter variant of fa_decode (CUDA-GRAPH-PLAN Phase 2). The sequence length is read from t_kv_dev[0] (resident device i32[1]) for the attention loop bound + per-split key range; the GRID n_splits is sized for bucket_max (the bucket’s max t_kv — baked at capture time). Empty splits (key range beyond the actual t_kv) write an empty partial (m=NEG_INF) so the shared combine skips them -> bit-correct for ANY actual t_kv <= bucket_max.

BIT-IDENTITY (the gate): pass bucket_max == actual_t_kv and this reproduces fa_decode EXACTLY (same n_splits, same per, same split boundaries, same combine) while reading t_kv from device. Bucketing (bucket_max > t_kv) is for the future captured path and changes split grouping (different but mathematically-equal log-sum-exp merge).

Source

pub fn fa_decode_dc_q8( &self, q: &CudaSlice<f32>, k: &CudaView<'_, u8>, v: &CudaView<'_, u8>, o: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, n_head_kv: usize, t_kv_dev: &CudaSlice<i32>, bucket_max: usize, scale: f32, k_tok_bytes: usize, v_tok_bytes: usize, g: bool, q8_out: Option<(&mut CudaSlice<i8>, &mut CudaSlice<f32>)>, ) -> Result<(), Box<dyn Error>>

fa_decode_dc with an optional q8_1 sink (wave 5b): when q8_out is given the combine emits (int8, per-32 scales) for the wo matmul_pre and skips the f32 O write.

Source

pub fn fa_geom_eager( &self, t_kv: usize, head_dim: usize, n_head_kv: usize, g: bool, ) -> (bool, usize)

EAGER fa_decode geometry for a given actual t_kv (CUDA-GRAPH-PLAN §3.3 bucketing). Returns (fa_vec, n_splits) EXACTLY as fa_decode computes them so the graph-capture path can key its bucket on the same (kernel, n_splits) pair and pass a bucket_max that reproduces eager’s n_splits bit-for-bit. (Per = ceil(t_kv/n_splits) is then recomputed from the DEVICE t_kv inside the kernel and matches eager when n_splits matches — the bit-identity contract.)

Source

pub fn fa_bucket_key( &self, t_kv: usize, head_dim: usize, n_head_kv: usize, g: bool, ) -> (bool, usize)

bucket_max (host t_kv to feed fa_decode_dc / full_attn_decode_dc) that makes the _dc kernel pick the SAME (fa_vec, n_splits) as eager would for actual t_kv. Because the dc launcher derives both from bucket_max via the same formulas, we just hand it t_kv itself: the n_splits is then identical, and the per-split boundaries (computed from the DEVICE t_kv in the kernel) match eager exactly. The bucket KEY (for the graph HashMap) is (fa_vec, n_splits).

Source

pub fn capture_graph_retained<F>( &self, step: F, ) -> Result<(CudaGraph, Vec<Box<dyn Any + Send>>), Box<dyn Error>>
where F: FnMut(&Engine) -> Result<(), Box<dyn Error>>,

CUDA-graph capture wrapper (CUDA-GRAPH-PLAN §3.2, llama.cpp warmup pattern). Runs step inline TWICE (warmup — lets the caching allocator settle to stable pointers and any one-time kernel attribute/JIT happen outside capture), then captures a THIRD invocation on the Engine’s decode stream (RELAXED mode) and instantiates it into a replayable CudaGraph. The closure must enqueue ONLY device work on e.stream() (no dtoh / no synchronize / no host branch on device data) — every per-step varying scalar must come from a device counter. Returns the instantiated graph; CudaGraph::launch() replays the whole step in one dispatch. capture_graph with CAPTURE-RETAIN: every Engine allocation made during the warmups and the capture is kept alive in the returned keeper — hold it as long as the graph replays (transients returning to the pool get reused by unrelated work and corrupt replays; the draft-graph root cause). Model-generic, next capture reuses it.

Source

pub fn capture_graph_retained_flags<F>( &self, flags: CUgraphInstantiate_flags, step: F, ) -> Result<(CudaGraph, Vec<Box<dyn Any + Send>>), Box<dyn Error>>
where F: FnMut(&Engine) -> Result<(), Box<dyn Error>>,

Retained capture with an explicit instantiate flag. ALLOC-FREE captured graphs (zero mem nodes — the gemma slotted door) should pass UPLOAD instead of AUTO_FREE_ON_LAUNCH: the auto-free flag’s launch-time mem-pool scan was measured at ~0.25us/node (205us on the 826-node step) even with nothing to free.

Source

pub fn capture_graph<F>(&self, step: F) -> Result<CudaGraph, Box<dyn Error>>
where F: FnMut(&Engine) -> Result<(), Box<dyn Error>>,

Source

pub fn gdn_scan_s128_view( &self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>, g: &CudaSlice<f32>, beta: &CudaSlice<f32>, state_in: &CudaView<'_, f32>, state_out: &mut CudaViewMut<'_, f32>, o: &mut CudaSlice<f32>, n_head: usize, t: usize, scale: f32, ) -> Result<(), Box<dyn Error>>

gdn_scan variant where state_in/out are CudaViews (resident SSM state, in-place per step).

Source

pub fn ssm_conv1d_view( &self, x: &CudaView<'_, f32>, w: &CudaSlice<f32>, y: &mut CudaSlice<f32>, conv_dim: usize, t: usize, d_conv: usize, silu: bool, ) -> Result<(), Box<dyn Error>>

conv1d where the input is a CudaView (resident conv state assembled in place).

Source

pub fn ssm_conv1d_tm( &self, qkv_tm: &CudaSlice<f32>, w: &CudaSlice<f32>, y: &mut CudaSlice<f32>, conv_dim: usize, t: usize, d_conv: usize, ) -> Result<(), Box<dyn Error>>

Depthwise causal conv1d + optional SiLU. x:[conv_dim, T+d_conv-1] channel-major (first d_conv-1 cols = carried state), w:[d_conv, conv_dim] kernel-major, y:[conv_dim, T] channel-major. FUSED prefill conv (token-major input, zero left-state): replaces transpose + zeros + conv_left_pad + ssm_conv1d with ONE launch reading the matmul output directly. Output channel-major [conv_dim, T], SiLU applied. BIT-IDENTICAL accumulation.

Source

pub fn ssm_conv1d_tm_state( &self, qkv_tm: &CudaSlice<f32>, conv_state: &mut CudaSlice<f32>, w: &CudaSlice<f32>, y: &mut CudaSlice<f32>, conv_dim: usize, t: usize, d_conv: usize, ) -> Result<(), Box<dyn Error>>

BATCHED verify conv (T>1, carried state): window reads the resident conv ring for negative rows; separate ring-update launch afterwards. BIT-IDENTICAL per value to the T=1 chain. T >= pad rides the pure input-column ring update (unchanged legacy path); T < pad (the MEMRA_SPEC_M2 t=2 verify arm) needs old-ring sources for the roll — the update kernel would race reading the ring it rewrites, so that arm clones the ring (dtod) and rolls via ssm_conv_ring_rebuild (PURE COPIES: the ring stores raw input columns; the final ring == what T sequential decode ring rolls leave).

Source

pub fn ssm_conv1d_tm_state_pad( &self, qkv_tm: &CudaSlice<f32>, conv_state: &mut CudaSlice<f32>, w: &CudaSlice<f32>, y: &mut CudaSlice<f32>, conv_dim: usize, t: usize, d_conv: usize, pad_len: Option<&CudaSlice<i32>>, ) -> Result<(), Box<dyn Error>>

task #14: pad_len = device true length for PADDED prime graphs — the ring update reads rows [len-pad, len) instead of the pad tail. None = the classic host-T path.

Source

pub fn ssm_conv1d_tm_state_pad_v( &self, qkv_tm: &CudaView<'_, f32>, conv_state: &mut CudaSlice<f32>, w: &CudaSlice<f32>, y: &mut CudaSlice<f32>, conv_dim: usize, t: usize, d_conv: usize, pad_len: Option<&CudaSlice<i32>>, ) -> Result<(), Box<dyn Error>>

qkv-view twin (task #16): batched prime reads the concat GEMM output directly.

Source

pub fn ssm_conv_ring_rebuild( &self, qkv_tm: &CudaSlice<f32>, ring_old: &CudaSlice<f32>, conv_state: &mut CudaSlice<f32>, conv_dim: usize, tc: usize, d_conv: usize, ) -> Result<(), Box<dyn Error>>

PREFIX conv-ring rebuild (spec REPLAY-FREE partial accept): overwrite the resident ring with the state a T=1 chain holds after only the FIRST tc columns of qkv_tm — the last pad entries of [ring_old | cols 0..tc-1]. PURE COPIES (the ring stores raw inputs; no arithmetic, cannot perturb FP order). ring_old = the pre-round snapshot ring.

Source

pub fn gdn_prep_decode( &self, conv_out: &CudaSlice<f32>, beta_raw: &CudaSlice<f32>, alpha: &CudaSlice<f32>, dt_bias: &CudaSlice<f32>, a: &CudaSlice<f32>, q_l2: &mut CudaSlice<f32>, k_l2: &mut CudaSlice<f32>, v_g: &mut CudaSlice<f32>, beta: &mut CudaSlice<f32>, g_log: &mut CudaSlice<f32>, d_state: usize, num_v: usize, num_k: usize, key_dim: usize, eps: f32, ) -> Result<(), Box<dyn Error>>

FUSED decode GDN prep (T=1): repack + q/k L2-norm + beta sigmoid + g_log in one launch. Replaces 5 tiny serialized kernels on the decode critical path. L2 reduce runs as a 32-lane warp tree (vs l2_norm_f32’s 256-thread two-level tree) — same math, different FP sum order; the argmax + run-spec gates are the authority.

Source

pub fn ssm_conv1d_gdn( &self, qkv_tm: &CudaSlice<f32>, w: &CudaSlice<f32>, q_g: &mut CudaSlice<f32>, k_g: &mut CudaSlice<f32>, v_g: &mut CudaSlice<f32>, conv_dim: usize, t: usize, d_conv: usize, d_state: usize, num_v: usize, num_k: usize, key_dim: usize, ) -> Result<(), Box<dyn Error>>

FUSED prefill conv + GDN repack: token-major qkv -> q_g/k_g/v_g in ONE launch (no conv_out materialization, no qkv_to_gdn_repack pass). BIT-IDENTICAL values; scatter matches qkv_to_gdn_repack’s modulo head-repeat mapping exactly.

Source

pub fn ssm_conv1d( &self, x: &CudaSlice<f32>, w: &CudaSlice<f32>, y: &mut CudaSlice<f32>, conv_dim: usize, t: usize, d_conv: usize, silu: bool, ) -> Result<(), Box<dyn Error>>

Source

pub fn gdn_scan_s128( &self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>, g: &CudaSlice<f32>, beta: &CudaSlice<f32>, state_in: &CudaSlice<f32>, state_out: &mut CudaSlice<f32>, o: &mut CudaSlice<f32>, n_head: usize, t: usize, scale: f32, ) -> Result<(), Box<dyn Error>>

Gated DeltaNet scan, S_v=128. q,k,v:[128,H,T]; g,beta:[H,T]; state:[128,128,H] transposed; o:[128,H,T]. Single sequence.

Source

pub fn ssm_conv1d_fused_decode_b( &self, qkv_cols: &CudaSlice<f32>, conv_state_ptrs: &CudaView<'_, u64>, w: &CudaSlice<f32>, conv_outs: &mut CudaSlice<f32>, conv_dim: usize, d_conv: usize, b_n: usize, ) -> Result<(), Box<dyn Error>>

Source

pub fn gdn_prep_decode_b( &self, conv_outs: &CudaSlice<f32>, beta_raws: &CudaSlice<f32>, alphas: &CudaSlice<f32>, dt_bias: &CudaSlice<f32>, a: &CudaSlice<f32>, q_l2: &mut CudaSlice<f32>, k_l2: &mut CudaSlice<f32>, v_g: &mut CudaSlice<f32>, beta: &mut CudaSlice<f32>, g_log: &mut CudaSlice<f32>, d_state: usize, num_v: usize, num_k: usize, key_dim: usize, eps: f32, conv_dim: usize, b_n: usize, ) -> Result<(), Box<dyn Error>>

Source

pub fn gdn_scan_s128_batched( &self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>, g: &CudaSlice<f32>, beta: &CudaSlice<f32>, state_in_ptrs: &CudaView<'_, u64>, state_out_ptrs: &CudaView<'_, u64>, o: &mut CudaSlice<f32>, n_head: usize, b_n: usize, scale: f32, ) -> Result<(), Box<dyn Error>>

Source

pub fn gdn_chunked_enabled() -> bool

A4 seam: chunked WY GDN prefill. DEFAULT ON (MEMRA_GDN_CHUNKED=0 = rollback to the sequential scan). Flipped 2026-07-04 with the full battery green: kernel-check ALL GREEN x {9B, 27B} incl the f64-truth chunk gates; run-gen argmax 82==82 both models on AND off (24/24 sweep runs); run-spec K={1,2,3,4,6,8} PASS x {9B synth, 9B text, 27B p2, 27B p3}; e2e first-16-token agreement 6/6 (full-256 drifts at index 47-125 on 5/6 prompts — accepted cache-state-FP class, batched-prime precedent). PREFILL-ONLY: decode + spec verify never route here (decode==verify dispatch identity law); prime_cache/forward/forward_last are the only callers.

Source

pub fn gdn_chunk_size() -> usize

A4 chunk size (MEMRA_GDN_CHUNK, default 32 — the sweep winner: the O(T*C) chunk matrices grow with C while the sequential state pass is C-flat, so smaller chunks win; C=32/64 also get the register-history solve template). Clamped to multiples of 32 in [32, 128] (kernel row mappings require it).

Source

pub fn gdn_chunk_k123( &self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>, g: &CudaSlice<f32>, beta: &CudaSlice<f32>, wb16: Option<&mut CudaSlice<u8>>, n_head: usize, t: usize, c: usize, hk: usize, k2w: Option<(&CudaSlice<u8>, &CudaSlice<u8>, &mut CudaSlice<u8>)>, ) -> Result<(CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>, CudaSlice<f32>), Box<dyn Error>>

A4: chunked WY / blockwise-inverse GDN prefill (see cu/hybrid.cu K1-K5 header for the math). Same contract as gdn_scan_s128 (layouts, state ping-pong) but chunk-parallel: NOT bit-identical to the sequential scan (chunked FP accumulation order); run-gen argmax + run-spec batteries are the accuracy authority. PREFILL callers only. task #18: K1-K3 of the chunked WY scan (shared by the per-seq path and the batched-prime varlen path). Returns (gcum, P, U, W); A is K3-internal.

Source

pub fn gdn_db_on() -> bool

task #21 de-broadcast seam: q/k stored at num_k distinct GQA heads instead of the num_v broadcast. MEMRA_GDN_DB=0 reverts. Only the chunked prefill path consumes the compact layout (hk plumbed; hk == H reproduces broadcast exactly).

Source

pub fn gdn_mma_enabled(&self, c: usize) -> bool

Whether the K4/K5 mma pair serves at chunk size c (mirrors gdn_scan_chunked’s seam read — env re-read per call ON PURPOSE, kernel-check pins both configs).

Source

pub fn gdn_wgmma_on(&self, c: usize) -> bool

task #22: whether the fused K4+K5 (+K2) wgmma path serves (nested inside the mma config; same per-call env read discipline).

Source

pub fn ssm_conv1d_gdn_state_pad( &self, qkv_tm: &CudaView<'_, f32>, conv_state: &mut CudaSlice<f32>, w: &CudaSlice<f32>, q_g: &mut CudaSlice<f32>, k_g: &mut CudaSlice<f32>, v_g: &mut CudaSlice<f32>, conv_dim: usize, t: usize, d_conv: usize, d_state: usize, num_v: usize, num_k: usize, key_dim: usize, hk: usize, pad_len: Option<&CudaSlice<i32>>, ) -> Result<(), Box<dyn Error>>

task #18 conv-fuse: carried-ring conv + SiLU + GDN repack in ONE pass (the conv_out intermediate and its transposed re-read disappear — 11.8ms of the T=2048 prime). Ring update stays the separate follow-up launch (pad-aware). BIT-IDENTICAL values to ssm_conv1d_tm_state_pad + qkv_to_gdn_repack.

Source

pub fn gdn_chunk_alloc( &self, n_head: usize, t: usize, c: usize, hk: usize, ) -> Result<GdnChunkBufs, Box<dyn Error>>

task #18 increment 2: allocate ONE sequence’s chunk buffers (no launches) — K1-K5 all run varlen afterwards. a/w become struct members so the varlen K2/K3 can write them.

Source

pub fn f32_to_bf16_v( &self, x: &CudaView<'_, f32>, dst: &mut CudaSlice<u8>, n: usize, ) -> Result<(), Box<dyn Error>>

view-source twin of f32_to_bf16 (the batched FA3 v mirror reads a concat view).

Source

pub fn f32_to_bf16_into( &self, x: &CudaSlice<f32>, dst: &mut CudaSlice<u8>, n: usize, ) -> Result<(), Box<dyn Error>>

f32 -> bf16 bulk mirror into a caller buffer (the K4/K5 operand mirrors).

Source

pub fn gdn_chunk_k123_vl8( &self, seqs: &[GdnSeqVl], n_head: usize, hk: usize, wq: Option<&GdnWVl8>, ) -> Result<(), Box<dyn Error>>

task #18 increment 2: varlen K1+K2+K3 — three launches run every sequence’s cumgate/attn/solve (per-block math identical to the per-seq kernels).

Source

pub fn gdn_prep_vl8( &self, seqs: &[GdnPrepVl], conv_w: &CudaSlice<f32>, dt_bias: &CudaSlice<f32>, a: &CudaSlice<f32>, conv_dim: usize, d_conv: usize, d_state: usize, num_v: usize, num_k: usize, key_dim: usize, hk: usize, eps: f32, ) -> Result<(), Box<dyn Error>>

task #18 increment 3: varlen PREP chain — conv(+ring) / repack / fused-l2 / fused gate-prep, 5 launches for every sequence (per-element math identical to the per-seq kernels; l2/gate fusions write disjoint outputs).

Source

pub fn gdn_mirror_vl8( &self, seqs: &[GdnSeqVl], n_head: usize, which: i32, hk: usize, ) -> Result<(), Box<dyn Error>>

varlen bf16 mirrors over the gdnseq_t table (which: 0 = k_l2 -> kb16, 1 = w -> wb16).

Source

pub fn gdn_tail_vl8( &self, seqs: &[GdnPrepVl], norm_w: &CudaSlice<f32>, d_state: usize, num_v: usize, eps: f32, ) -> Result<(), Box<dyn Error>>

varlen gated-norm tail (+f16out) — one launch replaces B gated_rmsnorm calls.

Source

pub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64

Raw device address helpers for the varlen by-value arg struct (single-stream launches; every buffer outlives the call — the f16 FFI discipline).

Source

pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64

Source

pub fn addr_f32v(&self, x: &CudaView<'_, f32>) -> u64

Source

pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64

Source

pub fn gdn_chunk_vl8( &self, seqs: &[GdnSeqVl], n_head: usize, scale: f32, hk: usize, wq: Option<&GdnWVl8>, ) -> Result<(), Box<dyn Error>>

task #18: the varlen K4+K5 pair — TWO launches run every sequence’s state pass and output pass (grid gains a seq dim; per-block math identical to the per-seq launches, so this is strictly bit-gateable against them).

Source

pub fn gdn_scan_chunked( &self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>, g: &CudaSlice<f32>, beta: &CudaSlice<f32>, kb16_pre: Option<&CudaSlice<u8>>, qb16_pre: Option<&CudaSlice<u8>>, state_in: &CudaSlice<f32>, state_out: &mut CudaSlice<f32>, o: &mut CudaSlice<f32>, n_head: usize, t: usize, scale: f32, c: usize, hk: usize, ) -> Result<(), Box<dyn Error>>

Source

pub fn gdn_scan_prefill( &self, q: &CudaSlice<f32>, k: &CudaSlice<f32>, v: &CudaSlice<f32>, g: &CudaSlice<f32>, beta: &CudaSlice<f32>, kb16_pre: Option<&CudaSlice<u8>>, qb16_pre: Option<&CudaSlice<u8>>, state_in: &CudaSlice<f32>, state_out: &mut CudaSlice<f32>, o: &mut CudaSlice<f32>, n_head: usize, t: usize, scale: f32, hk: usize, ) -> Result<(), Box<dyn Error>>

PREFILL GDN scan dispatch (the A4 seam): chunked WY form when enabled and T is in the batched-prefill regime, else the sequential scan. Callers: hybrid_forward::linear_attn (forward/forward_last) + linear_attn_prime (prime_cache). Decode (T=1) and the spec verify call gdn_scan_s128 DIRECTLY — the decode==verify dispatch identity is untouched.

MEMRA_GDN_DIFF=1: numerical-oracle mode — runs BOTH forms on the same inputs, prints the per-call (== per-layer, in call order) output/state error distribution, and keeps the SEQUENTIAL results so the run stays on the shipped path (stage-1 prototype evidence).

Source

pub fn gdn_glog( &self, alpha: &CudaSlice<f32>, dt_bias: &CudaSlice<f32>, a: &CudaSlice<f32>, g_log: &mut CudaSlice<f32>, n_head: usize, t: usize, ) -> Result<(), Box<dyn Error>>

softplus-based g_log: g_log[h,t] = a[h] * softplus(alpha[h,t] + dt_bias[h]). a pre-negated.

Source

pub fn sigmoid_v( &self, x: &CudaView<'_, f32>, y: &mut CudaSlice<f32>, n: usize, ) -> Result<(), Box<dyn Error>>

view twins (task #16): the batched prime’s GDN core reads the CONCAT projection buffers at row offsets (CudaView) — same kernels, same values, no split copies.

Source

pub fn gdn_glog_v( &self, alpha: &CudaView<'_, f32>, dt_bias: &CudaSlice<f32>, a: &CudaSlice<f32>, g_log: &mut CudaSlice<f32>, n_head: usize, t: usize, ) -> Result<(), Box<dyn Error>>

Source

pub fn sigmoid( &self, x: &CudaSlice<f32>, y: &mut CudaSlice<f32>, n: usize, ) -> Result<(), Box<dyn Error>>

Source

pub fn sig_mul_f16out( &self, a: &CudaSlice<f32>, g: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, dst16: &mut CudaSlice<u8>, n: usize, ) -> Result<(), Box<dyn Error>>

attn out-gate fused epilogue (task #17): dst = a * sigmoid(g) + fp16 twin, one launch (replaces sigmoid + mul + convert). Bit-identical class.

Source

pub fn gated_rmsnorm( &self, o: &CudaSlice<f32>, w: &CudaSlice<f32>, z: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(), Box<dyn Error>>

gated RMSNorm: dst = RMSNorm(o, w[ncols]) * silu(z), per row of ncols. nrows blocks.

Source

pub fn gated_rmsnorm_f16out( &self, o: &CudaSlice<f32>, w: &CudaSlice<f32>, z: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, dst16: &mut CudaSlice<u8>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(), Box<dyn Error>>

f16out twin of gated_rmsnorm (task #17): epilogue also emits the fp16 operand for the ssm_out GEMM. Bit-identical class (same floats + the cvt kernel’s __float2half).

Source

pub fn add_rms_norm_zq8( &self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, w: &CudaSlice<f32>, res: &mut CudaSlice<f32>, z: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn Error>>

add+RMSNorm emitting the f32 normed row AND its q8_1 quantization in one launch (the MoE layer input: z feeds the router matmul as f32, the expert dp4a as q8_1). BIT-IDENTICAL to add_rms_norm + quantize_q8_1. Returns (q, d) alongside the caller-provided res/z buffers.

Source

pub fn gated_rmsnorm_zv( &self, o: &CudaSlice<f32>, w: &CudaSlice<f32>, z: &CudaView<'_, f32>, dst: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(), Box<dyn Error>>

gated RMSNorm emitting q8_1 directly (fused quantize epilogue) — the ssm_out matvec input. BIT-IDENTICAL bytes to gated_rmsnorm + quantize_q8_1 (ncols % 32 == 0; blocks never straddle rows). Saves one launch per linear-attn layer (36/token on the 9B). z-view twins of gated_rmsnorm(+f16out) — task #16 batched-prime split removal.

Source

pub fn gated_rmsnorm_f16out_zv( &self, o: &CudaSlice<f32>, w: &CudaSlice<f32>, z: &CudaView<'_, f32>, dst: &mut CudaSlice<f32>, dst16: &mut CudaSlice<u8>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(), Box<dyn Error>>

Source

pub fn gated_rmsnorm_q8_1( &self, o: &CudaSlice<f32>, w: &CudaSlice<f32>, z: &CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn Error>>

Source

pub fn transpose( &self, inp: &CudaSlice<f32>, rows: usize, cols: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

transpose [rows,cols] row-major -> [cols,rows] row-major.

Source

pub fn repeat_heads( &self, inp: &CudaSlice<f32>, out: &mut CudaSlice<f32>, head_dim: usize, n_in: usize, n_out: usize, t: usize, ) -> Result<(), Box<dyn Error>>

repeat-interleave heads: in[head_dim,n_in,T] -> out[head_dim,n_out,T].

Source

pub fn q_gate_split( &self, qf: &CudaSlice<f32>, q_out: &mut CudaSlice<f32>, gate_out: &mut CudaSlice<f32>, head_dim: usize, n_head: usize, t: usize, ) -> Result<(), Box<dyn Error>>

q|gate split (on-device). qf:[T, n_head2head_dim] -> q_out,gate_out:[head_dim,n_head,T]. Replaces the dtoh->host-double-loop->htod in full_attn / full_attn_decode.

Source

pub fn qkv_to_gdn_repack( &self, conv_out: &CudaSlice<f32>, q_g: &mut CudaSlice<f32>, k_g: &mut CudaSlice<f32>, v_g: &mut CudaSlice<f32>, d_state: usize, num_v: usize, num_k: usize, key_dim: usize, t: usize, ) -> Result<(), Box<dyn Error>>

qkv->GDN repack (on-device). conv_out:[conv_dim,T] channel-major -> q_g/k_g/v_g:[d_state,num_v,T] with q/k head-repeat kh = vh % num_k (validated modulo mapping). Replaces the dtoh->host-q/k/v-repack->3x-htod in linear_attn / linear_attn_decode.

Source

pub fn conv_left_pad( &self, src: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, conv_dim: usize, t: usize, pad: usize, ) -> Result<(), Box<dyn Error>>

conv left zero-pad (prefill from zero state). src:[conv_dim,T] -> dst:[conv_dim,T+pad], cols 0..pad = 0, cols pad..pad+T = src. dst MUST be pre-zeroed. No dtoh/host-loop/htod.

Source

pub fn conv_assemble_and_roll( &self, qkv_col: &CudaSlice<f32>, conv_state: &mut CudaSlice<f32>, conv_in: &mut CudaSlice<f32>, conv_dim: usize, pad: usize, ) -> Result<(), Box<dyn Error>>

conv-state assemble + ring roll (decode T=1). conv_state:[conv_dim,pad] (resident), qkv_col:[conv_dim] -> conv_in:[conv_dim,pad+1]; AND rolls conv_state (keep last pad cols). Replaces the dtoh->host-conv-ring-assemble->ring-update->htod in linear_attn_decode.

Source

pub fn ssm_conv1d_fused_decode( &self, qkv_col: &CudaSlice<f32>, conv_state: &mut CudaSlice<f32>, w: &CudaSlice<f32>, conv_out: &mut CudaSlice<f32>, conv_dim: usize, d_conv: usize, ) -> Result<(), Box<dyn Error>>

RANK3 LEVER (conv fuse, T=1 DECODE): fused conv_assemble_and_roll + ssm_conv1d_silu in ONE launch. Assembles the conv window [conv_state | qkv_col] in registers, computes the depthwise causal conv + SiLU into conv_out, and rolls the ring — never materializing conv_in to HBM. Replaces e.conv_assemble_and_roll(…) + e.ssm_conv1d(…). BIT-IDENTICAL to that two-kernel sequence (same 8-wide accumulation order, same SiLU). conv_out is [conv_dim] (T=1).

Source

pub fn slice_range( &self, src: &CudaSlice<f32>, start: usize, len: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Copy a contiguous range [start, start+len) out of src into a fresh slice (device→device via host). Used for qkv split views. Small/rare; not perf-critical in Stage 1.

Trait Implementations§

Source§

impl KvDev for Engine

The memra-kv device seam (Phase D): the cache’s 7 ops delegate to the engine’s inherent methods (inherent methods win name resolution, so no recursion).

Source§

fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn Error>>

Source§

fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn Error>>

Source§

fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn Error>>

Source§

fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn Error>>

Source§

fn clone_dtod( &self, src: &CudaSlice<f32>, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Source§

fn copy_into( &self, dst: &mut CudaSlice<f32>, off: usize, src: &CudaSlice<f32>, len: usize, ) -> Result<(), Box<dyn Error>>

Source§

fn set_i32_one( &self, d: &mut CudaSlice<i32>, v: i32, ) -> Result<(), Box<dyn Error>>

Auto Trait Implementations§

Blanket Implementations§

Source§

impl<T> Any for T
where T: 'static + ?Sized,

Source§

fn type_id(&self) -> TypeId

Gets the TypeId of self. Read more
Source§

impl<T> Borrow<T> for T
where T: ?Sized,

Source§

fn borrow(&self) -> &T

Immutably borrows from an owned value. Read more
Source§

impl<T> BorrowMut<T> for T
where T: ?Sized,

Source§

fn borrow_mut(&mut self) -> &mut T

Mutably borrows from an owned value. Read more
Source§

impl<T> From<T> for T

Source§

fn from(t: T) -> T

Returns the argument unchanged.

Source§

impl<T, U> Into<U> for T
where U: From<T>,

Source§

fn into(self) -> U

Calls U::from(self).

That is, this conversion is whatever the implementation of From<T> for U chooses to do.

Source§

impl<T, U> TryFrom<U> for T
where U: Into<T>,

Source§

type Error = Infallible

The type returned in the event of a conversion error.
Source§

fn try_from(value: U) -> Result<T, <T as TryFrom<U>>::Error>

Performs the conversion.
Source§

impl<T, U> TryInto<U> for T
where U: TryFrom<T>,

Source§

type Error = <U as TryFrom<T>>::Error

The type returned in the event of a conversion error.
Source§

fn try_into(self) -> Result<U, <U as TryFrom<T>>::Error>

Performs the conversion.