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).

Modules§

batch16_decode
The 5..=32-row native-FP8 DECODE tier (#927) — same child-module reason as w8a8_prefill above: it reads this layer’s private kernel handles, and the arm’s rule, its opt-in and the measurements that made it opt-in do not fit in this file’s budget. The 5..=32-row native-FP8 dense-FFN DECODE tier — w8a16_gemv_batch16.
fp8_down
The native-FP8 M=1 decode DOWN projection (#928) — the arm rule. A CHILD module, not a sibling: this file is already at the CI size cap, and the nsys attribution that motivates the arm needs room this file does not have. The native-FP8 M=1 decode DOWN projection — which arm runs it (#928).
gateup_fused
The FUSED gate+up DECODE GEMM (ffn_gateup_fused, #927) — a child module because it reads this layer’s private kernel handles, and the round-13 receipt, the layout decision and the residency-neutrality argument need room this file does not have. The FUSED dense-FFN gate+up DECODE GEMM — one block-scaled FP8 cuBLASLt call at N = 2 * intermediate in place of two at N = intermediate.
m16_tc
The TENSOR-CORE 5..=32-row decode tier (ATLAS_FFN_M16_TC, #927) — same child-module reason as the two above: it reads this layer’s private kernel handles, and its rule, lever and the numerics seam it opens need room. The 5..=32-row native-FP8 dense-FFN decode tier on TENSOR CORES — w8a16_gemm_m16, behind ATLAS_FFN_M16_TC (#927).
w8a8_prefill
W8A8 block-scaled prefill branch (#917/#928). A CHILD module, not a sibling: it adds impl DenseFfnLayer methods that read this layer’s private kernel handles, and this file is already at the CI size cap. W8A8 block-scaled dense-FFN PREFILL — the gate/up/down GEMM arm that the native-FP8 dispatch in dense_ffn.rs reaches ahead of its W8A16 branches.

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).