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

Motivation

• On-device LLMs must balance accuracy against strict memory, bandwidth, power, and latency constraints, making cloud-optimized architectures often infeasible at the edge [4] [5] [6].This deployment tension motivates hardware–software co-design in embodied AI systems, including autonomous vehicles and mobile robots [1] [3].Source: §1 · p. 4

Hardware constraints

• Transformer inference spans bandwidth-bound attention, compute-bound feedforward layers, and KV-cache memory pressure, while depth and width scaling can shift execution into memory-bandwidth-limited regimes [7] [8] [11] [12].Under batch-1 edge inference, limited weight reuse and insufficient on-chip cache make repeated off-chip parameter loads dominate latency; Figure 2(a) illustrates the relevant roofline limits.Source: §1 · p. 4, Figure 2 · p. 4

Architecture search

• Because conventional NAS often optimizes a single objective under loose budgets [13], the paper uses Pareto optimization to select non-dominated architectures balancing accuracy and latency.This avoids exhaustive enumeration while supporting model selection for specific deployment constraints [7]. Figure 2(c) visualizes the Pareto frontier.Source: §1 · pp. 4–5, Figure 2 · p. 4

Contributions

• The paper develops a hardware co-design law combining loss scaling with roofline-based latency modelling to characterize accuracy–latency Pareto trade-offs under fixed hardware constraints.It frames architecture search as joint optimization over precision and performance, deriving feasible design regions under industrial hardware and application budgets.Source: §1 · p. 5

Empirical framework

• The study benchmarks approximately 1,942 architectures on NVIDIA Jetson Orin, selects 170 models, and trains each for 10B tokens to fit and validate the empirical loss-scaling and accuracy–latency Pareto structure.The methodology reduces architecture selection from months to a few days, with code, checkpoints, trained models, and evaluation protocols planned for release.Source: §1 · p. 5

2.1. Efficient LLM Architectures

• Efficient LLM architectures use sparse activation, cache reduction, sub-quadratic attention, hybrid SSM-attention designs, and gated attention to shift inference bottlenecks toward memory or routing.These approaches include MoE routing [23] [24], MLA [8], sliding-window attention [25], linear-complexity mechanisms [26], hybrid architectures [28], and gated attention [29] [30].Source: §2.1 · p. 6

2.2. On-Device LLM Deployment

• On-device LLM deployment co-optimizes models, compression, and systems, using efficient small models, quantization, inference engines, shared embeddings, and sparsity-aware mechanisms.Qwen [31], MiniCPM [32], and SmolLM use data-efficient training, while AWQ, GPTQ [35], vLLM [36], MLC-LLM, and PowerInfer [38] support heterogeneous hardware execution.Source: §2.2 · p. 6

2.3. Hardware-Aware Architecture Optimization

• NAS faces prohibitive search costs, weak latency and memory integration, and limited interpretability, while profiling and depth–width studies remain disconnected or inconsistent under deployment constraints.LLM-Viewer and LLMCompass provide hardware-aware analysis without integrating architecture search, and depth–width studies [42] [11] often neglect deployment constraints.Source: §2.3 · p. 6
• This work extends roofline analysis to map architectural choices, MoE sparsity, attention variants, and KV-cache strategies to compute- or bandwidth-limited edge-device regimes.The resulting architecture search is guided by a hardware co-design law rather than post-hoc benchmarking.Source: §2.3 · p. 6

3.1. Implicit Optimization Objective of Hardware Co-Design Law

• The hardware co-design objective selects an architecture θ* minimizing validation loss subject to fixed inference latency, memory, and precision constraints.The architecture is θ = (l, d, d_m, r, ρ), covering depth, width, KV-cache dimension, FFN expansion, and MoE activation rate.Source: §3 · p. 6, §3.1 · p. 6
• The architecture variables govern distinct efficiency mechanisms: KV-cache dimension controls memory and bandwidth, FFN expansion controls per-token compute, and MoE activation determines sparsity.For MoE models, r is defined as the total expansion ratio across activated experts so FFN compute remains comparable across sparsity levels.Source: §3.1 · p. 7
• Latency and memory surrogates depend jointly on architecture, hardware peak throughput and bandwidth, and workload batch size and input/output sequence lengths.Roofline behavior determines whether inference is compute- or bandwidth-bound, while workload configuration affects attention complexity and KV-cache access patterns.Source: §3.1 · p. 7
• Because exhaustive architecture enumeration is infeasible, the method replaces brute-force optimization with explicit surrogates for learning dynamics and system behavior.A fitted polynomial loss model and a roofline-based latency model enable scalable identification of Pareto-optimal architectures under fixed hardware constraints.Source: §3.1 · p. 7

