HELIX: HEterogeneous LLM Inference with FleXible and Fine-Grained Partitioning

HELIX: HE TEROGENEOUS LLM I NFERENCE WITH F LE X IBLE AND F INE -G RAINED PARTITIONING

arXiv:2611.06513v1 [cs.SE] 4 Nov 2026

Victoria Clerico 1 2 William Simon 1 Corey Liam Lammie 1 Orhun Görkem 1 Athanasios Vasilopoulos 1 Hsinyu Tsai 3 Jeronimo Castrillon 2 Abu Sebastian 1 Hadjer Benmeziane 1

A BSTRACT LLM inference is inherently heterogeneous, with performance bottlenecks shifting throughout inference phases down to individual operators. This has motivated heterogeneous serving systems that disaggregate execution across specialized accelerators, typically through fixed partitions such as prefill/decode or attention/FFN. However, these approaches commit to a partitioning strategy in advance, limiting their ability to adapt to different workloads and model architectures while forcing emerging accelerators into predefined roles. We present HELIX, a vLLM-based software stack that discovers, compiles and dispatches workload-aware partitions across heterogeneous accelerators at the operator granularity. At engine initialization, HELIX automatically determines hardware-specific operator placements that minimize execution latency, accounting for profiled execution costs, communication overheads, device capabilities, and memory constraints. It combines dynamic programming, piecewise compilation, and heterogeneous subgraph dispatch, enabling accelerators with partial operator coverage to participate in LLM inference without predefined model partitions. HELIX provides an open-source unified framework for exploring, evaluating, and integrating heterogeneous accelerators for LLM inference, lowering the barrier to emerging hardware by exposing them as subgraph executors. In systems combining H100 and B200 GPUs with emerging accelerators, HELIX resolves LLM partitions that multi-query decoding latency by up to 3× for dense and sparse transformer models and up to 3.5× for hybrid state-space models, relative to standard attention-FFN disaggregation. 1

I NTRODUCTION

The rapid adoption of LLMs for conversational assistants, code synthesis and complex reasoning tasks has made inference efficiency a first-order systems concern. Achieving efficient hardware mapping requires recognizing that LLM inference is far from monolithic, as the hardware demands it imposes vary substantially depending on which stage of the inference pipeline is examined. At the phase level, prefill processes the input prompt in parallel and is predominantly compute-bound, whereas decode generates tokens in an auto-regressive fashion and is typically bandwidth-bound. This asymmetry has motivated prefill/decode (P/D) disaggregation (Qin et al., 2025; Zhong et al., 2024), which assigns each phase to a dedicated hardware pool optimized for its respective bottleneck (see Figure 1a). Within the decode phase itself, a similar asymmetry emerges at the operator level: for highconcurrency multi-query serving, feed-forward networks 1 IBM Research–Zurich Research Lab, Rüschlikon, Switzerland Technische Universität Dresden, Dresden, Germany 3 IBM Research–Silicon Valley Lab, San Jose, CA, USA. Correspondence to: Victoria Clerico <mvc@zurich.ibm.com>. 2

(FFN) amortize the large set of weight parameters fetched from memory to generate multiple tokens. Instead, attention is strictly memory-bound as each new token must attend over the entire growing KV-cache, requiring repeated memory reads that scale with sequence length rather than reuse of loaded data. This operator-level asymmetry has encouraged Attention–FFN (A/F) disaggregated systems (Team, 2025; Liang et al., 2025; Zhu et al., 2025; Aubrey & Ghodsian, 2026), which map attention and FFN or expert layers onto separate resources (see Figure 1b). However, even these operator-level bottlenecks are not invariant. At the workload level, batch size modulates the arithmetic intensity of FFN kernels during decode (He et al., 2025), since larger batches amortize the loaded weights. Therefore, variations in the serving scenario can cause individual kernels to transition between memory- and compute-bound regimes, making them better suited to different hardware architectures. This heterogeneity has led to the development of disaggregated inference systems, with A/F emerging as a prominent direction. Beyond GPU-only deployments, the introduction of specialized accelerators has also encouraged heterogeneous disaggregation, since hardware properties favor specific operators, such as LPU-based FFN offloading from GPU (Aubrey & Ghodsian, 2026) and DRAM processin-memory (PIM) for memory-bound attention (Park et al.,

HW 2

HW 1

P

D

QKV

QKV

Attn

Attn

QKV

Out

Out

Attn

FFN

FFN

Out

a. Prefill-decode (P/D) disaggregation

HW 2

P D(Attn)

HW 1

HW 2

QKV

D(FFN) FFN

b. Attention-FFN disaggregation

HW 3

QKV

Attn

Attn

Out

Out

FFN

FFN

c. *This work: Dynamic partitioning on heterogeneous hardware

1 Dynamic Partitioning

4

5 Executor

Profilers

HW 1

HELIXvLLMBackend

HELIX: HEterogeneous LLM Inference with FleXible and Fine-Grained Partitioning

Profile à LUT à Search

2

Het. Backends

IMC|GPU|DRAMPIM

3

Compilation

Codegen|CUDAGraph

DRAM-PIM

GPU

CIM

HELIX

Figure 1. LLM inference disaggregation. (a) P/D assigns each phase to separate hardware. (b) A/F further splits decode attention and FFN. (c) HELIX maps operators across heterogeneous hardware (HW 1–3) at engine initialization via 1 profiling-based partition search, 2 heterogeneous backends, 3 piecewise compilation, and 4 - 5 heterogeneous dispatch with per-device profiling.

2024; He et al., 2025; Hong et al., 2026). These heterogeneous systems commonly use model-specific partition boundaries and focus the disaggregation decision on the sharding exploration rather than jointly deciding finegrained placements down to the individual operators. The disaggregation strategy is inherited from the A/F framing and never revisited against the particular model, hardware combination and serving scenario. As hardware architectures diversify, fixed partition boundaries cannot account for the trade-offs introduced by combining different hardware capabilities, adopting new model structures, and varying serving regimes. This paper argues that LLM serving must be rebuilt around dynamic, fine-grained partitioning. Here, dynamic denotes deployment-specific placement resolved during compilation time. What to offload and where to place it are decision variables searched over the LLM computational graph and resolved against the hardware’s capabilities, interconnect and computation cost for a model and a target serving regime. We formulate partitioning as a dynamic programming problem over the LLM computational graph. Using profiled per-operator execution and communication costs, our method finds latency-minimizing partitions for a given model, hardware configuration, and serving regime. We present HELIX, a heterogeneous execution stack built on vLLM (Kwon et al., 2023), which compiles and executes partitioned LLM graphs across available accelerators. Figure 1 places our approach in this landscape, moving from static P/D and A/F toward the dynamic partitioning that accommodates workloads, models, and hardware pool. In summary, this paper makes the following contributions: • We develop a profile-guided dynamic programming method for partitioning search to determine, at compilation time, how to split LLM inference stages across heterogeneous hardware, given workload, system, and model. • We design HELIX, a heterogeneous execution stack built on vLLM, that along with resolving partitions, runs LLM inference across the available devices,

backed by pluggable per-device profilers and backends. • We evaluate HELIX on dense Transformer, MoE, and hybrid SSM models across four hardware platforms and four heterogeneous configurations where found partitions can reduce multi-query decoding latency up to 3.5× compared to attention-FFN disaggregation.

2 BACKGROUND Disaggregated serving: P/D disaggregation assigns prefill and decode to separate compute resources connected by KV-cache transfers (Zhong et al., 2024) (see Figure 1a). Heterogeneous GPU cluster variants extend this idea by routing prefill to high-FLOPS devices and decode to highbandwidth ones (Patel et al., 2024; JIANG et al., 2025). A/F further separates attention and FFN or expert layers into dedicated GPU pools (Lin et al., 2024; Team, 2025; Zhu et al., 2025) (Figure 1b). At the orchestration layer, NVIDIA Dynamo (Pei et al., 2026) realizes these ideas by providing request routing, KV-cache-aware scheduling, and cross-engine coordination for disaggregated deployments. Frameworks such as AIConfigurator (Xu et al., 2026) optimize resource allocation and sharding settings for a given workload 1 . However, this remains an offline sizing problem over GPU resources rather than a search for operator placement across heterogeneous accelerator types. Heterogeneous hardware for inference: Beyond GPUonly disaggregation, recent systems pair compute-centric and memory-centric devices within a single inference pipeline for A/F. Attention-centric GPU-PIM designs place attention on PIM and retain other operators on the GPU (Park et al., 2024; Heo et al., 2024; Hong et al., 2026). PIM offloading also extends to state updates in SSMs (Kim et al., 2025). More recently, NVIDIA’s proposed a complementary design that keeps attention on GPUs (Vera Rubin) and offloads static FFN and MoE expert layers to 1

