LIA: A Single-GPU LLM Inference Acceleration with Cooperative AMX-Enabled CPU-GPU Computation and CXL Offloading

Reviewer: The Guardian

---

Summary

This paper introduces LIA, a framework for accelerating single-GPU LLM inference by leveraging cooperative computation between a GPU and a modern, AMX-enabled Intel CPU. The central thesis is that recent advancements in CPU matrix multiplication capabilities, specifically Intel's AMX, are significant enough to fundamentally alter the trade-offs in CPU-GPU offloading strategies. The authors propose an offloading algorithm that systematically determines which model sublayers to execute on the CPU versus the GPU to minimize latency or maximize throughput. Additionally, they introduce a CXL-based memory offloading policy to expand capacity for throughput-driven scenarios. The paper presents substantial performance improvements over existing frameworks like FlexGen and a CPU-only baseline (IPEX).

While the paper addresses a timely and important problem, its conclusions rest on a series of questionable methodological choices, most notably the comparison against baselines that are not optimized for the very hardware LIA is designed to exploit. The significant performance claims are therefore a conflation of algorithmic novelty and an underlying hardware advantage that is not fairly accounted for in the baseline comparison.

---

Strengths

1. Timely Exploration of New Hardware: The paper provides a valuable and, to my knowledge, one of the first comprehensive microbenchmark analyses of Intel's AMX performance for realistic LLM workloads (GEMM and GEMV). The characterization in Section 4 and Figure 5 is an important contribution to the community, quantifying the substantial leap from AVX512 and situating AMX performance relative to several generations of NVIDIA GPUs.

2. Systematic Offloading Policy: The formulation of the compute-offloading decision as an optimization problem (Section 5.1) is methodologically sound. The cost model, while simple, holistically considers key system parameters like compute throughput and memory/PCIe bandwidths to derive an optimal policy, moving beyond the heuristic-based approaches of prior work.

3. Well-Reasoned CXL Integration: The memory-offloading policy for CXL is logical and well-motivated. The authors correctly identify that for parameter transfers to the GPU, CXL bandwidth can be sufficient to hide behind the PCIe bottleneck (Observation-1, Section 6), while simultaneously recognizing that CXL's high latency is detrimental to CPU-bound computation, leading to their policy of keeping the KV cache in DDR.

---

Weaknesses

1. Fundamentally Unfair Baseline Comparison: The central weakness of this paper is the evaluation against inappropriate baselines. The primary competitor, FlexGen, is designed for CPUs with AVX instruction sets. The authors evaluate it on a Sapphire Rapids CPU but do not appear to have modified it to utilize AMX. Consequently, the reported speedups of up to 5.1x-19x are not solely attributable to LIA's "cooperative framework" but are heavily skewed by comparing an AMX-native framework to an AVX-based one. The paper fails to disentangle the gains from their scheduling algorithm versus the raw hardware advantage. A fair comparison would require an AMX-enabled version of FlexGen or another strong baseline that also leverages AMX. The comparison to IPEX (CPU-only) is a strawman for any model that could even partially fit on a GPU.

2. Over-reliance on Analytical Models and Simulation: A significant portion of the paper's results, particularly for large batch sizes and throughput-oriented scenarios, are derived from an analytical model rather than direct measurement (as indicated by stars in Figure 11, page 10). The authors state this model has an "average error of 12%" (Section 7, page 8), which is a non-trivial margin of error that casts doubt on the precision of the claimed speedups. Furthermore, the multi-GPU comparison in Section 7.8 relies on comparing their system (partially evaluated with their own model) against a simulation of a DGX-A100 system. Conclusions drawn from comparing a model to a simulation are speculative at best.

3. Unsupported Claims of Generalizability: In Section 7.7, the authors claim LIA's optimizations generalize to other model architectures like GPT, Llama2, and Bloom. This claim is substantiated only by results from their analytical model, not empirical evidence. Performance characteristics can vary significantly across model families due to differences in layer normalization, activation functions, and attention mechanisms. Without measured data on real hardware for at least one other major model family, this claim is unsubstantiated and constitutes an overstatement of the work's demonstrated contributions.

