WSC-LLM: Efficient LLM Service and Architecture Co-exploration for Wafer-scale Chips

Reviewer: The Guardian

---

Summary

The authors present WSC-LLM, a co-exploration framework designed to optimize Large Language Model (LLM) serving on wafer-scale chip architectures. The framework proposes to jointly explore hardware parameters (such as DRAM capacity and interconnect bandwidth) and software scheduling strategies (such as resource partitioning, placement, and memory management for prefill and decode phases). The core contributions include a Central Scheduler for resource allocation, a Memory Scheduler for KV cache management, and an evaluation methodology based on an extended version of the ASTRA-sim simulator. The paper claims that a wafer-scale architecture with moderate DRAM capacity yields the best performance and that their proposed framework significantly outperforms a state-of-the-art GPU-based serving system (Splitwise) by an average of 3.12x.

Strengths

1. The paper addresses a timely and significant problem at the intersection of LLM serving and next-generation hardware design. The challenges posed by LLM inference on wafer-scale systems are non-trivial, and the authors correctly identify key trade-offs.

2. The conceptual separation of the Central Scheduler and Memory Scheduler is logical, addressing distinct optimization problems (computation/placement vs. data management) that arise in disaggregated LLM inference.

3. The ablation study presented in Section 5.4 provides some insight into the relative contributions of the proposed scheduling components, suggesting that memory management becomes increasingly critical for larger models.

Weaknesses

The paper's conclusions are built upon a methodological foundation that lacks the rigor required for this venue. The claims, while significant, are not substantiated by the evidence provided.

1. Critically Flawed Evaluation Methodology: The entire quantitative analysis hinges on a simulator described in Section 4.6. The authors state they extend ASTRA-sim and, more critically, use a DNN to create a "mapping lookup table" to estimate performance metrics. The paper provides zero validation for this performance model. There is no information on the model's accuracy, error bounds, or how it was trained and tested against a ground truth (e.g., a cycle-accurate model or hardware measurement). The authors merely cite that other works [37, 89] have found this approach feasible. This is insufficient. Without this validation, every performance number, graph, and conclusion in Section 5 is speculative and cannot be trusted.

2. Confounded and Misleading Baseline Comparison: The central claim of a 3.12x performance improvement (Section 5.3, Figure 11) is derived from comparing the proposed simulated wafer-scale system to a real-world A100 GPU cluster. This is a fundamentally confounded comparison. The simulated wafer-scale chip is described with a total interconnect bandwidth of 6 TB/s per die (Section 5.1.1), orders of magnitude higher than the 400 GB/s inter-node bandwidth of the GPU cluster. The reported performance gains are almost certainly dominated by these vastly superior, and perhaps optimistic, hardware assumptions rather than the novelty of the scheduling framework itself. The paper fails to isolate the contribution of its scheduling algorithms from the contribution of the hypothetical hardware, making the headline claim unsubstantiated. The SW-Wafer experiment is a step in the right direction but is insufficient to fully de-confound these factors.

3. Overstated "Co-Exploration" Scope: The paper frames itself as an "architecture co-exploration framework." However, the actual exploration in Section 5.2 is a simple parameter sweep across just four pre-selected hardware configurations (Table 1). The fundamental architectural choices—a 2D-mesh topology, Dojo-style compute cores, and a specific die structure—are fixed. This is not a general exploration of the architectural design space but rather a tuning of a few parameters within a highly constrained template. The claims of generality are therefore not supported by the experiments.

4. Unjustified Heuristics and Parameters: The resource placement strategy in Section 4.2.2 relies on minimizing a TransferCost function which includes a hyperparameter α. The paper provides no details on how α is determined, its sensitivity, or its impact on the final placement. This introduces an element of arbitrariness into a key part of the methodology, potentially suggesting the results are cherry-picked based on a favorable but unexplained tuning.

Questions to Address In Rebuttal

The authors must address the following critical points to establish the validity of their work:

