PaperlayerView exact revision on arXiv

Source-linked AI summary

DynaFlow: Transparent and Flexible Intra-Device Parallelism via Programmable Operator Scheduling

Yi Pan, Yile Gu, Jinbin Luo, Yibo Wu, Ziren Wang, Hongtao Zhang, Ziyi Xu, Shengkai Lin, Baris Kasikci, Stephanie Wang

arXiv:2605.21603v1cs.DCcs.LG

TL;DR

Heterogeneous operator resource demands leave device resources idle, while intra-device parallelism is difficult to integrate with existing sequential models and sensitive to execution context. DynaFlow separates operator scheduling from model implementation and provides programmable scheduling with an execution backend. Across six ML systems, it achieves up to 1.29x throughput improvement with minimal code changes.

  • Problem

    Heterogeneous operator resource demands leave device resources idle, reducing inference and training throughput.

    Source: §1 · p. 1

  • Method

    DynaFlow decouples operator scheduling from model implementation, combining graph partitioning and programmable scheduling with a backend that manages execution dependencies.

    Source: §1 · p. 2

  • Results

    Up to 1.29x throughput improvement was achieved across six ML systems with minimal code changes.

    Source: §1 · p. 2

  • Takeaways & Limitations

    DynaFlow integrates representative intra-device parallelism strategies into state-of-the-art ML systems while preserving optimizations such as CUDA Graphs.

    Source: §7 · p. 12

    Naive batch splitting can worsen performance because each additional micro-batch requires another read of the model weights.

    Source: §2.1 · p. 4

Abstract

from arXiv · show

Intra-device parallelism addresses resource under-utilization in ML inference and training by overlapping the execution of operators with different resource usage. However, its wide adoption is hindered by a fundamental conflict with the static, sequential programming model of existing frameworks. Integrating these strategies requires invasive, model-specific code overhauls, representing an intractable engineering cost. This is further amplified by the high sensitivity of strategies to execution contexts (e.g., workload, model architecture, hardware), forcing developers to implement and maintain multiple specialized solutions. To address this, we propose DynaFlow, a framework that enables the transparent and flexible integration of intra-device parallelism by decoupling the logical model definition from the physical execution schedule. DynaFlow introduces a flexible frontend with annotations for graph partitioning and a programmable interface for defining custom intra-device parallelism strategies. Its efficient backend manages complex control/data-flow asynchronously, uses custom memory management to eliminate copy overheads, and preserves compatibility with optimizations like CUDA Graphs and TorchInductor. We demonstrate that DynaFlow can integrate representative parallelism strategies into 6 state-of-the-art ML systems with minimal code changes, achieving up to a 1.29x throughput improvement. DynaFlow is publicly available at https://github.com/uw-syfi/DynaFlow.

1 INTRODUCTION

Intra-device parallelism can improve resource utilization, but its non-sequential schedules conflict with existing sequential programming models and are costly to integrate. DynaFlow decouples operator scheduling from model implementation and provides a programmable frontend with an execution backend designed to manage the resulting complexity.

1 INTRODUCTION

  • Representative strategies overlap computation and communication, fuse kernels, or split input batches for concurrent execution.These approaches break sequential execution to overlap heterogeneous operations and improve hardware utilization.Source: Figure 1 · p. 1, §1 · p. 1
  • Sequential programming models make non-sequential parallelism costly to integrate, while strategy effectiveness varies by workload, model architecture, and hardware.The paper cites a dual-batch-overlap implementation in SGLang that took more than two months and 1.3K specialized lines for one model.Source: §1 · p. 2
  • DynaFlow decouples operator execution from model implementation, allowing users to define non-sequential execution plans without changing the model definition.The approach retains existing sequential and imperative programming models while adding a dynamic, programmable execution substrate.Source: §1 · p. 2
  • DynaFlow combines annotation-based graph partitioning and programmable scheduling with an asynchronous backend for control- and data-flow management.Its design must reconcile flexible scheduling with low overhead and compatibility with static-graph optimizations.Source: §1 · p. 2
  • Up to 1.29x throughput improvement was achieved across six ML systems with minimal code changes.DynaFlow also matched or outperformed specialized implementations by up to 1.1x.Source: §1 · p. 2

