Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

DASH: Deterministic Attention Scheduling for High-throughput Reproducible LLM Training

Xinwei Qiang, Hongmin Chen, Shixuan Sun, Jingwen Leng, Xin Liu, Minyi Guo

TL;DR

Determinism is essential for reproducibility in large-scale LLM training but imposes throughput penalties in attention backward passes due to serialized gradient accumulation. DASH reframes the backward pass as a DAG scheduling problem and introduces two complementary strategies—Descending Q-Tile Iteration and Shift Scheduling—to minimize the critical path. On NVIDIA H800 GPUs, these methods achieve up to $1.28\times$ throughput gains over the deterministic FlashAttention-3 baseline, with end-to-end transformer-block speedups averaging around $5\%$. Hardware realities such as inter-SM L2 latency and register pressure influence which strategy is preferable in a given setting, underscoring the need for adaptable scheduling. Overall, DASH narrows the determinism gap and enables more efficient reproducible training for large language models.

Abstract

Determinism is indispensable for reproducibility in large language model (LLM) training, yet it often exacts a steep performance cost. In widely used attention implementations such as FlashAttention-3, the deterministic backward pass can incur up to a 37.9% throughput reduction relative to its non-deterministic counterpart, primarily because gradient accumulation operations must be serialized to guarantee numerical consistency. This performance loss stems from suboptimal scheduling of compute and gradient-reduction phases, leading to significant hardware underutilization. To address this challenge, we formulate the backward pass of deterministic attention as a scheduling problem on a Directed Acyclic Graph (DAG) and derive schedules that minimize the critical path length. Building on this formulation, we present DASH (Deterministic Attention Scheduling for High-Throughput), which encapsulates two complementary scheduling strategies: (i) Descending Q-Tile Iteration, a reversed query-block traversal that shrinks pipeline stalls in causal attention, and (ii) Shift Scheduling, a theoretically optimal schedule within our DAG model that reduces pipeline stalls for both full and causal masks. Our empirical evaluations on NVIDIA H800 GPUs demonstrate that DASH narrows the performance gap of deterministic attention. The proposed strategies improve the throughput of the attention backward pass by up to 1.28$\times$ compared to the baseline, significantly advancing the efficiency of reproducible LLM training. Our code is open-sourced at https://github.com/SJTU-Liquid/deterministic-FA3.

DASH: Deterministic Attention Scheduling for High-throughput Reproducible LLM Training

TL;DR

Determinism is essential for reproducibility in large-scale LLM training but imposes throughput penalties in attention backward passes due to serialized gradient accumulation. DASH reframes the backward pass as a DAG scheduling problem and introduces two complementary strategies—Descending Q-Tile Iteration and Shift Scheduling—to minimize the critical path. On NVIDIA H800 GPUs, these methods achieve up to throughput gains over the deterministic FlashAttention-3 baseline, with end-to-end transformer-block speedups averaging around . Hardware realities such as inter-SM L2 latency and register pressure influence which strategy is preferable in a given setting, underscoring the need for adaptable scheduling. Overall, DASH narrows the determinism gap and enables more efficient reproducible training for large language models.

Abstract

Determinism is indispensable for reproducibility in large language model (LLM) training, yet it often exacts a steep performance cost. In widely used attention implementations such as FlashAttention-3, the deterministic backward pass can incur up to a 37.9% throughput reduction relative to its non-deterministic counterpart, primarily because gradient accumulation operations must be serialized to guarantee numerical consistency. This performance loss stems from suboptimal scheduling of compute and gradient-reduction phases, leading to significant hardware underutilization. To address this challenge, we formulate the backward pass of deterministic attention as a scheduling problem on a Directed Acyclic Graph (DAG) and derive schedules that minimize the critical path length. Building on this formulation, we present DASH (Deterministic Attention Scheduling for High-Throughput), which encapsulates two complementary scheduling strategies: (i) Descending Q-Tile Iteration, a reversed query-block traversal that shrinks pipeline stalls in causal attention, and (ii) Shift Scheduling, a theoretically optimal schedule within our DAG model that reduces pipeline stalls for both full and causal masks. Our empirical evaluations on NVIDIA H800 GPUs demonstrate that DASH narrows the performance gap of deterministic attention. The proposed strategies improve the throughput of the attention backward pass by up to 1.28 compared to the baseline, significantly advancing the efficiency of reproducible LLM training. Our code is open-sourced at https://github.com/SJTU-Liquid/deterministic-FA3.
Paper Structure (35 sections, 1 theorem, 4 equations, 10 figures, 1 table, 1 algorithm)

This paper contains 35 sections, 1 theorem, 4 equations, 10 figures, 1 table, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Deterministic FlashAttention Backward Pass
  4. GPU Architecture
  5. Determinism in Other Operations of the Transformer
  6. DASH: Scheduling Strategies for Deterministic Attention
  7. Problem Formulation
  8. Graph Construction.
  9. Optimization Constraint.
  10. Analysis of FlashAttention-3 Deterministic Backward Schedule
  11. Descending Q-Tile Iteration: A Robust Heuristic for Causal Masks
  12. Shift Scheduling
  13. Optimal Schedule for Full Masks
  14. Symmetric Shift Scheduling for Causal Masks
  15. Summary of Optimal Performance
  16. ...and 20 more sections

Key Result

lemma 1

Let $G_0=(V, E_0)$ be a DAG consisting of a single source node $s$, a single sink node $t$, and $n \geq 1$ parallel, isomorphic chains connecting $s$ to $t$. All edge weights in $E_0$ are strictly positive. Let the depth of a node $v$, denoted $depth(v)$, be the number of edges on the unique path fr

Figures (10)

  • Figure 1: Overview of the deterministic FlashAttention. Left: Tiled computation structure of the backward pass, highlighting the local and global reductions. Middle: Comparison between the non-deterministic (atomic-based) and deterministic (ordered) global reduction. Right: Performance degradation under causal and full attention masks, HD stands for head dimension.
  • Figure 2: Visualization of the Deterministic Scheduling Problem. The Gantt chart (left) shows a naive execution schedule for a problem with two KV-tiles ($i$-index) and two Q-tiles ($j$-index). Each task consists of a compute phase $C(i,j)$ and a reduction phase $R(i,j)$. Local reductions enforce contiguous execution on a single SM (e.g., all tasks for $i=0$ on SM0). A deterministic global reduction order introduces a cross-SM dependency (red arrow), forcing SM1 to idle and creating a pipeline bubble. The corresponding DAG (right) abstracts this schedule, where the critical path determines the end-to-end latency.
  • Figure 3: Backward scheduling of FlashAttention-3 for both mask shapes. Each colored segment denotes one block’s computation (cost $c$) followed by a reduction (cost $r$). Idle gaps correspond to pipeline bubbles. For clarity, since we assume the number of KV tiles equals the number of SMs, each SM processes exactly one KV tile; thus we omit the KV index in the visualization and show only the query index for each block.
  • Figure 4: Descending (reverse-order) query tile schedule for the causal mask. Reversing the $Q$-block traversal accelerates dependency resolution. Colors distinguish attention heads in the pipeline.
  • Figure 5: Illustrative example for Lemma \ref{['lemma:opt-schedule-revised']}. Left: Added dependency (zero-weight) edges preserve non-decreasing depth order and do not lengthen the critical path. Right: A backward (depth-decreasing) dependency edge violates the lemma’s condition and increases the critical path.
  • ...and 5 more figures

Theorems & Definitions (2)

  • lemma 1
  • proof