OLED-MoE: Accelerating MoE-Based dLLM Inference via Inter-Iteration Locality-Aware Expert Offloading

• Block-wise routing expands dLLM's active expert set, so intra-iteration prefetches can miss their timing and push H2D transfers or CPU compensation onto the decoding critical path.Existing AR systems rely on layer-wise prefetching that can be hidden by current-layer computation.Source: §1 · p. 2
• Adjacent denoising iterations have strongly overlapping same-layer expert sets, and high-confidence tokens tend to preserve routes, making their experts more likely to be reused.Token confidence is already produced during denoising, enabling reuse ranking without additional prefetch traffic.Source: §1 · p. 2
• OLED-MoE maps token confidence and gate scores to expert-retention priorities, adapts cache capacity by layer, and compensates misses using dynamic load and predicted reuse.The system targets actionable inter-iteration locality under limited GPU memory.Source: §1 · p. 2
• 1.23×–7.93× TPOT reduction and 1.44×–4.23× higher expert cache utilization over state-of-the-art offloading systems are reported for OLED-MoE.Source: §1 · p. 3

2.1 Semi-Autoregressive MoE-based dLLMs

• Semi-AR dLLMs denoise multiple masked tokens in parallel within a block, then proceed autoregressively across blocks and reuse completed-block KV states.An iteration is one model forward pass; AR processes one new token per iteration, while semi-AR processes a block of active masked tokens.Source: §2.1 · p. 3

2.2 Expert Offloading

• MoE expert weights can dominate model parameters, so offloading keeps a subset in GPU memory and places the rest in host memory for transfer or CPU execution when activated.GPU-resident experts are called the GPU expert cache; a miss is handled by H2D transfer or CPU computation.Source: §2.2 · p. 3
• Prior offloading systems use request-level caching, intra-iteration prefetching, or partial-CPU execution, differing in prediction timing and where expert computation runs.Prefetch systems predict experts for later layers during current-layer computation; partial-CPU systems execute some activated experts on the CPU.Source: §2.2 · pp. 3–4, Figure 3 · p. 4

3.1 Workload Characteristics of MoE-based dLLMs

• Semi-AR dLLMs activate around 8× more experts and incur around 16× higher expert FLOPs per MoE layer per iteration than AR decoding.This increases GPU cache pressure, H2D transfer demand, and CPU-side compensation cost.Source: §3.1 · p. 4

3.2 Limitations of Prior Offloading Systems

• Prior systems’ TPOTs are 1.70×–8.16× the ideal value, while cache utilization is 17.5%–73.2% lower under the same GPU expert-cache budget.Figure 5 compares TPOT and expert cache utilization at that shared budget.Source: Figure 5 · p. 4, §3.2 · p. 5
• Late prefetches and exposed CPU compensation put misses on the decoding critical path, while locality-unaware eviction and reuse-agnostic placement can undermine cache utilization.The larger block-wise active set makes prefetch slack scarce and raises CPU compensation load; current-cost-only placement can leave reusable experts on the CPU.Source: §3.2 · p. 5

3.3 Insight: Inter-Iteration Expert Locality

• Adjacent iterations share 85%–90% of activated experts, and block-boundary transitions retain more than 65% overlap; overlap declines as iteration gaps grow.Hidden-state cosine similarity reaches 98%, and similar locality trends appear in SDAR and RND1-Base.Source: §3.3 · p. 6

3.4 Implications and Challenges

• OLED-MoE’s challenges are predicting reusable experts under limited cache capacity and handling unavoidable misses without sacrificing future cache utility.Its proposed response is temporal-locality-aware retention and dynamic load- and reuse-aware cooperative compensation.Source: §3.4 · p. 6

4.1 System Overview

• OLED-MoE couples layer-wise cache retention with miss compensation: cache hits run on the GPU, while misses are transferred to the GPU or executed on the CPU.GPU-transferred misses can also become cache candidates for the next iteration.Source: §4.1 · p. 6

4.2 Problem Formulation

• The system makes retention, cache-allocation, and miss-compensation decisions online because activated experts evolve with each request and denoising iteration.Its objective balances current compensation latency against future cache utility; profiling supplies hardware costs, not request-specific reuse decisions.Source: §4.2 · pp. 6–7

4.3 Temporal-Locality-Aware Expert Retention

• Token confidence weights gate-derived routing information to prioritize experts whose routes are more likely to persist, while layer cache budgets favor both activation demand and cache effectiveness.Cache sizes are updated every five denoising iterations using smoothed runtime statistics.Source: §4.3 · pp. 7–8

4.4 Dynamic Load- and Reuse-Aware Cooperative Compensation