AIConfigurator added A/F support in a later commit (https: //github.com/ai-dynamo/aiconfigurator/pull/ 1597)

HELIX: HEterogeneous LLM Inference with FleXible and Fine-Grained Partitioning

Groq LPUs (Aubrey & Ghodsian, 2026), reporting substantial throughput gains over GPU-only execution. Similarly, emerging 3D compute-in-memory (3DCIM) hardware can heavily accelerate memory-bound matrix multiplications, and has been paired with DPUs for attention to avoid costly weight transfers (Jang et al., 2025; Liang et al., 2026; Büchel et al., 2025).

acts as a subgraph executor rather than hosting an independent serving engine, using a torch.compile-compatible backend to execute its supported operations. This design suits tightly interconnected single-node heterogeneous systems that lack the device pools and datacenter fabric that justify one engine per operator class as in existing methods (Liang et al., 2025; Xu et al., 2026; Team, 2025).

3

4

M OTIVATION

Challenges in Attention–FFN disaggregation: Figure 2 illustrates workload-dependent shifts in operator bottlenecks on an A100-80GB: increasing batch size moves decode general matrix multiplications (GEMMs) from memoryto compute-bound. Recent analytical work on A/F provisioning confirms this at the system level, showing that the optimal attention-to-FFN device ratio depends on stochastically evolving KV-cache occupancy and number of requests (Song et al., 2026). Whether an operator is computeor memory-bound does not alone determine where it should run. Optimal operator placement also depends on the hardware combination, and on which regimes each device performs best, since in heterogeneous systems, platforms are paired to complement each other and usually differ in the operators and regimes they execute efficiently. Other factors, including interconnect bandwidth, memory capacity, energy, and cost per token, further shape operator placement, which a fixed attention-FFN split cannot capture.

BS  32, ISL 64  4096

ISL  128, BS 1  256

Figure 2. Roofline analysis on an A100-80GB showing workloaddependent operator bottlenecks. Marker size increases with input sequence length (ISL, Left) and batch size (BS, Right).

Limitations of serving engines on heterogeneous platforms. Existing A/F solutions (Team, 2025; Aubrey & Ghodsian, 2026; Zhu et al., 2025; Liang et al., 2025) converge on a common design, where an orchestrator assigns hardware pools and sharding configurations to prefill, attention, and FFN, each served by an independent, networkconnected engine. Supporting a new accelerator therefore means porting the entire engine to it, which renders it impractical for emerging platforms and infeasible for specialized attention or GEMM units. HELIX instead intervenes one layer below, at the engine level, treating heterogeneous placement as a data-plane rather than an orchestration problem. Each accelerator

HELIX: P ROFILE - GUIDED PARTITIONING S EARCH

HELIX comprises two components that interface through the partitioning choice but that can be used independently. The first is a profile-guided partition search that, given a model, serving regime and hardware combination, maps the model’s main computational primitives (GEMMs, attention, state scan and update, etc.) to available hardware, minimizing end-to-end latency under memory capacity constraints. The second is a vLLM-based execution stack that runs the model under a given placement (Section 5). The search can therefore be replaced by a heuristic-based partition, and the stack can serve on its own as a heterogeneous model executor within a practical inference engine. The partition search operates on the model’s torch.fx graph, which contains on the order of a thousand nodes for models like Llama-8B (Figure 3(a-b)). Only a few of these nodes carry the non-negligible compute cost of the LLM computational graph (i.e., the computational primitives). We declare these placeable on a set of candidate devices defined by a compatibility map (e.g., QKV → HW1, HW2). The search omits the costs of the remaining element-wise and metadata operations (reshape, view, transpose, slice, etc.) as they typically account for a small percentage of the execution time (see Appendix Figure 10), while full graph execution and evaluation include these element-wise costs. These operations are absorbed into their producers only when supported by the assigned backend. Otherwise, they remain explicit placeable stages with their transfer dependencies but bear no cost (e.g., SiLU → HW1). Figure 3c illustrates the reduced directed acyclic graph (DAG). The reduced DAG has N stages, proportional to the number of layers and the number of placeable primitives per layer. Each primitive is profiled once, offline, over the shapes it takes under the evaluated workloads, so scoring a candidate placement costs a table lookup in microseconds rather than a deployment in minutes. Recent configuration-search tools adopt the same abstraction (Xu et al., 2026), which holds as long as a primitive’s isolated latency predicts its runtime latency. This is the case on GPUs, where consecutive kernels execute with little cross-kernel pipelining (see Appendix Figure 12). Our implementation can also use roofline-based cost estimates instead of profiled costs, but we omit this from our evaluation since roofline models are generally less accurate than profiled data (see Appendix Figure 11). Placing a computation graph onto heterogeneous devices is

HELIX: HEterogeneous LLM Inference with FleXible and Fine-Grained Partitioning

NP-complete in general (Topcuoglu et al., 2002), so practical methods exploit model structure. Following prior deviceplacement work (Hu et al., 2022), HELIX adopts a dynamic programming (DP) formulation and extends it to the residual dependencies and memory constraints in LLM stage graphs, keeping the search within the time budget of compilation time during the vLLM engine initialization. 4.1 Placement and objective We represent the reduced computation graph as a DAG G = (S, E), with stages S = {s0 , . . . , sn−1 } in topological order. In standard chain-like DAGs, every stage si (i > 0), consumes the output of si−1 . However, residual additions in LLMs create nonadjacent dependencies as shown in Figure 3, so a stage may have multiple predecessors, denoted pred(i). Each stage (si ) has an admissible device set (Ci ) a)

Transformer Block

b)

torch.fx.Graph

…

Up

arg7

linear

SiLU

Down

mul

+ K

slice

silu

V

linear

𝜀!,!&%

𝜀!"#,! 𝜀!,!&'

𝑆!"# : layers.1. mlp.gate_up_proj

𝜀!&%,!&'

slice

x

Q

𝑆! : layers.1. self_attn.o_proj

pow

RSMNorm

𝑆!&# : layers.1. self_attn

𝜀!"%,!

add

+

Gate

arg3

linear

Out

c) Directed Acyclic Graph

arg2

arg1

Attention

arg8

𝑆!"$ : layers.1. mlp.down_proj

𝜀!&',!&( 𝑆!"% : layers.2. self_attn.qkv_proj

add

Figure 3. Graph reduction from (a) a transformer block through (b) its torch.FX graph to (c) a placement DAG, preserving sequential backbone and residual dependencies.

and each device (d) has available memory (Md ). Assigning (si ) to d has a computational cost (Tcomp (si , d)) and a memory footprint (m(si )), comprising parameters, KV cache, or recurrent state. The budget Md is computed after reserving runtime headroom. A graph edge (sj , si ) ∈ E carries wij elements of B bytes. When sj and si are placed on d′ and d, the transfer cost Tcomm (w, d, d′ ) = w · B / βd,d′ + σ where B is the number of bytes per element and βd,d′ is effective directed-link bandwidth, plus σ a fixed startup latency, excluding buffer-capacity effects. Based on these definitions, a placement a : {0, . . . , n − 1} → D, a(i) ∈ Ci assigns each stage to an admissible device (i.e., a partition of the DAG up to stage si ). It determines where each operator (or set of operators) executes and which intermediate tensors cross a hardware link and, hence, the latency of the forward pass. We define the finish time of stage si as the arrival time Ai (a) of its latest-arriving input plus the execution time of the cost-bearing operation of the stage,
Li (a) = Ai (a) + Tcomp si , a(i) . (1) How arrivals from multiple inputs combine depends on

where their producers execute. Host inputs are treated as produced on dinit . For d = a(i), the output of a predecessor sj , j ∈ pred(i), is available at Lj (a) if also sj runs on d, and inputs already transferred to d for an earlier consumer are reused without a new transfer. The remaining inputs, produced on a remote device d′ ̸= d, contend for the link (d′ , d), whose link state Qd′ →d (initially 0) must record when it frees. Each link serves transfers in the order their data is consumed, and transfers feeding the same stage, in the order they become available. The k-th transfer completes at:

Ai,d′ (k) = max Ljk (a), Ai,d′ (k−1) +Tcomm wjk i , d, d′ , (2) Let Ai,d′ denote the time at which the transfer of the last input of si sent over link (d′ , d) finishes. The link then remains occupied until this time, so we set Qd′ →d ← Ai,d′ for the next stage to account for link contention. The first term of the max in Eq. 2 assumes a transfer starts once its input is released, so it cannot overlap with the producer’s own subgraph execution. However, execution of stages on the destination device d can overlap with transfers d′ → d. For example, a residual input to si can be transferred while si−1 computes on the same d. The second term makes the transfer wait for earlier transfers on the same link. Links between distinct device pairs proceed concurrently, so the arrival time of the last input is:   ′ Ai (a) = max max Lj (a), max A (3) i,d ′ j∈pred(i): a(j)=d

d ̸=d

Since arrivals combine by max, a transfer adds latency only if it delays the last input of some stage, either directly or by occupying the link for later transfers. Considering the previous equations, the best candidate placement (i.e. partition) is the one that minimizes the finish time of the last stage, a∗ = arg min Ln−1 (a) a

(4)

and is memory-feasible when the Pfootprint it accumulates on each accelerator stays below i : a(i)=d m(si ) ≤ Md 4.2 Partitioning search via dynamic programming The objective in Eq. 4 defines a space of possible placements of size DN for a model with N stages and up to D admissible devices, and every candidate must additionally be checked against the per-device memory budgets. Collecting profiling data makes each placement candidate cheap to score, a table lookup rather than a deployment, but does not make exhaustive search tractable. HELIX instead uses dynamic programming along the sequential backbone, retaining the earlier outputs needed by residual dependencies and the memory committed on capacity-constrained devices. We extend stage-indexed DP formulations (Hu et al., 2022) with these two state components, described below; Appendix C provides the full algorithm.

HELIX: HEterogeneous LLM Inference with FleXible and Fine-Grained Partitioning

