PaperlayerView exact revision on arXiv

Source-linked AI summary

Hardware Co-Design Scaling Laws via Roofline Modelling for On-Device LLMs

Luoyang Sun, Jiwen Jiang, Yifeng Ding, Fengfa Li, Yan Song, Haifeng Zhang, Jian Ying, Lei Ren, Kun Zhan, Wei Chen, Yan Xie, Cheng Deng

arXiv:2602.10377v1cs.LGcs.CL

TL;DR

On-device LLM deployment requires balancing quality with strict latency and hardware constraints, but this loss–latency relationship has not been systematically characterized. The paper combines an empirical architecture-to-loss scaling law with roofline latency modeling, evaluates architectures on NVIDIA Jetson Orin, and identifies Pareto-optimal designs. At equivalent latency to Qwen2.5-0.5B, the co-designed architecture achieves 19.42% lower WikiText-2 perplexity.

  • Problem

    On-device LLMs must satisfy strict latency and hardware constraints while maintaining accuracy, yet the loss–latency trade-off has not been systematically characterized.

  • Method

    The paper combines an empirical loss scaling law over architectural hyperparameters with roofline-based latency modeling and Pareto optimization.

  • Results

    19.42% lower perplexity on WikiText-2 than Qwen2.5-0.5B is achieved at the same inference latency on the target hardware.

  • Takeaways & Limitations

    The hardware co-design scaling law provides a practical guideline for selecting on-device LLM architectures under strict latency and resource budgets.

  • Takeaways & Limitations

    The framework is bounded by a loss law fitted on 170 architectures trained for 10B tokens, an idealized roofline latency model, and standard transformer components.

Abstract

from arXiv · show

Vision-Language-Action Models (VLAs) have emerged as a key paradigm of Physical AI and are increasingly deployed in autonomous vehicles, robots, and smart spaces. In these resource-constrained on-device settings, selecting an appropriate large language model (LLM) backbone is a critical challenge: models must balance accuracy with strict inference latency and hardware efficiency constraints. This makes hardware-software co-design a game-changing requirement for on-device LLM deployment, where each hardware platform demands a tailored architectural solution. We propose a hardware co-design law that jointly captures model accuracy and inference performance. Specifically, we model training loss as an explicit function of architectural hyperparameters and characterise inference latency via roofline modelling. We empirically evaluate 1,942 candidate architectures on NVIDIA Jetson Orin, training 170 selected models for 10B tokens each to fit a scaling law relating architecture to training loss. By coupling this scaling law with latency modelling, we establish a direct accuracy-latency correspondence and identify the Pareto frontier for hardware co-designed LLMs. We further formulate architecture search as a joint optimisation over precision and performance, deriving feasible design regions under industrial hardware and application budgets. Our approach reduces architecture selection from months to days. At the same latency as Qwen2.5-0.5B on the target hardware, our co-designed architecture achieves 19.42% lower perplexity on WikiText-2. To our knowledge, this is the first principled and operational framework for hardware co-design scaling laws in on-device LLM deployment. We will make the code and related checkpoints publicly available.

1. Introduction

On-device LLMs must balance model quality against strict hardware, memory, power, and latency constraints. The paper addresses this tension with hardware-guided architecture selection and Pareto analysis.

  • Motivation: On-device LLM deployment faces strict memory, bandwidth, power, and latency constraints that can make cloud-optimized architectures infeasible.High-accuracy models may exceed latency budgets, while latency-optimized pipelines can reduce accuracy.
  • Motivation: Transformer components occupy different compute–memory regimes, so arithmetic intensity, locality, KV-cache footprint, and routing patterns limit hardware utilization.Attention is bandwidth-bound, feedforward layers are compute-bound, and KV-cache stresses on-chip memory.
  • Research gap: NAS methods commonly optimize a single objective, but edge deployment requires explicit characterization of the loss–latency trade-off.The paper focuses on the interplay between model loss and inference latency in hardware-co-designed LLMs.
  • Approach: Pareto optimization identifies non-dominated architectures that balance accuracy and latency, allowing selection under specific deployment constraints.This avoids exhaustive enumeration when no single architecture is optimal for every deployment setting.
  • Contributions: The paper develops a hardware co-design law combining loss scaling with roofline-based latency modeling to characterize accuracy–latency trade-offs.It also evaluates approximately 1,942 architectures, trains 170 representative models for 10B tokens each, and formulates joint precision–performance optimization.
  • Paper organization: The paper organizes its framework around formulation, roofline-based Pareto architecture discovery, theoretical joint optimization, and practical deployment scenarios.These topics correspond to Sections 3–5 of the paper.