2 BACKGROUND AND MOTIVATION

Heterogeneous operators leave resources idle when executed sequentially, motivating intra-device strategies that overlap work but require careful scheduling. Their model-wide, runtime-sensitive integration creates substantial engineering and performance challenges, which DynaFlow addresses by decoupling execution from model implementation.

2.1 Intra-Device Parallelism Strategies

  • Sequential execution of compute-, memory-, and network-bound operators can stall one resource while another is active, leaving up to 15% performance on the table during network stalls.Intra-device strategies seek to improve utilization by overlapping heterogeneous operations.Source: §2 · p. 3, §2.1 · p. 3
  • Overlapping independent operators, fusing dependent operators in custom kernels, and splitting batches into micro-batches are three strategies for creating overlap.Their scheduling and implementation choices make deciding when to apply each strategy nontrivial.Source: §2.1 · p. 3
  • Overlap can increase peak memory and create resource contention; fusion requires maintaining expert-designed kernels across many operator combinations and optimizations.The described vLLM MoE-kernel abstraction is one-off and does not generalize to its other operators.Source: §2.1 · p. 3
  • Batch splitting adds weight reads per additional micro-batch, may worsen performance when applied naively, and requires control over micro-batch sizes and split/merge copies.It also inherits the challenges of overlapping and requires choosing when and at what granularity to split.Source: §2.1 · p. 4

2.2 Our Approach and Challenges

  • Intra-device parallelism conflicts with frameworks’ sequential execution order, requiring invasive global rewrites that can demand months of specialized code for one model.Its best strategy also varies with workload, model architecture, and hardware, so static configurations do not suffice.Source: §2.2 · p. 4
  • DynaFlow decouples operator execution from model implementation, adding an abstraction layer for transparent, flexible, context-aware schedules.The model can retain its sequential, imperative definition while users specify non-sequential execution plans.Source: §2.2 · p. 4
  • Scheduling granularity must avoid both the overhead and fusion obstruction of tiny operators and the lost intra-layer overlap opportunities of whole-layer scheduling.The section identifies granularity, a unified runtime-adaptive schedule interface, and performance-preserving execution as core design challenges.Source: §2.2 · p. 4
  • A flexible scheduler must manage control- and data-flow dependencies, avoid split/merge copy overheads, and remain compatible with static-graph optimizations such as CUDA Graphs and TorchInductor.Naive scheduling overhead or data movement can negate parallelism gains.Source: §2.2 · p. 4

3 DESIGN

DynaFlow separates model definition from execution scheduling through a frontend for graph partitioning and programmable schedules, paired with a backend for asynchronous execution and data-flow management. Its backend addresses resharding overheads and preserves compatibility with static-graph optimizations.

3.1 Overview

  • The frontend expresses parallelism strategies and scheduling granularities, while the backend manages dynamic dependencies, minimizes copies, and supports CUDA Graphs and TorchInductor.Integration proceeds through graph partitioning, a Python scheduler, and runtime dispatch to the backend.Source: §3.1 · pp. 4–5

3.2 Frontend Design

  • The frontend takes a TorchDynamo-traced graph as input, enabling developers to define custom parallelism strategies over extracted PyTorch graphs.The paper notes TorchDynamo is supported by systems including vLLM and SGLang.Source: §3.2 · p. 5

3.2.1 Graph Partition

  • Graph partition annotations target logical operators: SplitModule partitions at module boundaries, SplitFunc around function calls, and dynaflow.mark wraps other code sections.The design avoids scheduling tiny tensor-arithmetic operators, whose fine granularity can add more overhead than benefit.Source: §3.2.1 · p. 5

3.2.2 Programmable Scheduling

  • The Python-native scheduler balances high-level control with backend-managed data-flow, using split, get_ready_ops, and execute to initialize micro-batches, query ready subgraphs, and dispatch operators.The interface can express DBO’s conditional splitting and overlap, as well as TokenWeave-style fusion and overlap.Source: §3.2.2 · pp. 5–7

