Struct DenseFfnLayer 

pub struct DenseFfnLayer {
    pub weights: DenseFfnWeights,
    /* private fields */
}

Fields

weights: DenseFfnWeights

Implementations

impl DenseFfnLayer

pub fn new(weights: DenseFfnWeights, gpu: &dyn GpuBackend) -> Result<Self>

pub fn new_with_activation( weights: DenseFfnWeights, activation: FfnActivation, gpu: &dyn GpuBackend, ) -> Result<Self>

pub fn finalize_q4k_load( &mut self, gpu: &dyn GpuBackend, h: u32, inter: u32, stream: u64, ) -> Result<()>

Load-time finalize for the Q4_K MMQ prefill path (ATLAS_FFN_MMQ). MUST run at load, BEFORE the KV cache is sized, so the net FFN footprint is correct when the KV cache claims free memory. Order is critical: (1) eagerly materialize the Q4_K weights (+9.63 GB) so they are accounted for now rather than lazily on first prefill (which would over-subscribe AFTER the KV cache already grabbed the freed _t space → decode OOM-throttle); (2) free the transposed _proj_t copies (−9.63 GB, dead under Q4_K prefill — only the unreachable Some(wt) arms read them). Net FFN = baseline; decode untouched (NVFP4 gemv on the non-_t copies). No-op unless Q4_K is active.

pub fn finalize_nvfp4_mmq_load( &mut self, gpu: &dyn GpuBackend, h: u32, inter: u32, stream: u64, ) -> Result<()>

Eagerly materialize the block_nvfp4 gate/up copies for the ATLAS_FFN_NVFP4_MMQ W4A4 prefill arm at LOAD time (before KV sizing), then free the now-dead gate/up transposed _t copies so net FFN footprint stays at the NVFP4 baseline. Down is untouched (hybrid: it stays on the default t_m128 path for accuracy → keeps its _t copy). No-op unless the env + kernels are present.

pub fn set_fp8_weights( &mut self, gate: Fp8Weight, up: Fp8Weight, down: Fp8Weight, )

Install native block-scaled FP8 dense MLP weights. After this call the forward paths dispatch w8a16_gemv (decode) / w8a16_gemm (prefill) instead of w4a16 NVFP4. Caller must ensure those kernels are present in the target (they are for the qwen3_5/ornith nvfp4 bundle).

pub fn set_lora_weights(&mut self, w: LoraFfnWeights) -> Result<()>

Install the startup-static LoRA FFN overlay (gate/up/down deltas). Hard-rejects when FP8/BF16 weight overlays are installed — those decode branches early-return before the NVFP4 tail where the M1 delta insertions land, so a permissive install would silently skip deltas. holo is NVFP4, so it is unaffected.

pub fn set_q2_weights( &mut self, gate: PackedQ2Weight, up: PackedQ2Weight, down: PackedQ2Weight, gpu: &dyn GpuBackend, )

Install native keep-packed ternary Q2_0 dense MLP weights. After this call, decode forward dispatches q2_0_gemv per projection (weights stay 2-bit resident, no NVFP4 requant) as the highest-priority path. Caller must ensure the q2_0_gemv kernel is present in the target (checked at forward time; falls through to a clear error otherwise). Prefill for packed-Q2 is a deferred phase — see forward_prefill.

pub fn set_bf16_weights( &mut self, gate: DenseWeight, up: DenseWeight, down: DenseWeight, )

Install BF16 dense MLP weights. After this call, the forward paths dispatch to the BF16 GEMV/GEMM kernels instead of w4a16. The caller must ensure the BF16 kernels are loaded (see dense_gemv_bf16_k / dense_gemm_bf16_k checks). Small-batch paths reuse forward_prefill so they cannot enter NVFP4 kernels with the null placeholder weights used by BF16-native layers.

pub fn forward( &self, input: DevicePtr, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<DevicePtr>

Single-token decode: 2-3 kernel launches depending on activation. SiLU: dual GEMV + SiLU-fused down GEMV (2 launches). GELU: dual GEMV + gelu_mul + down GEMV (3 launches, no fused GELU down kernel).

pub fn forward_k2( &self, input: DevicePtr, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

K=2 speculative: batched GEMV for 2 tokens. 3 launches: dual batch2 (gate+up) + silu_mul + batch2 (down).

pub fn forward_k3( &self, input: DevicePtr, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

K=3 speculative: batched GEMV for 3 tokens. 3 launches: dual batch3 (gate+up) + silu_mul + batch3 (down).

pub fn can_forward_km(&self, m: u32) -> bool

Whether the M-row batched-GEMV verify path is available for m rows (batchm kernel present AND NVFP4 weights loaded — the batchm GEMV reads the non-transposed NVFP4 layout).

pub fn forward_km( &self, input: DevicePtr, m: u32, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

K=m (m<=8) speculative verify: batched GEMV for m tokens. 4 launches: batchm gate + batchm up + silu_mul + batchm down — each projection weight is read ONCE for all m rows at near-peak stream bandwidth. nsys (2026-07-18, M=4): the forward_prefill MMQ arm this replaces for the K=4 verify cost 54.8 ms/step across the 64-layer dense FFN stack (~156 GB/s effective at M=4); the batch GEMV family measures ~290 GB/s on the same shapes (w8a16_gemv_batch4 sibling), putting this path at the ~31 ms weight-traffic floor. m=5..8 uses w4a16_gemv_batch8 (batchm_bench: same weight-streaming bandwidth, removing the M>4 tile-GEMM cliff for chain-verify K=5..8).

pub fn forward_prefill( &self, input: DevicePtr, num_tokens: usize, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

Timed wrapper around the dense-FFN prefill.

★ THIS PATH HAD NO TIMERS AT ALL, and that hid the largest unexplained number on the board. Profiling nvidia/Gemma-4-31B-IT-NVFP4 at a 4096-token prompt: wall 28,180 ms, while EVERY profiled phase across ATTN prefill [...] and MoE prefill [...] summed to 3,269.8 ms. 88% of the prefill was invisible — not attributed to something slow, simply not instrumented. forward_prefill dispatches ~20 quantization arms and none of them reported elapsed time; only one-shot “which arm was chosen” INFO lines existed.

One coarse timer first, deliberately: it answers whether the missing time is here at all before anyone threads timers through twenty arms. Same <AREA> prefill [phase] N=<n>: <us>µs shape the attention and MoE paths already emit, so the existing log-summing one-liners pick it up unchanged.

pub fn forward_batched( &self, input: DevicePtr, num_tokens: usize, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

Batched forward (per-token loop). Used by forward_batched in model loop.