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
impl Engine
Sourcepub fn mmq_supports(&self, w: &GpuTensor) -> bool
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.
Sourcepub fn qmatvec_mmq(
&self,
w: &GpuTensor,
x: &CudaSlice<f32>,
m: usize,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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).
Sourcepub fn accprobe_act_bytes(&self, in_f: usize, m: usize) -> usize
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.
Sourcepub 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>>
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.
Sourcepub fn mmq_act_begin(&self)
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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.
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>>
Sourcepub fn fp8_blk_nan_count(
&self,
w_e4m3: &CudaSlice<u8>,
) -> Result<u32, Box<dyn Error>>
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).
Sourcepub fn mmq_iq_quantize_act(
&self,
x: &CudaSlice<f32>,
in_f: usize,
n_tokens: usize,
) -> Result<CudaSlice<u8>, Box<dyn Error>>
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).
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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
impl Engine
Sourcepub fn f16_scratch_swap(&self, new: Option<F16Scratch>) -> Option<F16Scratch>
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.
Sourcepub fn try_f16_gemm(
&self,
w: &GpuTensor,
x: &CudaSlice<f32>,
m: usize,
) -> Result<Option<CudaSlice<f32>>, Box<dyn Error>>
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).
Sourcepub 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>>
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.
Sourcepub fn f16_act(
&self,
x: &CudaSlice<f32>,
nelem: usize,
in_f: usize,
) -> Result<CudaSlice<u8>, Box<dyn Error>>
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).
Sourcepub fn try_f16_gemm_pre_into(
&self,
w: &GpuTensor,
xh: &CudaSlice<u8>,
m: usize,
y: &mut CudaSlice<f32>,
) -> Result<bool, Box<dyn Error>>
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).
Sourcepub 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>>
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).
Sourcepub fn try_f16_gemm_pre(
&self,
w: &GpuTensor,
xh: &CudaSlice<u8>,
m: usize,
) -> Result<Option<CudaSlice<f32>>, Box<dyn Error>>
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.
Sourcepub fn build_q8_f16_raw(
&self,
bytes: &CudaSlice<u8>,
in_f: usize,
out_f: usize,
) -> Result<CudaSlice<u8>, Box<dyn Error>>
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).
Sourcepub fn build_q8_f16(&self, t: &mut GpuTensor) -> Result<(), Box<dyn Error>>
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.
Sourcepub fn build_q4_f16_raw(
&self,
bytes: &CudaSlice<u8>,
in_f: usize,
out_f: usize,
) -> Result<CudaSlice<u8>, Box<dyn Error>>
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).
Sourcepub fn build_q5k_f16_raw(
&self,
bytes: &CudaSlice<u8>,
in_f: usize,
out_f: usize,
) -> Result<CudaSlice<u8>, Box<dyn Error>>
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.
Sourcepub fn build_q4k_f16_raw(
&self,
bytes: &CudaSlice<u8>,
in_f: usize,
out_f: usize,
) -> Result<CudaSlice<u8>, Box<dyn Error>>
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.
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
impl Engine
Sourcepub fn try_fp8_gemm(
&self,
w: &GpuTensor,
x: &CudaSlice<f32>,
m: usize,
) -> Result<Option<CudaSlice<f32>>, Box<dyn Error>>
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
impl Engine
Sourcepub fn try_fp8_blk_mmq(
&self,
w: &GpuTensor,
x: &CudaSlice<f32>,
m: usize,
) -> Result<Option<CudaSlice<f32>>, Box<dyn Error>>
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
impl Engine
Sourcepub fn fp8_blk_q8_0_bytes(out_f: usize, in_f: usize) -> usize
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.
Sourcepub fn fp8_blk_dequant_q8_0(
&self,
f8: &[u8],
grid: &[f32],
out_f: usize,
in_f: usize,
) -> Result<CudaSlice<u8>, Box<dyn Error>>
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).
Sourcepub 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>>
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
impl Engine
pub fn new(ordinal: usize) -> Result<Self, Box<dyn Error>>
pub fn ctx(&self) -> &Arc<CudaContext> ⓘ
Sourcepub fn pool_cached_bytes(&self) -> usize
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).
Sourcepub fn pool_reserved_used(&self) -> (usize, usize)
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.
Sourcepub fn stream(&self) -> Arc<CudaStream> ⓘ
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.
Sourcepub fn gkv_on() -> bool
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.
Sourcepub fn wkv_on() -> bool
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.
Sourcepub fn kv_fp8_on() -> bool
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub fn wpf_level() -> u32
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.
Sourcepub fn set_verify_exact(&self, on: bool)
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.
Sourcepub fn qkv_append_on() -> bool
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).
Sourcepub fn pdl_wb_on() -> bool
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.
Sourcepub fn pdl_mmvq_on() -> bool
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.
pub fn pdl_on() -> bool
Sourcepub fn prefetch_weight_l2(&self, w: &GpuTensor) -> Result<(), Box<dyn Error>>
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.
Sourcepub fn gather_row_bf16(
&self,
table: &CudaSlice<u8>,
tok: &CudaSlice<u32>,
idx: usize,
dst: &mut CudaSlice<f32>,
ncols: usize,
) -> Result<(), Box<dyn Error>>
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.
Sourcepub fn add_row_inplace(
&self,
logits: &mut CudaSlice<f32>,
bias: &CudaSlice<f32>,
n: usize,
row_off: usize,
) -> Result<(), Box<dyn Error>>
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).
Sourcepub fn prefetch_l2(
&self,
p: &CudaSlice<u8>,
n: usize,
) -> Result<(), Box<dyn Error>>
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).
Sourcepub fn router_gemv(
&self,
w: &CudaSlice<f32>,
x: &CudaSlice<f32>,
n_embd: usize,
n_experts: usize,
t: usize,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
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.
Sourcepub 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>>
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).
Sourcepub fn rows_permute(
&self,
src: &CudaSlice<f32>,
idx: &CudaSlice<i32>,
nrows: usize,
ncols: usize,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
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).
Sourcepub fn sigmoid_dot_rows(
&self,
x: &CudaSlice<f32>,
w: &CudaSlice<f32>,
n_embd: usize,
t: usize,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
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.
Sourcepub 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>>
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.
Sourcepub fn plain_tok_ring(
&self,
vam: &CudaSlice<u32>,
pos_start: &CudaSlice<i32>,
base: usize,
ring: &mut CudaSlice<u32>,
) -> Result<(), Box<dyn Error>>
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].
Sourcepub 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>>
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.
pub fn i32_copy_add( &self, src: &CudaSlice<i32>, dst: &mut CudaSlice<i32>, delta: i32, ) -> Result<(), Box<dyn Error>>
pub fn u32_copy( &self, src: &CudaSlice<u32>, dst: &mut CudaSlice<u32>, ) -> Result<(), Box<dyn Error>>
Sourcepub fn spec_adapt_k(
&self,
acc: &CudaSlice<u32>,
brk: &mut CudaSlice<u32>,
floor: usize,
cap: usize,
) -> Result<(), Box<dyn Error>>
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).
Sourcepub 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>>
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).
Sourcepub fn pos_iota(
&self,
pos0: &CudaSlice<i32>,
out: &mut CudaSlice<i32>,
t: usize,
) -> Result<(), Box<dyn Error>>
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.
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>>
Sourcepub 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>>
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).
Sourcepub fn pack_tok_p(
&self,
tok: &CudaSlice<u32>,
p: &CudaSlice<f32>,
out: &mut CudaSlice<u32>,
slot: usize,
) -> Result<(), Box<dyn Error>>
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).
pub fn tok_map_u32( &self, tok: &mut CudaSlice<u32>, map: &CudaSlice<u32>, ) -> Result<(), Box<dyn Error>>
Sourcepub 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>>
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.
Sourcepub 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>>
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).
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>>
Sourcepub fn spec_rollback_kv(
&self,
len_ptrs: &CudaSlice<u64>,
saved: &CudaSlice<i32>,
acc: &CudaSlice<u32>,
base: usize,
n_layer: usize,
) -> Result<(), Box<dyn Error>>
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).
Sourcepub 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>>
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).
Sourcepub 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>>
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).
Sourcepub fn gumbel_perturb(
&self,
x: &CudaSlice<f32>,
y: &mut CudaSlice<f32>,
n: usize,
seed: u64,
stream_pos: u32,
temp: f32,
) -> Result<(), Box<dyn Error>>
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).
Sourcepub fn mask_logits_col(
&self,
logits: &mut CudaSlice<f32>,
mask: &CudaSlice<u32>,
col: usize,
n: usize,
mask_words: usize,
) -> Result<(), Box<dyn Error>>
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.
Sourcepub 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>>
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.
Sourcepub fn sctr_inc(&self, ctr: &mut CudaSlice<u32>) -> Result<(), Box<dyn Error>>
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).
Sourcepub 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>>
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).
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub fn freeze_moe_cache(&self)
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.
Sourcepub fn export_moe_residency(&self) -> Option<Vec<(u16, u8, u16)>>
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.
Sourcepub fn frozen_cpu_experts_prefer_tokenwise_prime(&self) -> bool
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.)
Sourcepub fn moe_cache_enabled() -> bool
pub fn moe_cache_enabled() -> bool
True if the MoE residency cache is enabled (MEMRA_MOE_CACHE set).
Sourcepub fn moe_cache_stats(&self) -> Option<(u64, u64, u64, usize)>
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).
Sourcepub fn cpu_expert_stats(
&self,
) -> Option<(u64, u64, u64, u64, u64, u64, u64, u64, u64, u64, u64)>
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.
Sourcepub fn cpu_expert_predictor_stats(&self) -> (u64, u64)
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.
pub fn cpu_expert_exposed_wait_ns(&self) -> Option<u64>
Sourcepub fn cpu_expert_gpu_residency_stats(&self) -> Option<(u64, u64, u64)>
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.
Sourcepub fn moe_pread_stats(&self) -> Option<(u64, u64, u64, u64, u64, u64, u64)>
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).
Sourcepub fn moe_cache_reset_counters(&self)
pub fn moe_cache_reset_counters(&self)
Reset the MoE cache perf counters (to separate warmup from steady-state windows).
pub fn htod_bytes(&self, v: &[u8]) -> Result<CudaSlice<u8>, Box<dyn Error>>
Sourcepub fn htod_bytes_padded(
&self,
v: &[u8],
pad: usize,
) -> Result<CudaSlice<u8>, Box<dyn Error>>
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.
Sourcepub fn copy_into(
&self,
dst: &mut CudaSlice<f32>,
off: usize,
src: &CudaSlice<f32>,
len: usize,
) -> Result<(), Box<dyn Error>>
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.
Sourcepub fn copy_u8_into(
&self,
dst: &mut CudaSlice<u8>,
off: usize,
src: &CudaSlice<u8>,
len: usize,
) -> Result<(), Box<dyn Error>>
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).
Sourcepub fn copy_u8_range_into(
&self,
dst: &mut CudaSlice<u8>,
dst_off: usize,
src: &CudaSlice<u8>,
src_off: usize,
len: usize,
) -> Result<(), Box<dyn Error>>
pub fn copy_u8_range_into( &self, dst: &mut CudaSlice<u8>, dst_off: usize, src: &CudaSlice<u8>, src_off: usize, len: usize, ) -> Result<(), Box<dyn Error>>
D2D byte-range copy with explicit source and destination offsets.
Sourcepub fn prepare_kv_append(
&self,
kv: &mut KvLayer,
retain_from: usize,
append_rows: usize,
) -> Result<usize, Box<dyn Error>>
pub fn prepare_kv_append( &self, kv: &mut KvLayer, retain_from: usize, append_rows: usize, ) -> Result<usize, Box<dyn Error>>
Resolve an absolute append slot to the Step35 SWA layer’s physical rows. At wrap, copy only the aligned live prefix through temporary device storage and rebase it at row zero, keeping the audited attention range contiguous without changing its absolute start.
Sourcepub fn htod_u8_into(
&self,
dst: &mut CudaSlice<u8>,
off: usize,
src: &[u8],
) -> Result<(), Box<dyn Error>>
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.
pub fn view<'a>(&self, b: &'a CudaSlice<f32>, len: usize) -> CudaView<'a, f32>
Sourcepub fn view_u8_range<'a>(
&self,
b: &'a CudaSlice<u8>,
start: usize,
end: usize,
) -> CudaView<'a, u8>
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).
pub fn view_u8<'a>(&self, b: &'a CudaSlice<u8>, len: usize) -> CudaView<'a, u8>
Sourcepub 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>>
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].
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub fn inc_seqlen(&self, p: &mut CudaSlice<i32>) -> Result<(), Box<dyn Error>>
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.
Sourcepub 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>>
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).
Sourcepub fn copy_view_into(
&self,
dst: &mut CudaSlice<f32>,
off: usize,
src: &CudaView<'_, f32>,
len: usize,
) -> Result<(), Box<dyn Error>>
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).
Sourcepub fn clone_dtod(
&self,
src: &CudaSlice<f32>,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
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.
Sourcepub fn dtod_copy_view(
&self,
src: &CudaView<'_, f32>,
dst: &mut CudaSlice<f32>,
) -> Result<(), Box<dyn Error>>
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.
Sourcepub fn dtod_copy_view_i8(
&self,
src: &CudaView<'_, i8>,
dst: &mut CudaSlice<i8>,
) -> Result<(), Box<dyn Error>>
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).
Sourcepub fn dtod_copy_into(
&self,
src: &CudaSlice<f32>,
dst: &mut CudaSlice<f32>,
offset: usize,
) -> Result<(), Box<dyn Error>>
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()].
Sourcepub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn Error>>
pub fn uninit_i8(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn Error>>
Uninitialized i8 device buffer (decode_batch q8_1 row scratch).
Sourcepub 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>>
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.
Sourcepub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn Error>>
pub fn alloc_u8(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn Error>>
Allocate a reusable u8 GPU scratch buffer (for staged expert weights).
Sourcepub fn alloc_u8_uninit(&self, n: usize) -> Result<CudaSlice<u8>, Box<dyn Error>>
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).
Sourcepub fn memset_zeros_view(
&self,
dst: &mut CudaViewMut<'_, f32>,
) -> Result<(), Box<dyn Error>>
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).
Sourcepub fn stage_expert(
&self,
host_bytes: &[u8],
scratch: &mut CudaSlice<u8>,
off: usize,
) -> Result<(), Box<dyn Error>>
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).
Sourcepub fn moe_router_topk(
&self,
logits: &CudaSlice<f32>,
t: usize,
n_expert: usize,
n_used: usize,
) -> Result<(CudaSlice<i32>, CudaSlice<f32>), Box<dyn Error>>
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).
Sourcepub 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>>
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]).
Sourcepub 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>>
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).
Sourcepub fn stage_expert_async(
&self,
host_bytes: &[u8],
scratch: &mut CudaSlice<u8>,
off: usize,
) -> Result<CudaEvent, Box<dyn Error>>
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).
Sourcepub fn compute_wait(&self, ev: &CudaEvent) -> Result<(), Box<dyn Error>>
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).
Sourcepub 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>>
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
Sourcepub 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>>
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.
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>>
Sourcepub 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>>
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).
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>>
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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).
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>>
pub fn moe_pairs_gelu_mul( &self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, n: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>
pub fn moe_pairs_silu_mul( &self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, n: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>
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>>
Sourcepub 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>>
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).
Sourcepub 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>>
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).
Sourcepub 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>>
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).
Sourcepub 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>>
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).
Sourcepub fn rp_probe_q4(&self, m: usize) -> Result<(f64, f64), Box<dyn Error>>
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.
Sourcepub fn build_q4_rp4(&self, t: &mut GpuTensor) -> Result<(), Box<dyn Error>>
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.
Sourcepub fn build_q8_rp4(&self, t: &mut GpuTensor) -> Result<(), Box<dyn Error>>
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.
Sourcepub fn build_q8_rp4_raw(
&self,
bytes: &CudaSlice<u8>,
in_f: usize,
out_f: usize,
) -> Result<CudaSlice<u8>, Box<dyn Error>>
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).
Sourcepub fn build_q4k_rp4(&self, t: &mut GpuTensor) -> Result<(), Box<dyn Error>>
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.
pub fn build_q6k_rp4(&self, t: &mut GpuTensor) -> Result<(), Box<dyn Error>>
Sourcepub fn build_kq_rp4_raw(
&self,
bytes: &CudaSlice<u8>,
in_f: usize,
out_f: usize,
qtype: i32,
) -> Result<CudaSlice<u8>, Box<dyn Error>>
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).
Sourcepub fn kqrp_enabled() -> bool
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).
Sourcepub fn build_q4_rp_swap(
&self,
t: &mut GpuTensor,
) -> Result<bool, Box<dyn Error>>
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.
Sourcepub fn q4rp_enabled() -> bool
pub fn q4rp_enabled() -> bool
MEMRA_Q4RP seam (default ON): the Q4_0 split-plane decode mirror at model load.
Sourcepub 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>>
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).
Sourcepub fn u32_set_k(
&self,
dst: &mut CudaSlice<u32>,
v: u32,
idx: usize,
) -> Result<(), Box<dyn Error>>
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).
Sourcepub fn i32_add_k(
&self,
d: &mut CudaSlice<i32>,
v: i32,
) -> Result<(), Box<dyn Error>>
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).
Sourcepub fn i32_iota_from(
&self,
ctr: &CudaSlice<i32>,
dst: &mut CudaSlice<i32>,
n: usize,
) -> Result<(), Box<dyn Error>>
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).
Sourcepub fn u32_map_k(
&self,
buf: &mut CudaSlice<u32>,
map: &CudaSlice<u32>,
idx: usize,
) -> Result<(), Box<dyn Error>>
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).
Sourcepub 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>>
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).
Sourcepub fn moe_w_exscale(
&self,
w: &mut CudaSlice<f32>,
sel: &CudaSlice<i32>,
s: &CudaSlice<f32>,
n: usize,
) -> Result<(), Box<dyn Error>>
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.
Sourcepub 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>>
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).
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>>
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>>
Sourcepub 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>>
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].
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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”.
Sourcepub 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>>
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”.
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>>
Sourcepub 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>>
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.
Sourcepub fn axpy_into(
&self,
src: &CudaSlice<f32>,
alpha: f32,
dst: &mut CudaViewMut<'_, f32>,
n: usize,
) -> Result<(), Box<dyn Error>>
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).
Sourcepub fn add_scaled_rows(
&self,
src: &CudaSlice<f32>,
scale: &CudaSlice<f32>,
dst: &mut CudaSlice<f32>,
ncols: usize,
nrows: usize,
) -> Result<(), Box<dyn Error>>
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).
Sourcepub fn gather_rows(
&self,
src: &CudaSlice<f32>,
idx: &CudaSlice<i32>,
dst: &mut CudaSlice<f32>,
ncols: usize,
m_e: usize,
) -> Result<(), Box<dyn Error>>
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].
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub fn quantize_q8_1_view(
&self,
x: &CudaView<'_, f32>,
m: usize,
in_f: usize,
) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn Error>>
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.
pub fn quantize_q8_1( &self, x: &CudaSlice<f32>, m: usize, in_f: usize, ) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn Error>>
Sourcepub fn quantize_fp4_act(
&self,
x: &CudaSlice<f32>,
m: usize,
in_f: usize,
) -> Result<(CudaSlice<u32>, CudaSlice<u8>), Box<dyn Error>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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).
Sourcepub 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>>
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).
Sourcepub 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>>
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.
pub fn htod(&self, v: &[f32]) -> Result<CudaSlice<f32>, Box<dyn Error>>
pub fn htod_i32(&self, v: &[i32]) -> Result<CudaSlice<i32>, Box<dyn Error>>
Sourcepub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn Error>>
pub fn htod_i8(&self, v: &[i8]) -> Result<CudaSlice<i8>, Box<dyn Error>>
i8 upload (moe-devq8-check: synthetic q8_1 activation bytes).
pub fn htod_u64(&self, v: &[u64]) -> Result<CudaSlice<u64>, Box<dyn Error>>
Sourcepub fn dtoh_view(
&self,
d: &CudaView<'_, f32>,
) -> Result<Vec<f32>, Box<dyn Error>>
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).
pub fn dtoh(&self, d: &CudaSlice<f32>) -> Result<Vec<f32>, Box<dyn Error>>
Sourcepub fn dtoh_pair(
&self,
a: &CudaSlice<f32>,
b: &CudaSlice<f32>,
) -> Result<(Vec<f32>, Vec<f32>), Box<dyn Error>>
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.
Sourcepub fn dtoh_i32(&self, d: &CudaSlice<i32>) -> Result<Vec<i32>, Box<dyn Error>>
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).
Sourcepub fn dtoh_u8(&self, d: &CudaSlice<u8>) -> Result<Vec<u8>, Box<dyn Error>>
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).
pub fn zeros(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn Error>>
Sourcepub fn prob_of_token_device(
&self,
logits: &CudaSlice<f32>,
tok: &CudaSlice<u32>,
n_vocab: usize,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
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.
Sourcepub 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>>
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).
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>>
pub fn argmax_token_device( &self, logits: &CudaSlice<f32>, n_vocab: usize, ) -> Result<CudaSlice<u32>, Box<dyn Error>>
Sourcepub fn argmax_token_device_into(
&self,
logits: &CudaSlice<f32>,
tok: &mut CudaSlice<u32>,
n_vocab: usize,
) -> Result<(), Box<dyn Error>>
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.
Sourcepub 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>>
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.
Sourcepub fn htod_u32_v(&self, v: &[u32]) -> Result<CudaSlice<u32>, Box<dyn Error>>
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).
pub fn dtoh_u32(&self, d: &CudaSlice<u32>) -> Result<Vec<u32>, Box<dyn Error>>
Sourcepub fn htod_u32_into(
&self,
dst: &mut CudaSlice<u32>,
src: &[u32],
) -> Result<(), Box<dyn Error>>
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).
pub fn alloc_u32_zeroed( &self, n: usize, ) -> Result<CudaSlice<u32>, Box<dyn Error>>
Sourcepub 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>>
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).
Sourcepub fn dtoh_i32_one(&self, d: &CudaSlice<i32>) -> Result<i32, Box<dyn Error>>
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.
Sourcepub fn i32_set_k(
&self,
dst: &mut CudaSlice<i32>,
v: i32,
) -> Result<(), Box<dyn Error>>
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).
pub fn set_i32_one( &self, d: &mut CudaSlice<i32>, v: i32, ) -> Result<(), Box<dyn Error>>
Sourcepub fn set_u32_one(
&self,
d: &mut CudaSlice<u32>,
v: u32,
) -> Result<(), Box<dyn Error>>
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.
Sourcepub fn dtoh_u32_one(&self, d: &CudaSlice<u32>) -> Result<u32, Box<dyn Error>>
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.
Sourcepub fn upload_u8(&self, bytes: &[u8]) -> Result<CudaSlice<u8>, Box<dyn Error>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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).
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>>
Sourcepub fn uninit_q8_pair(
&self,
n: usize,
) -> Result<(CudaSlice<i8>, CudaSlice<f32>), Box<dyn Error>>
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).
pub fn uninit(&self, n: usize) -> Result<CudaSlice<f32>, Box<dyn Error>>
Sourcepub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn Error>>
pub fn alloc_i8_uninit(&self, n: usize) -> Result<CudaSlice<i8>, Box<dyn Error>>
i8 uninitialized scratch (same contract as uninit).
Sourcepub 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>>
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).
Sourcepub fn qkvnorm_w_on_prefill(rows: usize, ncols: usize) -> bool
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.
Sourcepub 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>>
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.
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>>
Sourcepub 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>>
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).
Sourcepub fn softcap(
&self,
y: &mut CudaSlice<f32>,
cap: f32,
n: usize,
) -> Result<(), Box<dyn Error>>
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).
Sourcepub fn mask_ids_rows(
&self,
y: &mut CudaSlice<f32>,
ids: &CudaSlice<i32>,
n_ids: usize,
n_vocab: usize,
t: usize,
) -> Result<(), Box<dyn Error>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub fn add_scale(
&self,
a: &CudaSlice<f32>,
b_in: &CudaSlice<f32>,
c: f32,
dst: &mut CudaSlice<f32>,
n: usize,
) -> Result<(), Box<dyn Error>>
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).
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>>
Sourcepub 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>>
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]).
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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).
Sourcepub 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>>
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).
Sourcepub fn build_q4_out_concat3(
&self,
w0: &GpuTensor,
w1: &GpuTensor,
w2: &GpuTensor,
) -> Result<Option<GpuTensor>, Box<dyn Error>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub fn l2_v2_on(ncols: usize) -> bool
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.
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>>
pub fn l2_norm( &self, x: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, ncols: usize, nrows: usize, eps: f32, ) -> Result<(), Box<dyn Error>>
Sourcepub fn l2_norm_decode(
&self,
x: &CudaSlice<f32>,
dst: &mut CudaSlice<f32>,
ncols: usize,
nrows: usize,
eps: f32,
) -> Result<(), Box<dyn Error>>
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.
Sourcepub 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>>
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].
Sourcepub 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>>
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]).
Sourcepub 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>>
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).
Sourcepub fn gelu_tanh_mul(
&self,
gate: &CudaSlice<f32>,
up: &CudaSlice<f32>,
dst: &mut CudaSlice<f32>,
n: usize,
) -> Result<(), Box<dyn Error>>
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.
pub fn silu_mul( &self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, n: usize, ) -> Result<(), Box<dyn Error>>
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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).
pub fn add( &self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, n: usize, ) -> Result<(), Box<dyn Error>>
pub fn mul( &self, a: &CudaSlice<f32>, b_in: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, n: usize, ) -> Result<(), Box<dyn Error>>
Sourcepub fn matmul(
&self,
w: &GpuTensor,
x: &CudaSlice<f32>,
m: usize,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
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.
Sourcepub fn uses_q8_1_fast(&self, w: &GpuTensor) -> bool
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).
Sourcepub fn matmul_pre(
&self,
w: &GpuTensor,
aq: &CudaSlice<i8>,
ad: &CudaSlice<f32>,
x_fallback: &CudaSlice<f32>,
m: usize,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
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).
Sourcepub fn matmul_decode_exact(
&self,
w: &GpuTensor,
x: &CudaSlice<f32>,
m: usize,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
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.
Sourcepub fn matmul_decode_exact_pre(
&self,
w: &GpuTensor,
aq: &CudaSlice<i8>,
ad: &CudaSlice<f32>,
m: usize,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
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).
Sourcepub 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>>
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.
Sourcepub fn matmul_decode_exact_dual(
&self,
w0: &GpuTensor,
w1: &GpuTensor,
x: &CudaSlice<f32>,
m: usize,
) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn Error>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub fn matmul_q8_fused2(
&self,
w0: &GpuTensor,
w1: &GpuTensor,
aq: &CudaSlice<i8>,
ad: &CudaSlice<f32>,
) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn Error>>
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.
Sourcepub fn matmul_q8_fused2_x(
&self,
w0: &GpuTensor,
w1: &GpuTensor,
x: &CudaSlice<f32>,
) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn Error>>
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).
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub fn matmul_q4_fused2(
&self,
w0: &GpuTensor,
w1: &GpuTensor,
aq: &CudaSlice<i8>,
ad: &CudaSlice<f32>,
) -> Result<Option<(CudaSlice<f32>, CudaSlice<f32>)>, Box<dyn Error>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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.
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>>
Sourcepub 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>>
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).
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub fn q8_ffn_fuse2_on(&self) -> bool
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.
Sourcepub fn e4m3_dual_on(&self) -> bool
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub fn qmatvec_e4m3_blk_mmvq_batched(
&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,
mcols: usize,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
pub fn qmatvec_e4m3_blk_mmvq_batched( &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, mcols: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>
BLOCK-128 e4m3 BATCHED matvec (lane/rp-on-st, 2026-08-06): the weight-read-once twin of
qmatvec_e4m3_blk_mmvq for m=2..16. Per (token,row) BIT-IDENTICAL to the grid.y=m launch
(same fmaf chain, same per-k32 s * ad fold, same warp reduce), so it inherits the
decode-exactness contract while reading the weight ONCE for up to mcols columns instead
of m times. mcols must be one of {2,4,8,16} and satisfy mcols >= m.
Sourcepub fn qmatvec_e4m3_blk_batched_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,
mcols: usize,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
pub fn qmatvec_e4m3_blk_batched_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, mcols: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>
Test entry for the kernel_check exactness gate: the block-128 e4m3 batched MMVQ from raw
bytes with an internal q8_1 quantize (mirrors qmatvec_batched_raw).
Sourcepub 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>>
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).
Sourcepub 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>>
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).
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>>
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>>
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>>
pub fn matmul_pre_noscale( &self, w: &GpuTensor, aq: &CudaSlice<i8>, ad: &CudaSlice<f32>, m: usize, ) -> Result<Option<(CudaSlice<f32>, f32)>, Box<dyn Error>>
Sourcepub fn mmvq_supports(&self, qtype: i32) -> bool
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).
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub fn batched_supports(&self, qtype: i32) -> bool
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.
Sourcepub fn iq_fast_enabled() -> bool
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.
Sourcepub fn b8_enabled() -> bool
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.
Sourcepub fn batched_mcols(m: usize) -> usize
pub fn batched_mcols(m: usize) -> usize
Compile-time column batch for a runtime m: 2 -> b2, 3..4 -> b4, 5..8 -> b8.
Sourcepub fn sm_count(&self) -> i32
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.
pub fn batched_variant( &self, _m: usize, in_f: usize, out_f: usize, qtype: i32, row_bytes: usize, mcols: usize, rp: bool, ) -> &'static str
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>>
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub fn matmul_group_xh(
&self,
ws: &[&GpuTensor],
x: &CudaSlice<f32>,
xh: &CudaSlice<u8>,
m: usize,
) -> Result<Vec<CudaSlice<f32>>, Box<dyn Error>>
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.
Sourcepub 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>>
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].
Sourcepub fn row_gather_dev(
&self,
src: &CudaSlice<f32>,
dst: &mut CudaSlice<f32>,
len_d: &CudaSlice<i32>,
ncols: usize,
) -> Result<(), Box<dyn Error>>
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.
Sourcepub fn matmul_group(
&self,
ws: &[&GpuTensor],
x: &CudaSlice<f32>,
m: usize,
) -> Result<Vec<CudaSlice<f32>>, Box<dyn Error>>
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.
Sourcepub fn matmul_group_multi(
&self,
ws: &[&GpuTensor],
xs: &[&CudaSlice<f32>],
ms: &[usize],
) -> Result<Vec<Vec<CudaSlice<f32>>>, Box<dyn Error>>
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.
Sourcepub fn gemm_supports(&self, w: &GpuTensor) -> bool
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).
Sourcepub fn qmatvec_gemm(
&self,
w: &GpuTensor,
aq: &CudaSlice<i8>,
ad: &CudaSlice<f32>,
m: usize,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub fn scale_inplace(
&self,
y: &mut CudaSlice<f32>,
s: f32,
n: usize,
) -> Result<(), Box<dyn Error>>
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.
Sourcepub fn bf16_to_f32(
&self,
data: &CudaView<'_, u8>,
n: usize,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
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).
Sourcepub 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>>
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.
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>>
Sourcepub 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>>
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].
Sourcepub 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>>
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).
Sourcepub 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>>
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).
Sourcepub 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>>
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.
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>>
Sourcepub fn sdpa_naive_w_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,
window: usize,
k_tok_bytes: usize,
v_tok_bytes: usize,
) -> Result<(), Box<dyn Error>>
pub fn sdpa_naive_w_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, window: usize, k_tok_bytes: usize, v_tok_bytes: usize, ) -> Result<(), Box<dyn Error>>
WINDOWED twin of sdpa_naive_quantized_view (step35 SWA prefill): dequant the KV byte
view into f32 workspaces with the SAME fa_dequant_kv_ws_f32 launch, then run
sdpa_naive_w instead of sdpa_naive. window == 0 is the unwindowed form (the kernel
treats a non-positive window as “no window mask”), so this is a strict superset of the
unwindowed function above and produces bit-identical output at window == 0.
Why this exists: EVERY windowed FlashAttention stamp in flash_attn.cu is head_dim-256
only (fa_prefill_w_f32 == fa_prefill_f32_body<256>, and the quantized-view windowed
twins likewise), while step35 is head_dim 128. Its SWA layers therefore have no windowed
FA path and take this f32 floor in v0 — same cache bytes, same numeric class as the
unwindowed quantized-view fallback, so the chunk-invariance contract holds on both.
Sourcepub 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>>
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].
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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).
Sourcepub 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>>
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).
Sourcepub fn f32_to_bf16(
&self,
x: &CudaSlice<f32>,
n: usize,
) -> Result<CudaSlice<u8>, Box<dyn Error>>
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.
pub fn f32_to_f16( &self, x: &CudaSlice<f32>, n: usize, ) -> Result<CudaSlice<u8>, Box<dyn Error>>
Sourcepub fn bf16_to_f16(
&self,
xb: &CudaSlice<u8>,
n: usize,
) -> Result<CudaSlice<u8>, Box<dyn Error>>
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).
Sourcepub fn bf16_to_f16_into(
&self,
xb: &CudaSlice<u8>,
n: usize,
y: &mut CudaSlice<u8>,
) -> Result<(), Box<dyn Error>>
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.
Sourcepub fn fa_prefill_vl8(
&self,
seqs: &[FaSeqVl],
head_dim: usize,
n_head: usize,
n_head_kv: usize,
scale: f32,
) -> Result<(), Box<dyn Error>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub fn fa_prefill_view_ws_w_hd128(
&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,
window: usize,
k_tok_bytes: usize,
v_tok_bytes: usize,
) -> Result<(), Box<dyn Error>>
pub fn fa_prefill_view_ws_w_hd128( &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, window: usize, k_tok_bytes: usize, v_tok_bytes: usize, ) -> Result<(), Box<dyn Error>>
WINDOWED fa_prefill_view_ws twin at head_dim 128 (lane/pp-prefill 2026-08-07):
step35’s SWA prefill (win=512, 33 of 45 layers) previously had NO windowed FA prefill
stamp — every windowed twin was hd256-only — and took sdpa_naive_w_quantized_view,
the f32 floor, at 565 ms/layer on a pp4096 where the hd128 FA family does the harder
causal-4096 in 3.3 ms (41% of the whole prime; research/pp-prefill-20260807 anatomy).
Same two-pass shape as the unwindowed function: dequant K/V ONCE into the resident
bf16 workspace, then the windowed qw kernel (fa_prefill_qw_db_w_hd128, cp.async
double-buffered; MEMRA_PRIME_DEQW_DB=0 selects the single-buffer twin). The window
mask is fa_prefill_f32_body’s exact predicate; window == 0 is bit-identical to
fa_prefill_view_ws by construction (default-arg body). NEW NUMERIC CLASS vs the
f32 floor on SWA rows (bf16 MMA online-softmax vs f32 serial softmax) — adoption is
gated by the full battery, and the class must change UNIFORMLY for a whole request
(kernel selection keys on seq_end, never per chunk — the chunkfix law).
hd128-only deliberately: the only windowed-prefill consumer at another head_dim is
gemma4 (hd256), which already has fa_prefill_w_f32.
Sourcepub 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>>
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).
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>>
Sourcepub 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>>
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).
Sourcepub 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>>
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).
Sourcepub fn fa_rows_eligible(&self, base_len: usize, head_dim: usize) -> bool
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).
Sourcepub 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>>
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].
Sourcepub 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>>
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).
Sourcepub 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>>
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).
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub fn fa_geom_eager(
&self,
t_kv: usize,
head_dim: usize,
n_head_kv: usize,
g: bool,
) -> (bool, usize)
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.)
Sourcepub fn fa_bucket_key(
&self,
t_kv: usize,
head_dim: usize,
n_head_kv: usize,
g: bool,
) -> (bool, usize)
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).
Sourcepub fn capture_graph_retained<F>(
&self,
step: F,
) -> Result<(CudaGraph, Vec<Box<dyn Any + Send>>), Box<dyn Error>>
pub fn capture_graph_retained<F>( &self, step: F, ) -> Result<(CudaGraph, Vec<Box<dyn Any + Send>>), 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.
Sourcepub fn capture_graph_retained_flags<F>(
&self,
flags: CUgraphInstantiate_flags,
step: F,
) -> Result<(CudaGraph, Vec<Box<dyn Any + Send>>), Box<dyn Error>>
pub fn capture_graph_retained_flags<F>( &self, flags: CUgraphInstantiate_flags, step: F, ) -> Result<(CudaGraph, Vec<Box<dyn Any + Send>>), 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.
pub fn capture_graph<F>(&self, step: F) -> Result<CudaGraph, Box<dyn Error>>
Sourcepub 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>>
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).
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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.
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>>
Sourcepub 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>>
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.
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>>
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>>
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>>
Sourcepub fn gdn_chunked_enabled() -> bool
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.
Sourcepub fn gdn_chunk_size() -> usize
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).
Sourcepub 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>>
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.
Sourcepub fn gdn_db_on() -> bool
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).
Sourcepub fn gdn_mma_enabled(&self, c: usize) -> bool
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).
Sourcepub fn gdn_wgmma_on(&self, c: usize) -> bool
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).
Sourcepub 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>>
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.
Sourcepub fn gdn_chunk_alloc(
&self,
n_head: usize,
t: usize,
c: usize,
hk: usize,
) -> Result<GdnChunkBufs, Box<dyn Error>>
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.
Sourcepub fn f32_to_bf16_v(
&self,
x: &CudaView<'_, f32>,
dst: &mut CudaSlice<u8>,
n: usize,
) -> Result<(), Box<dyn Error>>
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).
Sourcepub fn f32_to_bf16_into(
&self,
x: &CudaSlice<f32>,
dst: &mut CudaSlice<u8>,
n: usize,
) -> Result<(), Box<dyn Error>>
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).
Sourcepub fn gdn_chunk_k123_vl8(
&self,
seqs: &[GdnSeqVl],
n_head: usize,
hk: usize,
wq: Option<&GdnWVl8>,
) -> Result<(), Box<dyn Error>>
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).
Sourcepub 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>>
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).
Sourcepub fn gdn_mirror_vl8(
&self,
seqs: &[GdnSeqVl],
n_head: usize,
which: i32,
hk: usize,
) -> Result<(), Box<dyn Error>>
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).
Sourcepub fn gdn_tail_vl8(
&self,
seqs: &[GdnPrepVl],
norm_w: &CudaSlice<f32>,
d_state: usize,
num_v: usize,
eps: f32,
) -> Result<(), Box<dyn Error>>
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.
Sourcepub fn addr_f32(&self, x: &CudaSlice<f32>) -> u64
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).
pub fn addr_f32_mut(&self, x: &mut CudaSlice<f32>) -> u64
pub fn addr_f32v(&self, x: &CudaView<'_, f32>) -> u64
pub fn addr_u8(&self, x: &CudaSlice<u8>) -> u64
Sourcepub fn gdn_chunk_vl8(
&self,
seqs: &[GdnSeqVl],
n_head: usize,
scale: f32,
hk: usize,
wq: Option<&GdnWVl8>,
) -> Result<(), Box<dyn Error>>
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).
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>>
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub fn sigmoid_v(
&self,
x: &CudaView<'_, f32>,
y: &mut CudaSlice<f32>,
n: usize,
) -> Result<(), Box<dyn Error>>
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.
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>>
pub fn sigmoid( &self, x: &CudaSlice<f32>, y: &mut CudaSlice<f32>, n: usize, ) -> Result<(), Box<dyn Error>>
Sourcepub 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>>
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.
Sourcepub fn attn_head_gate(
&self,
a: &CudaSlice<f32>,
g: &CudaSlice<f32>,
dst: &mut CudaSlice<f32>,
dst16: Option<&mut CudaSlice<u8>>,
head_dim: usize,
n_head: usize,
t: usize,
) -> Result<(), Box<dyn Error>>
pub fn attn_head_gate( &self, a: &CudaSlice<f32>, g: &CudaSlice<f32>, dst: &mut CudaSlice<f32>, dst16: Option<&mut CudaSlice<u8>>, head_dim: usize, n_head: usize, t: usize, ) -> Result<(), Box<dyn Error>>
step35 (Step-3.7-Flash) SEPARATE head-wise attention gate: one scalar per query head,
broadcast over head_dim. dst = a * sigmoid(g) where a/dst are [head_dim, n_head, T]
(the q_gate_split layout) and g is the PRE-sigmoid attn_gate projection output in
token-major [T, n_head]. dst16 is the optional fp16 operand for wo (None -> skipped).
NOT interchangeable with sig_mul_f16out, which gates FULL WIDTH (qwen35 packs one gate
value per (head, dim) element inside wq). Using this for that, or that for this, silently
applies the wrong number of distinct gate values.
Sourcepub fn swiglu_clamped_mul_scaled(
&self,
gate: &CudaSlice<f32>,
up: &CudaSlice<f32>,
gs: f32,
us: f32,
limit: f32,
dst: &mut CudaSlice<f32>,
n: usize,
) -> Result<(), Box<dyn Error>>
pub fn swiglu_clamped_mul_scaled( &self, gate: &CudaSlice<f32>, up: &CudaSlice<f32>, gs: f32, us: f32, limit: f32, dst: &mut CudaSlice<f32>, n: usize, ) -> Result<(), Box<dyn Error>>
step35 CLAMPED SwiGLU: dst = min(silu(gate*gs), limit) * clamp(up*us, +-limit).
Verbatim from llama.cpp llama-graph.cpp:2146-2165 (routed, swiglu_clamp_exp) and
:1751-1770 (shared, swiglu_clamp_shexp), non-DEEPSEEK4 branch.
This is NOT swigluoai_mul_scaled: that one clamps the gate BEFORE swish and multiplies by
(1 + clamp(up)). Caller MUST check limit > 1e-6 (upstream’s eps gate) and use the plain
silu_mul_scaled path otherwise — at limit=0 this kernel would clamp every positive
activation to zero. On Step-3.7-Flash only layers 43 (7.0) and 44 (16.0) have a live limit.
Sourcepub 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>>
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.
Sourcepub 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>>
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).
Sourcepub 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>>
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.
Sourcepub 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>>
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.
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>>
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>>
Sourcepub fn transpose(
&self,
inp: &CudaSlice<f32>,
rows: usize,
cols: usize,
) -> Result<CudaSlice<f32>, Box<dyn Error>>
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.
Sourcepub 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>>
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].
Sourcepub 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>>
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.
Sourcepub 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>>
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.
Sourcepub fn conv_left_pad(
&self,
src: &CudaSlice<f32>,
dst: &mut CudaSlice<f32>,
conv_dim: usize,
t: usize,
pad: usize,
) -> Result<(), Box<dyn Error>>
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.
Sourcepub 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>>
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.
Sourcepub 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>>
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).
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).
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).