Module dense_ffn

Module dense_ffn 

Source
Expand description

Dense SwiGLU FFN component for non-MoE models.

Forward: gate = gate_proj(x), up = up_proj(x), out = down_proj(SiLU(gate) * up) 2 fused kernel launches per decode token (dual GEMV + SiLU-fused down GEMV).

Structs§

DenseFfnLayer
DenseFfnWeights
DenseFfnWeightsBf16
BF16 dense MLP weights — alternative to NVFP4 for precision-sensitive models (Gemma-4-31B). Each is [N, K] row-major BF16. When installed on a DenseFfnLayer via set_bf16_weights, the forward paths dispatch to dense_gemv_bf16 / dense_gemm_bf16 instead of the w4a16 NVFP4 kernels. Costs ~3.4 GB extra GPU memory on Gemma-4-31B (3 × hidden×intermediate × 2 bytes) vs NVFP4’s 0.5 bytes/weight.
DenseFfnWeightsFp8
Native block-scaled FP8 dense MLP weights — loaded directly from an FP8 checkpoint (no NVFP4 requant). When installed via set_fp8_weights, decode dispatches w8a16_gemv and prefill w8a16_gemm per projection (BF16 act × FP8 E4M3 weight with 2D block scales), mirroring the SSM/attention FP8 path.
DenseFfnWeightsQ2
Native keep-packed ternary Q2_0 dense MLP weights — loaded directly from a PrismML Q2_0 GGUF (ATLAS_GGUF_NATIVE_Q2=1) with NO dequant / NVFP4 requant. Each projection is a raw block_q2_0 buffer (2-bit codes + inline fp16 scale per group). When installed via set_q2_weights, decode dispatches q2_0_gemv (BF16 act × 2-bit weight, dequant-in-dot-product), mirroring the FP8 path but with the weights ~4× smaller resident.

Enums§

FfnActivation
Activation function for gated FFN (SiLU for Qwen/Llama, GELU for Gemma-4).