FlashAttention Semantic Contract

FlashAttention is an IO-aware implementation strategy for scaled dot-product attention: it tiles the attention computation and maintains online softmax summaries so the full n × n attention matrix does not need to be materialized. TorchLean separates the semantic contract from execution:

• this file gives the proof layer semantic contract for a fused FlashAttention operator;
• the current CUDA runtime composes attention in Lean over LibTorch numerical primitives (NN/Runtime/Autograd/Engine/LibTorch/Ops/Attention.lean), retaining the capsule identity libtorch.direct_attention; it materializes full score matrices and saves probabilities;
• the numerical FFI boundary is documented separately because Lean does not verify LibTorch or CUDA code. This runtime does not select a fused FlashAttention implementation.

The fused operation has the same denotation as standard masked scaled dot-product attention over the spec scalar. Different tile sizes are runtime scheduling choices, not semantic choices.

Refinement target

The forward and backward contracts reuse standard attention directly. Their equality theorems let a downstream refinement certificate use either name without duplicating the attention calculation. Tile sizes remain implementation metadata; this file defines no online softmax recurrence, tile loops, or memory-traffic model.

The current LibTorch composition is tested operationally. Its foreign numerical calls remain outside Lean's kernel; a fused implementation needs its own refinement argument.

References:

• Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré, "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness", arXiv:2205.14135.
• Tri Dao, "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning", arXiv:2307.08691.
• Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao, "FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision", arXiv:2407.08608.

structure Spec.FlashAttentionConfig : Type

Runtime tiling metadata for a FlashAttention-style fused implementation.

The spec-level denotation below intentionally ignores these fields: they describe how an implementation schedules work, not what mathematical function the operator computes.

• blockQ : ℕ

  Query block size used by a tiled implementation. 0 means "backend default".

• blockK : ℕ

  Key/value block size used by a tiled implementation. 0 means "backend default".

@[instance_reducible]

instance Spec.instDecidableEqFlashAttentionConfig : DecidableEq FlashAttentionConfig

def Spec.instDecidableEqFlashAttentionConfig.decEq (x✝ x✝¹ : FlashAttentionConfig) : Decidable (x✝ = x✝¹)

def Spec.instReprFlashAttentionConfig.repr : FlashAttentionConfig → ℕ → Std.Format

@[instance_reducible]

instance Spec.instReprFlashAttentionConfig : Repr FlashAttentionConfig

def Spec.FlashAttentionConfig.default : FlashAttentionConfig

Backend-default tiling.

def Spec.flashAttention {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (config : FlashAttentionConfig) {nQ nK dModel : ℕ} {h1 : nQ ≠ 0} {h2 : nK ≠ 0} (ctx : AttentionContext α nQ nK dModel h1 h2) : TorchLean.Tensor α [nQ, dModel]

Semantic FlashAttention forward operator.

Tile sizes do not change the result. A fused implementation must refine the same masked scaled-dot-product attention calculation as an unfused implementation.

def Spec.flashAttentionBackward {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (config : FlashAttentionConfig) {nQ nK dModel : ℕ} {h1 : nQ ≠ 0} {h2 : nK ≠ 0} (ctx : AttentionContext α nQ nK dModel h1 h2) (dOut : TorchLean.Tensor α [nQ, dModel]) : TorchLean.Tensor α [nQ, dModel] × TorchLean.Tensor α [nK, dModel] × TorchLean.Tensor α [nK, dModel]

Semantic FlashAttention backward/VJP operator.

This is the local derivative contract a fused backward kernel should refine. The actual CUDA kernel may recompute attention probabilities from row statistics rather than storing the full attention matrix, but the returned adjoints must match this spec-level VJP up to the chosen floating-point error envelope.

@[simp]

theorem Spec.flashAttention_eq_scaledDotProductAttention {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (config : FlashAttentionConfig) {nQ nK dModel : ℕ} {h1 : nQ ≠ 0} {h2 : nK ≠ 0} (ctx : AttentionContext α nQ nK dModel h1 h2) : flashAttention config ctx = scaledDotProductAttention ctx

Forward semantic correctness of the fused FlashAttention spec.

@[simp]

theorem Spec.flashAttentionBackward_eq_scaledDotProductAttentionBackward {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (config : FlashAttentionConfig) {nQ nK dModel : ℕ} {h1 : nQ ≠ 0} {h2 : nK ≠ 0} (ctx : AttentionContext α nQ nK dModel h1 h2) (dOut : TorchLean.Tensor α [nQ, dModel]) : flashAttentionBackward config ctx dOut = scaledDotProductAttentionBackward ctx dOut

Backward/VJP semantic correctness of the fused FlashAttention spec.