3.3 Backend Design

  • The backend executes dynamic schedules generated by the frontend.Its design addresses the performance challenges of dynamic scheduling.Source: §3.3 · p. 7

3.3.1 Data-Flow and Memory Management

  • Static graph analysis records tensor reference counts and merge preallocation needs; runtime hooks direct outputs into contiguous buffers, enabling zero-copy resharding.Reference counts support lifecycle management and garbage collection, while preallocation avoids the copy overhead of naive splitting or concatenation.Source: §3.3.1 · p. 7

3.3.2 Compatibility with Low-Level Optimizations

  • DynaFlow applies TorchInductor compilation per subgraph and captures separate CUDA graphs per micro-batch and subgraph, reusing graph pools across subgraphs and batch sizes.These static optimizations operate at subgraph level while the runtime Python scheduler dispatches preprocessed subgraphs dynamically.Source: §3.3.2 · p. 7

4 IMPLEMENTATION

DynaFlow is implemented as a torch.compile backend that captures PyTorch operator graphs and supports asynchronous execution of scheduled subgraphs.

4 IMPLEMENTATION

  • DynaFlow is a 4.1K-line Python torch.compile backend that can be plugged into PyTorch systems by intercepting the model forward call.The cited text notes that the dynaflow.mark wrapper depends on newer PyTorch versions; the examples are compatible with PyTorch 2.6+.Source: §4 · pp. 7–8
  • The supplied table captions describe integration and strategy implementation costs in lines of code, but provide no cell values.Table 1 distinguishes core-framework, per-model, and attention-backend changes; Table 2 covers different parallelism strategies.Source: Table 1 · p. 8, Table 2 · p. 8
  • TorchDynamo captures an operator-level graph from Python bytecode, supporting compatibility with PyTorch systems such as vLLM and SGLang.The graph is the frontend input to DynaFlow's scheduling process.Source: §4 · p. 8
  • An asynchronous callback-driven backend tracks graph dependencies and returns ready subgraphs to the frontend scheduler.Dispatch requests are enqueued by the frontend; the engine updates dependency state and determines which subgraphs are ready.Source: §4 · p. 8

5 EVALUATION

The evaluation tests representative intra-device parallelism strategies across six ML systems, on B200 and H100 testbeds, using open-source models and real and synthetic workloads. It measures integration effort, backend overhead, throughput, initialization cost, and ablation effects.

5.2.1 Frontend Effectiveness

  • 75 LoC integrated DynaFlow into vLLM, MoE model annotations averaged 8 lines per model, and other system changes were usually under 100 lines.Implementing six strategies averaged 11 LoC for graph partitioning rules and 31 LoC for dynamic scheduling logic.Source: §5.2.1 · p. 9

5.2.2 Backend Efficiency

  • 6.4x lower CPU execution time was achieved with low-level optimizations than with a non-optimized variant; sequential fallback took 4.7 ms versus vLLM's 4.4 ms.Dynamic scheduling took 10.8 ms, below GPU execution time for larger batches, the primary scenarios where intra-device parallelism is beneficial.Source: Figure 8 · p. 8, §5.2.2 · p. 9

5.3.2 Dual-Batch Overlap

  • 1.14x throughput improvement was achieved by DynaFlow-enabled dual-batch overlap, with up to 1.1x better performance than vLLM's hand-built implementation.On lightweight real workloads, aggressive batch splitting can be detrimental; DynaFlow's dynamic scheduling avoids the degradation observed with vLLM's static threshold.Source: Figure 10 · p. 10, §5.3.2 · p. 10

5.3.1 NanoFlow

  • 1.29x throughput improvement over baseline vLLM was achieved by DynaFlow-based NanoFlow on Llama-3-70B and Qwen-2.5-72B.For Llama-3-8B, the improvement was 1.17x; the speedup mainly came from overlapping network- and memory-bound operations during batch splitting.Source: §5.3.1 · pp. 9–10, Figure 9 · p. 9