1. Simulator Validation: Please provide rigorous validation data for the DNN-based performance model described in Section 4.6. What is its prediction accuracy (e.g., Mean Absolute Percentage Error) against a ground truth for latency, DRAM access, and communication overhead? Without this, the results are not credible.

2. De-confounding Performance Gains: How can the authors de-confound the performance gains attributed to their scheduling algorithm from the gains stemming from the assumed hardware's massive bandwidth advantage? A more convincing experiment would be to implement a simplified version of the WSC-LLM scheduler on the GPU baseline to demonstrate its benefits on a fixed hardware platform.

3. Hyperparameter Justification: Please provide a justification and sensitivity analysis for the hyperparameter α used in the resource placement strategy (Section 4.2.2). How does the system's performance change with different values of α, and how was the value used in the experiments chosen?

4. Clarification of Constraints: Algorithm 1 (Section 4.2.1) iterates through "all feasible instance sizes." What are the constraints that determine feasibility? The text mentions a die limit and memory capacity, but is the search space truly exhaustive, and what is the typical runtime of this "offline" algorithm for the configurations tested?

WSC-LLM: Efficient LLM Service and Architecture Co-exploration for Wafer-scale Chips

Review Form: The Synthesizer

---

Summary

This paper presents WSC-LLM, a novel co-exploration framework designed to optimize both the architecture of wafer-scale chips and the scheduling of Large Language Model (LLM) inference workloads. The authors correctly identify that the unique characteristics of wafer-scale hardware—namely, vast on-wafer bandwidth but a fixed area that forces trade-offs between compute, memory, and communication resources—create a complex, coupled optimization problem. The core contribution is a holistic framework that systematically explores this design space. It features a sophisticated Central Scheduler that partitions and places prefill/decode instances with awareness of the 2D-mesh topology, and a novel Memory Scheduler that leverages high Die-to-Die (D2D) bandwidth to manage KV cache storage across the entire wafer. Through a comprehensive design space exploration and comparison against state-of-the-art LLM serving systems, the paper demonstrates that a balanced wafer-scale architecture with moderate DRAM capacity yields the best performance, and that its scheduling strategies significantly outperform existing methods.

---

Strengths

The true strength of this paper lies in its ambitious and well-executed synthesis of two traditionally separate domains: computer architecture and distributed systems software.

1. Excellent Problem Formulation and Contextualization: The authors have done an outstanding job articulating the core tensions in designing wafer-scale systems for LLMs. Figure 1 (page 2) provides a remarkably clear and concise visualization of the fundamental architectural trade-off (DRAM capacity vs. DRAM costs) and the key inefficiencies in existing disaggregated scheduling approaches. This demonstrates a deep understanding of the problem that goes beyond a superficial application of known techniques. The work correctly positions itself at the intersection of wafer-scale computing (e.g., Cerebras, Tesla Dojo), chiplet integration, and disaggregated LLM inference systems (e.g., Splitwise).

2. Holistic Co-exploration Framework: The central idea of a unified framework for co-exploration is the paper's most significant contribution. Instead of proposing a fixed hardware architecture and then designing a scheduler for it, WSC-LLM provides a methodology for discovering a near-optimal hardware/software pairing. This is a crucial step forward for the field. As we move into an era of specialized hardware, such co-design frameworks will become indispensable, and this paper provides a strong early example in a highly relevant domain.

3. Actionable Architectural Insights: The study produces results that are not merely academic but provide genuine guidance for hardware architects. The conclusion in Section 5.2 (page 10) that a moderate DRAM capacity (Case 3) is superior to both lower (Cases 1, 2) and higher (Case 4) capacity designs is a powerful, non-obvious finding. It beautifully illustrates the law of diminishing returns, showing that beyond a certain point, the D2D bandwidth lost to accommodating more DRAM becomes the primary performance bottleneck. This is a critical insight for the future design of AI accelerators.