• OLED-MoE dispatches cache misses using both current token load and predicted reuse: a hardware-aware split initializes transfers, then bounded swaps refine the GPU/CPU assignment.The heuristic avoids enumerating all partitions, which is too expensive in the decoding loop.Source: §4.4 · pp. 8–9

4.5 Implementation

• OLED-MoE is implemented on dInfer v0.1 with vLLM v0.10.2 as the generation backend, using Python-level hooks to intercept routing metadata.The intercepted expert IDs, gate scores, and token confidence support priority updates, miss dispatch, and layer-wise cache management.Source: §4.5 · p. 9

5.1 Experimental Setup

• The evaluation uses LLaDA2.0-mini and LLaDA2.0-flash across GSM8K, HumanEval, and GPQA, with two GPU configurations and cache capacities of 40 or 100 experts per layer.Unless stated otherwise, experiments use batch size 1 and block length 16; prior systems are ported to a common backend for policy comparisons.Source: §5.1 · p. 10, §5.2 · p. 11

5.2 Overall Performance

• OLED-MoE reduces TPOT by 1.95× on average and up to 3.73× across settings, while raising expert cache utilization from 51.5% to 77.7% versus the strongest traditional baseline.It consistently has the lowest TPOT among practical systems and approaches the ideal cache bound.Source: §5.2 · p. 11
• Long prefill sequences yield high cache utilization for both OLED-MoE and FineMoE, while H2D transfers and computation of missed experts dominate TTFT.Prefill occurs once per request and is not OLED-MoE’s optimization target.Source: Figure 14 · p. 10, §5.2 · p. 11

5.3 Ablation Study

• Adding Priority, LayerAdapt, and Coop reduces TPOT by a further 20.56%, 9.09%, and 9.14%, respectively, on average over the inter-iteration caching baseline.The full system’s reductions versus Naive-LRU are larger at Cache=40 than Cache=100, and on LLaDA2.0-flash than LLaDA2.0-mini.Source: Figure 15 · p. 11, §5.3 · p. 12
• Priority improves cache utilization over LRU and LFU at small capacities by combining routing information with token confidence; the gap narrows as capacity grows.LayerAdapt helps most when average per-layer cache capacity is comparable to the number of activated experts, while its advantage narrows at very low or ample capacity.Source: §5.3 · p. 12
• Reuse-aware dispatch increases the effective transfer ratio over load-only dispatch, retaining and reusing a larger fraction of H2D-transferred experts in the next iteration.It directs limited transfer capacity toward experts with greater near-term reuse value while balancing CPU and H2D workloads.Source: §5.3 · pp. 12–13

5.4 Sensitivity and Overhead

• Larger batches enlarge the activated-expert union and reduce expert-level sparsity, while the TPOT-minimizing block length shifts with cache capacity.Offloading is most effective at single-request and small-batch sizes; smaller blocks suit tight cache budgets, whereas larger caches can support longer blocks.Source: Figure 16 · p. 12, §5.4 · p. 13
• At Cache=100, OLED-MoE uses about 40% of full-residency expert GPU memory, with TPOT 32% higher for LLaDA2.0-mini and 23% higher for LLaDA2.0-flash.TPOT falls as PCIe bandwidth rises, especially at Cache=40; LLaDA2.1 shows trends comparable to LLaDA2.0, and policy overhead is 3.69–6.03 ms per forward (0.59%–4.53% of forward latency).Source: Figure 17 · p. 12, Figure 18 · p. 12, §5.4 · p. 13

Expert offloading for MoE inference

• AR-oriented MoE offloading uses predictive caching, prefetching, routing-aware execution, or expert management, whereas OLED-MoE targets the larger concurrent activations of block-wise dLLM decoding.OLED-MoE predicts and retains reusable experts across denoising iterations.Source: §6 · p. 14

CPU-GPU hybrid inference for MoE models

• Existing CPU-GPU hybrid MoE systems primarily optimize placement for current execution or transfer cost, while OLED-MoE also accounts for future reuse in semi-AR dLLM decoding.Its compensation engine jointly considers current load balance and predicted inter-iteration reuse.Source: §6 · p. 14

Diffusion LLM inference acceleration

• Diffusion-LLM acceleration techniques reduce computation or latency but do not address large MoE expert-weight memory pressure; OLED-MoE focuses on that offloading problem.KV-cache offloading and weight quantization are complementary memory-saving techniques that OLED-MoE could combine with expert offloading.Source: §6 · p. 14

7 Conclusion

• OLED-MoE combines temporal-locality-aware retention, layer-adaptive cache allocation, and load- and reuse-aware cooperative compensation for MoE-based dLLM inference.Across evaluated workloads, it reduces TPOT by 1.23×–7.93× and improves expert cache utilization by 1.44×–4.23× over state-of-the-art baselines.Source: §7 · p. 14