3.2. Precision Modeling via Loss

• The precision model fits validation loss as a separable polynomial function of architecture, using 170 dense and MoE architectures to predict loss across the design space.The formulation captures different width-scaling exponents for sparsity-driven and base-capacity terms and reparameterizes KV-cache dimension through the GQA group ratio.Source: §3.2 · p. 7

3.3. Performance Modeling via Latency

• The performance model uses roofline analysis, expressing latency through arithmetic intensity and hardware peak compute throughput and sustained memory bandwidth.The resulting end-to-end formulation covers both context prefill and autoregressive decoding, while the current model targets conventional memory-centric accelerators.Source: §3.3 · p. 8

3.4. Pareto-optimal architectures

• Combining the loss and latency models defines Pareto-optimal architectures as those for which no alternative improves both objectives, yielding a hardware co-design scaling law for deployment under strict budgets.The resulting principle guides architecture selection by trading validation loss against end-to-end inference latency on a given hardware platform.Source: §3.4 · p. 8

4. Pareto-Optimal Architecture Search

• PLAS combines empirical scaling-law loss prediction, roofline-based latency estimation, and Pareto-frontier construction to select architectures under application-specific latency budgets.The framework evaluates loss and latency jointly rather than relying on exhaustive architecture measurement.Source: Figure 3 · p. 9, §4 · p. 9, §4.3 · p. 11

4.1.2. Scaling Law Fitting

• The fitted scaling law achieves training R^2 = 0.975 and validation R^2 = 0.952, supporting predictable architecture-level loss modeling across heterogeneous dense and sparse designs.Its exponents remain stable across depth, width, sparsity, and FFN expansion when training compute, data budget, and optimization protocol are fixed.Source: Figure 4 · p. 10, §4.1.2 · p. 10

4.2. Latency Modeling

• Roofline modeling classifies operators by arithmetic intensity and evaluates 50,000+ configurations in approximately 20 minutes, with top candidates empirically validated for latency ranking.For VLA deployment, the analysis uses batch size 1, 1,024 input tokens, and 16 output tokens; prefill is mainly compute-bound while decode is weight-memory-bound.Source: §4.2.1 · p. 11, §4.2.2 · p. 11

4.3.2. Precision–Performance Trade-off

• INT8 consistently shifts Pareto frontiers toward lower latency at equivalent loss, but gains remain sub-linear because higher-precision nonlinear operations and conversion overhead persist.The relevant optimization target differs across prefill, decode, and total latency, so optimal architectures depend on workload characteristics.Source: §4.2.2 · p. 11, §4.3.2 · p. 12, Figure 5 · p. 12

4.3.4. Architecture Parameter Evolution

• Sparse MoE architectures constitute 100% of Pareto-optimal configurations and are generally wide-and-shallow; optimal expert counts vary by phase, while balanced workloads typically favor around 8 experts.Both phases consistently prefer Top-K=1 routing.Source: §4.3.4 · pp. 13–14

4.3.5. Empirical Validation

• 19.42% lower WikiText-2 perplexity (50.88 vs. 63.14) is achieved by the co-designed architecture at Qwen2.5-0.5B’s measured NVIDIA Jetson Orin latency.Both models use identical data mixtures and optimization protocols, and Figure 8(b) shows consistently lower training loss for the co-designed architecture.Source: Figure 8 · p. 15, §4.3.5 · p. 15

5.1. From Empirical Search to Principled Optimization

• Architecture selection is recast from exhaustive empirical search over 1,942 candidates as explicit constrained optimization with closed-form solutions for different hardware regimes.The formulation explains structural principles governing Pareto-frontier solutions and their potential transfer to new hardware platforms.Source: §5.1 · p. 16

5.2. Problem Formulation and Constraint Types

• Sparsity affects model storage but not per-token computation, creating the asymmetry that drives the theoretical results across prefill, decode, and memory constraints.The constraints jointly account for compute-bound prefill, bandwidth-bound decode with KV-cache access, and storage-bound parameter memory.Source: §5.2 · p. 16