4.2.1 The frontier state In LLM DAGs, fan-in and fan-out stages arise from the residual connections within and between layers (see Figure 3). Therefore, the optimal placement of a fan-in stage i does not depend only on the immediate predecessor stage i − 1 (as DP solution could track) but must be resolved based on all possible arrivals from all of its producers. We therefore extend the standard DP state to retain all stages within a fan-out→fan-in structure. For instance, if si is a fan-in and consumes both si−1 and si−2 , and, si−1 also consumes si−2 , a partial placement that minimizes the completion time of si−1 may not minimize that of si , so states of si−2 and si−1 shall not be resolved until si is evaluated. Under Equations 2–3, a partial placement now depends on predecessor assignments and finish times together with the directed-link availability Q. Outside the fan-out and fan-in, every stage consumes at most the output of its predecessor, where a single transfer is in flight at a time per link. For a fan-in stage si , link contention can arise between the residual skip and a transfer feeding an intermediate stage si−1 . Link becomes available by the time si−1 starts executing, which lets us simplify Qd′ →d = Ai−1 when si−1 runs on d, and 0 otherwise. This simplification relies on our LLM DAG abstraction and general DAGs, since they can have many transfers in flight at once, would require tracking Q explicitly in the state. A residual input to si is then sent after the inputs of si−1 (since they are consumed earlier) and its transfer overlaps with si−1 computation. We call the set of stages between a fan-out and its fan-in the frontier: it opens when the residual producer is placed and closes once the fan-in consumes its output. The DP is no longer indexed by stage and device (i, d) but by the joint placement of the stages within the frontier. Under each frontier assignment, the DP state retains one or more labels ℓ, i.e., partial placements that share the frontier devices aℓ (j) but differ in the finish times Lℓj and in the devices to which their outputs have already been transferred to. Assigning si to an admissible device extends a label as

Li (ℓ, d) = Tcomp (si , d) + Ai d, {(aℓ (j), Lℓj )}j∈pred(i) , (5) where Ai follows Eqs. 2–3 with Qd′ →d = Ai−1 when a(i − 1) = d, and 0 otherwise. In the fan-in example above, a frontier assignment fixes the devices of si−2 and si−1 , and each label under it stores their finish times. Placing si on d combines these finish times through Eq. 5, and the new label stores Li (ℓ, d) together with the finish times of the stages that remain in the frontier. The number of nodes that enter the frontier before being resolved is bounded by the width (ω) of the residual connections in the graph (how many stages are skipped), so the number of tracked states due to the frontier is at most |D|ω .

4.2.2 Memory as a resource dimension The frontier-state DP makes the latency of a placement computable but does not ensure that it fits in memory. HELIX extends the DP state with an occupancy vector u, similar to solutions to the multi-dimensional multiple-choice knapsack problem (Hifi et al., 2004), whose entries carry the aggregate footprint on each device. Extending a partial placement ℓ with si 7→ d updates only that device’s enℓ⊕(s 7→d) try, ud i = uℓd + ⌈m(si , d)/b⌉, and partitions where ′ uℓd > Md are discarded. However, not every device needs an entry in the occupancy vector for all serving regimes and modes. We consider T the set of memory-tracked devices, which contains only the overflowable ones, i.e., those whose capacity would be exceeded if every stage listing d among its candidates Ci were assigned to it. Therefore, the tracked occupancy vectors are limited to devices d ∈ T . A GPU, for instance, supports every operator and becomes overflowable once parameters and KV cache/SSM exceed its capacity. Reducing the memorytracked devices is crucial as search time grows quickly with T. Adding occupancy to the state enlarges the search space, as partial placements that would be interchangeable under latency considerations, now can differ in how they occupy memory. To keep the states tractable, HELIX prunes states by Pareto dominance. A label is dominated by another if the latter has no greater running latency, no later finish time for every live frontier stage, and no greater tracked occupancy. The effectiveness of pruning in keeping the search time small depends on how often occupancies are comparable, so byte-level precision makes cases where dominance occurs rare, keeping apart states whose occupancies differ too little to affect any later feasibility decision. Therefore, occupancy is tracked in integer bucket units of size b. Each stage occupies then ⌈m(si )/b⌉ and each device has capacity ⌊Md /b⌋. Further implementation details are in Appendix B. 4.3 Complexity and optimality Optimality: DP finds an optimal placement for the reduced-graph objective in Equation 4. Pareto pruning discards a partial-placement label only if another label with the same frontier label has lower or equal running latency, finish time on every live frontier stage, and occupancy on every device. We treat pruning of the memory representation as a heuristic to reduce the amount of tracked states and evaluate the resulting placement quality empirically. Additionally, accounting for memory occupancy in buckets may be a source of approximation. A larger b leaves fewer possible occupancy values per device, so more memory states match and get pruned, reducing the search space and time. However, faster search time comes at the cost of accuracy in how the occupancy is being tracked, as coarser buckets can conservatively reject feasible placements (see Appendix Figure 14a). We set b to the greatest common divisor of

HELIX: HEterogeneous LLM Inference with FleXible and Fine-Grained Partitioning

all stage footprints. Capacity rounding is exact since all footprints are multiples of b. Complexity: Without memory overflow, search over the DAG visits each stage just once, in addition to extending each frontier state over admissible devices. With ω stages in the frontier, there are at most |D|ω frontier assignments, resulting in O(N |D|ω+1 ) transitions when one sufficient state is kept per assignment, against the O(|D|N ) complexity of exhaustive search. In transformer-like models, ω does not depend on the model’s depth and usually stays between 2 and 4 stages, so |D|ω stays small and the search is linear in N · D. Instead, tracking memory increases the amount of labels retained per stage. Considering P as the maximum number of nondominated labels retained per key. Labels sharing a frontier key and an occupancy vector are compared on arrival time alone,Q so at most one is kept per occupancy vector, giving P ≤ d∈T (⌊Md /b⌋ + 1), where ⌊Md /b⌋ refers to the amount of values each occupancy vector can take, which depends jointly on the memory capacity and the bucket size. The search therefore considers O(N P |D|ω+1 ) label extensions, with additional cost for arrival evaluation and dominance checks. P grows polynomially with N for a fixed T until it saturates the memory constraint, and then becomes linear in N , which can happen in early stages of the graph under high memory constraints. In practice, HELIX targets regimes where each accelerator hosts most of the parameters of its supported operators, since heterogeneous platforms pair devices to take advantage of their compute properties rather than memory aggregation, so T stays small. If model placement is dictated by memory constraints instead of complementary hardware strengths, search should lean towards resolving sharding configuration as well.

5

HELIX: THE EXECUTION PATH IN V LLM

The placement resolved in the previous section defines which device each stage belongs to, and leaves open the mechanism by which each fragment is compiled for its target device, and how activations cross between them. Figure 4 illustrates the HELIX software stack. At the frontend, HELIX integrates with vLLM, inheriting its scheduling and KV-cache management. Model loading is intercepted by HELIXvLLMBackend, which hooks into TorchDynamo’s tracing pipeline and splits the traced graph for a given workload (derived from the model’s context length and vLLM’s maximum number of concurrent sequences), model and available hardware, resolving operator-to-backend assignments from kernel support, profiled runtime cost and memory capacity 1 with the methodology described in 4. The HELIXModelRunner reuses the existing piecewise compilation which is used to compile CUDAGraph safe and unsafe regions, to convert each subgraph into a standalone artifact specialized to its target and stitching them with

Frontend High-level scheduler

Graph Tracing

1

HELIXModelRunner

2

Piecewise Compilation Profiler Aggregator Model executor

GPU

CIM

…

HELIXvLLMBackend

Partitioning search (DP)

Partitioning map 3

Heterogeneous Backends

4

Profilers

5

Device transfers

PIM

Figure 4. HELIX software stack: graph partitioning based on DP 1 , piecewise compilation 2 on heterogeneous backends 3 , and profiler feedback 4 through execution 5 .

designated transfer modules 2 . Compute artifacts are dispatched sequentially, while transfer modules are issued asynchronously so that a transfer starts as soon as its producer finishes and overlaps with compute on other devices, as assumed by the cost model (Eqs. 2–3) 3 . Each hardware is exposed to the stack through a vLLM-compatible compilation interface, so that graph tracing can verify operator support. Adding new hardware only requires implementing this interface and profiling the costs used for partitioning, so a vendor can evaluate an emerging accelerator within a full serving stack without porting an entire model. Compute costs are collected per device from either hardware profilers on physical accelerators or software simulators modeling emerging hardware, allowing HELIX to mix both sources within a single deployment. These costs are then aggregated in the model runner to report user metrics 4 during execution 5 . 5.1 Interception point: piecewise compilation vLLM V1 integrates torch.compile as a core pipeline component: TorchDynamo captures tensor operations into a computation graph, and TorchInductor lowers it to kernels. Since some operations are size-variable (e.g., attention over a KV cache of unpredictable length), vLLM uses piecewise CUDA Graphs, splitting the graph into CUDA-Graph-safe and -unsafe regions executed as independently dispatchable artifacts, with attention detached from the rest. HELIX intercepts at this attention/non-attention boundary and extends it, partitioning the graph at any boundary where target hardware changes and turns each piece into hardware-specific subgraphs, leaving vLLM’s GPU-side optimizations (kernel fusion, CUDA Graph replay, shape specialization) untouched. vLLM suits HELIX because its graph tracing is device-independent, with boundaries defined over generic tensor operations, so any torch.compile-compatible accelerator can become a dispatch target. Engines like TensorRT-LLM (NVIDIA Corporation, 2026) or SGLang (SGLang Team, 2026), whose segmentation is bound to the GPU execution model, lack a

HELIX: HEterogeneous LLM Inference with FleXible and Fine-Grained Partitioning

comparable interception point for non-CUDA substrates. 5.2 Heterogeneous graph partitioning and dispatch The resolved placement from Section 4 enters the stack here. HELIX intercepts the pre-dispatch torch.fx.Graph, traverses it node by node, and assigns each node to the backend that the DP partition places its DAG stage to. Contiguous runs of nodes sharing a target backend are then fused into maximal subgraphs, minimizing backend transitions and preserving kernel-fusion opportunities within each partition. At every boundary where the target backend changes, HELIX inserts a dedicated transfer module, resolved by data flow rather than by the positional adjacency of subgraphs, so that transfers carry only the tensors the target backend actually consumes. Registering the module inserts a node into the stitching graph, redirects all affected consumers to the transferred copies and re-compiles. Each transfer becomes its own subgraph in the dispatch list, interleaved with the compute subgraphs. This insertion point may also serve as the integration surface for communication backend development (such as NIXL (ovidiusm et al., 2026)) as heterogeneous platform support matures. Each resulting subgraph is then compiled by its respective backend into standalone compiled artifacts and dispatched independently at run time.