2. Related Work

Prior work improves LLM efficiency through sparse, memory-aware, and hybrid architectures, but existing NAS and profiling approaches do not fully integrate architecture search with hardware-aware performance modeling. This paper positions its roofline-based framework as that integration.

  • Efficient architectures: Recent efficiency-oriented LLM designs use MoE routing, KV-cache reduction, sub-quadratic attention, and SSM–attention hybrids.Gated attention further modulates token-dependent computation, shifting bottlenecks toward memory or routing sensitivity.
  • Edge deployment: Quantization methods and inference engines support efficient execution of compact language models on heterogeneous edge hardware.Examples include AWQ, GPTQ, vLLM, MLC-LLM, and PowerInfer.
  • Architectural trade-offs: Deeper designs with shared embeddings and sparsity-aware mechanisms can offer favorable accuracy–efficiency trade-offs under resource constraints.These findings concern architectural efficiency rather than a unified hardware–architecture search procedure.
  • Research gap: NAS automates architecture design but faces prohibitive search costs, limited latency and memory integration, and weak interpretability for LLMs.Profiling tools analyze hardware behavior but lack integration with architecture search, while depth–width studies often neglect deployment constraints.
  • Positioning: The paper extends roofline analysis to map architectural choices, MoE sparsity, attention variants, and KV-cache strategies to hardware-limited regimes.This enables architecture search guided by a hardware co-design law instead of post-hoc benchmarking.

3. Formulating Hardware Co-Design Law for on-Device LLM

The hardware co-design law formulates architecture selection as minimizing validation loss under latency and memory constraints. It combines empirical loss scaling with roofline latency modeling to identify Pareto-optimal designs for specified hardware and workloads.

  • 3. Formulating Hardware Co-Design Law for on-Device LLM: The deployment objective is to identify an architecture under fixed inference latency and precision constraints, unlike classical training-focused scaling laws.The formulation targets the deployment regime rather than fixed training budgets.
  • 3.1. Implicit Optimization Objective of Hardware Co-Design Law: The optimization minimizes validation loss subject to latency and memory constraints over architecture θ.The architecture includes depth, width, KV-cache dimension, FFN expansion ratio, and expert activation rate.
  • 3.1. Implicit Optimization Objective of Hardware Co-Design Law: The KV-cache dimension governs decoding memory and bandwidth, while the FFN expansion ratio controls intermediate-layer size and per-token compute.For MoE models, the expert activation rate is ρ = K/E, and total expansion accounts for activated experts.
  • 3.1. Implicit Optimization Objective of Hardware Co-Design Law: The latency surrogate models end-to-end context encoding and autoregressive generation using hardware throughput, memory bandwidth, and workload configuration.Roofline analysis determines whether inference is compute-bound or bandwidth-bound.
  • 3.1. Implicit Optimization Objective of Hardware Co-Design Law: Because brute-force enumeration is infeasible, the framework uses surrogate models to approximate stable trends in learning dynamics and system behavior.The loss surrogate is fitted from 170 dense and MoE architectures to predict validation loss from architectural parameters.
  • 3.2. Precision Modeling via Loss: The empirical loss model treats architecture components through separable scaling behavior, including distinct width-scaling exponents for sparsity-driven and base-capacity terms.The KV-cache dimension is reparameterized using the GQA group ratio for theoretical analysis.
  • 3.3. Performance Modeling via Latency: The latency model derives an approximate roofline expression from arithmetic intensity, peak compute throughput, sustained memory bandwidth, and total FLOP and memory traffic.Total inference latency includes both prefill and decode for the specified layer count and token lengths.
  • 3.4. Pareto-optimal architectures: Pareto-optimal architectures jointly optimize validation loss and end-to-end latency, with the frontier containing designs that cannot improve one objective without worsening the other.The resulting architecture selection principle and parameter scaling behavior under hardware constraints are called the hardware co-design scaling law.

4. Pareto-Optimal Architecture Search

