Deploy efficient MoE models on resource-constrained edge devices by learning chunk-level activation sparsity that achieves 3.67× speedup. Use when you need to compress LLMs for on-device inference while maintaining reasoning quality and supporting speculative decoding acceleration.
Deploy efficient MoE models on resource-constrained edge devices by learning chunk-level activation sparsity that achieves 3.67× speedup. Use when you need to compress LLMs for on-device inference while maintaining reasoning quality and supporting speculative decoding acceleration.
BlockFFN: Sparse Mixture-of-Experts Optimized for End-Device Deployment
Standard Mixture-of-Experts architectures suffer from two fundamental constraints that prevent efficient on-device deployment: non-differentiable routing mechanisms that limit training convergence, and low chunk-level sparsity where consecutive tokens activate overlapping expert sets, blocking parallelization opportunities. BlockFFN addresses both problems through differentiable routers and specialized loss functions that enforce structured sparsity patterns compatible with hardware acceleration.
The core insight is that chunk-level sparsity (whether the same expert is active across consecutive tokens) is what enables acceleration kernels, not just token-level sparsity. By learning when to deactivate experts across sequences of tokens, the model becomes compatible with speculative decoding and GEMM-based acceleration while remaining highly parameter-efficient.
Core Concept
BlockFFN replaces TopK routing with a differentiable linear router combined with ReLU gating, enabling end-to-end optimization. Critically, it introduces two novel loss functions: Activation Locality Loss encourages neighboring tokens to activate similar expert sets, and Chunk Sparsification Loss directly minimizes the probability that an expert activates within L-token windows.
This architecture design enables hardware accelerators to compute only the union of activated experts across n tokens in parallel, masking unnecessary activations post-computation. The result is genuine speedup on edge devices rather than theoretical parameter reduction.
Architecture Overview
Router Module: Linear transformation with ReLU activation followed by RMSNorm, enabling differentiable expert selection unlike non-differentiable TopK
Expert Modules: Non-gated MLPs with Swish activation, sized smaller than typical (d_e << d_h) for fine-grained sparsity control
Activation Pattern Learning: Sigmoid-sharpened activation patterns for smooth gradients during Activation Locality Loss computation
Chunk-Level Sparsity Mechanism: Probabilistic computation of expert activation within L-token windows using binary models
CUTLASS Acceleration Kernel: GEMM-based implementation processing n=32 tokens simultaneously, computing only the union of activated experts
Implementation
Router Architecture Setup
The router combines simplicity with differentiability, avoiding TopK bottlenecks while maintaining effective expert selection.
import torch
import torch.nn as nn
import torch.nn.functional F
(nn.Module):
():
().__init__()
.hidden_dim = hidden_dim
.num_experts = num_experts
.sharpness = sharpness
.router_linear = nn.Linear(hidden_dim, num_experts, bias=)
.rms_norm = nn.RMSNorm(num_experts)
():
scores = .router_linear(x)
scores = F.relu(scores)
scores = .rms_norm(scores)
probs = torch.sigmoid(.sharpness * scores)
probs, scores
(nn.Module):
():
().__init__()
.fc1 = nn.Linear(hidden_dim, expert_dim)
.fc2 = nn.Linear(expert_dim, hidden_dim)
.dropout = nn.Dropout(dropout)
.activation = nn.SiLU()
():
x = .fc1(x)
x = .activation(x)
x = .dropout(x)
x = .fc2(x)
x
as
class
BlockFFNRouter
def
__init__
self, hidden_dim, num_experts, sharpness=1.0
super
self
self
self
# Linear router without bias (simpler gradients)
self
False
self
def
forward
self, x
# x: (batch, seq_len, hidden_dim)
# Linear routing scores
self
# (batch, seq_len, num_experts)
# ReLU activation for sparsity
# RMSNorm for numerical stability
self
# Compute activation probabilities via sigmoid with sharpness
self
return
class
BlockFFNExpert
def
__init__
self, hidden_dim, expert_dim, dropout=0.1
super
# Smaller expert dimensions for finer granularity
self
self
self
self
# Swish
def
forward
self, x
self
self
self
self
return
The router avoids gradient blocking by using ReLU and normalization instead of TopK, enabling reliable training.
Activation Locality Loss
This loss encourages neighboring tokens to select similar expert sets, which creates structured sparsity compatible with acceleration.
defactivation_locality_loss(activation_probs, alpha=2.0):
"""
Minimize divergence between neighboring token activation patterns.
Args:
activation_probs: (batch, seq_len, num_experts) probabilities from sigmoid
alpha: sharpness parameter for approximation (higher = sharper transition)
Returns:
locality loss encouraging similar patterns in adjacent tokens
"""
batch_size, seq_len, num_experts = activation_probs.shape
# Sharpen probabilities for binary-like patterns
sharpened = torch.sigmoid(alpha * (2 * activation_probs - 1))
# Compute BCE between adjacent tokens
locality_loss = 0.0for i inrange(seq_len - 1):
# Binary cross entropy between token i and i+1
curr = sharpened[:, i, :] # (batch, num_experts)
next_token = sharpened[:, i + 1, :]
# BCE: -[p*log(q) + (1-p)*log(1-q)]
bce = -(curr * torch.log(next_token + 1e-8) +
(1 - curr) * torch.log(1 - next_token + 1e-8))
locality_loss += bce.mean()
return locality_loss / max(seq_len - 1, 1)
Chunk Sparsification Loss
This loss directly minimizes the probability that an expert activates anywhere within an L-token window, enabling acceleration kernel compatibility.
defchunk_sparsification_loss(activation_probs, chunk_length=8):
"""
Minimize probability that expert i activates within L-token chunk.
Probability that expert i does NOT activate in any token: ∏(1 - p_ik)
Probability that expert i DOES activate: 1 - ∏(1 - p_ik)
Args:
activation_probs: (batch, seq_len, num_experts)
chunk_length: L in the paper (typically 8 for speculative decoding)
Returns:
sparsification loss encouraging deactivation across chunks
"""
batch_size, seq_len, num_experts = activation_probs.shape
num_chunks = (seq_len + chunk_length - 1) // chunk_length
sparsif_loss = 0.0for chunk_idx inrange(num_chunks):
start = chunk_idx * chunk_length
end = min(start + chunk_length, seq_len)
# Get probabilities for this chunk
chunk_probs = activation_probs[:, start:end, :] # (batch, L, num_experts)# Compute probability of NOT activating: ∏(1 - p_ik)
prob_no_activation = torch.prod(1 - chunk_probs, dim=1) # (batch, num_experts)# Probability of activation: 1 - ∏(1 - p_ik)
prob_activation = 1 - prob_no_activation
# Loss: minimize activation probability (encourage deactivation)
sparsif_loss += prob_activation.mean()
return sparsif_loss / num_chunks
Full BlockFFN Layer with Combined Losses
Integrate the router, experts, and loss functions into a complete layer.
As training progresses, gradually increase λ_cs to enforce stronger chunk-level sparsity.
defschedule_sparsity_factor(step, total_steps, lambda_cs_max=0.5):
"""
Linearly increase sparsification weight from 0 to lambda_cs_max.
Prevents early-stage training instability.
"""return lambda_cs_max * (step / total_steps)
# In training loopfor step, batch inenumerate(train_loader):
lambda_cs = schedule_sparsity_factor(step, total_steps, lambda_cs_max=0.5)
loss = train_block_ffn_step(model, batch, optimizer, lambda_cs=lambda_cs)
Practical Guidance
Key Hyperparameters
Parameter
Value
Notes
Router Sharpness α
2.0
Increases with training; controls activation discretization
λ_al (Activation Locality)
0.1-0.3
Weight for locality loss; higher = more aligned neighboring patterns
λ_cs (Chunk Sparsification)
0.1-0.5
Weight for chunk loss; schedule to increase over training
Chunk Length L
8 tokens
Aligned with speculative decoding batches (32 tokens = 4 chunks)
Number of Experts
40 (optimal)
Empirically optimal for sparsity-performance tradeoff
Expert Dimension
d_e << d_h
256-512 for 2048-dim hidden states
Token-Level Sparsity Target
80-84%
Achieved via router without explicit TopK
Chunk-Level Sparsity Target
70%+
Primary metric for acceleration compatibility
When to Use
Deploying LLMs to edge devices (Jetson Orin, mobile hardware) with strict memory constraints
Requiring real-time inference speedup on consumer-grade GPUs (RTX 4080, A100)
Building systems that need both parameter efficiency AND hardware-friendly sparsity patterns
Implementing speculative decoding where the model handles 32-token batches simultaneously
Creating energy-efficient inference pipelines for mobile or embedded applications
When NOT to Use
Training from scratch with small models (<0.1B parameters); sparsity overhead exceeds benefits
Applications with no hardware acceleration support (pure CPU inference gets minimal speedup)
Scenarios where absolute accuracy is critical and you cannot tolerate 1-2% performance drops
Highly specialized domains where model capacity is already tight (sparsity compounds scaling down)
Common Pitfalls
Over-aggressive chunk length: Using L=4 instead of L=8 causes unbalanced expert loads and training instability; match your hardware's speculative batch size
Neglecting locality warmup: Applying full λ_al from step 0 causes optimization conflicts; start with λ_al=0, gradually increase to 0.1 over first 1000 steps
Mismatched expert dimensionality: Setting expert_dim too large (>d_h/2) negates sparsity benefits; keep d_e < 1/4 of hidden dimension
Ignoring chunk boundary effects: Edge chunks (last partial chunk) behave differently; pad seq_len to multiples of chunk_length for consistency
Static routing weights: Using fixed λ_cs throughout training leads to degraded initial convergence; implement adaptive scheduling to increase gradually
Reference
Tian, Y., Wang, Z., Liu, S., et al. (2024). BlockFFN: Towards End-Side Acceleration-Friendly Mixture-of-Experts with Chunk-Level Activation Sparsity. arXiv preprint arXiv:2507.08771.