6 E VALUATION M ETHODOLOGY 6.1 Models and Metrics Llama3.1-8B is a dense transformer model, while Mixtral8x7B replaces the FFN with a sparse mixture of experts, both with quadratic grouped-query attention (GQA). Bamba-9B is a hybrid that interleaves a few GQA layers among state-space mixer layers, exhibiting two profiles on the weightless computations. Because partitioning decisions are driven by the mix of computational primitives rather than by any specific release, each model stands for its architectural family. Newer models within these families largely reuse the same primitives, varying mainly the attention mechanism and proportions between dynamic and static operators, while producing DAGs of comparable size (Appendix Table C). We focus on the decode phase, which dominates serving time and is where heterogeneous accelerators target memory-bound operators; prefill can be handled by existing P/D disaggregation, although the partitioning search applies to it unchanged with prefill-shaped profiles. Partitions are judged by the latency a single decoding step and the generation throughput over multiple steps. 6.2 Hardware platforms and profiling methodology GPU (H100-80GB SXM, B200-192GB SXM). Latencies come primarily from the public PerfDatabase (Xu et al., 2026), which reports single-kernel latencies profiled on

vLLM’s execution path. We follow its profiling protocol2 to add in-house measurements for ISLs and batch sizes it does not cover (GEMMs, GQA) and for the Mamba mixer, validating our data against the released data where they overlap. DRAM-PIM. We adopt the AttAcc (Park et al., 2024) accelerator, built on 8-Hi HBM3 stacks, which targets attention. In batched decoding, the score (QK ⊤ ) and context (αV ) operations reduce to per-request GEMVs over the KV cache. AttAcc places GEMV units at the bank level so that they execute at DRAM internal bandwidth rather than over the external HBM I/O. We instantiate a single-GPU-scale system of 5 AttAcc stacks (equivalent to the HBM of an A10080GB SXM) and obtain attention latencies from the authors’ open-source simulator3 . 3D compute-in-memory (3DCIM). 3DCIM is an analog in-memory computing accelerator built around vertically stacked non-volatile memory, which maps linear layers to tier-level matrix–vector multiplications executed within the array (Jang et al., 2025; Liang et al., 2026; Büchel et al., 2025). Stationary weights eliminate parameter-fetch traffic, making it well suited to compute-bound fully connected and MoE expert layers. We model an accelerator inspired by (Büchel et al., 2025) with a capacity of ∼ 86B parameters, and a 200 ns MVM integration time. The 3DCIM simulator serves as a backend for HELIX-vLLM execution, serving to demonstrate how a simulator and hardware profiling backends can be implemented simultaneously in HELIX-vLLM. Further details on the 3DCIM system simulator are available in Appendix D. Finally, connections between accelerators are assumed peerto-peer using NVLink 3.0 with 300 GB/s and a startup latency of 2 microseconds (Ma et al., 2026). 6.3 Baselines We test three heterogeneous system configurations: 3DCIM + GPU, AttAcc + GPU, and 3DCIM + GPU + AttAcc. The pairwise configurations reflect the GPU-plus-accelerator pairings discussed in Section 2, where a fixed partitioning strategy is prescribed by the design itself. Concretely, for 3DCIM + GPU, the baseline partition follows mainstream A/F as proposed in (Aubrey & Ghodsian, 2026; Zhu et al., 2025) where either the FFN block or the experts are offloaded from the GPU. For DRAM-PIM + GPU, the baseline follows the strategy adopted in AttAcc (Park et al., 2024) where only the attention is offloaded from the GPU; this is referred to in the rest of the paper as the AttAcc baseline. All methods use identical profiles, operator-support constraints, communication assumptions and memory budgets. 2 https://github.com/ai-dynamo/aiconfigur ator 3 https://github.com/scale-snu/attacc_simu lator

HELIX: HEterogeneous LLM Inference with FleXible and Fine-Grained Partitioning

7 R ESULTS 7.1 DP-based Partitioning search We first assess the efficiency and correctness of HELIX’s DP-based partitioning search. Table 7.1 shows execution times for HELIX DP-based partitioning search against an exhaustive-like baseline over DAGS of different structures. The first four rows show search times for a small synthetic DAG of 16 stages isolating the effects that add complexity to HELIX DP. The first two rows show the effect of the frontier, raising slightly the search time from 1.4 milliseconds to 3 milliseconds when adding residual skips of width 3 every four stages (similar to the pattern in transformers). The third and fourth rows isolate the effect of tracking memory across the search for three P equally-constrained P devices in terms of capacity where m(si ) ≈ 1.4 × Md , and shows how state pruning, based on Pareto dominance, keeps the search time short (785 ms vs. 267 s) without compromising the final result. On all synthetic DAGs, the DP search returns the same solution as the exhaustive search. For a realistic graph, enumerating the entire search space is intractable for 3162 combinations, so we instead conducted Monte Carlo simulations, sampling 4,000 initial random placements and timed them out after 30 minutes of search, and comparing their cost against the DP solution. None of the random searches returned a partition with lower latency than DP. Further ablations over the algorithms can be found in Appx. Section B. Placing the full set of stages of Llama takes roughly 16 milliseconds, which makes the partition search seamless at compilation time. DAG

D/T

HELIX

Chain(16) Skip(16)

3/0 3/0 3/3 3/3

1.4 ms 3.1 ms 267 s* 785 ms

3/0

15.967 ms

Skip(16) Llama 8B (162)

Baseline Exh. 3342.6 s 3129.8s 2018.9 Rand. < 30 min

Est. Latency = = = = 3.1±0.45×

Table 1. HELIX (DP) search time is compared against exhaustive search on small graphs (top four rows; * = DP without state pruning), with the last column showing whether estimated latency matches (=). On the full LLM graph (bottom row), it is compared against the average of 4,000 random searches (30-minute timeout), whose found partitions have higher latency than DP’s.

Having validated the DP-based search, we now apply it to different heterogeneous systems to show the advantage of tuned partitions over standard disaggregation techniques when different accelerators are combined. Subsections 7.2– 7.4 each target a different heterogeneous platforms, resolving the partitions and reporting the decode latency the search estimates from pre-profiled data across a range of models, query counts, and context lengths.

7.2 A/F-style system: GPU + 3DCIM Here we test a heterogeneous system composing an H100 and 3DCIM accelerator, which maps directly onto the A/F disaggregation trend and therefore is the selected baseline. Figure 5 reports the latency of a single decoding step for three strategies: H100-only execution, the standard A/F split (FFN or expert layers on 3DCIM, all other operations on GPU), and the optimal partition found by the DP search, to which the other two are normalized. At small batch sizes, where decoding is dominated by weight fetching, the HELIX partition is ∼3.14–6× faster than single-H100 for Llama 8B and ∼4.48–7.3× faster for Bamba 9B, for ISL of 8K and 2K, respectively. For Mixtral 8x7B, whose parameters do not fit on a single GPU, the baseline is 2xH100, and HELIX is ∼5–9.8× faster over the same ISL range. Pairing (HELIX)

ISL2048 TP2

ISL8192 TP2

Figure 5. Single-step decoding latency of the H100-only and A/F baselines, normalized to HELIX. For Mixtral 8x7B, which does not fit on a signle H100, the baseline uses tensor parallelism (TP) = 2. Partitions can be found in Appendix Figure 16.

a GPU with CIM core that eliminates memory-bottleneck is already advantageous, and offloading only the FFN/MoE (A/F) captures much of the gain, though it is not the optimal split in all scenarios. With up to 32 concurrent queries the optimal partition differs the most from A/F, especially for Bamba 9B. Bamba’s hybrid structure makes standard A/F hard to apply, and our search discovers the better split automatically by assigning all GEMMs to 3DCIM, including those around the mixer and attention, rather than only the FFN. Its mixer projections have larger input and output dimensions than the attention projections in Llama 8B and Mixtral 8x7B and account for most of its execution time (see Appx. Figure 10), so offloading them yields speedups of roughly 1.7–3.2× over A/F depending on batch size. These gains roughly hold at longer contexts, since the SSM state update, unlike attention, does not grow with context length4 . For Llama 8B and Mixtral 8x7B, KV-cache fetching increasingly dominates decode latency as context grows, so the optimal partition lands gains only 2.2–1.1× over A/F for batch sizes up to 16 and context lengths of 2K to 8K, respectively. Around batch size 64, the HELIX partition 4

Bamba-9B-v1 has a maximum model length of 4096 tokens, so longer contexts are shown to be representative of how a model less bound by attention would perform at higher context lengths.

HELIX: HEterogeneous LLM Inference with FleXible and Fine-Grained Partitioning

converges to A/F, meaning the optimal split offloads only the FFN/MoE to the 3DCIM accelerator. Beyond this batch count, the best partition assigns all GEMMS back to the GPU for all models except for the Mixtral, because transfer latencies increase with batch size and mask the per-GEMM speedups 3DCIM can have over H100. For Mixtral, the experts remain in 3DCIM since loading all eight experts’ weights keeps its computation memory-bound even at higher concurrency. This results in ×1 w.r.t. GPU-only baseline, but for this serving scenario, attention dominates execution, so all partitions perform similarly. 7.3 GPU + DRAM-PIM The limited gains of any HELIX partition at large batch sizes and context lengths stem from attention remaining resident on the GPU. We now relax this constraint and study a system where the second accelerator can also execute attention (i.e. DRAM-PIM), allowing the partition search to split attention as well as the projections across both devices. For DRAM-PIM, GEMM latency is estimated with a roofline model. Mixtral 8x7B is excluded from this study because mapping its experts to DRAM-PIM is non-trivial, and its expert weights (∼90,GB in FP16) exceed the GPU’s 80,GB memory. The AttAcc simulator models multi-head atten(HELIX)