4. Ambiguity Regarding Software Stack Maturity: Footnote 4 on page 5 states that "the recently-introduced AMX libraries are less optimized." This is a critical caveat that is buried and not discussed in the main text. It raises serious questions about the replicability and robustness of the foundational microbenchmarks in Figure 5. If the libraries are immature, the reported performance may not be representative of the hardware's true capability, or conversely, the gains over AVX512 could be even larger. This uncertainty undermines the foundation upon which the entire paper is built.

---

Questions to Address In Rebuttal

1. Please justify the choice of an AVX-based FlexGen as the state-of-the-art baseline on an AMX-enabled processor. Did you make any attempt to build a stronger, AMX-aware baseline to ensure a fair comparison? If not, how can you deconvolve the performance gains of your algorithm from the underlying hardware performance differential between AMX and AVX?

2. The paper relies heavily on an analytical model for many key results. Can you provide a sensitivity analysis for this model? How do the model's predictions change with variations in its core assumptions (e.g., PCIe bandwidth, CPU memory latency)? Why were direct measurements not feasible for the starred configurations in Figure 11?

3. To support the claim of generalizability (Section 7.7), please provide measured, end-to-end performance data for LIA on at least one non-OPT model, such as Llama2-70B, and compare it against FlexGen on the same model.

4. Regarding Footnote 4, please elaborate on the maturity of the AMX software stack used in your evaluation. How might the presented results and the optimal offloading policies change as these libraries mature and AMX performance presumably improves? Does this not risk making your derived "optimal" policies obsolete?

Reviewer Persona: The Synthesizer (Contextual Analyst)

---

Summary

This paper presents LIA, a framework for accelerating Large Language Model (LLM) inference on a single GPU by leveraging cooperative computation with modern, powerful CPUs. The work is built upon a crucial and timely insight: the introduction of specialized matrix multiplication hardware in recent Intel CPUs (Advanced Matrix Extensions, or AMX) has fundamentally altered the performance landscape, making the CPU a viable computational partner rather than just a passive memory host.

The authors first provide a rigorous performance characterization of AMX-enabled processors (Sapphire Rapids and Granite Rapids), demonstrating that their matrix math throughput is competitive with some GPUs and a significant fraction of modern ones. Building on this, LIA introduces a systematic compute-offloading algorithm that partitions LLM sublayers between the CPU and GPU to optimize for end-to-end latency or throughput. Finally, the paper explores the use of CXL memory to cost-effectively expand system memory for large-batch, throughput-driven scenarios, proposing a simple but effective memory tiering policy. The experimental results show substantial improvements in latency, throughput, and energy efficiency over existing CPU-only and CPU-GPU collaborative frameworks.

Strengths

1. Central Thesis is Timely and Insightful: The core contribution of this work lies in identifying and exploiting a major shift in the hardware landscape. For years, the systems community has treated CPUs in ML workloads as slow, general-purpose engines, primarily useful for their large memory capacity. This paper compellingly argues that this assumption is now obsolete due to accelerators like AMX. The microbenchmarks presented in Section 4 (pages 4-5) provide the crucial evidence for this claim and serve as a strong foundation for the entire paper. This is an excellent example of systems research that is deeply connected to hardware evolution.

2. Holistic and Well-Designed System: LIA is more than just a proof-of-concept for AMX. The authors have designed a complete system that addresses the problem holistically. The formulation of the offloading decision as an optimization problem (Section 5.1, page 6) is principled and provides a generalizable framework. The inclusion of system-level optimizations like efficient GPU memory usage and overlapping communication (Section 5.2, page 6) shows attention to practical implementation details. The integration of CXL (Section 6, page 7) is forward-looking and addresses the next major bottleneck: memory capacity for large-batch inference.

3. Connects Disparate Hardware Trends: A key strength of this work is its ability to synthesize two independent but complementary hardware trends: the rise of on-CPU acceleration (AMX) and the emergence of disaggregated, tiered memory (CXL). The paper demonstrates a synergistic relationship where AMX makes the CPU a more valuable compute resource, and CXL provides the memory capacity needed to feed it in high-throughput scenarios. This provides a powerful architectural blueprint for future cost-effective inference servers.