5.3.3 Communication Overlap

  • 1.15x speedup was achieved in most evaluated communication-overlap cases across HuggingFace Transformers, Megatron-LM, xDiT, and FastVideo.Megatron-LM backward overlap was less effective because its fused TP backward operator could not be captured with TorchDynamo's limited compiled-autograd support; longer FastVideo sequences were attention-bound.Source: §5.3.3 · p. 10, Figure 11 · p. 11

5.3.4 TokenWeave

  • 1.22x throughput improvement over HuggingFace Transformers was achieved with DynaFlow-integrated TokenWeave, while batch-size-based CTA selection improved throughput by up to 12% over the original framework.The integration also achieved up to 1.21x over vLLM and similar performance to the original TokenWeave on optimal workloads.Source: §5.3.4 · p. 10, Figure 12 · p. 11

5.3.5 Flux / 5.3.6 Comet

  • Integrating Comet's fused expert-parallel communication kernels into Megatron-LM via DynaFlow achieved up to 1.25x speedup on Qwen2-MoE and 1.27x on Mixtral, matching the original fork with orthogonal optimizations disabled.Separately, integrating Tritondistributed fused compute-communication kernels into vLLM targeting Linear and AllReduce subgraphs caused up to 20% degradation versus the original baseline; fused-kernel latency was 1.6x higher than separate GEMM and AllReduce.Source: §5.3.6 · p. 11, §5.3.5 · p. 11

5.4 Overhead Analysis / 5.5 Ablation Study

  • 1.14x speedup was achieved with all ablation-study optimizations enabled; disabling CUDA Graph capture and replay reduced throughput to 0.96x, while disabling zero-copy memory pre-allocation gave 1.10x.Static splitting reduced throughput to 1.00x; initialization incurred 4.3 s of CUDA graph capture versus vLLM's 2.4 s and used 1.80 GiB versus 0.98 GiB of graph memory.Source: §5.5 · pp. 11–12, §5.4 · p. 11, Figure 14 · p. 12

6 RELATED WORK

Related work spans generic intra-device parallelism systems, targeted compiler approaches, and general-purpose runtimes, each with constraints on integration or scheduling flexibility.

Intra-Device Parallelism

  • DynaFlow complements fixed-strategy systems by providing a framework to integrate and dynamically select or combine approaches, since no single strategy is optimal across all configurations.The text cites workload sensitivity and the need to intercept operator execution as obstacles for generic systems, while other systems focus on overlap or kernel fusion.Source: §6 · p. 12
  • Megakernels target kernel-launch overhead and fine-grained operator overlap, while dynamicity for conditional MoE computation remains an open problem.Source: §6 · p. 12

General-Purpose ML Runtimes and Compilers

  • XLA's limited exposure of implementation details and PyTorch's manual scheduling burden make expressing workload-dependent overlap and splitting difficult.Both runtimes support some compiler-based overlapping for targeted operators, but covering all operators, schedules, and implementation options is challenging.Source: §6 · p. 12

7 CONCLUSION

DynaFlow decouples model definition from physical execution scheduling to address the engineering difficulty of integrating intra-device parallelism. The evaluation reports minimal-code integrations with speedups over original systems and existing native implementations.

7 CONCLUSION

  • DynaFlow combines annotation-based graph partitioning and a programmable scheduler with an asynchronous backend that manages data flow and preserves CUDA Graph compatibility.Source: §7 · p. 12
  • The evaluation reports up to 1.29x speedup over original systems and up to 1.1x over existing native implementations.Source: §7 · p. 12

A.1 Abstract

  • The artifact includes source code, evaluation scripts, and instructions to reproduce end-to-end throughput evaluations for Figures 9–12.Source: §A.1 · p. 16

A.2 Artifact check-list (meta-information)

The artifact checklist describes DynaFlow’s public release, evaluation setup, and resource requirements. Its repository includes the core package, framework patches, and benchmarking scripts.

