Agentic AI for Heterogeneous System Design Across Chip, Scale-Up, and Scale-Out Domains

AI Analysis

Abstract

Key Highlights

• ArchAgent is a hierarchical multi-agent framework for heterogeneous system design, leveraging LLM-based planning and specialist agents.
• A deterministic analytical evaluator provides all performance estimates, ensuring auditability and preventing LLM-generated performance claims.
• ArchAgent achieved a 1.93x weighted geometric mean speedup over a homogeneous GPU baseline for mixed AI+HPC workloads under a 1 MW power budget.
• The framework demonstrated superior search efficiency, discovering high-quality designs in 18 evaluations compared to 100 for hill climbing.
• Discovered heterogeneous designs often feature a large GPU-like partition complemented by smaller, role-specialized dataflow and FP64 partitions.
• The design space explored includes RISC-V accelerators, alongside GPUs and dataflow processors.
• Key performance bottlenecks identified in winning configurations are all-to-all communication and FP64 memory bandwidth.

Technical Details

• ArchAgent employs a hierarchical multi-agent framework with a planner agent that decomposes system design into chip, scale-up, and scale-out subproblems, handled by specialist sub-agents.
• System designs are represented as JSON documents specifying partitions, chip archetypes (including RISC-V accelerators like Tenstorrent's Blackhole), scale-up/scale-out fabrics, power allocation, and workload roles.
• The deterministic analytical evaluator uses a roofline formulation for conventional accelerators, pipelined layer sharding for SRAM-resident dataflow architectures, and an α–β formulation for communication costs.
• The model supports LLM training (tensor, pipeline, data parallelism, ZeRO-style memory checks), disaggregated serving (prefill/decode), and four HPC kernels (FP64 stencil, HPCG-like conjugate-gradient, 3D FFT, coupled AI-surrogate–solver loop).
• Model calibration was performed against published MLPerf results and vendor-reported measurements.
• Evaluation methodology uses an iso-power constraint of 1 MW and optimizes a weighted geometric mean of per-workload speedups relative to a homogeneous H100+NVLink4+InfiniBand-NDR reference system.
• A same-generation control experiment was used to isolate the benefits of architectural heterogeneity from generational hardware improvements.

Implications

• Future mixed AI+HPC systems could benefit from scale-out fabrics with higher injection bandwidth (e.g., ≥200GB/s per endpoint) and hardware support for in-network collective and all-to-all communication.
• Memory systems with HBM4-class bandwidth (e.g., ≥10TB/s per socket) or accelerators combining high memory bandwidth with substantial FP64 capability are suggested to address bandwidth-bound HPC workloads.
• System procurement could be organized around workload roles, with specialized heterogeneous components complementing predominantly AI-oriented systems.
• Allocating a small fraction (2–10%) of the power budget to specialized decode or solver partitions can significantly improve overall system performance.
• The ArchAgent framework can serve as a benchmark for studying the impact of open-weight LLM capabilities on hardware-design search tasks due to its auditable and deterministic evaluation process.

Limitations

• The analytical evaluator does not capture switch-level congestion, uses fixed factors for compute/communication overlap, and does not account for KV-cache transfer in disaggregated serving, which may affect communication-intensive configurations.
• The wafer-scale archetype is approximated as composable slices, potentially underestimating intra-wafer bandwidth and overestimating inter-slice communication performance.
• The iso-cost analysis inherits uncertainty from publicly reported list prices.
• LLM stochasticity affects search outcomes, suggesting the need for reporting results across multiple random seeds.
• Future work includes validating selected top-ranked configurations with higher-fidelity simulator-in-the-loop evaluation (e.g., ASTRA-sim, Timeloop).
• Ablation studies isolating the contributions of the specialist-agent hierarchy and bottleneck-attribution feedback are considered future work.