PLAS combines empirical loss scaling, roofline-based latency estimation, and Pareto analysis to select hardware-aware LLM architectures. Its analysis identifies precision, sparsity, width, depth, expert count, and FFN expansion choices that shape accuracy–latency trade-offs on edge hardware.

  • Framework: PLAS fits an empirical loss model, estimates inference latency with roofline analysis, and combines both to construct Pareto frontiers for hardware-aware selection.The framework uses trained architectures to approximate validation loss and guides selection under application-specific latency budgets.
  • Scaling Law Fitting: Stable scaling-law exponents across heterogeneous dense and sparse architectures support explicit modeling of architecture-level loss for Pareto-optimal selection.The fitted law remains consistent across depth, width, sparsity, and FFN expansion, with stronger held-out generalization than architecture-specific fitting.
  • Precision–Performance Trade-off: INT8 consistently shifts Pareto frontiers toward lower latency at equivalent loss, but its speedup remains below the theoretical 2× improvement.Higher-precision nonlinear components and quantization/dequantization overhead limit the realized gain.
  • Architecture Patterns: Sparse MoE architectures comprise 100% of Pareto-optimal configurations across latency regimes under batch-one on-device inference.They provide greater total capacity while maintaining comparable activated parameters per token, improving loss-per-FLOP trade-offs over dense counterparts.
  • Architecture Patterns: Pareto-optimal designs favor wide-and-shallow architectures, with depth generally below 20 layers and width delivering more efficient loss reduction under tight latency constraints.Width saturates at the search-space limit before additional capacity is allocated to depth.
  • Architecture Patterns: Optimal expert counts depend on latency objective: prefill favors fewer experts, decode favors more, and balanced total latency typically favors around 8 experts.The difference follows from prefill memory pressure versus decode capacity expansion at batch size one.
  • Architecture Patterns: Pareto-optimal configurations often use FFN expansion ratios below 1× rather than the conventional 4×, reallocating parameters toward experts or model width.This compact FFN design is reported as viable under memory-constrained inference.

5. Theoretical Framework for Hardware-Aware Architecture Optimization

The framework formalizes hardware-aware architecture selection as constrained optimization over loss, latency, memory, and architectural parameters. Its theoretical results derive regime-specific solutions, including width–sparsity scaling and phase-aware design guidance.

  • Framework: The framework replaces exhaustive architecture search with constrained optimization under latency and memory regimes.It formalizes Pareto-optimal design rather than treating it only as an empirical outcome.
  • Latency-constrained regime: Under latency-only constraints, the optimal strategy minimizes MoE activation rate because fewer active experts reduce loss without increasing per-token latency.The paper connects this result to top-1 routing and larger expert pools when memory permits.
  • Memory-constrained regime: Under memory constraints, optimal depth saturates the budget and follows l* ∝ d^-2, establishing a depth–width trade-off.This inverse scaling is reported as consistent with empirical Pareto-frontier behavior.
  • Width–sparsity scaling: Wider models should use sparser MoE, with ρ* ∝ d^-1.19 and approximately 2.3× lower activation rate when width doubles.The paper gives practical examples of ρ≈0.15 for a 2B model with d≈2048 and ρ≈0.25 for a 500M model with d≈1024.
  • Phase-aware tuning: Prefill- and decode-dominant workloads require different FFN and GQA choices because decode latency includes sequence-length-dependent KV-cache coupling.Prefill-oriented models use smaller FFN ratios and larger GQA values, whereas decode-oriented models use larger FFN ratios and balance GQA against KV-cache bandwidth.
  • Practical deployment: Architecture selection falls from months to under one week through theoretical prediction plus small-scale validation.The deployment example considers 10 TOPS compute, 50 GB/s bandwidth, 4 GB memory, and a decode-latency target below 100ms.
  • Scope and limitations: The framework assumes a 170-architecture, 10B-token scaling fit and idealized roofline latency, with real-system deviations of 10–20%.It also assumes standard transformer components; hybrid architectures require re-derived constraints.

6. Conclusion

The paper presents a hardware-aligned framework connecting transformer architecture, model quality, and end-to-end inference efficiency. Evaluating 1,942 architectures reveals Pareto-optimal regimes and structural principles for hardware–model co-design on edge accelerators.

  • The framework connects transformer architecture to model quality and end-to-end inference efficiency through equivalent-parameter scaling and hardware-aware latency modeling.
  • Evaluating 1,942 architectures identifies structural principles and Pareto-optimal regimes governing loss, latency, and roofline efficiency.

A. Architecture Search Space

The architecture search space spans dense and sparse transformer designs across depth, width, MoE configuration, and grouped-query attention settings.

  • The search space covers 4–32 layers and hidden dimensions from 768–3072.
  • MoE configurations range from dense models to 16 experts with Top-1 or Top-2 routing.
  • Grouped-query attention settings are included alongside depth, width, and MoE variables.

