Struct DenseFfnLayer 

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

Fields

weights: DenseFfnWeights

Implementations

impl DenseFfnLayer

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

impl DenseFfnLayer

pub fn set_w8a8_decode_weights( &mut self, w: W8a8Ffn, hidden: u32, inter: u32, ) -> Result<()>

2026-09-28: Install the W8A8 gate, up and down. Refuses a GELU FFN and shapes other than [inter, hidden] (gate, up) and [hidden, inter] (down).

impl DenseFfnLayer

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

2026-09-25: Install native block-scaled FP8 dense MLP weights. The FP8 kernel handles are optional lookups (zero when absent), so the caller installs these only on a target that carries the FP8 kernels.

pub fn set_fp8_gate_up_fused(&mut self, fused: Fp8Weight)

2026-09-25: Install the [2*inter, hidden] FP8 gate+up weight that gateup_fused_plan can select. Contract: the gate and up passed to Self::set_fp8_weights are views into fused (gate.weight == fused.weight, up.weight == fused.weight + inter*hidden, and likewise for the scale grid). Only the first equality is checked, by debug_assert!. Otherwise the fused and unfused arms of one layer compute from different bytes.

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

2026-09-25: Install the LoRA overlay for gate/up/down. Errors when an FP8, BF16 or packed-Q2 overlay is installed (their branches return before the NVFP4 code that applies the deltas), or when the layer cannot take the split-SiLU decode path.

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

2026-09-25: Install packed Q2_0 dense MLP weights and look up the Q2_0 MMQ kernels. A missing q2_0_gemv_vec, q2_0_gemv_vec_batchm or dequant kernel makes the path that needs it return an error.

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

2026-09-25: Install BF16 dense MLP weights. The BF16 kernel handles are optional lookups, so the caller installs these only on a target that carries them. forward_k2, forward_k3 and forward_km then hand the layer to forward_prefill, so no NVFP4 kernel reads its NVFP4 weights.

impl DenseFfnLayer

pub fn fixed_ok(&self, m: usize, ctx: &ForwardContext<'_>) -> bool

2026-09-30: Whether m rows run a fixed format on this layer.

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

2026-09-30: The FFN of m rows at input under the fixed format, into moe_output.

impl DenseFfnLayer

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

2026-09-25: Decode for one token; returns moe_output holding [1, hidden]. Packed-Q2, FP8 and BF16 layers take their own branches. NVFP4 that single_row_w4a4 admits runs forward_km at one row (W4A4, 2026-09-28). Other NVFP4 runs the fused gate+up GEMV, then either act_mul plus the down GEMV (split SiLU: the default, and always with a LoRA adapter or GELU) or the SiLU-fused down GEMV. METRALE_DECODE_FFN_VIA_GEMM=1 runs NVFP4 through w4a16_prefill_gemm instead.

impl DenseFfnLayer

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

2026-09-25: FFN for 2 rows. Packed-Q2 goes to forward_km_q2, FP8/BF16 to forward_prefill, and to forward_km when small_batch_w4a4 and can_forward_km(2); otherwise NVFP4 runs w4a16_gemv_dual_batch2, act_mul and w4a16_gemv_batch2, plus the LoRA deltas.

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

2026-09-25: FFN for 3 rows; the same routing as forward_k2, with the batch3 kernels.

pub fn small_batch_w4a4(&self) -> bool

2026-09-28: Whether the 2- and 3-row steps go through forward_km, whose projections run W4A4 where the tier admits them (WeightQuantTier::ffn_small_batch_w4a4).

pub fn single_row_w4a4(&self) -> bool

2026-09-28: Whether the single decode row goes through forward_km (WeightQuantTier::ffn_single_row_w4a4).

pub fn narrow_rows(&self) -> u32

2026-09-28: Row edge of the dense FFN’s narrow arms (W4a16BatchmTiers::ffn_edge).

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

2026-09-25: Whether forward_km can serve m rows: a batchm tier resolved, and the layer has NVFP4 gate weights or an FP8 overlay.

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

2026-09-25: FFN for m verify rows. FP8/BF16 layers go to forward_prefill. NVFP4 runs gate, up and down through ops::w4a4_proj::nvfp4_proj_small_m with the batchm_kernel(m) tier (W4A4 where the weight-quantization tier admits it), plus act_mul and the LoRA deltas. There is no packed-Q2 arm; can_forward_km is false for a layer without NVFP4 or FP8 weights.

impl DenseFfnLayer

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

2026-09-25: Dense-FFN prefill of num_tokens rows of input; the output lands in moe_output. Under ctx.profile it synchronizes stream and logs FFN prefill [dense_total] N=<n>: <us>µs.

impl DenseFfnLayer

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

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

2026-09-25: Load-time setup for the Q4_K MMQ prefill arm. Returns at once for packed-Q2 or FP8 layers, when one of the four Q4_K kernels is missing, or when METRALE_FFN_MMQ is unset. Otherwise it builds the gate/up Q4_K copies and the down copy (int8 when int8_gemm_faith2 and requant_a_bf16_int8 resolved and METRALE_FFN_MMQ_DOWN_Q4K is unset, Q4_K otherwise), synchronizes stream, and frees the transposed _t copies.

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

2026-09-25: Load-time setup for the NVFP4 MMQ prefill arm. Returns at once for packed-Q2 or FP8 layers, GELU layers, when a kernel it needs is missing, or when METRALE_NO_FFN_NVFP4_MMQ is set. Otherwise it repacks gate and up, and down unless METRALE_NO_FFN_NVFP4_MMQ_DOWN is set, synchronizes stream, and frees the transposed _t copies of the repacked projections.

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

2026-09-25: Same as forward_prefill.