5.3. Optimal Activation Rate Across Constraint Regimes

• The optimal activation rate depends on the regime: maximize sparsity under latency constraints, use ρ* ∝ d^-1.19 under memory constraints, and apply the corresponding formula under dual constraints.For memory-limited devices, a 2B-parameter model with d≈2048 should use ρ≈0.15.Source: §5.3 · pp. 17–18, §5.5.1 · p. 19

5.4. Optimal Depth, FFN Ratio, and GQA

• Optimal depth saturates the active constraint with l* ∝ d^-2, while FFN and GQA optima have regime-dependent coefficients and KV-cache corrections.Prefill-optimized models use smaller FFN ratios and larger GQA values than decode-optimized models at equivalent performance levels.Source: §5.4 · pp. 18–19

5.5.1. Key Structural Insights

• The analysis shows memory-optimal parameters are phase-independent, whereas prefill and decode differ by 2× coefficients and decode additionally depends on sequence length through KV-cache coupling.The width-sparsity law implies that doubling width should reduce activation rate by approximately 2.3×.Source: §5.5.1 · p. 19

5.5.2. Actionable Design Guidelines

• Practitioners should select sparsity from the active constraint, allocate depth after choosing width, and tune FFN/GQA differently for prefill-dominant, decode-dominant, and balanced workloads.On new hardware, measuring compute, bandwidth, and memory enables regime identification and theorem-based prediction followed by 3–5 small-scale validation runs.Source: §5.5.2 · pp. 19–20

5.5.3. Limitations and Future Extensions

• The framework is limited by loss-law extrapolation beyond 170 architectures trained for 10B tokens, idealized roofline assumptions with 10–20% deviations, and standard-transformer scope.Future work targets training dynamics, refined latency models, hybrid architectures, and validation across TPUs and specialized accelerators.Source: §5.5.3 · p. 20

5.6. Summary

• The framework characterizes optimal activation rates, structural parameter asymmetries, and hardware-specific design rules, reducing architecture selection from months to under one week.Its predictions align closely with empirical Pareto frontiers while explaining why different hardware constraints produce distinct optimal architectures.Source: §5.6 · p. 21, §5.5.3 · p. 21

• The framework connects transformer architecture to model quality and end-to-end inference efficiency using equivalent-parameter scaling and hardware-aware latency modelling.Source: §6 · p. 21
• Across 1,942 evaluated architectures, the analysis identifies structural principles and Pareto-optimal regimes governing loss, latency, and roofline efficiency.Source: §6 · p. 21
• The resulting hardware co-design scaling law provides a stable, predictable architecture–hardware relationship for extrapolating optimal designs across deployment regimes.Source: §6 · p. 21

A. Architecture Search Space

• The search spans 4–32 layers, 768–3072 hidden dimensions, dense-to-16-expert MoE configurations with Top-1/Top-2 routing, and grouped-query attention.These hyperparameters define the architecture search space used for scaling-law fitting.Source: A. Architecture Search Space · p. 26
• The design covers both dense and sparse architectures while avoiding degenerate or ill-conditioned regimes.Source: A. Architecture Search Space · p. 26

Training Data

• All 170 configurations use 10B tokens from general, mathematics, and code data, a budget chosen to reveal scaling behavior while remaining computationally tractable.The training corpus is planned for public release.Source: B. Pre-training Details · p. 26

Optimization

• Optimization uses AdamW with β1 = 0.9, β2 = 0.95, weight decay 0.01, cosine learning-rate decay from 1 × 10^-4 to 1 × 10^-6, 0.2% warmup, QK-Norm, and global batch size 256.QK-Norm is applied to stabilize MoE pre-training.Source: B. Pre-training Details · p. 26

Evaluation

• Performance is measured by held-out validation loss averaged over the final 1,000 optimization steps, with WikiText-2 perplexity additionally reported for downstream evaluation.The validation set contains approximately 1B tokens, and final-step averaging reduces variance.Source: B. Pre-training Details · p. 26

Computational Cost

• The full study consumes 13,600 GPU-hours by training each configuration on 8 NVIDIA H200 GPUs for approximately 10 hours.The total is calculated as 170 × 8 × 10 GPU-hours.Source: B. Pre-training Details · p. 26