A.2 Artifact check-list (meta-information)

  • DynaFlow, vLLM, SGLang, and HuggingFace Transformers are listed as programs, and the artifact is publicly available under an MIT license.The artifact checklist says a DOI will be provided after artifact evaluation finishes.Source: §A.2 · p. 16
  • The evaluation environment uses Linux, Python 3.10+, PyTorch 2.8+, CUDA 12.8+, and an eight-GPU NVIDIA DGX B200 system.The checklist specifies two Intel Xeon 8570 CPUs; exact evaluation replication requires a DGX B200 node.Source: §A.2 · p. 16, §A.3.2 · p. 16
  • Experiments evaluate offline inference and training throughput across batch sizes and sequence lengths, reporting tokens/s or seq/s in JSON files with visualization charts.The listed disk requirement is approximately 400 GB, with 1–2 hours to prepare and 12 hours to complete experiments.Source: §A.2 · p. 16

A.3.1 How delivered

  • The public GitHub repository contains the DynaFlow core package, evaluated vLLM and SGLang patches, and automated benchmarking scripts.The checklist names HuggingFace Transformers as a program; the repository contents passage specifically lists patches for vLLM and SGLang.Source: §A.3.1 · p. 16, §A.2 · p. 16

A.3.3 Software dependencies

The software dependencies specify minimum framework, CUDA, and driver versions, with a higher PyTorch requirement for the Figure 5 API mark.

A.3.3 Software dependencies

  • The environment requires PyTorch 2.8+, CUDA 12.8+, and NVIDIA driver 580+, while the API mark in Figure 5 requires PyTorch 2.10+.These are stated software requirements for running the environment and using the marked API.Source: §A.3.3 · p. 16

A.3.4 Data sets

The data setup uses ShareGPT for varied input and output lengths and generated data for fixed-size cases. The artifact instructions point users to the repository’s ae branch for installation.

A.3.4 Data sets

  • ShareGPT tests diverse input and output lengths, while randomly generated datasets supply fixed-size inputs and outputs.Figure 15 is identified as serving throughput for DynaFlow-based DBO integration in vLLM under PCIe interconnect; the caption gives no result values.Source: §A.3.4 · p. 16, Figure 15 · p. 16

A.4 Installation

  • Installation instructions for DynaFlow and the modified systems are provided under examples/ae in the GitHub repository’s ae branch.Source: §A.4 · p. 16

A.5 Evaluation and expected result

Evaluation scripts produce throughput metrics and provide a plotting workflow that saves visualizations as PDF figures.

A.5 Evaluation and expected result

  • Evaluation scripts save throughput metrics in tokens/s or seq/s as JSON under examples/ae/results, and plotting scripts save PDF figures.The plotting scripts are located under examples/ae/plot.Source: §A.5 · p. 16

B LOW-BANDWIDTH DBO PERFORMANCE

DynaFlow-enabled DBO is evaluated under simulated low-bandwidth communication by routing peer-to-peer traffic over PCIe rather than NVLink. In this setting, it improves throughput over baseline vLLM, while small LMSYS-Chat batches limit the benefit of batch splitting.

B LOW-BANDWIDTH DBO PERFORMANCE

  • The experiment simulated a low-bandwidth interconnect on two GPUs by forcing peer-to-peer communication over PCIe instead of NVLink.The setup used NCCL environment variables because hardware constraints prevented direct multi-node evaluation.Source: B LOW-BANDWIDTH DBO PERFORMANCE · p. 16
  • Up to 2.06x throughput improvement over baseline vLLM was achieved by DynaFlow-enabled DBO with communication time significantly increased.Native vLLM DBO also gained up to a 1.37x speedup, but remained constrained by its static splitting policy.Source: B LOW-BANDWIDTH DBO PERFORMANCE · p. 16
  • On LMSYS-Chat, speedups remained low for both implementations because its small batch sizes are not large enough to benefit from batch splitting.Source: B LOW-BANDWIDTH DBO PERFORMANCE · p. 16

Related papers

Loading 2605.21603v1…