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

Abstract: Hybrid Transformer-Mamba LLMs enhance long-context efficiency, but their heterogeneous computation and communication patterns complicate efficient hardware acceleration. Chiplet-based architectures offer a scalable solution by integrating specialized compute and memory units. However, the design space spanning static architectural configurations and dynamic runtime policies is prohibitively large to explore exhaustively. To address this challenge, we present HYDRA, a comprehensive design space exploration framework for hybrid LLM serving on heterogeneous chiplet systems. HYDRA jointly explores chiplet composition, placement, inter-chiplet bandwidth provisioning, dynamic batching, and runtime scheduling. It integrates communication-aware placement, dynamic batching, elastic task scheduling, and a fast Markov-based performance estimator that captures multi-tenant runtime dynamics for efficient and accurate exploration. Across all workloads, HYDRA delivers 1.55x the throughput and 43.7 percent lower time-to-first-token on average, with throughput gains reaching up to 2.3x compared to state-of-the-art baselines. These results highlight that co-designing architecture and runtime policies is critical for efficient large-scale LLM serving on heterogeneous chiplet systems.

• The paper introduces HYDRA, a framework that jointly optimizes chiplet composition, communication-aware placement, dynamic batching, and elastic scheduling, achieving 1.55× higher throughput and 43.7% lower TTFT on average.
• HYDRA’s CTMC-based estimator recovers 93% of the maximum throughput found by exhaustive simulation while reducing design-space exploration from days to 4–15 minutes, although detailed simulation remains necessary for final validation.
• The evaluation shows that mixture-aware configurations generalize better across evolving hybrid models, raising worst-case normalized performance to 0.86–0.91 compared with as low as 0.31 for single-model-specialized designs.

Motivation and problem statement

Hybrid Transformer–Mamba LLMs such as Jamba, Zamba, and Nemotron-H combine attention blocks, whose KV-cache traffic grows linearly with context length, with Mamba blocks that maintain a compact recurrent state and sustain nearly constant compute and memory demand during decoding. This heterogeneity spans both model components and serving phases: prefill is compute-intensive for both block types, while decode imposes fundamentally different memory and bandwidth profiles on Attention versus SSM layers. Chiplet-based integration on a silicon interposer is a natural fit for such heterogeneous workloads, but designing these systems requires jointly exploring macro-architecture (chiplet composition, placement, NoI bandwidth) and runtime policies (batching, scheduling), which prior DSE frameworks do not address. Existing frameworks—Gemini for homogeneous CNN/MLP chiplets (Sisini, 2023), Cascade for edge workloads, and WSC-LLM for wafer-scale Transformer serving—assume static workload traces and fixed scheduling policies, and therefore miss the dynamic, multi-tenant behavior of LLM serving.

The paper identifies two obstacles that make naïve exploration impractical. First, the design space is broad: exhaustive simulation of large heterogeneous configurations can take days to weeks of wall-clock time on a 64-core CPU even with fixed placement and scheduling. Second, runtime behavior is unpredictable: asynchronous request arrivals with input-dependent prefill and decode lengths cause resource demands across compute, memory, and inter-chiplet communication to vary continuously, invalidating static analytical models and offline task mapping.

Framework overview

HYDRA takes as inputs a workload specification (models, datasets, request traces) and system constraints (chiplet library, NoI topology, interposer size, die-area budget), and jointly explores four dimensions: chiplet composition, communication-aware placement, dynamic batching, and elastic task scheduling. The design space includes four accelerator chiplet types—prefill- and decode-optimized variants of both Attention ($A_p$, $A_d$) and Mamba ($M_p$, $M_d$) accelerators—alongside HBM3 chiplets, under interposer areas of 2700–3000 mm² comparable to NVIDIA B200 and AMD MI350 packages, with individual chiplets capped at roughly 121 mm². A notable coupling the framework captures explicitly is that higher D2D bandwidth provisioning consumes die area for PHYs, micro-bumps, and routing logic, reducing area available for compute units and SRAM; bandwidth, computation, and buffering must therefore be co-optimized rather than treated independently.

Communication-aware placement

Placement proceeds in two stages. HBM chiplets are placed along the interposer's outer ring and partitioned between Mamba and Attention groups; the partition $(N_M, N_A)$ maximizing $\min(r_M, r_A)$—the minimum of the two groups' maximum supported request concurrency—is selected to avoid memory bottlenecks in either component. Compute chiplets are then placed to minimize weighted Manhattan-distance communication cost to their dominant HBM partners, with larger groups claiming preferred locations first to resolve conflicts. The two-stage decomposition trades global optimality for tractability; the paper does not evaluate how far this greedy sequential resolution departs from a joint optimum.

Dynamic batching and elastic scheduling

HYDRA generalizes continuous batching (in the spirit of vLLM and Orca) to heterogeneous chiplets. Requests are admitted when batch size is below a user-defined maximum and estimated token demand stays within a system-level token budget derived from available resources; virtualized memory blocks decouple admission from physical allocation, and request-level preemption evicts the request with the fewest generated tokens under memory pressure, prioritizing near-completion requests to reduce wasted work.

Task dispatch begins with an offline spatial pipeline mapping that minimizes the bottleneck chiplet latency using profiled task–chiplet affinity tables ($t_{comp}$ from MARCA characterizations for Mamba kernels, TSTC for Attention kernels, vendor specifications for commercial parts, or cycle-level RTL simulation for custom designs). At runtime, an elastic scheduler monitors per-chiplet queue occupancy and selectively reassigns ready tasks to underutilized compatible chiplets while preserving communication locality—a deliberate middle ground between static affinity mapping and aggressive load-driven policies.