ISL2048

ISL8192

Figure 6. Single-decoding-step output token latency of the H100only and AttAcc-style disaggregation baselines, both normalized to the latency given by the HELIX partition.

tion (MHA) instead of GQA, but since the authors report no latency benefit of GQA on PIM (Park et al., 2024), we consider the simulated MHA latencies with a GQA memory footprint. Figure 6 compares GPU-only and AttAcc-style disaggregation against the HELIX partition. For 4 queries, HELIX finds partitions that place both GEMMs and attention on DRAM-PIM for Llama 8B, achieving 2–2.7× speedups over AttAcc. However, GEMMs on DRAM-PIM are beneficial only at low batch sizes as its compute throughput (∼30 TFLOPS) is far below the H100’s. Beyond 16 queries, AttAcc-style placement is the one found to be optimal. Notably, at 64 queries with 8K tokens in context, memory pressure hits, and execution under standard AttAcc disaggregation or GPU-only would be infeasible. HELIX still finds a feasible partition by moving part of the attention heads from DRAM-PIM to H100, which now would host

the GEMMs along with 13 attention layers. Instead, the chosen partitioning for Bamba 9B is exactly AttAcc-style, as the large input and output projections are faster on the H100 regardless of the batch count. Detailed partitions for Llama 8B are available in Appx. Figure 17. Therefore, gains over GPU-only execution shrink across all workloads for Bamba 9B, since only 3 of its 32 layers contain GQA and the rest are SSM blocks. Models like Bamba do not benefit from AttAcc-style partitions. Prior work has explored accelerating SSM blocks with PIM architectures (Kim et al., 2025), which we leave for future study. Yet, the previous results in 7.2 showed that HELIX still finds partitions that outperform more suitable baselines such as A/F for this model. This highlights how the best strategy depends jointly on the model architecture and the system configuration. With newer models adopting different attention variants (e.g., multi-head latent attention, SSM, gated delta attention) and different ratios of dynamic to static operations, a strategy designed for one model family, such as AttAcc, does not directly transfer to another. 7.4 Three-way heterogeneity In this section, we study a more complex system of three hardware components to illustrate how HELIX can guide heterogeneous system design by simulating as-yet unrealized hardware combinations on a given model and workload. The two pairwise studies 7.2 and 7.3 show that 3DCIM is a strong accelerator for most of the GEMMs, moving towards H100 for larger batch count as computation gains are masked by transfer latencies, while DRAM-PIM remains the fastest attention substrate across all evaluated workloads. However, no existing disaggregation strategy covers how to partition a model across three components with such different properties, as both A/F and AttAcc, are GPU-centered and only serve as pairwise references. Figure 7 shows the partitions found for each workload and their single decodingstep speedup achieved over the two baselines. Below 64 concurrent queries, all projections stay on 3DCIM, and attention is consistently assigned to DRAM-PIM, a placement neither reference or existing system in literature can express. In this range, the optimal partition is 1.8–3.5× faster than A/F and 1–8.7× faster than AttAcc for Llama 8B, and 1.8– 3.6× faster than A/F for Mixtral 8x7B depending on batch size and context. Notably, neither baseline is optimal across the full range of workloads, since A/F is the better option below 16 concurrent queries at every evaluated context length and the AttAcc one beyond that point, while HELIX partitions outperform both in all cases. The GPU could still be involved in running operations such as Rotary Position Embedding (RoPE). Since our DAG abstraction absorbs RoPE into the attention stage, we assume it can run on the DRAM-PIM base die alongside Softmax, avoiding costly data transfers. Bamba again shows the largest gains over A/F. Because its Mamba mixer is tied to the GPU, every

HELIX: HEterogeneous LLM Inference with FleXible and Fine-Grained Partitioning H100

Llama 8B

DRAMPIM

3DCIM

Mixtral 87B

Bamba 9B

BS

Figure 7. Partitions found for the three-way system (H100 + 3DCIM + DRAM-PIM) across batch sizes and sequence lengths, and the single-decoding-step speedup over the A/F and AttAcc-style baselines. × marks memory-infeasible configurations.

partition spans all three accelerators, which opens opportunities for finer-grained pipelined scheduling as future work, using sub-batching to increase resource utilization. Replacing the H100 with a B200 To examine whether these placements persist when the GPU changes, Figure 8 repeats the study with a B200 replacing the H100. The B200’s higher compute throughput and memory bandwidth makes it more competitive for both GEMMs and attention, shifting where each is best placed. At batch size 4 and context length 2K tokens, for instance, HELIX partition offloads all GEMMS to 3DCIM rendering a 1.9× speedup against A/F, compared to 3× speedup over A/F when applying the same partition with the H100-based three-way system, since now GEMMs in the B200 are slower than 3DCIM but faster than H100. As batch size grows, partitions for Llama 8B and Bamba 9B converge to A/F at 32 queries (instead of 64), and, to B200-only execution at 64. The B200 matches and starts to outperforms DRAM-PIM for attention execution around batch count 16. This suggests that B200’s higher memory bandwidth based on HBM3e, combined with its compute cores, is enough to compete with the benefits of near-memory computing on an HBM3 memory system. Changing only the GPU scale shifts partition boundaries, confirming that each hardware combination demands its own partition search rather than a fixed strategy. (HELIX)

ISL2048

ISL8192

Figure 8. Single-decoding-step latency for B200 + 3DCIM + DRAM-PIM and speedup over pair-wise baselines. Partitions in Appendix Figure 19

7.5 Offline serving with HELIX-vLLM Figure 9 reports end-to-end generation throughput with the found partitions deployed in vLLM over multiple decoding steps. Unlike the single-step estimates above, which add up

the costs of the profiled primitives only, these runs account for every operation in the parent graph, including elementwise and normalization operations, except for RoPE when executed in DRAM-PIM.

HELIX

Figure 9. Generation throughput over complete decoding graphs in HELIX-vLLM for the three-way heterogeneous partitions found by HELIX, compared to the A/F and AttAcc-style baselines.

8

R ELATED W ORK

Device placement of DNN graphs on distributed environments has been widely studied. Tarnawski et al. (Tarnawski et al., 2020) formulate operator placement across identical accelerators and CPUs as an integer program for latency and a dynamic program for pipelined throughput, where split decisions are triggered by memory capacity and graph parallelism. PipeEdge (Hu et al., 2022) splits chain-like DNN graphs at the layer granularity across heterogeneous edge devices to minimize the slowest pipeline stage within memory constraints. HELIX instead targets residual LLM graphs with hundreds of stages across multiple accelerator types, where splits are driven by per-operator performance rather than capacity alone, and resolves placement within compilation-time budgets. For LLM inference specifically, methods in literature have started to move beyond fixed A/F splits towards finer-grained execution. PAPI (He et al., 2025) and DOPS (Yang et al., 2026) co-design runtime operator scheduling with tailored hardware to dispatch kernels between DRAM-PIM and a GPU or an NPU, respectively, deciding among predefined kernel classes for a fixed hardware pairing. In contrast, HELIX is hardware-agnostic, searching placements over the full model graph for arbitrary accelerator combinations within a serving engine. OpWeave (Li

HELIX: HEterogeneous LLM Inference with FleXible and Fine-Grained Partitioning

et al., 2026) also argues A/F boundaries are suboptimal for newer models and proposes general operator partitions within vLLM to minimize the dollar cost per million tokens. Its heterogeneity is limited to GPUs of different capability rather than accelerators with complementary strengths where each excells different operators and serving regimes. The two approaches are therefore orthogonal, addressing the same motivation under different objectives and hardware assumptions.

9

C ONCLUSION

HELIX proposes method to resolve workload-specific operator partitions and dispatch them across heterogeneous hardware within a unified vLLM engine. By treating hardware as graph executors, we enable emerging accelerators (e.g., PIM, 3DCIM) with partial operator support to participate in LLM serving. Future work can extend HELIX to optimize energy consumption and resource utilization alongside latency, by targeting pipelined execution when multiple accelerators are available. Pre-compiling partitions for several serving regimes would further enable online adaptation as well as auto-tuning through runtime profiling to inform partitioning decisions.

10

ACKNOWLEDGMENTS

We thank Thomas Parnell, Alberto Mannari, Garima Singh, and Irem Boybat for their valuable technical discussions. We also acknowledge Vijay Narayanan and Jeff Burns for their managerial support.

