Source-linked AI summary
Hardware Acceleration of Block-Diffusion LLM for Edge Devices
Wei-Hsing Huang, Kiseok Lee, Ming-Yen Lee, Weiyu Sun, Cheng-Jhih Shih, Gayatri Tanksali, Arpit Khandelwal, Pin-Jun Chen, Yingyan Celine Lin, Shimeng Yu
TL;DR
Batch-one edge inference cannot amortize weight traffic, and block-diffusion refinement still repeatedly reads prefix KV and FFN weights despite immutable completed blocks. The paper co-designs WIFiV-LPDDR with BRQ-KV and DAT-FFN for precision-tagged persistent-state delivery. On modeled Jetson-class platforms, the full stack improves energy and latency while keeping every corresponding compressed benchmark score within one absolute percentage point of baseline.
Problem
Batch-one edge inference lacks cross-request weight reuse, while repeated block-diffusion refinement still streams prefix KV and FFN weights.
Method
The paper co-designs WIFiV-LPDDR, BRQ-KV, and DAT-FFN to deliver precision-tiered KV and FFN state for native block diffusion.
Results
3.79x/3.96x energy-reduction factors and 2.88x/4.44x latency speedups are reported for evaluated 1.5B/7B models, with compressed benchmark score drops below one absolute percentage point.
Takeaways & Limitations
The full stack reduces modeled batch-one block-diffusion energy and latency while retaining near-baseline benchmark scores under the reported settings.
Abstract
from arXiv · showhide
Single-stream (batch-one) edge inference cannot amortize weight traffic across requests. Full-attention diffusion LLMs recompute the entire sequence at every step; native block diffusion makes completed blocks immutable and exactly cacheable, yet refinement still streams prefix KV and FFN weights. We co-design WIFiV-LPDDR, a wide-I/O LPDDR system for precision-tagged reads, BRQ-KV for a canonical low-rank-plus-INT8-residual prefix with query-dependent per-entry precision, and DAT-FFN for drift-mapped canonical replacement, adjacent-stage-corrected low-bit delta, or cached-state carry while keeping live activations unquantized. Both map to an input-stationary mixed-precision systolic array. For the evaluated 1.5B/7B models on modeled Jetson-class platforms, the full stack provides arithmetic-mean energy-reduction factors of 3.79x/3.96x and arithmetic-mean latency speedups of 2.88x/4.44x at the reported DAT-FFN settings; every corresponding compressed model-benchmark score drops by less than one absolute percentage point from its baseline.
I. INTRODUCTION
The paper targets batch-one edge inference, where native block diffusion enables immutable-prefix caching but refinement still repeatedly streams prefix KV and FFN weights. It co-designs memory, precision policies, and accelerator delivery to reduce this traffic while preserving benchmark quality.
- Motivation: Batch-one edge inference lacks cross-request batching, while block diffusion refines only an active block and makes completed blocks immutable and cacheable.Full-attention diffusion keeps all response positions mutable; representative workloads generate multi-token responses.
- Motivation: Refinement repeatedly streams immutable prefix KV and static FFN weights from DRAM, motivating the proposed memory–algorithm co-design.The system targets repeated block-diffusion refinement rather than prefill, which contributes at most 3% of measured latency and energy.
- Contributions: BRQ-KV preserves a recoverable canonical prefix while reversibly reassigning per-entry residual precision using block queries.Every entry remains represented through its low-rank base, and the lowest tier is base-only.
- Contributions: DAT-FFN maps drift to canonical replacement, adjacent low-bit correction, or cached-state carry while reusing a block-entry channel order.The design keeps live Q/K/V and FFN activations in BF16 and maps both policies to an input-stationary mixed-precision array.
- Results: 3.79x/3.96x energy-reduction factors and 2.88x/4.44x latency speedups are achieved for evaluated 1.5B/7B models, with score drops below one absolute percentage point.These are arithmetic means at the reported DAT-FFN settings across five workloads.
B. Memory Systems for Future Edge AI
The paper positions WIFiV-LPDDR and its precision policies against edge-memory and prior KV/FFN approaches. Its key distinction is reversible, multi-level persistent-state delivery coupled with drift-aware FFN updates.
- Memory Systems for Future Edge AI: WIFiV-LPDDR combines dense 4F2 VCT memory, CBA FinFET peripheral logic, and fine-pitch 2.5D integration to scale LPDDR-family bandwidth.The architecture separates array and peripheral process budgets and increases effective DQ width through the interposer interface.
- Prior Policies: MAGE reduces KV reads through binary top-k participation but leaves FFN unchanged, so each step still streams all FFN weights under the evaluated model.Entries outside its selected subset contribute nothing to sparse-attention steps.
- Prior Policies: Chipmunk periodically refreshes active-neuron masks using a full first-layer projection, while unselected neurons receive no current-step correction until refresh.This differs from persistent FFN-weight delivery under the paper’s weight-streaming model.
- Proposed KV Policy: BRQ-KV assigns reversible q8/q4/q2/q0 residual precision while retaining every historical entry through its low-rank base.No entry is evicted, and lower-tier entries may later be promoted.
- Proposed FFN Policy: DAT-FFN fixes one channel ordering per block and moves replacement, view-delta, and carry boundaries with drift, avoiding reranking or an extra full-matrix projection.The recurrent-state refresh occurs after each interior step.
A. Notation and Decoding Setting
Native block diffusion organizes decoding around immutable completed context and a mutable active block, with compatible query spans determining reuse. The notation distinguishes persistent q8 storage from read-time lower-precision views.
- Notation and Decoding Setting: Each decoding step updates either an S-token sub-block or the whole B-token block, defining the query span used for tier-map compatibility.The phases are block entry, interior refinement, and block completion.
- Notation and Decoding Setting: Completed context is immutable while the active block remains mutable, allowing compatible steps on the same span to reuse its tier map.Sub-block-free backbones set T = B.
- Precision Tiers: q8 is the only persistent integer code, whereas q4 and q2 are read-time views; q0 retains a KV low-rank base or carries FFN state without a current update.The tier budget applies to KV entries or FFN channels.
- Hardware Model: 4x bandwidth is the conservative evaluation point for the LPDDR5-based modeled WIFiV-LPDDR instantiation.Figure 2 presents a generation-agnostic organization and an LPDDR5-based modeled implementation.
B. WIFiV-LPDDR Architecture and Mixed-Precision Data
WIFiV-LPDDR co-designs memory, peripheral circuits, packaging, and precision-tagged delivery to scale LPDDR bandwidth for mixed-precision edge inference.
- Device, circuit, and package co-design: WIFiV-LPDDR combines a dense 4F2 VCT array, CBA FinFET peripheral circuits, and fine-pitch 2.5D integration with the compute die.The CBA separates array and peripheral process and thermal budgets, while a silicon interposer provides the die-to-die interface.
- Device, circuit, and package co-design: Package-enabled DQ width and stronger peripheral logic raise aggregate bandwidth through wider I/O and higher per-pin transfer rates.The LPDDR-generation-agnostic organization is parameterized by N_DQ, R_pin, and η_link.
- Device, circuit, and package co-design: Fine-pitch 2.5D integration reduces interconnect energy per bit while enabling bandwidth scaling without conventional package-on-package transfer costs.The cited VCT/CBA and 2.5D-interconnect results serve as technology anchors for the design.
- Memory-side precision conversion: Precision-tagged reads store canonical q8 residual or weight codes, then emit only the requested q8, q4, or q2 payload after local conversion.The conversion path uses clamp, bit selection, rounding, saturation, comparisons, and multiplexing depending on the payload type.
- Memory-side precision conversion: 8e_core + p e_down + E_xcode(p) models movement energy for p∈{8,4,2}, while q0 incurs zero modeled movement energy.The dominant downstream-transfer term scales with transmitted bits, so lower precision reduces bandwidth demand, transfer latency, and movement energy.
C. Block-Refreshed Query-Ranked KV Transport
BRQ-KV preserves a canonical low-rank-plus-q8 residual prefix while adapting residual precision per entry from block-query rankings and reusing compatible tier maps across refinement steps.
- Canonical prefix representation: Fixed rank-two basis matrices factorize keys before RoPE and values directly, producing coefficient matrices and signed-INT8 residual codes in the persistent master.The canonical representation stores shared BF16 bases, per-entry coefficients, and residual codes in DRAM.
- Canonical prefix representation: BRQ-KV retains every historical KV entry through a low-rank base while reversibly assigning q8, q4, q2, or q0 residual precision.q4 and q2 are read-time views, while q0 omits only residual traffic and MACs rather than evicting the entry.
- Block-refreshed query ranking: The first step of each new decode block uses q8 warm-up, while later spans build query-ranked maps and compatible steps reuse them.Map construction ranks the prefix with current query groups; block or span completion invalidates the tier map.
- Block-refreshed query ranking: A fixed aggregate budget partitions the prefix ordering into q8/q4/q2/q0 tiers, with reuse-step residual width b̄_kv = 8ρ8 + 4ρ4 + 2ρ2 bits per entry.The tier map begins at the next compatible step after construction, while warm-up and map-building steps request canonical q8 residuals.
- Mixed-precision attention execution: Mixed-precision attention keeps BF16 queries, low-rank branches, and probability tiles on the input-stationary array while streaming residuals at delivered tier precision.Keys are reconstructed to BF16 before RoPE, and finalized block K/V rows are encoded with fixed bases before persistent append.
D. Drift-Adaptive Tiered FFN
DAT-FFN reuses a fixed channel order while mapping drift to replacement, adjacent-view corrections, or cached-state carry. Its canonical q8 representation supports derived q4/q2 views and input-stationary execution with unquantized BF16 activations.
- DAT-FFN maps per-step input drift to replacement, weight-view delta, or carry over cached gated-FFN state.Live activations remain BF16; approximation comes from selected views, carry, and one-step tier crossings.
- Canonical q8 FFN weights generate q4 and q2 read-time views plus adjacent corrections, avoiding stored low-bit weight copies.Offline calibration fixes the code maps, clamp and shift parameters, and low-rank correction factors.
- Factorized runtime corrections compute adjacent-view updates from low-rank code differences and grouped activation projections.The method applies GPTQ scaling through grouped correction partials.
- Fixed channel ordering lets drift move only tier boundaries, while q8, q4/q2, and q0 branches select replacement, correction, or cached-state reuse.Same-tier updates are exact relative to the selected view; tier crossings use a prior anchor for one step before refresh.
- The input-stationary array keeps BF16 activation tiles stationary while tiered weight rows stream and correction branches accumulate separately.Gate/Up use X_t or δX_t; Down uses the corresponding gated-activation tile.
- Refresh materializes nonzero-tier state into the next anchor, carries unchanged q0 state, and patches either one span or the full block.The refresh pass commits recurrent state, q8 partials, and current tier boundaries.
A. Evaluation Setup
The evaluation models compute, memory, conversion, packing, metadata, and utilization for 1.5B/7B Fast-dLLM v2 on Jetson-class platforms. Four normalized configurations isolate LPDDR, compression, and their combination under fixed calibration and DAT-FFN settings.
- LLMET models compute, on/off-chip memory, conversion, KV traffic, packing, metadata, and utilization for Jetson Orin Nano 4GB and NX 16GB platforms.LPDDR5 baselines provide 51.2/102.4 GB/s; WIFiV-LPDDR is modeled at four times 204.8/409.6 GB/s.
- Baseline, Compression, WIFiV-LPDDR, and Full isolate the algorithmic and memory contributions, with every added compute and data-movement cost included.Each pair is normalized to Baseline, and model factors are arithmetic means of five configuration ratios.
- A fixed GSM8K-training calibration subset sets the target KV residual width without using test examples.The target is b̄_kv = 4.6861 residual bits/entry.
- The fixed KV tier shares are q8/q4/q2/base-only = (0.32326, 0.43082, 0.18837, 0.05755) across models and benchmarks.The same calibration subset fixes KV clip/shift, threshold maps, and DAT-FFN amplitude factors.
- Reported DAT-FFN settings vary by model and benchmark while keeping KV fixed and using BF16 activations with fixed offline corrections.Settings are 0/3/2/3/4 for GSM8K/MATH/IFEval/HumanEval/MBPP on 1.5B and 4/2/6/4/5 on 7B.
B. Results
Across ten model-benchmark pairs, Full improves every pair and outperforms either component alone, with the largest arithmetic-mean gains for the combined stack.
- Full achieves arithmetic-mean energy-reduction factors of 3.79×/3.96× and latency speedups of 2.88×/4.44× for 1.5B/7B.These values use the reported DAT-FFN settings across five workloads.
- Full improves every evaluated model-benchmark pair and exceeds either Compression or WIFiV-LPDDR alone for every pair.Compression atop WIFiV-LPDDR lowers energy more than latency because transferred-bit savings persist while compute and on-chip work limit latency.
- Prefill contributes at most 3% of measured latency or energy across the ten evaluated pairs.
V. CONCLUSION
The paper co-designs WIFiV-LPDDR, BRQ-KV, and DAT-FFN for batch-one block-diffusion inference. On modeled Jetson-class systems, the Full stack delivers substantial mean energy and latency gains while keeping benchmark degradation below one percentage point.
- The co-designed stack combines WIFiV-LPDDR, BRQ-KV, and DAT-FFN for batch-one block-diffusion inference.
- 3.79×/3.96× energy reduction and 2.88×/4.44× latency speedup are reported for 1.5B/7B models on modeled Jetson-class systems.
- Every corresponding compressed benchmark score drops by less than one absolute percentage point from its baseline.