Markov-based performance estimator

Full event-driven simulation costs hours per design point, while roofline-style models are fast but ignore runtime dynamics. HYDRA's estimator occupies the middle ground by reducing the high-dimensional scheduling problem to a low-dimensional state space: a continuous-time Markov chain whose state $s \in \{0,\dots,N\}$ denotes the number of compute chiplets assigned to prefill versus decode. State-dependent aggregate service capacities $G_p(s)$ and $G_d(s)$ capture elastic cross-phase reassignment, and a fluid-queue approximation tracks state occupancy fractions $A_d$0 without analytic solution of the CTMC. Expected throughput is the occupancy-weighted decoding capacity, and TTFT is derived from the same trajectory, capturing the fundamental prefill-latency versus decode-throughput trade-off.

The authors are explicit about the estimator's role and limits: it is a pruning and ranking mechanism, not a replacement for detailed simulation, which remains necessary to capture transient contention, NoI effects, memory management, batching decisions, and pipeline imbalance. Empirically it achieves an average cosine similarity of 0.9 with full simulation, though per-point gaps reach 8–12% throughput on some workloads (e.g., Nemotron-H–LW), attributable to abstracted queueing and contention effects.

Evaluation

The evaluation covers three models (Nemotron-H-4B hybrid, LLaMA3-7B Transformer, Mamba-2.8B SSM) across four datasets spanning long-prefill/short-decode through short-prefill/long-decode regimes (ArXiv-4K, BWB, LongWriter-6K, LMSYS-Chat-1M), on a 24-chiplet system with 2–16 HBM3 stacks and D2D bandwidth options of 256–640 GB/s, all normalized to 22 nm via DeepScale scaling.

Headline results: Across all 12 model–dataset pairs, HYDRA's combined optimizations deliver 1.55× throughput and 43.7% lower TTFT on average relative to the static baseline, with throughput gains up to 2.3×. The Markov estimator recovers 93% of maximum throughput found by exhaustive simulation while reducing exploration time from days (e.g., 8 days 8 hours for Nemotron-H) to 4–15 minutes, matching roofline-model runtime but finding designs that achieve only 74% of achievable throughput under the roofline approach.

Component-level findings:

Component	Key result
Communication-aware placement	3%–100% throughput gain over round-robin; 1.29× TP/TTFT vs. RR, 1.04×–1.80× vs. random
Elastic scheduling	3%–150% throughput gain, 1.02×–3.14× TTFT reduction vs. static; beats FCFS by 1.04×–1.30× and work-stealing in 6 of 8 cases
Dynamic batching	+15% throughput, −23% TTFT on average

NoI utilization heatmaps show that round-robin placement produces widespread bandwidth pressure from long-distance HBM–compute traffic, whereas communication-aware placement localizes data movement. The authors concede that placement cannot eliminate hotspots arising from the intrinsic bandwidth demand of steady-state decoding under fixed NoI provisioning—an honest limitation of locality optimization alone. The consistent picture is that communication locality, not raw compute provisioning, is often the dominant bottleneck at this scale.

Specialization–generality trade-off

A cross-model study using Jamba-tiny (2 Attention / 14 Mamba blocks), Zamba2-7B (13/81), and Nemotron-H-4B (4/24) quantifies the risk of over-specialization. Single-model-specialized balanced configurations degrade severely off-target—for example, the Jamba-specialized design drops to 0.42 normalized performance on Zamba and 0.31 on Nemotron-H. The all-model configuration raises these worst cases to 0.86 and 0.91 respectively (2.77× and 2.84× improvements) while limiting worst-case degradation to 14% relative to each workload's specialized optimum. Throughput-oriented ($A_d$1) designs transfer somewhat better because they are dominated by sustained decode capacity, but the same pattern holds. The practical implication is clear: platforms expecting multiple or evolving hybrid LLMs should adopt mixture-aware DSE, whereas single-model services can retain specialized configurations.

Limitations and open questions

The paper scopes its objective function to throughput and TTFT; power, thermal, and reliability are explicitly deferred, despite growing importance as chiplet counts and bandwidth scale. The Markov estimator's accuracy degrades where fine-grained queueing and contention dominate (up to ~12% throughput error), so detailed simulation remains mandatory for final validation. The evaluation relies on kernel characterizations from published MARCA and TSTC results rather than fabricated hardware, and the placement decomposition's distance from a jointly optimal floorplan is unquantified. An open question the specialization study raises but does not resolve is how configurations generalize to hybrid architectures outside the evaluated Attention-to-Mamba ratio range, or to architectures introducing new operator classes beyond Attention and SSM blocks.

Conclusion

HYDRA demonstrates that efficient serving of dynamic hybrid Transformer–Mamba workloads on heterogeneous chiplet systems requires co-exploring macro-architecture and runtime policy rather than optimizing either in isolation. Its combination of communication-aware placement, continuous batching adapted to heterogeneous resources, locality-preserving elastic scheduling, and a CTMC-based fast estimator yields substantial average gains (1.55× throughput, 43.7% TTFT reduction) while compressing DSE time from days to minutes with 93% throughput recovery. The framework's main contribution is methodological: a tractable formulation that makes architecture–runtime co-design searchable at package scale, with the specialization–generality analysis providing concrete guidance for deployment scenarios involving evolving model mixes.