Struct MoeLayer 

pub struct MoeLayer {
    pub weights: MoeWeights,
    pub pre_expert_norm: Option<DenseWeight>,
    pub is_dflash_capture_layer: bool,
    /* private fields */
}

MoE feed-forward network component.

Not a TransformerLayer — used as a component inside layers for the FFN/MoE block after post-attention norm.

Fields

weights: MoeWeights

pre_expert_norm: Option<DenseWeight>

Pre-expert norm: applied to input AFTER routing but BEFORE expert dispatch. Gemma-4 26B: router sees raw residual, experts see pre_feedforward_layernorm_2(residual).

is_dflash_capture_layer: bool

Implementations

impl MoeLayer

pub fn fp32_routing_active(&self, levers: &ModelLevers) -> bool

True when the ATLAS_FP32_ROUTING path is active: the SSM-side MoE-input norm should emit an FP32 router_in (residual_add_rms_norm_gatef32) which the gate GEMM then consumes at full precision. Requires the f32 kernels to be present and the softmax-routed dense-gate config (NVFP4 gate / sigmoid+bias stay BF16). Default off → BF16 routing unchanged. The lever is the LAST term on purpose: the four preconditions are properties of this layer’s weights and kernels, and only the final one is configuration. levers is passed rather than read because this is called once per layer per DECODE TOKEN from six sites.