4. Novelty in Scheduling for Wafer-Scale Topologies: The paper successfully adapts and extends concepts from disaggregated inference to the unique 2D-mesh topology of wafer-scale systems. The topology-aware placement strategy (Section 4.2.2, page 6) and the wafer-wide Memory Scheduler (Section 4.4, page 7) are key innovations that directly exploit the hardware's strengths (high D2D bandwidth) to mitigate its constraints (communication locality). The ablation study in Section 5.4 (page 11) convincingly shows that the Memory Scheduler's contribution becomes increasingly vital for larger models, which is another significant finding.

---

Weaknesses

While the core ideas are strong, the paper could be improved by addressing the implications of some of its methodological abstractions and exploring the boundaries of its proposed solution.

1. Fidelity of the Performance Evaluator: The entire co-exploration process hinges on the accuracy of the Evaluator module (Section 4.6, page 9), which uses a DNN to predict performance and avoid full, slow simulations. While this is a standard and necessary technique for tractable design space exploration, its implications are not fully discussed. The validity of the paper's central claims rests on this model's fidelity. A brief discussion of the model's validation, its potential sources of error (e.g., modeling network contention), and the sensitivity of the final architectural choice to this error would significantly strengthen the work's foundations.

2. Static Nature of Resource Allocation: The resource partitioning and placement strategies (Section 4.2.1 and 4.2.2) are performed offline. This is a reasonable simplification for demonstrating the framework's potential. However, real-world serving environments are highly dynamic, with fluctuating request loads and potentially evolving model popularity. The current framework does not address how the system might adapt to such long-term changes, which could lead to suboptimal static partitioning. Acknowledging this as a scope limitation and an avenue for future work (e.g., dynamic re-partitioning) would provide a more complete picture.

3. Limited Scope of Explored Architectural Parameters: The design space exploration in Section 5.1.2 (page 9) focuses primarily on the number of DRAM chiplets, which in turn affects DRAM capacity/bandwidth and D2D bandwidth. While this is arguably the most critical trade-off, a true co-exploration could extend to other parameters like the size and number of compute cores per die, the on-die NoC bandwidth, or the SRAM capacity per core. The current work serves as an excellent proof-of-concept, but its conclusions are conditioned on the fixed compute die design.

---

Questions to Address In Rebuttal

1. On the Evaluator's Accuracy: Could the authors provide more detail on the validation of the DNN-based performance model used in the Evaluator (Section 4.6)? Specifically, what is the reported prediction error against a cycle-accurate or full-system simulation for key metrics like latency and throughput, and how does this potential error margin impact the confidence in identifying Case 3 as the definitive optimal architecture?

2. On the Generality of the Optimal Architecture: The conclusion that the balanced Case 3 architecture is optimal is a key result. How sensitive is this finding to the workload characteristics? For example, if faced with a workload composed exclusively of very long-context requests (e.g., summarizing entire books), which are heavily prefill-bound and generate massive KV caches, would the optimal point in the design space shift towards Case 4 (higher DRAM bandwidth/capacity)?

3. On Network Contention from the Memory Scheduler: The Memory Scheduler (Section 4.4) is a powerful concept that leverages the entire wafer's DRAM pool for KV cache. As multiple, distributed decoding instances access these remote KV cache blocks, this could create "hotspots" or significant contention on the 2D-mesh network. Is this cross-instance network traffic and potential congestion fully modeled within the Evaluator? And how does the system arbitrate or manage this contention in practice?

Title: WSC-LLM: Efficient LLM Service and Architecture Co-exploration for Wafer-scale Chips
Reviewer: The Innovator

---

Summary

The authors present WSC-LLM, a framework designed to co-explore hardware architecture parameters and software scheduling strategies for serving Large Language Models (LLMs) on wafer-scale chips. The framework aims to navigate the complex trade-offs between computation, memory capacity, and communication bandwidth inherent in such systems. It proposes strategies for resource partitioning, placement, and memory management tailored to the 2D-mesh topology of wafer-scale systems. The core claim of novelty rests on this co-exploration framework and its constituent algorithms for optimizing disaggregated LLM inference in this specific hardware context.