R EFERENCES Almaci, T., Lammie, C., Benmeziane, H., Fournier, C., Görkem, O., Boybat, I., and Simon, W. A. Ghl: An extensible c++ library for flexible and performant graph pattern matching and rewriting. In Archibald, B. and Semeráth, O. (eds.), Graph Transformation, pp. 187–198, Cham, 2026. Springer Nature Switzerland. ISBN 978-3032-29730-3. Aubrey, K. and Ghodsian, F. Inside NVIDIA Groq 3 LPX: The low-latency inference accelerator for the NVIDIA Vera Rubin platform. NVIDIA Technical Blog, March 2026. URL https://developer.nvidia.com /blog/inside-nvidia-groq-3-lpx-the-l ow-latency-inference-accelerator-for -the-nvidia-vera-rubin-platform/. [Accessed 2026-10-16]. Büchel, J., Vasilopoulos, A., Simon, W. A., Boybat, I., Tsai, H., Burr, G. W., Castro, H., Filipiak, B., Le Gallo, M., Rahimi, A., Narayanan, V., and Sebastian, A. Efficient scaling of large language models with mixture of experts and 3d analog in-memory computing. Nature Computational Science, 5(1):13–26, 2025. doi: 10.1038/s43588-024-00753-x. URL https://do i.org/10.1038/s43588-024-00753-x. He, Y., Mao, H., Giannoula, C., Sadrosadati, M., GómezLuna, J., Li, H., Li, X., Wang, Y., and Mutlu, O. Papi: Exploiting dynamic parallelism in large language model decoding with a processing-in-memory-enabled computing system. In Proceedings of the 30th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2, ASPLOS ’25, pp. 766–782, New York, NY, USA, 2025. Association for Computing Machinery. ISBN 9798400710797. doi: 10.1145/3676641.3716009. URL https: //doi.org/10.1145/3676641.3716009. Heo, G., Lee, S., Cho, J., Choi, H., Lee, S., Ham, H., Kim, G., Mahajan, D., and Park, J. Neupims: Npupim heterogeneous acceleration for batched llm inferencing. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3, ASPLOS ’24, pp. 722–737, New York, NY, USA, 2024. Association for Computing Machinery. ISBN 9798400703867. doi: 10.1145/3620666.3651380. URL https: //doi.org/10.1145/3620666.3651380. Hifi, M., Michrafy, M., and Sbihi, A. Heuristic algorithms for the multiple-choice multidimensional knapsack problem. Journal of the Operational Research Society, 55(12):1323–1332, Dec 2004. ISSN 1476-9360. doi: 10.1057/palgrave.jors.2601796. URL https://doi. org/10.1057/palgrave.jors.2601796.

HELIX: HEterogeneous LLM Inference with FleXible and Fine-Grained Partitioning

Hong, Y., Yang, J., Peng, B., and Yao, J. Repa: Reconfigurable pim for the joint acceleration of kv cache offloading and processing. In Proceedings of the 31st ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2, ASPLOS ’26, pp. 1622–1639, New York, NY, USA, 2026. Association for Computing Machinery. ISBN 9798400723599. doi: 10.1145/3779212.3790212. URL https://doi.org/10.1145/3779212.3790 212. Hu, Y., Imes, C., Zhao, X., Kundu, S., Beerel, P. A., Crago, S. P., and Walters, J. P. Pipeedge: Pipeline parallelism for large-scale model inference on heterogeneous edge devices. In 2022 25th Euromicro Conference on Digital System Design (DSD), pp. 298–307, 2022. doi: 10.1109/ DSD57027.2022.00048. Jang, Y., Hwang, S., Lee, H., Jung, S., Lee, D., Shim, W., and Kung, J. Dissecting and re-architecting 3d nand flash pim arrays for efficient single-batch token generation in llms. In 2025 IEEE 43rd International Conference on Computer Design (ICCD), pp. 237–244, 2025. doi: 10.1109/ICCD65941.2025.00040. JIANG, Y., Yan, R., and Yuan, B. Hexgen-2: Disaggregated generative inference of LLMs in heterogeneous environment. In The Thirteenth International Conference on Learning Representations, 2025. URL https: //openreview.net/forum?id=Cs6MrbFuMq. Kim, W., Lee, Y., Kim, Y., Hwang, J., Oh, S., Jung, J., Huseynov, A., Park, W. G., Park, C. H., Mahajan, D., and Park, J. Pimba: A processing-in-memory acceleration for post-transformer large language model serving. In Proceedings of the 58th IEEE/ACM International Symposium on Microarchitecture, MICRO 2025, pp. 292–307. ACM, October 2025. doi: 10.1145/3725843.3756121. URL http://dx.doi.org/10.1145/3725843 .3756121. Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., and Stoica, I. Efficient memory management for large language model serving with pagedattention, 2023. URL https://arxiv.or g/abs/2309.06180. Li, Z., Mei, Y., Pan, S., Chen, Z., Zhang, X., Wu, M., Lin, S., Yang, Y., Zhang, Z., Miao, X., Vinayak, R., and Jia, Z. Opweave: Flexible operator disaggregation for heterogeneous llm serving, 2026. URL https: //arxiv.org/abs/2609.14237. Liang, L., Luo, T., Zhong, S., Zhao, D., Ma, Q., Wei, R., Wang, J., Li, M., Sun, G., Wang, Z., and Cai, Y. Diamond: Dynamic inference for adaptive edge moe with heterogeneous in-nand and near-dram compute architecture. In

2026 ACM/IEEE 53rd Annual International Symposium on Computer Architecture (ISCA), pp. 985–1000, 2026. doi: 10.1109/ISCA66397.2026.00078. Liang, Y., Chen, Z., Zuo, P., Zhou, Z., Chen, X., and Yu, Z. Injecting adrenaline into llm serving: Boosting resource utilization and throughput via attention disaggregation, 2025. URL https://arxiv.org/abs/2503.2 0552. Lin, B., Zhang, C., Peng, T., Zhao, H., Xiao, W., Sun, M., Liu, A., Zhang, Z., Li, L., Qiu, X., Li, S., Ji, Z., Xie, T., Li, Y., and Lin, W. Infinite-llm: Efficient llm service for long context with distattention and distributed kvcache, 2024. URL https://arxiv.org/abs/2401.02669. Ma, Y., Shen, S., Chen, T., Langer, A., Kraus, J., Glick, B., Belusar, C., Hammond, J., and Hoefler, T. Demystifying nvshmem: A system-level analysis on symmetric memory and device-initiated operations in gpu communication, 2026. URL https://arxiv.org/abs/2606.0 5951. NVIDIA Corporation. TensorRT-LLM, 2026. URL https: //github.com/NVIDIA/TensorRT-LLM. Opensource library for high-performance LLM inference on NVIDIA GPUs. ovidiusm, Ranadive, A., Khazraee, M., Rivkin, A., Hirsch, C., Stamler, T., Mailhot, N., Pressler, D., Akhmetshin, R., Shalev, M., e-eygin, Yastrebov, I., Buslov, V., Karthik, A., Leksikov, E., Vegas, T., Alroy, I., Brinskiy, M., Danino, R., NirWolfer, Vijayakrishnan, P., Xiang, Y., fengjica, Morag, G. E., Nayshtut, A., Itigin, Y., yaroslavs, Zegerson, Y., Garaud, O., and dfarge. Nvidia inference xfer library. https://github.com/ai-dynamo/nixl, jul 24 2026. URL https://github.com/ai-dynamo/ nixl. Park, J., Choi, J., Kyung, K., Kim, M. J., Kwon, Y., Kim, N. S., and Ahn, J. H. Attacc! unleashing the power of pim for batched transformer-based generative model inference. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2, ASPLOS ’24, pp. 103–119, New York, NY, USA, 2024. Association for Computing Machinery. ISBN 9798400703850. doi: 10.1145/3620665.3640422. URL https://doi.or g/10.1145/3620665.3640422. Patel, P., Choukse, E., Zhang, C., Shah, A., Goiri, I., Maleki, S., and Bianchini, R. Splitwise: Efficient generative llm inference using phase splitting. In 2024 ACM/IEEE 51st Annual International Symposium on Computer Architecture (ISCA), pp. 118–132, 2024. doi: 10.1109/ISCA59077.2024.00019.

HELIX: HEterogeneous LLM Inference with FleXible and Fine-Grained Partitioning

Pei, Y. R., C, K., King, G., Zhou, H., Mancuso, J., Sharma, A., ishandhanani, Panda, B., Gil, D., McCormick, R., Alec, hhzhang16, jthomson04, Tokarev, D., Prashanth, K., Shah, N., Verma, T., Sharma, T., atchernych, SaturleyHall, H., Saereesitthipitak, S., Kan, T.-L., Hao, J., Montfort, T., Hung, K., mohammedabdulwahhab, Agarwal, A., Olson, R., GuanLuo, and Jacky. ai-dynamo/dynamo, 2026. URL https://docs.nvidia.com/dyna mo/getting-started/quickstart. Qin, R., Li, Z., He, W., Zhang, M., Wu, Y., Zheng, W., and Xu, X. Mooncake: A kvcache-centric disaggregated architecture for llm serving, 2025. URL https://ar xiv.org/abs/2407.00079. SGLang Team. SGLang, 2026. URL https://github .com/sgl-project/sglang. High-performance serving framework for large language models and multimodal models. Song, C., Wang, M., Zhou, H., Liang, H., Lyu, Y., Chen, Z., Fan, Y., and Zhou, Z. Analytical provisioning for attention-ffn disaggregated llm serving under stochastic workloads, 2026. URL https://arxiv.org/abs/ 2601.21351. Tarnawski, J., Phanishayee, A., Devanur, N., Mahajan, D., and Paravecino, F. N. Efficient algorithms for device placement of dnn graph operators. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS ’20, Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN 9781713829546. Team, S. Step-3 is large yet affordable: Model-system codesign for cost-effective decoding, 2025. URL https: //arxiv.org/abs/2507.19427. Topcuoglu, H., Hariri, S., and Wu, M.-Y. Performanceeffective and low-complexity task scheduling for heterogeneous computing. IEEE Transactions on Parallel and Distributed Systems, 13(3):260–274, 2002. doi: 10.1109/71.993206. Xu, T., Liu, Y., Lu, X., Zhao, Y., Zhou, X., Feng, A., Chen, Y., Shen, Y., Zhou, Q., Chen, X., Sherstyuk, I., Li, H., Thakkar, R., Hamm, B., Li, Y., Huang, X., Wu, W., Shanbhag, A., Kim, H., Chen, C., and Lai, J. Aiconfigurator: Lightning-fast configuration optimization for multi-framework llm serving, 2026. URL https://arxiv.org/abs/2601.06288. Yang, J., Li, J., Fu, Y., Zhao, H., Chen, Z., Wu, Q., Yang, Y., and Yan, B. Beyond prefill-decode disaggregation: Dissecting llm inference for heterogeneous platforms via dynamic operator scheduling, 2026. URL https:// arxiv.org/abs/2607.25498.

