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

AI summary

Overview

Research area: GPU kernel optimization and LLM training systems — specifically, the backward pass of attention kernels (FlashAttention-3) and the performance cost of enforcing bitwise reproducibility.

Technical level: Advanced. The paper combines a formal DAG scheduling formulation with a proof (Lemma 1, Appendix B) and detailed GPU architecture reasoning (register pressure, L2 cache segmentation, inter-SM synchronization latency).

Scope in one sentence: The paper formalizes deterministic attention backward scheduling as a critical-path minimization problem on a DAG, proposes three scheduling strategies, and evaluates them on NVIDIA H800 GPUs against the FlashAttention-3 deterministic baseline.

What This Paper Is About

Floating-point addition is not associative, so when gradient contributions are combined on GPUs via unordered atomic operations, the results differ slightly from run to run. FlashAttention-3 offers a deterministic mode that forces a fixed accumulation order using synchronization barriers, but this costs up to a 37.9% throughput reduction relative to the non-deterministic version. DASH's goal is to close that gap by treating the deterministic backward pass as an explicit scheduling problem and re-arranging how compute and reduction phases are interleaved across Streaming Multiprocessors (SMs), rather than accepting serialization as an unavoidable tax.

Key Contributions

1. Diagnosis of the root cause. The authors identify the misalignment between tile execution order and the required accumulation ordering as the principal source of the deterministic throughput penalty — not serialization itself.
2. First DAG-based formalization. They model the deterministic attention backward pass as a directed acyclic graph whose critical-path length is the objective to minimize, with zero-weight dependency edges encoding legal accumulation orders and all operations for one KV tile constrained to run contiguously on a single SM (to keep dK and dV accumulation register-resident).
3. Three complementary scheduling strategies. Descending Q-Tile Iteration (a reversed query-block traversal heuristic for causal masks), Shift Scheduling (a cyclic per-SM visiting order claimed to be theoretically optimal under the DAG model for full masks), and Symmetric Shift Scheduling (a two-phase workload-folding scheme with symmetric pairing of KV blocks for causal masks).
4. Empirical validation with a practical caveat. Up to a 1.28× speedup of the deterministic attention backward pass on NVIDIA H800 GPUs, plus the finding that theoretical optimality can lose to simpler heuristics under real hardware constraints. Code is open-sourced at https://github.com/SJTU-Liquid/deterministic-FA3.

Main Findings

• The size of the problem: Enabling deterministic mode in FlashAttention-3 can lower throughput by up to 37.9%, which the authors frame as a severe cost when scaling to hundreds of thousands of GPUs.
• Full mask, Shift Scheduling: Shift Scheduling outperforms the FlashAttention-3 baseline across most tested sequence lengths, but slightly degrades relative to the baseline at the maximum sequence length of 16,384.
• Why the 16,384 anomaly happens: The DAG model assumes zero-cost dependency edges, but inter-SM synchronization actually travels through the L2 cache, costing roughly 200 cycles for a local L2 segment access and over 500 cycles for a remote segment access on H800-class GPUs. At sequence length 16,384 with a KV block size of 128, one head spans 128 blocks, often mapped to 128 SMs, so many synchronization signals cross to remote L2 segments. Shift Scheduling's more intricate dependency graph makes it more sensitive to this than the baseline's linear dependency chain.
• Causal mask, headdim = 64: Symmetric Shift Scheduling achieves the highest performance, validating its workload balancing. Descending Q-Tile Iteration performs poorly here because FlashAttention-3's L2-aware LPT scheduler interleaves multiple heads across SMs; with a small L2 footprint per head, the causal stalls that descending iteration targets are largely masked.
• Causal mask, headdim = 128, the performance inversion: Descending Q-Tile Iteration surpasses the theoretically optimal Symmetric Shift Scheduling. The folded task space requires extra loop counters and intermediate state, adding roughly 10 registers per thread over the base requirement. At headdim = 128 the base register pressure is already high, and per Nsight Compute this pushes threads past the hardware register limit, forcing compiler-generated register spills to local memory whose latency negates the algorithmic benefit.
• Complementary schedules: The authors conclude Symmetric Shift is theoretically optimal under the DAG model while Descending is the practically preferred choice for large head dimensions on current GPUs.
• End-to-end transformer block impact: Causal-mask models show 2% to 10% improvement; full-mask models show approximately 4%; the average is around 5%, which the authors say aligns with their internal training experience on thousands of GPUs.
• Numerical stability (Table 1): Max gradient deviation averaged over 10 identical backward passes. Full mask: 2.4×10⁻⁴ non-deterministic versus 0 deterministic. Causal mask: 4.9×10⁻⁴ non-deterministic versus 0 deterministic. Non-deterministic kernels thus cause run-to-run gradient deviations of O(10⁻⁴), while deterministic ones produce bitwise identical outcomes.
• Scope of the determinism problem: The paper argues the cost of enforcing determinism is minimal for other transformer operations — GEMMs only exhibit nondeterminism under split-K or stream-K partitioning (unnecessary at large batch), attention forward passes and normalizations reduce within a single block, and elementwise operations are inherently deterministic.
• Not reported: The paper text presents relative speedups rather than absolute throughput values (e.g., TFLOPS) for the backward pass.

Methodology in Plain English