pub fn apply_zero_expert( &self, out: DevicePtr, x: DevicePtr, n: u32, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<()>

Forward pass: gate → top-K routing → batched expert FFN → blend.

All expert dispatch stays on device — zero D2H synchronization. 9 kernel launches per MoE layer (down from 58).

When gelu_activation is true, falls back to the sorted prefill path (which uses separate activation kernel) to avoid fused SiLU decode kernels. LongCat zero-computation experts: out[t,:] += zero_accum[t] * x[t,:] where zero_accum was written by the softmax+bias router kernels (the folded weights of selected identity experts). MUST run after the routed blend for the SAME tokens whose routing wrote zero_accum. No-op (no launch) when the model has no zero-experts.

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

impl MoeLayer

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

C=4 NVFP4 routed MoE decode with FP32 atomic accumulation.

Gate/top-K remain batched. Gate+up reuses the token-major kernel, then routed down projections atomic-add weighted FP32 contributions into a tiny [4,H] scratch accumulator. Finalization casts routed output to BF16 and optionally blends shared expert output.

impl MoeLayer

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

Batched forward: GEMM gate for N tokens, per-token expert dispatch.

Gate projection reads weights once for N tokens (GEMM M=N). Expert dispatch remains per-token (data-dependent routing).

impl MoeLayer

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

4. Computes local experts on local + received tokens
5. Sends results back (combine)
6. Weighted sum into output

Currently scaffolding only — builds routing table and logs statistics. Expert compute and actual dispatch use the existing per-token path. The all-reduce fallback is used for the actual output until dispatch kernels are implemented.

impl MoeLayer

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

Fused K=2 forward: process 2 tokens through MoE in 5 kernel launches.

Gate GEMV batch2 → batched topK → fused expert gate+up → fused silu+down → fused wsum+blend. Expert buffers sized for 2*top_k slots. Shared expert buffers reuse logits/ssm_qkvz (sized for 2 tokens). Output at moe_output() [2, H].

impl MoeLayer

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

Fused K=3 forward: process 3 tokens through MoE in 5 kernel launches.

Gate GEMV batch3 → batched topK → fused expert gate+up → fused silu+down → fused wsum+blend. Expert buffers sized for 3*top_k slots. Output at moe_output() [3, H].

impl MoeLayer

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

N-token prefill via grouped GEMM: sort-by-expert → tensor-core GEMM per expert.

Each expert’s weight matrix is loaded once (not per-token), cutting LPDDR5X reads from ~6 GB (GEMV) to ~150 MB (grouped GEMM) at N=1024.

Pipeline: gate → topK → sort → grouped gate/up GEMM → SiLU → grouped down GEMM → unpermute + weighted reduce → shared expert blend. Shared expert uses checkpoint-native BF16 when installed, otherwise W4A16.

impl MoeLayer

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

Token-major fused decode for small N>=4.

This reuses the generic moe_prefill kernels without the sorted/grouped GEMM path. It batches gate/top-k and processes all (token, expert-slot) routes in three token-major kernels:

gate GEMM -> batched topK -> gate+up -> silu+down -> wsum/blend.

First pass is NVFP4 + shared-expert only, matching Holo’s current decode path. FP8/BF16/unified-layout variants deliberately fall back to the existing implementation until they have equivalent generic kernels.

impl MoeLayer

pub fn set_pre_expert_norm(&mut self, norm: DenseWeight)

Transpose MoE weights for coalesced prefill GEMM reads.

Transposes per-expert routed weights [N, K/2] → [K/2, N] to enable the cp.async pipelined FP8-MMA K64 kernels. This doubles expert memory (~17 GB for 35B, ~30 GB for 122B) but eliminates the catastrophic uncoalesced B reads in the fallback grouped GEMM, cutting MoE prefill time by ~2x. Set pre-expert norm (Gemma-4 26B: pre_feedforward_layernorm_2). Applied to input AFTER routing but BEFORE expert dispatch.

pub fn set_gelu_activation(&mut self, gpu: &dyn GpuBackend) -> Result<()>

Set GeGLU activation for MoE experts (Gemma-4 26B). Replaces SiLU with GELU in the sorted/unfused path and forces decode to use the sorted path (avoiding fused SiLU kernels).

pub fn transpose_for_prefill( &mut self, gpu: &dyn GpuBackend, config: &ModelConfig, ) -> Result<()>

pub fn transpose_gate_up_for_prefill( &mut self, gpu: &dyn GpuBackend, config: &ModelConfig, ) -> Result<()>

Transpose only the gate+up routed weights, leaving the down projection in its original layout. Cuts the transpose memory cost from ~3× (gate+up+down) to ~2× per expert. Used by MiniMax M2.7-NVFP4 EP=2 when the full transpose doesn’t fit but gate+up does — the fused moe_w4a16_fused_gate_up_k64_n128 kernel still runs (capturing the dominant gate+up bandwidth savings), while down stays on the uncoalesced grouped-GEMM path.

pub fn transpose_for_prefill_unified( &mut self, gpu: &dyn GpuBackend, config: &ModelConfig, ) -> Result<()>

Phase 8a unified-layout transpose pass: build persistent transposed gate/up/down for all experts, freeing the untransposed copies between phases so the entire pass fits in tight memory budgets that the non-unified transpose_for_prefill_impl(true) would reject.

Phased flow (memory math for MiniMax M2.7-NVFP4 EP=2 ≈ 47 GB free): A. Transpose gate+up (allocs +39 GB; free ≈ 8 GB) B. Free gate+up untransposed (frees 39 GB; free ≈ 47 GB) C. Transpose down (allocs +20 GB; free ≈ 27 GB) D. Free down untransposed (frees 20 GB; free ≈ 47 GB)

Net memory: same as starting point, but layout is now unified (transposed-only) — the [N, K/2] decode kernels can no longer run; dispatch must use the _t decode kernels (which do).

Caller responsibilities:

1. Set ATLAS_UNIFIED_MOE_LAYOUT=1 so MoeLayer::use_t_layout_for_decode() returns true at dispatch time.
2. Call this method INSTEAD of transpose_for_prefill / transpose_gate_up_for_prefill.

pub fn transpose_for_prefill_hybrid( &mut self, gpu: &dyn GpuBackend, config: &ModelConfig, ) -> Result<()>

Hybrid-layout transpose pass — analogue of transpose_for_prefill_unified that keeps the untransposed originals so decode + MTP verify dispatch can continue using the warp-reduction kernels. Allocates ~58 GB transposed alongside the existing ~58 GB originals on MiniMax M2.7-NVFP4 EP=2; fits in 122 GB GB10 with KV-cache headroom up to ~32K context. Caller is responsible for memory-fit gating (factory checks free memory before invoking this).

pub fn build_cutlass_grouped_sfb( &mut self, gpu: &dyn GpuBackend, config: &ModelConfig, stream: u64, ) -> Result<()>

Build per-expert swizzled SFB weight-scale tables for the CUTLASS grouped NVFP4 path (ATLAS_HOLO_MOE_GROUPED_CUTLASS). For each expert, swizzle the [K/16,N] gate_ptrs_t/up_ptrs_t scale into the CUTLASS SFB atom via pack_weight_sfb, then upload the per-expert pointer arrays. The grouped kernel pairs these with gate_ptrs.packed ([N,K/2]) + the real per-expert scale2. Requires FAST_MOE=full (gate_ptrs_t/up_ptrs_t present); no-op else.

impl MoeLayer

pub fn set_down_transpose_scratch( &mut self, scratch_packed: DevicePtr, scratch_scale: DevicePtr, packed_ptrs_t: DevicePtr, scale_ptrs_t: DevicePtr, )

Wire a shared per-prefill down_proj scratch + transposed pointer table.

Called by the factory after the persistent MoE transpose pass falls back to gate+up only. The scratch and pointer tables are shared across all MoE layers — one allocation reused layer-by-layer during the sequential forward. The same scale2_vals buffer is reused from the existing untransposed down_ptrs (transpose preserves per-tensor scales).

impl MoeLayer

pub fn predequant_for_prefill( &mut self, gpu: &dyn GpuBackend, config: &ModelConfig, stream: u64, ) -> Result<()>

Pre-dequant dense (non-expert) NVFP4 weights to FP8 for zero-overhead prefill.

Only affects gate GEMM and shared expert GEMMs. Expert weights stay NVFP4 (they’re bandwidth-bound so FP8 wouldn’t help).

pub fn set_fp8_experts( &mut self, experts: &[Fp8ExpertWeight], shared_expert: Fp8ExpertWeight, gpu: &dyn GpuBackend, ) -> Result<()>

Set FP8 expert weights for native FP8 dispatch.

Builds device-side pointer tables from FP8 expert weights so the fused FP8 MoE kernel can index by expert_id at dispatch time. Also stores the shared expert FP8 weights for direct pointer passing.

pub fn set_bf16_experts( &mut self, gate_experts: &[DenseWeight], up_experts: &[DenseWeight], down_experts: &[DenseWeight], shared_gate: DevicePtr, shared_up: DevicePtr, shared_down: DevicePtr, gpu: &dyn GpuBackend, ) -> Result<()>

Set BF16 expert weights for the FP8-dequant-on-load MoE path.

Activated by ATLAS_FP8_DEQUANT_MOE_TO_BF16=1. Eliminates the per-layer 0.989 FP8 cosine ceiling (measured in bench/fp8_dgx2_drift/cosine_run.py) by serving experts as BF16 throughout, matching vLLM-BF16 reference numerics. Memory cost: 2× expert weights vs native FP8.

shared_* are the shared expert’s BF16 gate/up/down DevicePtrs (or DevicePtr::NULL when the model has no shared expert).

pub fn set_bf16_shared_expert( &mut self, gate_proj: DenseWeight, up_proj: DenseWeight, down_proj: DenseWeight, ) -> Result<()>

Install checkpoint-native BF16 shared-expert weights independently of routed-expert precision.

impl MoeLayer

pub fn new( weights: MoeWeights, num_experts: usize, gate_nvfp4: Option<QuantizedWeight>, gpu: &dyn GpuBackend, config: &ModelConfig, ) -> Result<Self>

pub fn new_with_hash( weights: MoeWeights, num_experts: usize, gate_nvfp4: Option<QuantizedWeight>, tid2eid_dev: Option<DevicePtr>, gpu: &dyn GpuBackend, config: &ModelConfig, ) -> Result<Self>

Like MoeLayer::new but with an optional DeepSeek-V4 hash-routing tid2eid table ([vocab_size, top_k] i64). Some marks this as a hash-routed layer.