Paper Notes: DYNAFLOW
DYNAFLOW
The paper named ‘DYNAFLOW: TRANSPARENT AND FLEXIBLE INTRA-DEVICE PARALLELISM VIA PROGRAMMABLE OPERATOR SCHEDULING’. And the link is https://arxiv.org/abs/2605.21603. The work is done by people from UW and SJTU.
Motivation
GPUs waste lots of resources when running LLMs — one operator uses compute cores while memory bus sits idle. Intra-device parallelism fixes this by running different operators at the same time to fully use GPU hardware. But existing ML frameworks are built for sequential code. To add intra-device parallelism, developers have to rewrite huge chunks of model code, which takes months and costs massive engineering work.
What’s more, No single parallel strategy works everywhere. The best way to overlap operators changes based on model size, input workload, and GPU hardware. Developers must build and maintain many separate custom codes for different scenarios.
Background
Models are shifting from purely compute-bound to a sequence of operators with highly diverse resource requirements. Such as compute-bound operators(matrix multiplications), memory-bound operators like decode attention, network-bound like communication.
Recent research has explored intra-device parallelism, a class of strategies that aims to maximize resource utilization within a single device. Techniques such as overlapping computation with communication, fine-grained kernel fusion, or further splitting the input batch for concurrent execution have shown significant throughput improvement.
Modern fast LLM inference systems such as vLLM and SGLang are designed to run operations one after another. But single-GPU parallel computing needs to run multiple tasks at the same time, which clashes with their original design. So adding intra-device parallel support means massive, disruptive code changes.
On top of the huge coding work already needed, there’s another big problem: there’s no one-size-fits-all parallel strategy. How well a strategy works depends entirely on your model, input batch size, and GPU hardware. Since every scenario needs its own custom implementation, developers have to build and maintain many separate pieces of code, making the engineering workload even heavier.
Key Ideas
Key idea for solving these challenges is to decouple operator execution from the model implementation. The paper proposes a new programmable execution substrate. It supports flexible parallel operator scheduling without changing model definitions.
DynaFlow is implemented as a torch.compile backend to support transparent integration with all PyTorch-based systems. Torch.compile is a built-in PyTorch compiler. It takes original model code, converts it into optimized GPU code.
How It Works
Overview
Integrating DynaFlow into preexisting machine learning systems proceeds in three distinct phases. First, at model initialization time, developers leverage the frontend’s annotation APIs to partition the model’s computation graph into independently schedulable subgraph fragments. Second, practitioners implement a bespoke scheduling policy inside a Python-first scheduler function; they utilize the frontend’s high-level programming interfaces to specify target execution sequences and computation-communication overlapping patterns. Third, during inference runtime, DynaFlow intercepts the model’s forward pass invocation and triggers this user-defined scheduler. The custom scheduler dynamically constructs a full execution plan and dispatches partitioned subgraphs to an asynchronous backend runtime, which encapsulates and handles all underlying low-level GPU execution mechanics transparently.
See Figure 4.
前端调度器 → execute() → 后端引擎（更新依赖、筛选就绪子图）
后端引擎 → 就绪子图列表 → get_ready_ops() → 前端调度器frontend
The frontend augments vanilla Python code with a unified, dynamic programming abstraction for expressing diverse parallelism strategies and scheduling granularities.
During model initialization, it takes a TorchDynamo-traced computational graph as input. TorchDynamo extracts PyTorch graphs from Python bytecode.
DynaFlow’s design strikes a balance: it provides a set of high-level APIs to abstract this complexity, but embeds them within a fully Python-native frontend to preserve flexibility. To implement a custom strategy, a developer inherits from a base class, OpSchedulerBase, and overrides its schedule method. Inside this method, the developer interacts with the DynaFlow backend using a set of high-level APIs.
backend
The backend focuses on efficient execution, managing the complex control and data-flow dependencies introduced by dynamic scheduling, minimizing data-flow overheads like memory copies, and ensuring compatibility with low-level optimizations such as CUDA graphs and TorchInductor.
The DynaFlow backend is responsible for executing the dynamic schedules generated by the frontend.
During initialization, the StaticAnalysis routine walks the full computational graph to precompute per-tensor metadata: reference counters governing tensor lifetime tracking, alongside a prealloc flag that flags tensors scheduled for later tensor-merging operations. At inference runtime, the RuntimeExecute routine consumes the precomputed metadata. It decrements tensor reference counters to drive automatic garbage collection. For operator outputs marked with the prealloc tag, the runtime reserves a contiguous memory buffer upfront. An embedded runtime hook routes operator outputs straight into the corresponding buffer segment, delivering zero‑copy data resharding for downstream consuming operators.
Compatibility with Low-Level Optimizations
Another challenge in dynamic scheduling is ensuring compatibility with other performance optimizations that intercept execution at the operator graph level, such as TorchInductor and CUDA Graphs.
DynaFlow addresses this conflict by applying static optimization techniques at the subgraph level.
Evaluation
Use lines of code(LOC) to evaluate frontend effectiveness.
We evaluate the benefits brought by its backend design of low-level optimization compatibility by measuring the CPU execution time of different DynaFlow configurations.
Takeaways
1）无 DynaFlow（原生 vLLM/SGLang）
TorchDynamo 整图 → TorchInductor 串行编译 → 小块静态子图单独录 CUDA Graph → 串行顺序执行
短板：无法并行重叠计算 / 通信，动态调度能力为 0，多批次并发显存开销高。
2）有 DynaFlow
TorchDynamo 整图 → 前端按 Module/mark 切分逻辑子图 → 后端静态分析（依赖 + 内存元数据）→ 每个子图独立 Inductor 编译 / 独立 CUDA Graph 录制（分微批次隔离）→ 异步乱序并行执行（后端自动筛选就绪子图并发跑）Key Techniques
In this section, I will describe techniques the paper mentioned.
Mixture-of-Experts (MoE)
MoE splits FFN into multiple independent sub-networks called experts, each token is routed to only a small subset of experts. Shared Expert is a universal FFN branch that runs for all tokens, no routing needed. It is always executed alongside the routed unique experts to improve model quality.
Input Token
|
v
┌─────────┴─────────┐
| |
v v
┌─────────┐ ┌──────────────┐
│ Shared │ │ Router │ ← Runs concurrently
│ Expert │ └──────┬───────┘
└────┬────┘ |
| ┌────┼────┐
| ▼ ▼ ▼
| ┌──┐ ┌──┐ ┌──┐
| │E0│ │E1│ │E2│ ← Top-K routed experts
| └┬─┘ └┬─┘ └┬─┘
| └────┼────┘
| |
v v
┌───────────────────┐
│ y = Shared(x) + │
│ Σ Routed(x) │
└───────────────────┘ML System Structure
┌───────────────────────────────────────────────┐
│ Level 6: ML Service Layer │
│ │
│ Inference Serving System │
│ │
│ vLLM / SGLang / Triton Server │
│ │
│ Responsibilities: │
│ - Request batching │
│ - KV Cache management │
│ - User request scheduling │
└───────────────────────┬───────────────────────┘
│
▼
┌───────────────────────────────────────────────┐
│ Level 5: Model Graph Layer │
│ │
│ Neural Network Representation │
│ │
│ Transformer Graph │
│ │
│ Example: │
│ │
│ Attention → FFN → Norm → Attention │
│ │
│ Framework: PyTorch / JAX │
└───────────────────────┬───────────────────────┘
│
▼
┌───────────────────────────────────────────────┐
│ Level 4: Graph Optimization / Compiler Layer │
│ │
│ Transform Model Graph │
│ │
│ Examples: │
│ - PyTorch Inductor │
│ - XLA │
│ - TensorRT │
│ │
│ Tasks: │
│ - Operator fusion │
│ - Kernel generation │
│ - Memory optimization │
└───────────────────────┬───────────────────────┘
│
▼
┌───────────────────────────────────────────────┐
│ Level 3: Runtime / Scheduler Layer │
│ │
│ Decide HOW operations execute │
│ │
│ Example: │
│ │
│ GEMM │
│ ↓ │
│ AllReduce │
│ ↓ │
│ LayerNorm │
│ │
│ Scheduler controls: │
│ - Execution order │
│ - Parallelism │
│ - Overlap │
└───────────────────────┬───────────────────────┘
│
▼
┌───────────────────────────────────────────────┐
│ Level 2: Distributed Communication Layer │
│ │
│ Multi-GPU Coordination │
│ │
│ Examples: │
│ - NCCL │
│ - MPI │
│ │
│ Operations: │
│ │
│ GPU0 ───── AllReduce ───── GPU1 │
└───────────────────────┬───────────────────────┘
│
▼
┌───────────────────────────────────────────────┐
│ Level 1: Hardware Execution Layer │
│ │
│ GPU / TPU / Accelerator │
│ │
│ CUDA Cores │
│ Tensor Cores │
│ HBM Memory │
│ │
└───────────────────────────────────────────────┘Attention后端
当模型不大、Batch 也不大时，还有一个更隐蔽的瓶颈会冒出来：CPU 发射 Kernel 的开销。GPU 自己不会主动干活，它每做一个操作（一次矩阵乘、一次 LayerNorm、一次激活……），都要由 CPU 通过一次”启动（launch）“命令告诉它。跑一遍 Transformer，有成百上千个这样的小操作，就对应成百上千次 CPU→GPU 的启动调用。每次启动本身有固定的 CPU 开销（构造参数、驱动调度、提交队列），量级在微秒级。
平时这点开销不算什么——因为 Prefill 一步要处理成千 Token，每个 GPU Kernel 一跑就是几百微秒甚至毫秒，CPU 那点启动开销被淹没了。但 Decode 每步只处理 1 个（或很少几个）Token，每个 Kernel 的实际计算可能就几微秒，结果CPU 发射这个 Kernel 的时间比 GPU 执行它的时间还长。于是 GPU 干几微秒就停下来，眼巴巴等 CPU 发下一条命令——算力再次闲置，只不过这次不是等数据，是等指令。
Attention 是 Transformer 里最重、最讲究的算子，它的 Kernel 实现直接决定推理速度。vLLM 的设计哲学是不绑死一种实现，而是提供可插拔的 Attention 后端，根据硬件、模型结构和场景自动或手动选择最优的那一个。Attention 后端是同一个数学操作的不同工程实现。它们算出的结果在数学上等价，区别在于访存模式、并行策略、对特定硬件指令和模型结构的适配程度。选对后端，同样的卡能跑出明显不同的吞吐。
vLLM V1 的统一调度器会把 Prefill 的 Chunk 和 Decode 请求塞进同一步——那么这一步的 Attention 该怎么算？毕竟 Prefill 部分要算”一整段新 Token 之间 + 对历史的注意力”，Decode 部分只算”1 个新 Token 对全部历史的注意力”，两者的计算形态不一样。
这就要求 Attention 后端原生支持混合批次（mixed prefill/decode batch）：在一次 Kernel 调用里，同时正确处理批中既有 Prefill 又有 Decode 的请求。FlashAttention 3 等现代后端正是为此设计的——它能接受变长的、混合阶段的输入，用统一的 Kernel 高效算完，而不需要把 Prefill 和 Decode 拆成两次 Kernel 调用。
CUDA Graph
CUDA Graph 是 NVIDIA 提供的解药：把一连串固定的 GPU 操作录制（capture）成一张”图”，之后只需一条命令重放（replay）整张图，GPU 就会依次执行录好的所有 Kernel——成百上千次 CPU 启动被压缩成一次。但 CUDA Graph 有个硬性前提：录制的操作序列和张量形状必须固定。这对 Decode 是天作之合——每步都是”处理固定 Batch 个 Token”，形状稳定，可以针对不同的 Batch Size 分别录制好图，运行时按当前 Batch Size 选对应的图重放。
torch.compile
能不能把很多小 Kernel 本身合并成更少、更大的 Kernel？ 这就是 torch.compile 干的事——它把模型的前向计算编译成优化后的图，通过算子融合（fusion）等手段减少 Kernel 数量。