Zhong, Y., Liu, S., Chen, J., Hu, J., Zhu, Y., Liu, X., Jin, X., and Zhang, H. Distserve: disaggregating prefill and decoding for goodput-optimized large language model serving. In Proceedings of the 18th USENIX Conference on Operating Systems Design and Implementation, OSDI’24, USA, 2024. USENIX Association. ISBN 9781-939133-40-3. Zhu, R., Jiang, Z., Jin, C., Wu, P., Stuardo, C. A., Wang, D., Zhang, X., Zhou, H., Wei, H., Cheng, Y., Xiao, J., Zhang, X., Liu, L., Lin, H., Chang, L.-W., Ye, J., Yu, X., Liu, X., Jin, X., and Liu, X. Megascale-infer: Efficient mixtureof-experts model serving with disaggregated expert parallelism. In Proceedings of the ACM SIGCOMM 2025 Conference, SIGCOMM ’25, pp. 592–608, New York, NY, USA, 2025. Association for Computing Machinery. ISBN 9798400715242. doi: 10.1145/3718958.3750506. URL https://doi.org/10.1145/3718958. 3750506.

11

AI U SAGE D ISCLOSURE

Generative AI was used to assist with code development, including plotting code for visualizing our experimental results, and with grammatical and typographical proofreading, not for writing.

HELIX: HEterogeneous LLM Inference with FleXible and Fine-Grained Partitioning

VALIDATION OF THE C OST-M ODEL A SSUMPTIONS

Dominance of the main computational primitives on the execution time: Figure 10 breaks down end-to-end execution time on an A100 GPU. GEMMs, attention, and state scan/update operations account for over [95%] of the total runtime across all evaluated models, while the remaining operations (e.g., normalization, activations, element-wise ops) contribute only in a small fraction.

Latency Ratio (%)

SSM

100 80 60 40 20 0

Attention

4k 8k 16k Llama-3.1-8B

GEMM

4k 8k 16k Mixtral-8x7B

Others

10

B200

fi

H100

fl

Pro led

Roo ine

103 Decode attention latency [µs]

Our placement problem formulation relies on three assumptions: (i) end-to-end latency is dominated by a small set of computational primitives, (ii) a roofline-based cost model does not give a sufficiently accurate latency estimate for these primitives, requiring per kernel profiling, and (iii) kernels profiled in isolation are representative of their behavior within full model execution. This section provides empirical evidence for each.

A100 2

GEMM latency [µs]

A

102 101 100

N = K = 4096 2

4

8

16 32 Batch size

64

128

10−1

ISL 128 → 16K 2

4

8

16 32 Batch size

64

128

Figure 11. Latency estimation based on roofline cost-model vs. obtained from profiling for A100, H100 and B200 GPUs.

Isolated kernel profiling vs. end-to-end runs Finally, we verify that profiling kernels individually does not distort their latency compared to execution within the full model, where effects such as caching or fusion can lead to differences in the timings. Figure 12 shows that latencies obtained from isolated profiling closely match those extracted from end-to-end traces.

4k

Bamba-9B

Figure 10. Breakdown of end-to-end execution time on an A100 between the main computational primitives (GEMMs, attention, state scan and update) and the remaining operations.

Accuracy of roofline cost models: Figure 11 compares the latencies predicted by the roofline model with those measured through profiling on A100, H100, and B200 GPUs. For GEMMs, the estimates only loosely track the measurements, since achieved resource utilization varies across operation shapes. In contrast, strictly memory-bound attention (ISL of 16K tokens, shown as large dots in the right plot) aligns most closely with the roofline estimates. The roofline models for the three GPUs are defined in A, while profiling numbers were extracted from AIConfigurator source code(Xu et al., 2026). GPU

Peak compute (TOPS)

HBM bandwidth (TB/s)

A100 H100 B200

624 2000 4500

2.0 3.35 8.0

Table 2. Peak compute throughput and memory bandwidth used in the roofline cost model.

Figure 12. Individual kernel profiling versus end-to-end latency kernel extraction

B PARTITION SEARCH ALGORITHM Algorithm 1 gives the pseudocode of the placement search based on dynamic programming described in Section 4. Preprocessing: Memory is tracked only on the set T of devices whose capacity can be exceeded if every stage that lists them as a candidate were placed there. Stage footprints and capacities are expressed in units of the greatest common divisor (gcd) of the stages b = gcdi m(si ), rounding stage sizes up and capacities down, so each device’s usage is a small integer. Labels: A label ℓ is a partial placement of s0 , . . . , si . It stores the memory usage uℓ on tracked devices, (ii) the placement of the frontier, i.e. the stages whose outputs are still needed by unplaced consumers, together with the time each output within is available on each device, and (iii) the running latency Lℓmax of the placed subgraph until the current stage.

HELIX: HEterogeneous LLM Inference with FleXible and Fine-Grained Partitioning

Algorithm 1 HELIX partitioning search Require: DAG G = (S, E), Ci , Tcomp , Tcomm , m(si ), Md , dinit ∗ Ensure: Placement P a 1: T ← {d : i: d∈Ci m(si ) > Md } 2: b ← gcdi m(si ) 0 3: L ← {ℓ0 } with empty frontier, uℓ0 = 0, Lℓmax =0 4: for i = 0, . . . , n − 1 do 5: L′ ← ∅ 6: for all ℓ ∈ L and d ∈ Ci do 7: if d ∈ T and uℓd + ⌈m(si )/b⌉ > ⌊Md /b⌋ then 8: continue 9: end if 10: ℓ′ ← ℓ ⊕ (si 7→ d) ′ 11: uℓd ← uℓd + ⌈m(si )/b⌉ if d ∈ T 12: Qd′ →d ← Ai−1 if aℓ (i − 1) = d else 0 ′ 13: Lℓi (d) ← Eq. 5 ′ 14: Lℓj (d) ← Ai,aℓ (j) for each j ∈ pred(i) transferred to d
′ ′ 15: Lℓmax ← max Lℓmax , Lℓi (d) 16: Remove from the frontier of ℓ′ every sj with no consumer after si 17: L′ ← L′ ∪ {ℓ′ } 18: end for 19: if L′ = ∅ then 20: return infeasible 21: end if 22: L ← labels of L′ not Pareto-dominated within their frontier assignment 23: end for 24: ℓ∗ ← arg minℓ∈L Lℓmax 25: return a∗ = aℓ∗

Label expansion and pruning: For each stage si , every label ℓ is expanded with every candidate device d ∈ Ci . If d ∈ T and placing si on d would exceed its capacity, the expanded label is discarded. Otherwise, the start time of si is the latest arrival over its predecessors, each incurring a transfer cost when placed on a different device (Eq. 5), ′ and its finish time updates Lℓmax . Stages with no remaining consumers are removed from the frontier. Among labels with the same frontier placement, a label is discarded if another has no larger memory usage, running latency, and frontier arrival times. If no label survives at some stage, the assignment is infeasible. After the last stage, the label with minimal Lℓmax gives the placement a∗ . Graph structure and number of devices. Figure 13 reports the search time of HELIX-DP without tracking any memory devices (|T | = 0) on three synthetic DAGs: (a) a sequential chain, (b) a chain with a single residual skip from the first to the last stage, and (c) a transformer-like DAG with residual skips across every block. All devices

can execute every stage, but they differ in compute throughput and memory bandwidth, so stage costs differ across devices and the search remains heterogeneous. We additionally vary the number of devices D from 2 to 4. For the chain and the transformer-like DAG, we sweep the number of stages N from 32 to 1024. Search time grows roughly linearly with N , as expected from the O(N |D|ω+1 ) bound, and stays below 0.4 s at 1024 stages with four devices. Although the transformer-like DAG contains many residual skips, each one closes at the next fan-in, within a couple of stages. Therefore ω is bounded by dependencies rather than by the depth of the model. The long skip in (b) is the adversarial case. The frontier spans every stage between a residual producer and its fan-in, so a skip from the first to the last stage makes it cover the entire graph, and ω grows with N . The number of frontier keys then grows as |D|N , matching exhaustive search, so we restrict this sweep to N ≤ 12. Search time increases by roughly a factor of D per added stage, reaching ∼0.1 s for D = 2 and ∼9 s for D = 3 at 12 stages, and ∼60 s for D = 4 at only 11 stages. LLM DAGs avoid this regime, since their residual connections span only a few stages. Memory tracking. Figure 14a evaluates the approximation introduced by tracking occupancy in buckets instead of bytes. We take the first five layers of Llama 8B with 5 stages each of different footprints (QKV, Attention, Out, Gate-Up, Down) and simulate the search for a batch of 32 queries at a context length of 16K tokens, and place them on three devices: an H100, 3DCIM and DRAM-PIM. Parameters and KV cache, both in 16-bit precision, create a DAG footprint of ≈ 12 GB of memory, and therefore both DRAM-PIM and H100 are set with a memory of 10 GB to force the tracking of their capacities. Exhaustive search (or brute force BF) runs over the same stages and conditions till termination, and therefore rovides the optimal placement. We then run HELIX-DP with a bucket size b ranging from 0.01% to 10% of device capacity. This means the memory is treacked in units of between 10 MB to 1 GB, whereas the smallest stage of this DAG (Out projection of 4096×4096 parameters) has a ≈ 33 MB of footprint. We compare the search times and their reported latency of the found placement and report them both relative to BF. At b = 0.01%, HELIX-DP returns the same placement as BF with a search time over three orders of magnitude lower. Coarser buckets can reduce search time by roughly another order of magnitude, since more occupancy vectors coincide, which reduces the term P from the complexity, meaning more labels are pruned by dominance. This comes at the cost of rounding each footprint up to ⌈m(si )/b⌉, which overestimates occupancy and the DP search ends up rejects feasible placements that could lead to the minimum latency. The latency of the found placement is to 1.08× slower than optimal solution retruned by exhaustive search at b = 0.1%, 1.46× at b = 1%, and close

