---
title: 'BlockFFN: Efficient MoE for LLMs'
url: https://www.emergentmind.com/topics/blockffn
type: topic
---

# BlockFFN: Efficient MoE for LLMs

BlockFFN is a mixture-of-experts (MoE) feed-forward network architecture designed to maximize hardware efficiency and activation sparsity in large language models (LLMs), specifically targeting practical end-side (single-device) deployment. It introduces fully differentiable, flexible routing and jointly optimizes for both token-level and chunk-level activation sparsity, deploying these innovations with acceleration kernels that integrate speculative decoding and sparse computation. BlockFFN is notable for its ability to achieve over 80% token-level sparsity (TLS), over 70% chunk-level sparsity (CLS) for 8-token batches, and end-to-end inference speedups up to 3.67× over dense baselines, all while matching or exceeding the perplexity and downstream metrics of prior MoE approaches [2507.08771].

## 1. Motivation and Core Principles

BlockFFN addresses two central bottlenecks in MoE-based LLMs: the non-differentiable/inflexible routing typical of vanilla MoE and the incompatibility of conventional sparsity patterns with high-performance single-device inference. Existing MoE designs often result in low CLS—meaning that, although each token uses few experts, the union of experts across a chunk of tokens remains large, precluding efficient chunk-wise computation. Furthermore, non-differentiable routing hinders model optimization, and activation patterns unaligned across tokens undermine chunk-based acceleration (e.g., via speculative decoding). The design paradigm of BlockFFN is to satisfy three criteria simultaneously: (1) fully differentiable, flexible routing, (2) maximization of both TLS and CLS, and (3) hardware-friendly kernels fusing activation sparsity with speculative decoding [2507.08771].

## 2. Architecture: Routing and Expert Design

BlockFFN layers are built from sparse MoE blocks of the form
$$
\mathrm{FFN}(x) = \sum_{i=1}^{N_e} A_i(x) E_i(x),
$$
where $N_e$ is the number of experts, $E_i(x)$ is the output of expert $i$, and $A_i(x)$ is its routing weight. The routing of tokens is made fully differentiable: given $x \in \mathbb{R}^{d_h}$,
$$
A^0(x) = W_{\mathrm{router}}^T x,\quad
A^1(x) = \mathrm{ReLU}(A^0(x)),\quad
A(x) = \mathrm{RMSNorm}(A^1(x)),
$$
where $W_{\mathrm{router}} \in \mathbb{R}^{d_h \times N_e}$ are learnable parameters. ReLU enforces token-specific, data-dependent sparsity by zeroing out unused experts, while RMSNorm rescales active weights, preventing attenuation of non-zero routing signals even under strong regularization [2507.08771].

Each expert is implemented as a compact two-layer MLP (“block”) with Swish activation:
$$
E_i(x) = W_{\mathrm{down}}^{(i)T} \sigma(W^{(i)T}_{\mathrm{up}} x),\quad \sigma = \mathrm{Swish},
$$
where $W^{(i)}_{\mathrm{up}} \in \mathbb{R}^{d_h \times d_e}$ and $W^{(i)}_{\mathrm{down}} \in \mathbb{R}^{d_e \times d_h}$.

## 3. Activation Sparsity: TLS, CLS, and Training Objectives

BlockFFN optimizes for both token-level and chunk-level activation sparsity, crucial for practical speedups in chunked and speculative inference:

- **Token-Level Sparsity (TLS)** quantifies the fraction of experts not used by an individual token $x_k$. Defining $\mathcal{S}_k = \{i : A_i(x_k) > 0\}$,
$$
\mathrm{TLS} = 1 - \frac{1}{K} \sum_{k=1}^K \frac{|\mathcal{S}_k|}{N_e},
$$
where $K$ is the number of tokens.

- **Chunk-Level Sparsity (CLS)** measures the fraction of experts unused across a chunk of $L$ tokens. For chunk $\{x_k\}_{k=1}^L$, define $\mathcal{S}_{\mathrm{union}} = \bigcup_{k=1}^L \mathcal{S}_k$,
$$
\mathrm{CLS}_L = 1 - \frac{|\mathcal{S}_\mathrm{union}|}{N_e}.
$$
BlockFFN achieves $\mathrm{TLS} \geq 80\%$ and $\mathrm{CLS}_8 \geq 70\%$, outperforming prior MoEs which typically collapse to $\mathrm{CLS}_8 < 60\%$.

To induce these properties, BlockFFN augments the standard language modeling loss with two auxiliary terms:

| Training Objective      | Definition | Purpose                                  |
|------------------------|------------|-------------------------------------------|
| $\mathcal{L}_{\mathrm{al}}$   | $\frac{1}{K}\sum_{k=1}^{K} \mathrm{BCE}(\sigma_\alpha(A^0(x_k)),\sigma_\alpha(\mathrm{LeftShift}(A^0(x_k))))$ | Encourages activation locality among adjacent tokens |
| $\mathcal{L}_{\mathrm{cs}}$   | $\frac{1}{N_e} \sum_{i=1}^{N_e} \mathcal{P}^i_{\mathrm{act}}$ | Penalizes union of active experts across $L$-token chunk |

The total objective is
$$
\mathcal{L}_\mathrm{total} = \mathcal{L}_\mathrm{lm} + \lambda_\mathrm{al} \mathcal{L}_\mathrm{al} + \lambda_\mathrm{cs} \mathcal{L}_\mathrm{cs},
$$
with an adaptive scheduler on $\lambda_\mathrm{cs}$ to steadily tighten chunk-level sparsity without destabilizing training [2507.08771].

## 4. Inference Kernels and End-Side Acceleration

BlockFFN provides custom CUDA (CUTLASS) kernels to exploit activation sparsity and speculative decoding for multi-token inference. The key steps are:

1. Compute routing activations $A^0 \rightarrow A^1 \rightarrow A$ for $n$ tokens.
2. Form $\mathcal{U} = \bigcup_{j=1}^n \{i : A_i(x^{(j)}) > 0\}$; $|\mathcal{U}| \ll N_e$ by high CLS.
3. Perform up-projection for all tokens and experts in $\mathcal{U}$ with one GEMM.
4. Apply masking to zero unused experts per token.
5. Run Swish and corresponding down-projection; mask and reduce.

For batch size $n$, the block is dense in $\mathcal{U}$, leveraging GPU tensor core efficiency. The expected FFN speedup is
$$
\mathrm{Speedup}_{\mathrm{FFN}} \approx \frac{N_e}{|\mathcal{U}|} = \frac{1}{1-\mathrm{CLS}_{\mathrm{spec}}}
$$
or $\frac{1}{1-\mathrm{TLS}}$ for $n=1$, and empirical results approach these upper bounds [2507.08771].

## 5. Empirical Results and Hardware Impact

Benchmarking on NVIDIA Jetson Orin NX (16 GB), BlockFFN (2.8B parameters) matches or exceeds previous MoE baselines (Switch/Mixtral TopK, GRIN, ReMoE, DeepSeek-MoE) on both sparsity and language modeling quality.

- *TLS*: consistently over 80%
- *CLS$_8$*: consistently over 70%
- *Inference speed*: up to 47.17 tokens/sec (3.67× dense baseline with speculative decoding)
- *Perplexity*: For a 1.2B model, BlockFFN achieves PPL = 8.69 (dense PPL = 8.49, TopK MoE = 8.87, ReMoE = 8.78)
- *Reuse ratio*: over 85% of active experts across token chunks, facilitating SRAM reuse

In practical deployment, this allows experts to be loaded into SRAM a single time and reused across tokens, reducing memory access inefficiency and further increasing throughput.

## 6. Relationship to Existing MoE Architectures

BlockFFN contrasts with prior MoE schemes in several respects. Vanilla MoE and TopK-based routing deliver high TLS but low CLS, leading to inefficient hardware utilization for chunked or speculative workloads. Heuristic or non-differentiable routers typical of earlier MoE models are replaced by a fully trainable ReLU+RMSNorm router, which jointly optimizes for local and chunked activation sparsity. Previous efforts frequently face performance or training stability trade-offs under strong sparsity constraints; BlockFFN employs auxiliary objectives with adaptive scheduling to mitigate these challenges [2507.08771].

## 7. Implications and Future Directions

BlockFFN establishes a new standard for end-side, inference-efficient MoE design, making large LLMs viable on single-device or memory-constrained hardware. The integration of chunk-level sparsity into both architectural and training design enables direct compatibility with mainstream acceleration techniques such as speculative decoding and chunked GEMM. A plausible implication is broader application of high-CLS MoE blocks in domains requiring fast, memory-efficient inference, including edge AI and real-time NLP systems. Future work may explore scaling BlockFFN to larger model sizes, tighter integration with more sophisticated speculative strategies, and adaptation to non-NLP modalities [2507.08771].

Source: https://www.emergentmind.com/topics/blockffn