C. Scaling Law Coefficients

• Depth has the strongest scaling effect on loss, with α_l = 1.63, while FFN expansion contributes modestly at α_r = 0.17.These fitted exponents quantify sensitivity to layer count and FFN expansion ratio.Source: C. Scaling Law Coefficients · p. 27
• The width–sparsity coupling β_1 = −0.33 implies that wider models require greater sparsity, or lower ρ, to remain competitive.The sparsity-related loss penalty grows with width.Source: C. Scaling Law Coefficients · p. 27
• KV-cache configuration has minimal loss impact at α_m = 0.05 but significantly affects inference efficiency.This separates KV-cache choices from the loss-sensitive architectural factors while preserving their importance for latency.Source: C. Scaling Law Coefficients · p. 27

D. Latency Modeling Details

• Roofline modeling classifies operators as compute- or memory-bound using arithmetic intensity relative to hardware’s compute-to-bandwidth ratio.Latency is estimated analytically from FLOPs, peak compute, and memory access.Source: D. Latency Modeling Details · p. 27
• Over 50,000 configurations can be evaluated in minutes, enabling large-scale architecture exploration with the analytical latency model.Top Pareto candidates are empirically measured with vLLM and subprocess isolation to validate GPU memory accounting and optimality.Source: D. Latency Modeling Details · p. 27
• For batch size 1, prefill latency scales as O(S^2), decode latency as O(S), and long generations are dominated by decode.Larger batches make prefill compute-bound while decode remains memory-bound because of weight loading, with diminishing batching returns.Source: D. Latency Modeling Details · p. 27
• Latency optimization targets depend on workload: decode for interactive streaming, prefill for long-context processing, and total latency for balanced summarization.These targets align optimization with the dominant phase in each deployment scenario.Source: D. Latency Modeling Details · p. 28

E.2. Transformer Architecture; E.3.1. Attention Projection Layers

• GQA reduces K/V computation, weight memory, and KV-cache storage by factor gqa, with gqa = 8 yielding approximately 44% lower attention-projection cost.The attention model uses nh query heads and nkv key-value heads, with gqa = nh/nkv and reduced K/V dimension d/gqa.Source: §E.2 · p. 28, §E.3.1 · pp. 29–30

E.3.2. Attention Score Computation

• Prefill attention scales as O(S^2), whereas decode attention scales as O(S); softmax is usually negligible relative to matrix multiplications when d is much larger than nh.Projection FLOPs can dominate attention-score FLOPs, while KV-cache traffic remains O(BSdm) in both phases.Source: §E.3.2 · pp. 30–31

E.3.3. FFN Layers

• MoE decouples capacity from computation: E = 8 and K = 2 provides 4× more parameters at the same compute, but increases decode memory pressure by 1/rho.For effective expansion r, MoE FLOPs match dense FFN costs while all E experts must be stored; FFNs contribute approximately 60–70% of total FLOPs when r ≈ 4.Source: §E.3.3 · pp. 31–32

E.4. Per-Layer Coefficient Summary; E.5. Inference Phases and Roofline Model; E.6.4. Decode Latency: Unified Form

• The roofline formulation distinguishes compute and bandwidth limits using FLOPs F, memory traffic W, peak compute π_H, and bandwidth β_H, yielding phase-specific latency constraints.The normalized decode coefficient is W = 2 + 2gqa + 3r, while sparsity affects storage but not per-token computation.Source: §E.4 · pp. 32–33, §E.5 · p. 33, §E.6.4 · p. 34

E.5. Inference Phases and Roofline Model

• Decode is typically memory-bandwidth-bound at small batch sizes, with KV-cache access becoming a critical bottleneck as sequence length grows.Prefill is typically compute-bound, whereas decode combines weight loading with KV-cache traffic and scales linearly in sequence length for attention scores.Source: §E.3.2 · p. 31, §E.5 · p. 33, §E.6.2 · p. 34

F.4. KKT Conditions; G.1. Problem Statement

• Under latency constraints, optimal MoE design minimizes rho by increasing expert-pool size E while keeping active experts K fixed, because per-token cost is invariant to rho and loss improves as rho decreases.The decode latency constraint satisfies ∂g_T/∂rho = 0, while memory constraints still limit the feasible expert pool.Source: §F.3 · p. 36, §F.4 · p. 37, §G.1 · p. 39