HELIX: HEterogeneous LLM Inference with FleXible and Fine-Grained Partitioning … (a)

Chain of stages

…

… (b)

Skip = N1

(c)

Transformer-like

Figure 13. HELIX-DP search time without memory tracking for (a) a chain of stages, (b) a chain with a residual skip from the first to the last stage (note the smaller range of N ), and (c) a transformer-like DAG, with D ∈ {2, 3, 4} heterogeneous devices. (a)

(b)

(c) M2 D3

Figure 14. (a) Search time and accuracy of the found partition in terms of latency for HELIX DP versus exhaustive search (brute-force BF) depending on the bucket size, i.e. the unit on which the memory capacity is tracked for all devices. (b) Dependency of the search time for D=3 devices and T =2, varying on which of the devices remains out of the memory tracking (not overflowable or unconstrained) and (c) Dependency of the search time depending on how constrained the devices are with respect to the footprint of the DAG stages

to 2× at b = 10%. This motivates setting b to the greatest common divisor of the stage footprints (Section 4), which represents every footprint exactly. We follow the same testing settings as for Figure 14a, but while varying the size of the DAG in terms of stages, memory capacities are adjusted proportionally to keep memory feasible placements. Figure 14b fixes D = 3 and |T | = 2 and varies which device is left untracked out of the three considered ones. Leaving the fastest device untracked keeps search time around 10−2 s up to 128 stages, while leaving the slowest device outside the memory tracking is indistinguishable from tracking all three, reaching ∼2500 s at 128 stages. The device with the lowest computational cost (i.e the fastest) for the stages dominates the placement, so its occupancy takes the largest number of distinct values and dominates the number of retained labels P . The slowest device receives few stages, so tracking it barely adds states. Search cost therefore highly depends on whether the devices that excel at the majority of stages in terms of execution cost, must be tracked. Figure 14c fixes N = 64 and |T | = 3 and varies the memory pressure, defined as the aggregate stage footprint over the aggregate device capacity. This shows the dependency of the search time with how fast devices hit their memory capacity when searching through the graph in topological

order. If a DAG has a footprint, for instance of 12 GB, and an admissible device d has only 1 GB of capacity, the amount of labels that will involve d will significanlty reduce compared to a case on which Md is 11.99 GB. Search time drops from ∼130 s at 37% pressure to ∼8 s at 96%. Tighter capacities leave fewer reachable occupancy vectors feasible, so P saturates earlier and more labels are discarded. The expensive regime is thus the loosely constrained one, where capacities are large relative to the footprint. In practice, this regime often makes tracking unnecessary altogether, since devices that cannot overflow are removed from T .

C

C OMPUTATIONAL PRIMITIVES OF

EVALUATED MODELS Table C shows the number of different computational primitives each of the evaluated models from the results section have, and therefore the ones that need pre-profiled costs in order to perform the partition search. Llama 8B and Mixtral 8x7B have 5 distinct primitives, i.e. placeable stages for the DAG construction, that repeat across all layers. Additionally, we also consider the first and final linear layers of the network (embedding and lm head). Bamba 9B instead alternates between mamba-like (in, mixer and out) and

HELIX: HEterogeneous LLM Inference with FleXible and Fine-Grained Partitioning

attention-like layers (QKV, attn and out), both followed by the standard FFN sequence of projections (Gate-Up and Down). Many newer models are still based on standard transformer-like patterns. They vary in their hidden dimensions, number of layers and number of experts, but can usually be summarized with the same computational primitives. A current trend is to replace attention blocks with more computationally efficient versions. For example, Qwen3-Next-80B and Kimi-Linear-48B extend the MoE structure and interleave linear attention layers with standard (GQA) or latent attention (MLA) layers, respectively. These newer attention blocks take a higher number of input tensors than the usual queries, keys and values, and therefore vary in the number of projections around the attention primitive. This represents further opportunities for HELIX to merge or split such parallel projections across hardware.

D

places the experts of an FFN layer vertically across different tiers of 3DCIM tiles, our experiments for MoE-based workloads model serialized expert computation.

E

D ETAILED HELIX PARTITIONS

This section shows the full partitions found by HELIX for each heterogeneous system, with each stage coloured by the device it is mapped to. Their discussion can be found in the results section. Llama 8B

3DCIM SYSTEM SIMULATOR AND

BACKEND The numbers for the 3DCIM platform are obtained via a simulator whose frontend is adapted to be registered as a custom backend for HELIX. With this frontend, the simulator can directly consume the graph captured by TorchDynamo when HELIX partitions the workload and allocates a partition to 3DCIM. To match individual operators to accelerator-level operators, the simulator transforms the given FX-level graph with a series of transformations based on pattern matching and replacement. For instance, linear layers involving individual slices of the Q, K, and V matrices are fused into a single linear layer with the combined QKV matrix, which can be efficiently computed on analog 3DCIM tiles. The transformed torch.fx graph is then lowered to a C++-based IR (Almaci et al., 2026), which manages lower-level operations such as tiling, batch splitting, data transfer modeling, and mapping. The resulting computational graph IR is scheduled, and the end time of the last operation is reported as the latency. The 3DCIM platform system architecture is adopted from (Büchel et al., 2025), which utilizes analog 3DCIM tiles of size 512 × 512, connected to a landing pad with I/O buffers via a common bus. This bus conservatively blocks concurrent broadcasts and gathers across multiple activations. The bus transmits a tile-width’s worth of inputs (512 bytes) in a bit serial manner with a startup latency of 50 cycles, according to internal circuit design findings. A total of 1024 3DCIM tiles are used to accommodate the weight capacity requirements of frontier LLM models. The MVM integration time on a 3DCIM tile is 200 ns. The greedy in-order weight mapping strategy from (Büchel et al., 2025) is used, which maximizes tensor parallelism by distributing a single-stage MVM computation across as many tiles as possible. 3DCIM tiles are constrained by the one-tier-ata-time constraint, which allows only one memory tier to be activated at a time. Since the weight mapping strategy

BS Mixtral 87B

BS

Figure 15. Partitions for 3DCIM + H100 for Llama 8B and Mixtral 8x7B

Bamba 9B

BS

Figure 16. Partitions for 3DCIM + H100 for Bamba 9B

Partitions for H100 and + DRAM-PIM : The AttAcc simulator (Park et al., 2024) only models attention, so we estimate projection latencies on DRAM-PIM with a roofline model using AttAcc’s bank-level, power-constrained configuration. Each HBM3 stack provides an internal bandwidth of 9 × 670.4 GB/s, and the PIM units are provisioned at 1 FLOP per byte, giving BWPIM ≈ 30 TB/s and FLOPSPIM ≈ 30 TFLOPS for 5 stacks. A projection with

HELIX: HEterogeneous LLM Inference with FleXible and Fine-Grained Partitioning

Model

# of nodes

# of Layers

# of Primitives

No. of Stages

Llama-3.1 8B Bamba 9B Mixtral 8x7B GPT-OSS Qwen3-Next-80B

2995 1784 3043 3455 4525

32 32 32 36 48

5 (QKV, Out, Attn, Gate-Up, Down) 8 (In, Mixer, Out — QKV, Attn, Out + Gate-Up, Down) 5 (QKV, Out, Attn, Router, Experts) 5 (QKV, Out, Attn, Router, Experts) 9 (QKVZ, BA, LinAttn, Out — QKV, Attn, Out

162 162 162 182 278

Kimi-Linear-48B

2203

27

+ Router, Experts) 12 (QKVZ, BA, LinAttn, F B, G B, Out —

204

KVA, QZ, MLA, Out + Router, Experts)

an R × C weight matrix applied to V input vectors moves B = B (RC + V C + V R) bytes, with B bytes per element, and performs F = 2V RC operations. Its latency is   F B , . TPIM = max BWPIM FLOPSPIM In FP16, the arithmetic intensity of a projection is roughly V FLOP/byte, so it reaches the compute roof already at V ≈ 1 and its latency then grows linearly with V . In contrast, the H100 remains memory-bound up to a few hundred vectors. DRAM-PIM is therefore faster only while its ∼9× bandwidth advantage outweighs this linear growth, and the H100 outperforms it after a few concurrent queries (for our studies after batch size 8). Llama 8B

H100

DRAMPIM

Llama 8B

DRAMPIM

BS

Mixtral 87B

BS

Figure 17. HELIX partitions for H100 + DRAM-PIM

Bamba-9B

DRAMPIM

BS

Figure 19. HELIX partitions for B200 + 3DCIM + DRAM-PIM for Llama 8B and Mixtral 8x7B

BS

Figure 18. HELIX partitions for B200 + 3DCIM + DRAM-PIM for Bamba

Related documents

• PipeMax: Enhancing Offline LLM Inference on Commodity GPU Servers#155244
• LeanStream: A Speculate-and-Refine Streaming Framework for Efficient on-Device LLM Inference#925446
• LeanStream: A Speculate-and-Refine Streaming Framework for Efficient on-Device LLM Inference#921210
• Beyond Greedy Chunking: SLO-Aware Sliding-Window Scheduling for LLM Inference#259421
• DolphinRouting: Cluster-Level Binary Routing for Disaggregated LLM Inference Serving#1491900