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

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.