Enum FfnComponent 

pub enum FfnComponent {
    Moe(MoeLayer),
    Dense(DenseFfnLayer),
    None,
}

FFN component: MoE (expert routing), dense SwiGLU, or None (standalone attention).

Variants

Moe(MoeLayer)

Dense(DenseFfnLayer)

None

No FFN — used by Nemotron-H standalone attention layers.

Implementations

impl FfnComponent

pub fn is_none(&self) -> bool

pub fn is_dense(&self) -> bool

True for a plain dense (SwiGLU) FFN. Wide-batch verify paths gate their forward_prefill fast path on this: batching reads dense weights once (big win at N=17), but on a 256-expert MoE the grouped-GEMM is a net loss at small batch (per-expert M~1 + sort/permute overhead), so MoE keeps its per-token loop.

pub fn moe_grouped_decode_ok(&self) -> bool

True when this MoE FFN can serve DECODE through the grouped read-once GEMM (forward_prefill) instead of the pairwise per-slot loop. The is_dense() comment above asserts grouped is “a net loss at small batch” on a 256-expert MoE, but that was never measured for decode CONCURRENCY (n=4) where the pairwise path re-reads ~14-20 distinct experts as 40 per-slot CTAs. Native-NVFP4-routed only (forward_prefill’s unconditional grouped path); dense/none are false.

pub fn fp32_routing_active(&self) -> bool

ATLAS_FP32_ROUTING active for this FFN (MoE only; false otherwise).

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

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

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

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

Whether the K=m (m<=8) batched-GEMV verify FFN is available (dense only — MoE / missing batch4/batch8 kernel / non-NVFP4 weights → false). Lets callers gate branch entry BEFORE computing the pre-FFN norm, so there is no half-done fallthrough to forward_prefill.

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

K=m (m=4..8) verify FFN via batched GEMV (dense only). Returns false when the path is unavailable (MoE / missing batchm kernel / non-NVFP4 weights) so the caller can fall back to forward_prefill.

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

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

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

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