B. Pre-training Details

The study trains all configurations under controlled data, optimization, and evaluation procedures to support fair comparison of architectural scaling behavior.

  • Each model is trained for 10B tokens using a mixture of general-corpus, mathematics, and code data.
  • Optimization uses AdamW, cosine learning-rate decay from 1 × 10^-4 to 1 × 10^-6, warmup, QK-Norm, and global batch size 256.
  • Performance is measured with held-out validation loss averaged over the final 1,000 optimization steps, plus WikiText-2 perplexity.
  • The full study costs 13,600 GPU-hours on 8 NVIDIA H200 GPUs per configuration for approximately 10 hours.

C. Scaling Law Coefficients

This section defines fitted loss and latency models for architecture exploration, using scaling-law coefficients and roofline analysis to evaluate many configurations efficiently.

  • The fitted parametric loss model is reported with concrete coefficients in Table 6.
  • Depth exhibits the strongest scaling, with α_l = 1.63, while width and sparsity are coupled through β1 = −0.33.
  • KV-cache configuration has minimal impact on loss, with α_m = 0.05, but significantly affects inference efficiency.
  • Roofline analysis classifies operators as compute-bound or memory-bound using arithmetic intensity relative to hardware compute-to-bandwidth ratio.
  • The approach enables rapid evaluation of over 50,000 configurations in minutes for large-scale architecture exploration.
  • Latency analysis distinguishes sequence-length, output-length, and batch-size effects across prefill and decode workloads.
  • Optimization targets depend on deployment scenario: decode for interactive applications, prefill for long-context processing, and total latency for balanced workloads.

E. Problem Formulation and Roofline Analysis

The roofline formulation decomposes decoder-only Transformer computation, memory traffic, and latency across prefill and decode, exposing how architecture and workload determine hardware bottlenecks.

  • E. Problem Formulation and Roofline Analysis: The analysis models decoder-only Transformers with per-layer attention, feed-forward, and mixture-of-experts components.
  • E.3.1. Attention Projection Layers: Grouped-Query Attention reduces K/V dimensions and KV-cache storage by the factor gqa.
  • E.3.1. Attention Projection Layers: For gqa = 1, projections require 8BSD^2 FLOPs and 4d^2b_w memory; for gqa = 8, the analysis reports an approximately 44% reduction.
  • E.3.2. Attention Score Computation: Softmax is typically negligible relative to matrix multiplications, contributing O(n_hS^2) versus O(Sd^2) when d≫n_h.
  • E.3.2. Attention Score Computation: Prefill attention scales as O(S^2), whereas decode scales as O(S); KV-cache access remains a critical decode bottleneck.
  • E.3.3. FFN Layers: MoE decouples capacity from computation: E = 8 and K = 2 provide 4× more parameters at the same compute, but increase decode memory pressure.
  • E.4. Per-Layer Coefficient Summary: The roofline model uses FLOPs, memory traffic, peak compute, and memory bandwidth to estimate latency for prefill and decode.
  • F. KKT Conditions: Under latency constraints, computation and bandwidth are invariant to ρ while loss increases with ρ, favoring maximum sparsity within memory limits.

K.3. Derivation of Activation Rate

This derivation characterizes optimal activation rates under prefill and decode constraints, with decode requiring additional treatment because KV-cache access couples the variables.

  • Validity of the derived result requires η_pb_w < 2.
  • Table 19 contrasts prefill and decode, with decode including KV-cache memory access and requiring a quadratic equation for the optimal depth-width ratio ρ*.
  • The coupled solutions r*, gqa*, and d* must be obtained by solving the KKT system.
  • Constraint forms, activation-rate solutions, depth solutions, and coefficient comparisons are organized in Tables 20–23.

L.5. Key Results

The key results describe how optimal sparsity and latency coefficients vary by constraint regime, including scenario independence under memory constraints and asymmetry between prefill and decode.

  • Under a memory constraint, the section derives a width-sparsity scaling result.
  • Memory-constrained ρ* is identical for prefill and decode.
  • Latency-constrained coefficients differ between prefill and decode because their derivatives with respect to r differ.
  • The decode latency constraint includes the KV-cache term 2l̄S_db_kv/gqa, increasing constraint tightness and coupling gqa more strongly.
  • Appendix notation extends the symbols used in the main tables.
Loading 2602.10377v1…