F.6. Solution Summary; G.7. Solution Summary; H.4. Special Case: δ→0; H.5. Remaining Solutions; H.6. Solution Summary

• With both latency and memory constraints active, r*, gqa*, and d* are coupled through the KKT system and must be obtained numerically, while the negligible-KV-cache limit recovers the simplified weight-dominated formula.The activation-rate solutions in the single-constraint cases correspond to Theorems 5.1 and 5.2, and the dual-constrained case corresponds to Theorem 5.3(b).Source: §F.6 · p. 38, §G.7 · p. 41, §H.5 · p. 43

I.6. Comparison with Case D1

• Prefill and decode produce different optimal architectures: prefill favors smaller r and larger GQA groups, whereas decode doubles the r coefficient from 1/6 to 1/3 and halves the GQA coefficient from 4 to 2.These differences reflect compute-bound prefill versus memory-bandwidth-bound decode with KV-cache access.Source: §I.6 · p. 45, Table 18 · p. 45, §I.4 · p. 45

J.2. Equivalence to Case D2

• All KKT conditions and solutions are identical to Case D2.Source: §J.2 · p. 46

J.3. Solution Summary

• All solutions are identical to Case D2.The activation-rate result corresponds to Theorem 5.2 in the main text.Source: §J.3 · p. 47
• The activation-rate result corresponds to Theorem 5.2 in the main text.Source: §J.3 · p. 47

K.1. Problem Statement

• The prefill constraint is compute-bound.It does not involve KV-cache.Source: §K.1 · p. 47
• The prefill constraint does not involve KV-cache.Source: §K.1 · p. 47

K.2. Constraint Compatibility

• Constraint compatibility is analyzed by dividing the two constraints.Source: §K.2 · p. 47

K.3. Derivation of Activation Rate

• ηpbw < 2 is required for validity, corresponding to Theorem 5.3(a) in the main text.Source: §K.3 · p. 48

K.4. Comparison with Case D3

• Prefill optimization depends only on compute ξFd^2 and admits a closed-form ρ∗, whereas decode optimization includes weight loading and KV-cache access and requires solving a quadratic equation.The KV-cache term 2Sdbkvgqa becomes dominant for large sequence lengths.Source: §K.4 · p. 48, Table 19 · p. 48

K.5. Remaining Solutions

• The remaining variables r∗, gqa∗, and d∗ are obtained by solving a coupled KKT system.Source: §K.5 · p. 48

K.6. Solution Summary

• The solution summary defines αattn = 2 + 2/gqa∗ and treats r∗ and gqa∗ as coupled solutions, while reaffirming the activation-rate result’s correspondence to Theorem 5.3(a).Source: §K.6 · p. 48

L.1. Constraint Forms

• The constraint formulations are presented in Table 20.Source: Table 20 · p. 49, §L.1 · p. 49

L.2. Solution Matrix: Activation Rate ρ∗

• The activation-rate solution matrix ρ∗ is presented in Table 21.Source: Table 21 · p. 49, §L.2 · p. 49

L.3. Solution Matrix: Depth l∗

• The depth solution matrix l∗ is presented in Table 22.Source: Table 22 · p. 49, §L.3 · p. 49

L.4. Coefficient Comparison

• The coefficient comparison is presented in Table 23.Source: Table 23 · p. 49, §L.4 · p. 49

L.5. Key Results

• Width-sparsity scaling is characterized under a memory constraint.Source: §L.5 · p. 50
• The memory-constrained ρ∗ is identical for Prefill and Decode, showing scenario independence.Source: §L.5 · p. 50
• Latency-constrained coefficients differ between Prefill and Decode because ∂ξF/∂r = 6 differs from the decode-side derivative.Source: §L.5 · p. 50
• The decode latency constraint includes the KV-cache term 2l̄Sdbkv/gqa, increasing effective constraint tightness and coupling gqa more strongly to the constraint.Source: §L.5 · p. 50
• In the D3 dual-constrained case, the KV-cache contribution leads to a quadratic form.Source: §L.5 · p. 50

L.6. Notation

• Table 24 extends Table 7 with notation used in the Appendices.Source: Table 24 · p. 50, §L.6 · p. 50