4. Significant and Well-Documented Performance Gains: The results are not merely incremental. Achieving up to 5.1x higher throughput and 19x lower latency compared to the state-of-the-art single-GPU offloading framework (FlexGen) is a very strong result. This clearly validates that rethinking the role of the CPU is not just an academic exercise but a source of major real-world performance improvements.

Weaknesses

While the work is strong, its framing could be broadened to better contextualize its contributions within the wider landscape of heterogeneous computing.

1. Implicit Dependency on High-End CPUs: The paper's success hinges on the availability of top-tier Intel Xeon processors with a high core count. While this is the correct platform to demonstrate the maximum potential, it leaves open the question of where the "break-even" point is. A discussion on the performance sensitivity to CPU core count and AMX capabilities would help readers understand the cost-performance trade-offs more broadly. The current work presents an architecture that pairs a high-end GPU with a high-end CPU, which is still a significant investment.

2. Limited Discussion on Architectural Generality: The paper is naturally focused on the Intel AMX + NVIDIA GPU ecosystem. However, the core idea of a powerful CPU partner is not unique to Intel. AMD is developing similar capabilities, and the ARM world has its own vector and matrix extensions (SVE, SME). The discussion of Grace Hopper in Section 8 (page 12) is a good start, but the paper would be even more impactful if it framed its optimization framework (Equation 1, page 6) as a general model for heterogeneous systems, where AMX is just one instance of a powerful host processor.

3. Static Memory Tiering Policy: The proposed CXL policy—storing model parameters on CXL and intermediate values (like the KV cache) on DDR—is pragmatic and effective. However, it is a static policy. This is a missed opportunity to discuss the potential for more dynamic, access-pattern-aware data placement, which is a rich area for future research that this work directly enables.

Questions to Address In Rebuttal

1. Could the authors comment on the sensitivity of LIA's performance to the CPU's core count? The experiments use high-end 40-core and 128-core CPUs. How would the optimal offloading policy and resulting performance change on a lower-end, say 16-core, AMX-enabled server CPU? This is critical for understanding the true cost-efficiency of this approach.

2. The optimization framework in Section 5.1 is based on performance parameters of the specific CPU and GPU. How portable is this framework conceptually? If one were to build a similar system using an AMD CPU with its own on-chip AI engine, would the framework apply directly with only new performance measurements, or would it require fundamental changes?

3. The paper makes a strong case for cost-efficiency against multi-GPU setups. Could you provide a brief TCO (Total Cost of Ownership) sketch? A modern H100 GPU and a Granite Rapids CPU still represent a very expensive server. How does the estimated cost per token served compare to, for example, a server with two last-generation A100 GPUs, which might have a similar acquisition cost?

Review Form: The Innovator (Novelty Specialist)

---

Summary

This paper introduces LIA, a framework for accelerating single-GPU LLM inference for models that exceed GPU memory. The authors' contribution is composed of three primary elements. First, they provide a timely performance characterization of Intel's Advanced Matrix Extensions (AMX) on recent CPUs, establishing that modern CPUs are computationally powerful enough to be viable partners in LLM inference, not just data hosts. Second, leveraging this insight, they propose a novel, systematic compute-offloading algorithm. This algorithm formulates the decision of which Transformer sub-layers to execute on the CPU versus the GPU as an optimization problem, creating a dynamic policy based on batch size and sequence length. This contrasts with prior art that used static, heuristic-based offloading. Third, the paper introduces a specific memory-offloading policy for systems with both DDR and CXL memory, proposing to store model parameters on CXL while keeping latency-sensitive intermediate data (like the KV cache) on DDR to mitigate the performance penalty of CXL on CPU-side computation.

---

Strengths

The primary strength of this paper lies in its novel reframing of the CPU's role in CPU-GPU cooperative inference. My evaluation identifies the following genuinely new contributions:

1. Shift from Static to Dynamic Offloading Policy: The central novel idea is the move away from the rigid offloading policies seen in prior work. Frameworks like FlexGen [43] and FastDecode [23] identified a single, structurally-determined "best" candidate for offloading (the attention scoring sublayer) based on a general principle (low arithmetic intensity). LIA's contribution is the formulation of this choice as a formal optimization problem (Equation 1, page 6). The policy vector p and the cost model that considers system parameters, B, and L to find an optimal configuration is, to my knowledge, a new and more sophisticated approach for this problem domain. It correctly intuits that the "best" sublayer to offload is not fixed, but is a function of the workload.

2. A Motivated, Non-Trivial CXL Memory Policy: The use of CXL for memory expansion is not in itself novel. However, the proposed policy is not a naive extension. The novelty lies in the insight derived from "Observation-2" (page 8), which correctly identifies that the CPU's own compute performance is severely hampered by CXL's higher latency. The resulting policy—store parameters on CXL (primarily for GPU transfer via DMA, which is less latency-sensitive) and keep CPU-accessed data (KV cache) on DDR—is a simple, elegant, and novel solution tailored specifically to the performance characteristics of their cooperative computing model. This is a clear advancement over simply treating CXL as a generic memory pool.

3. Foundational Insight on Modern CPU Capability: While performance studies are common, the specific analysis in Section 4 (page 4) serves as the novel insight that motivates the entire work. Prior art has operated under the assumption that CPUs are orders of magnitude slower than GPUs, thus limiting offloading to only the most trivial compute tasks. By demonstrating that AMX-enabled CPUs achieve throughput that is competitive with previous-generation GPUs (e.g., SPR-AMX vs. P100/V100 for certain workloads in Figure 5, page 5), the authors provide the foundational evidence needed to justify their more ambitious and dynamic offloading algorithm. This re-evaluation of the hardware baseline is a key part of the paper's novelty.

---

Weaknesses

While the core ideas are novel in their application, their conceptual origins are not without precedent in adjacent fields. The paper's claims could be strengthened by acknowledging this and positioning the work more precisely.

1. Framing of the Optimization Model: The concept of using a cost model to partition a workload across heterogeneous processors is a classic problem in systems research. The authors present their algorithm in Section 5.1 (page 6) as a core contribution, which it is. However, the novelty is not in the invention of cost-based scheduling, but in its specific formulation for the unique data flow and computational characteristics of Transformer sub-layers. The paper would be more intellectually honest if it contextualized its model within the broader history of heterogeneous scheduling to more sharply define its specific contribution.

2. Simplicity of the CXL Policy: The CXL policy, while effective and well-motivated, is ultimately a binary, static decision. It represents a single, clever data placement rule rather than a comprehensive memory management framework. This is not a fatal flaw, but it does limit the scope of the novelty. The contribution is a specific, useful heuristic, not a general theory for tiered memory management in LLM inference.

---

Questions to Address In Rebuttal

1. On the Novelty of the Cost Model: Please explicitly differentiate the novelty of your formulation in Section 5.1 (page 6) from the general, well-established field of cost-based task scheduling on heterogeneous systems. What specific characteristics of the LLM inference problem make your model a non-trivial application of these classic ideas?

2. On the Generalizability of the Offloading Policy: The optimization is performed over a fixed set of six sub-layers derived from the OPT architecture. The novelty of the framework appears tied to this specific structure. How does your method generalize to architectures with fundamentally different layer structures, such as Mixture-of-Experts (MoE) models where the feed-forward network is replaced by a sparse routing mechanism? Is the novel contribution the specific model for OPT, or a more general meta-framework for constructing such models?

3. On the Limits of the CXL Policy: The proposed CXL policy is a static partitioning of data types (parameters vs. KV cache). Did the authors consider or evaluate more dynamic policies? For example, in a scenario with an extremely long context length, the KV cache itself could become enormous. Would a tiered policy, where older KV cache tokens are migrated from DDR to CXL, be feasible? What are the conceptual boundaries of your proposed policy's novelty?