The authors start from how the FlashAttention backward pass is organized. Gradients dK and dV are reduced along the query axis and can be finished locally inside one SM, but dQ requires reducing partial contributions scattered across many SMs. The fast, non-deterministic approach lets those SMs update dQ in global memory via concurrent atomicAdd, so the final value depends on which CTA finishes first. The deterministic approach instead dictates a fixed order and uses barriers, which is where performance is lost.

To fix this, the authors build an abstract graph. Each tile-processing task becomes a short chain of two stages — a compute phase of cost c and a reduction phase of cost r — connected by a zero-weight edge wherever the accumulation order imposes a dependency. All work for a single KV tile must stay on one SM because register-resident accumulation of dK and dV depends on it. Latency then equals the longest path through the graph, so the whole problem becomes: add the fewest-path-lengthening dependencies possible while keeping the SMs balanced.

They prove a lemma (Lemma 1) stating that a zero-weight dependency edge from node u to node v preserves the original critical path if and only if depth(u) ≤ depth(v). Translated to hardware: two tiles that reduce into the same dQ block cannot be allowed to reach their reduction phase simultaneously on different SMs — that would force a depth-decreasing edge and lengthen the critical path.

Three schedules follow from this. Reversing the query-tile traversal order (Descending Q-Tile Iteration) resolves causal dependencies sooner and lets the next head start filling freed SMs. For full masks, giving SM i the cyclic KV order (i, i+1, …, n−1, 0, …, i−1) produces balanced work and, because the timestamps differ per row, a conflict-free reduction sequence for every dQ block. For causal masks, where workloads decrease linearly across the sequence, the authors pair the i-th and (n−1−i)-th KV blocks so long and short tasks cancel out, implemented as a two-phase schedule: a cyclic shift over the dense lower-left rectangle, then an analytically modeled "folding" of the leftover triangles into a logical square traversed from the main diagonal.

Experiments follow the FlashAttention-3 methodology: total tokens fixed at 16,384 with sequence length varying from 512 to 16,384, hidden dimension fixed at 2,048, head dimensions of 64 and 128, BF16 random inputs, CUDA 12.6 and Triton 3.4 on NVIDIA H800 GPUs. The baseline is the deterministic backward pass of FlashAttention-3, with the Triton tutorial's causal implementation also benchmarked. FlashAttention-2 is omitted because published Hopper-class benchmarks show FlashAttention-3 consistently outperforming it.

Why This Matters

Impact on research: The paper's central lesson is methodological. Deterministic training is usually treated as a fixed overhead, but this work shows the cost is largely a scheduling artifact, and that a provably optimal schedule can still lose to a simpler one when register pressure or inter-SM latency dominate. That trade-off — theoretical optimality versus hardware reality — is directly useful to anyone designing GPU kernels. It also replaces the intuition-based approach to determinism with a formal critical-path framework that others can extend.

Real-world applications:

• Large-scale LLM pretraining on thousands of GPUs, where the paper says deterministic training is increasingly adopted as standard industry practice.
• Debugging and validating training. Reproducibility lets practitioners diagnose instabilities such as loss divergence and evaluate the impact of architectural modifications without run-to-run noise confounding the comparison.
• Multimodal and vision workloads that use full attention masks — the paper evaluates on the vision model SAM-huge.
• Diffusion and diffusion-based language models — the paper evaluates on StableDiffusion3.5 (medium and large) and LLaDA-1b.

Industry relevance: The evaluation includes production-scale models (LLaMA3-8b, Qwen2.5-7b, Mistral-8×7b under causal masks) with LLM batch size 1 and sequence lengths of 8k, 16k, and 32k, and full-mask models at batch size 16 and sequence length 4k. The roughly 5% average end-to-end speedup on an entire transformer block is modest in isolation but compounds over multi-week training runs, and one author group is affiliated with ByteDance Seed, suggesting direct applicability to industrial training pipelines.

Future Directions

1. Realizing Symmetric Shift's theoretical advantage. The authors expect its benefits to be fully realized on newer architectures with greater on-chip resources, such as Blackwell GPUs with TMEM, or devices with larger register files, or under kernel designs less constrained by register allocation than the current FlashAttention-3 implementation.
2. Reducing the register overhead of the folded schedule. About 10 extra registers per thread caused the spilling that reversed the ranking at headdim = 128; restructuring the folded task space to use fewer live counters is an open engineering problem.
3. Incorporating realistic inter-SM communication cost into the DAG model. The current model assumes zero-cost dependency edges, which the authors identify as the reason Shift Scheduling degrades at sequence length 16,384; modeling L2-segment latency (roughly 200 versus over 500 cycles) could yield schedules that also win in extreme-parallelism, long-sequence regimes.
4. Bridging the model-reality gap generally. The authors state explicitly that their DAG is a simplified abstraction whose purpose is insight rather than accurate prediction of execution time, leaving open the question of how to build scheduling models that are predictive without becoming intractable.

Target Audience

This paper is most valuable to GPU kernel engineers and ML systems researchers working on attention implementations, and to infrastructure engineers responsible for large-scale deterministic training pipelines. It also suits researchers interested in the formal scheduling of parallel computation, and anyone who needs concrete evidence that algorithmically optimal schedules can be defeated by register pressure and memory-hierarchy effects on real hardware. Readers should be comfortable with GPU execution concepts such as SMs, CTAs, shared memory versus L2 cache, and register spilling; the DAG formalization itself is presented accessibly, with the lemma proof deferred to Appendix B.

Authors’ 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.