Module dense_ffn

Module dense_ffn 

Source
Expand description

2026-09-25: Dense gated FFN layer, down(act(gate(x)) * up(x)) with SiLU or GELU, for single-token decode, small verify batches and prefill.

Owner: model-layers (dense FFN). Invariants:

  • On success, every forward entry point leaves its [rows, hidden] BF16 output in ctx.buffers.moe_output(); decode forward also returns that pointer.
  • forward, forward_k2, forward_k3 and forward_prefill serve an installed packed-Q2 overlay first, then native FP8, then BF16, and use the NVFP4 weights only when none is installed.

Modules§

batch16_decode
2026-09-25: The opt-in 5..=32-row native-FP8 decode tier. 2026-09-26: The opt-in 5..=32-row native-FP8 dense-FFN tier, w8a16_gemv_batch16: one weight pass for up to 16 rows, and two launches on row halves for 17..=32.
fp8_down
2026-09-25: Arm choice for the native-FP8 decode FFN. 2026-09-25: Chooses the arm that runs the native-FP8 decode FFN for one token (Fp8DownArm), as a pure function DenseFfnLayer::forward calls.
gateup_fused
2026-09-25: The fused FP8 gate+up GEMM (ffn_gateup_fused). 2026-09-25: The fused dense-FFN gate+up decode arm: one W8A8 block-scaled GEMM at N = 2 * inter, not two.
m16_tc
2026-09-25: The tensor-core 5..=32-row FP8 tier (m16_tc). 2026-09-25: The tensor-core w8a16_gemm_m16 arm of the FP8 dense-FFN ladder, 5..=32 rows.
w8a8_prefill
2026-09-25: The W8A8 block-scaled FP8 prefill arm. 2026-09-25: The W8A8 block-scaled arm of the native-FP8 dense FFN: which projections take it, and the GEMM it runs.

Structs§

DenseFfnLayer
DenseFfnWeights
DenseFfnWeightsBf16
2026-09-25: BF16 dense MLP weights. Once installed with set_bf16_weights, decode runs dense_gemv_bf16 and prefill runs cuBLASLt, dense_gemm_tc or dense_gemm_bf16.
DenseFfnWeightsFp8
2026-09-26: Native block-scaled FP8 E4M3 dense MLP weights. Once installed with set_fp8_weights, decode runs the arm fp8_down::fp8_down_arm picks and prefill runs the w8_gemm! ladder in prefill_fp8.
DenseFfnWeightsQ2
2026-09-25: Packed Q2_0 dense MLP weights (block_q2_0), kept 2-bit on the device. Once installed with set_q2_weights, decode runs q2_0_gemv_vec per projection and prefill runs the Q2_0 MMQ GEMM or a BF16 dequant (forward_prefill_inner).

Enums§

FfnActivation
2026-09-25: Activation of the gated FFN: SiLU runs moe_silu_mul, GeLU runs gelu_mul.