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_prefillabove: 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 atN = 2 * intermediatein place of two atN = 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, behindATLAS_FFN_M16_TC(#927). - w8a8_
prefill - W8A8 block-scaled prefill branch (#917/#928). A CHILD module, not a
sibling: it adds
impl DenseFfnLayermethods 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 indense_ffn.rsreaches ahead of its W8A16 branches.
Structs§
- Dense
FfnLayer - Dense
FfnWeights - Dense
FfnWeights Bf16 - BF16 dense MLP weights — alternative to NVFP4 for precision-sensitive
models (Gemma-4-31B). Each is
[N, K]row-major BF16. When installed on aDenseFfnLayerviaset_bf16_weights, the forward paths dispatch todense_gemv_bf16/dense_gemm_bf16instead 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. - Dense
FfnWeights Fp8 - Native block-scaled FP8 dense MLP weights — loaded directly from an FP8
checkpoint (no NVFP4 requant). When installed via
set_fp8_weights, decode dispatchesw8a16_gemvand prefillw8a16_gemmper projection (BF16 act × FP8 E4M3 weight with 2D block scales), mirroring the SSM/attention FP8 path. - Dense
FfnWeights Q2 - 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 rawblock_q2_0buffer (2-bit codes + inline fp16 scale per group). When installed viaset_q2_weights, decode dispatchesq2_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).