My analysis concludes that while the paper addresses a timely problem with a comprehensive engineering effort, its foundational novelty is limited. The primary contribution is the application-specific synthesis and refinement of existing concepts from distributed systems, scheduling, and hardware/software co-design, rather than the introduction of a fundamentally new algorithmic or architectural paradigm. The most notable novel component is a greedy algorithm for remote KV cache placement.

Strengths

1. Problem-Specific Algorithmic Contribution: The memory scheduling algorithm for KV cache placement (Algorithm 2, Page 8) is a concrete and novel contribution. While the concept of memory disaggregation is not new, this algorithm provides a specific, topology-aware greedy strategy to utilize distributed idle DRAM for KV cache, which is a key challenge in disaggregated LLM serving.

2. Holistic System Integration: The framework's strength lies in its comprehensive integration of multiple optimization layers (architectural DSE, instance partitioning, physical placement, memory management) into a single toolchain. While the concept of a co-design framework is not novel, creating a functional one for this complex domain is a non-trivial engineering achievement.

3. Context-Specific Heuristic: The "decoding-centered" resource placement strategy (Section 4.2.2, Page 6) is a simple and logical heuristic tailored to the producer-consumer dataflow of prefill and decoding phases on a 2D mesh. It is a novel, albeit incremental, heuristic for this specific problem.

Weaknesses

1. Limited Conceptual Novelty of the Core Framework: The central claim of a "co-exploration framework" is an overstatement of novelty. The field of hardware/software co-design has long used automated Design Space Exploration (DSE) to find optimal system configurations. WSC-LLM appears to be a well-executed but conceptually standard DSE framework applied to a new and important domain. It does not introduce a new theory or methodology for co-design itself.

2. Heavy Reliance on Assembling Prior Art: The framework's scheduling engine is built upon a foundation of well-established techniques from recent LLM serving literature, which the authors correctly cite. These include:

   • Disaggregated Inference: The core idea of separating prefill and decoding resources was popularized by systems like Splitwise [62].
   • Continuous Batching: A key technique from systems like Orca [91] and vLLM [45].
   • Chunked Prefill: A strategy to handle long prompts, proposed in works like DeepSpeed-FastGen [10].
     The novelty in WSC-LLM is the integration of these techniques for a wafer-scale target, not the techniques themselves.

3. Algorithmic Contributions are Largely Heuristic Search: The "Optimal Resource Partition Algorithm" (Algorithm 1, Page 7) is essentially a structured grid search over instance sizes and pre-defined Tensor Parallelism (TP) strategies. While systematic, this is a standard methodology for performance exploration and does not represent a novel algorithmic paradigm for optimization. It formalizes an exhaustive search within a constrained space.

Questions to Address In Rebuttal

1. The term "co-exploration framework" suggests a novel methodology. Can the authors precisely articulate what makes the WSC-LLM framework conceptually different from standard hardware/software Design Space Exploration (DSE) frameworks, beyond its application to the specific domain of wafer-scale LLM inference?

2. The decoding-centered placement strategy (Section 4.2.2) is a key contribution for handling the 2D-mesh topology. This producer-consumer placement problem (prefill instances producing KV cache for decoding consumers) is analogous to problems in NoC design and general parallel task mapping. Can the authors contrast their heuristic with prior work on topology-aware task placement for producer-consumer patterns and clarify the novelty?

3. The KV cache placement algorithm (Algorithm 2) is a greedy strategy. While this is the paper's strongest novel component, its simplicity warrants discussion. Could you elaborate on the conditions under which this greedy approach is sufficient and identify scenarios (e.g., highly fragmented memory, complex communication patterns) where it might lead to suboptimal placements compared to more complex graph-based or ILP formulations?

4. If you had to isolate the single most significant and novel conceptual contribution of this work, what would it be? Is it the framework itself, one of the specific algorithms, or the insight that moderate DRAM is optimal? Please be specific.