Technical Summary: HYDRA – A Heterogeneous Chiplet DSE Framework for Serving Dynamic Hybrid LLM Workloads

1. Problem Statement

The rapid evolution of Large Language Models (LLMs) has introduced significant challenges for hardware acceleration, particularly with the emergence of hybrid Transformer–Mamba architectures (e.g., Jamba, Nemotron-H). These models combine the scalable long-context processing of State Space Models (SSMs) with the strong dependency modeling of Transformers, creating heterogeneous computation and communication patterns.

Existing hardware systems face two primary limitations when serving these hybrid workloads:

1. Architectural Heterogeneity: SSM blocks increasingly rely on element-wise operations and recurrent states, while Transformer blocks are dominated by matrix multiplications and KV-cache memory traffic. Furthermore, the prefill and decode phases exhibit fundamentally different resource demands. For instance, prefill is compute-intensive for both, but during decoding, Transformer blocks incur linearly increasing memory bandwidth overhead due to KV-cache access, whereas Mamba blocks maintain constant demand.
2. Design Space Complexity: The integration of diverse processing units onto a single monolithic chip is becoming cost-prohibitive. While chiplet-based architectures offer a scalable solution by modularizing compute, memory, and networking, the design space for these systems is prohibitively large. It spans static macro-architectural configurations (chiplet composition, placement, inter-chiplet bandwidth) and dynamic runtime policies (batching, scheduling). Exhaustive exploration of this space is infeasible, often requiring weeks of simulation time.
3. Runtime Unpredictability: Multi-tenant LLM serving involves asynchronous request arrivals and input-dependent behaviors (varying prefill/decode lengths). This creates highly dynamic resource demands that static performance models and offline task mapping fail to capture accurately.

2. Methodology: The HYDRA Framework

HYDRA is a comprehensive Design Space Exploration (DSE) framework designed to jointly optimize architecture and runtime policies for hybrid LLMs on heterogeneous 2.5D chiplet systems. It integrates four key dimensions:

A. System Modeling and Constraints

HYDRA takes workload specifications (models, datasets, request traces) and system constraints (chiplet library, Network-on-Interposer (NoI) topology, die-area budget) as inputs.

• Chiplet Library: Includes specialized accelerators for prefill/decode phases of both Attention (Transformer) and Mamba (SSM) blocks, as well as HBM memory chiplets.
• NoI Modeling: Accounts for topology (e.g., mesh), die-to-die (D2D) bandwidth, and latency, recognizing that higher bandwidth consumes die area that could otherwise be used for compute or SRAM.

B. Optimization Strategies

HYDRA employs a three-stage optimization process:

1. Communication-Aware Chiplet Placement:
   • Step 1 (HBM Placement): HBM chiplets are placed on the interposer's outer ring. They are partitioned into groups for Mamba and Attention blocks to balance memory capacity and minimize interference.
   • Step 2 (Compute Placement): Compute chiplets are placed to minimize Manhattan distance to their dominant memory partners (HBM groups), reducing inter-chiplet communication costs and NoI traffic.
2. Dynamic Request Batching:
   • Replaces static batching with a continuous admission policy. Requests are admitted to an active batch as soon as resources (compute, memory, token budget) are available, without waiting for batch synchronization.
   • Utilizes virtualized memory blocks (similar to vLLM) to decouple memory allocation from admission, allowing dynamic growth of KV-cache or recurrent states.
   • Implements request-level preemption to prevent out-of-memory conditions, prioritizing requests closer to completion.
3. Elastic Task Scheduling:
   • Generates an initial spatial pipeline mapping based on task-chiplet affinity (prefill vs. decode optimization).
   • A runtime elastic scheduler monitors chiplet utilization and dynamically rebalances the load. If a chiplet is underutilized due to transient pipeline stalls or bursty arrivals, tasks are migrated to it, provided communication overhead remains bounded.

C. Fast Performance Estimation (Markov-Based)

To overcome the prohibitive cost of full event-driven simulation for every candidate configuration, HYDRA introduces a Markov-based performance estimator:

• Abstraction: It models the system as a Continuous-Time Markov Chain (CTMC) where states represent the number of chiplets allocated to prefill vs. decode phases.
• Fluid Queue Model: Instead of tracking individual events, it uses a fluid model to approximate the aggregate work and state transitions over time.
• Output: It predicts system throughput and Time-to-First-Token (TTFT) by calculating steady-state probabilities and service capacities. This estimator serves as a fast pruning mechanism to identify promising regions of the design space, which are then validated by detailed simulation.

3. Key Contributions

1. HYDRA Framework: A DSE framework specifically for hybrid LLM workloads on 2.5D heterogeneous chiplet systems, available as an open-source tool.
2. Co-Design Strategy: A unified approach combining communication-aware placement, dynamic batching, and elastic scheduling to improve compute and memory utilization in multi-tenant environments.
3. Markov-Based Estimator: A novel performance modeling technique that captures runtime dynamics (elastic scheduling, dynamic batching) with high accuracy (average cosine similarity of 0.9 with full simulation) while reducing exploration time from days to minutes.
4. Comprehensive Evaluation: Extensive experiments across hybrid (Nemotron-H), pure Transformer (LLaMA3), and pure SSM (Mamba2) models on diverse datasets (ArXiv, Chat, LongWriter, etc.).

4. Experimental Results

The authors evaluated HYDRA against state-of-the-art baselines and static configurations across 12 model-dataset pairs.

• Performance Gains: Compared to baselines, HYDRA achieves an average 1.55× throughput improvement and 43.7% lower Time-to-First-Token (TTFT). In optimal cases, throughput gains reached up to 2.3×.
• Efficiency of Estimator: The Markov-based estimator reduced DSE time from days/weeks to minutes (e.g., 8 days 8 hours for exhaustive simulation vs. 15 minutes for HYDRA). It recovered 93% of the maximum throughput found by exhaustive simulation on average.
• Placement Impact: Communication-aware placement improved throughput by 3%–100% and the throughput-to-TTFT ratio compared to round-robin or random placement, primarily by reducing NoI link utilization and latency.
• Runtime Adaptivity: The combination of elastic scheduling and dynamic batching was critical. Elastic scheduling improved throughput by 3%–150% over static mapping, while dynamic batching contributed an additional 15% throughput gain and 23% TTFT reduction on average.
• Specialization vs. Generality: The study found that while single-model specialized configurations yield the best performance for that specific model, "mixture-aware" DSE (optimizing for multiple models) significantly improves worst-case performance on unseen models (e.g., improving transfer performance by up to 2.84×) with minimal loss to the target model's optimum.

5. Significance and Claims

The paper argues that co-designing architecture and runtime policies is critical for efficient large-scale LLM serving on heterogeneous chiplet systems.

• Beyond Static Models: The work demonstrates that static architectural designs and offline scheduling are insufficient for the dynamic, multi-tenant nature of modern LLM serving. Runtime adaptivity (elastic scheduling and dynamic batching) is essential to exploit optimized macro-architectural configurations.
• Communication Locality: The results highlight that communication locality (minimizing inter-chiplet distance) is often a more dominant bottleneck than raw compute provisioning in chiplet-based systems.
• Scalable Exploration: HYDRA provides a practical pathway to explore the vast design space of heterogeneous chiplet systems, which was previously intractable due to simulation costs. By accurately modeling the coupling between batching, scheduling, and resource contention, the framework enables the identification of Pareto-optimal designs that balance throughput and latency.

The authors conclude that HYDRA offers a scalable framework for exploring next-generation multi-tenant LLM systems, noting that future work could extend the framework to include power, thermal, and reliability models.