feat(attention): complete Flash Attention VJP implementation

[Brooooooklyn]

Fused Flash Attention Backward (VJP) Kernels for Metal

Summary

This PR adds fused Flash Attention backward pass (VJP) kernels for Apple Silicon GPUs, implementing the two-kernel architecture from Flash Attention 2 (Dao, 2023) using MLX's STEEL tiling framework. The fused backward eliminates the O(L^2) attention matrix materialization, reducing peak memory by 70-95% at the cost of additional recomputation. An auto-dispatch policy routes between fused and unfused paths based on sequence length and memory pressure.

Key additions:

• Two new Metal kernels: steel_attention_vjp_dq (dQ gradients) and steel_attention_vjp_dkv (dK/dV gradients)
• JIT compilation with baked constants (gqa_factor, scale, scale_log2, alignment flags, causal, block mask) for dead-code elimination
• Delta precomputation as lazy MLX graph ops (MFA pattern)
• Threadgroup memory aliasing: red_smem over Q_smem+dO_smem (temporally disjoint), reducing D=128 dKV from 23 KB to 14.8 KB (enables 2 TGs/core)
• Sparse block mask support in both backward kernels (zero overhead when unused via function constant gating)
• Auto-dispatch with per-config L thresholds and memory ceiling
• Support for D={64, 96, 128}, float16/bfloat16, causal masking, GQA

Background

The Flash Attention Backward Recomputation Trade-off

Flash Attention (Dao et al., NeurIPS 2022) achieves O(N) memory for attention by never materializing the full N x N attention matrix. The forward pass tiles the computation with no extra FLOPs — it's the same work, just blocked differently. But the backward pass fundamentally changes the compute profile.

Forward pass (2 matmuls per block-pair):

S = Q @ K^T // attention scores
P = softmax(S) // attention weights
O = P @ V // output

Backward pass (5 matmuls per block-pair):

// Must recompute S and P from saved LSE (logsumexp) — this is the cost of not saving O(N^2)
S = Q @ K^T // recompute (1)
P = exp2(S * scale_log2 - LSE) // recompute from saved LSE (2)
dP = dO @ V^T // (3)
dS = scale * P * (dP - delta) // elementwise
dQ = dS @ K // (4)
dK = dS^T @ Q // (5)
dV = P^T @ dO // (5, in dKV kernel)

The key insight from Flash Attention 2 (Dao, 2023, Section 3.1): the backward pass requires ~2.5x the FLOPs of the forward pass because the attention matrix S must be recomputed from saved LSE in both the dQ kernel and the dKV kernel. This is the fundamental cost of O(N) memory.

Published backward/forward ratios

These ratios are well-documented across implementations:

Implementation	Hardware	Backward / Forward	Source
Flash Attention 2	A100 (CUDA)	~2.0-2.5x (dense), ~1.7-2.0x (causal)	Dao 2023, arXiv:2307.08691
Flash Attention 2	A100	Forward: 73% peak TFLOPS, Backward: 63% peak	FA2 paper, Table 2
Metal Flash Attention	M1/M3 (Metal)	~2-3x vs unfused backward for long sequences	MFA v0.9.1 benchmarks

The backward is inherently more expensive than forward in Flash Attention. This is not a bug — it's the price of O(N) memory.

Reference implementation: Metal Flash Attention (MFA)

Our implementation follows the architectural patterns from Metal Flash Attention by Philip Turner, the only prior fused attention backward for Apple Silicon:

• Two-kernel split: Separate dQ and dKV kernels with no atomics (MFA pattern)
• Delta precomputation: delta = rowsum(dO * O) computed once outside kernels (MFA pattern)
• Log2 domain: P = exp2(S * scale_log2 - LSE) instead of exp(S * scale - LSE) for Metal exp2 efficiency (MFA pattern)
• Register pressure management: WM=2 (64 threads) for D=128 dKV with BQ=16 tiles to avoid register spilling

The main difference from MFA: we use MLX's JIT compilation to bake runtime constants (gqa_factor, scale, alignment, causal) as #define literals, enabling the Metal compiler to eliminate dead code paths. MFA achieves this via pre-compiled metallib with #define constants baked at build time.

Architecture

Two-Kernel Design

The backward is split into two independent kernels that run concurrently:

 ┌─────────────────────────────────┐
│ Delta Precomputation │
│ delta = rowsum(dO * O) [f32] │
└──────────┬──────────────────────┘
│
┌──────────┴──────────────────────--┐
│ │
┌─────┴─────-┐ ┌─────┴──────┐
│ dQ Kernel │ │ dKV Kernel │
│ │ │ │
│ Grid: │ │ Grid: │
│ [NQ,H,B] │ │ [NK,kvH,B] │
│ │ │ │
│ WM=4 │ │ WM=1 (D64) │
│ 128 threads│ │ WM=2 (D96+)│
│ 4 simdgrps │ │ 32-64 thrd │
│ │ │ │
│ Loop: KV │ │ Loop: GQA │
│ blocks │ │ then Q blks│
└─────┬──────┘ └─────┬──────┘
│ │
│ dQ dK, dV
└───────────────┬───────────────────┘
│
Output grads

Per-Head-Dimension Configuration

D	Kernel	BQ	BK	WM	Threads	Regs/thread	Notes
64	dQ	32	32	4	128	~180	Full occupancy
64	dKV	32	32	1	32	~220	Single simdgroup, no reduction
96	dQ	32	32 (M3+) / 16	4	128	~200	BK=32 on M3+ for bandwidth
96	dKV	32	16	2	64	~280	Threadgroup reduction for dK/dV
128	dQ	32	32 (M3+) / 16	4	128	~220	BK=32 on M3+
128	dKV	16	16	2	64	~202	BQ=16 halves register tiles, avoids spilling

D=128 dKV is the most register-constrained configuration. At BQ=32 WM=2, it would need ~338 regs/thread (spills at >256). BQ=16 reduces TQ from 2 to 1 per simdgroup, cutting register usage to ~202 at the cost of 2x more Q-tile iterations.

Threadgroup Memory Aliasing

The dKV kernel uses four smem buffers: Q_smem, dO_smem (iteration phase), KV_smem (iteration phase), and red_smem (reduction phase, WM>1 only). Since Q/dO and red are temporally disjoint (iterations end before reduction begins), red_smem is aliased over the Q_smem+dO_smem region:

Before aliasing (D=128 BQ=16 WM=2):
Q_smem(4,352) + dO_smem(4,352) + KV_smem(6,144) + red_smem(8,192) = 23,040 bytes
→ 1 threadgroup/core (32KB limit)
After aliasing:
max(Q+dO(8,704), red(8,192)) + KV(6,144) = 14,848 bytes
→ 2 threadgroups/core (latency hiding doubles)

A static_assert verifies the Q+dO region is large enough to alias red, and an explicit threadgroup_barrier before the reduction phase ensures correctness.

Sparse Block Masks

Both kernels support an optional block_mask buffer (uint8_t[NQ_tiles * NK_tiles]) for skipping tile pairs in sparse attention patterns (sliding window, local attention, block-sparse). The mask is gated by a function constant (has_block_mask, index 302) / JIT define (VJP_HAS_BLOCK_MASK):

• When has_block_mask = false (default): all mask checks are dead-code eliminated — zero overhead
• When has_block_mask = true: tile-skip check at the top of each loop iteration, before any smem loads
• dQ kernel: skips K-blocks where block_mask[qb * NK_tiles + kb] == 0
• dKV kernel: skips Q-blocks where block_mask[qb * NK_tiles + kb] == 0

JIT Constant Baking

Both kernels are JIT-compiled with #define constants:

#defineVJP_GQA_FACTOR4// eliminates GQA loop for GQA=1
#defineVJP_SCALE0.0883883461f// constant-folds into FMA
#defineVJP_SCALE_LOG20.127552539f
#defineVJP_BAKED_FC1// signals JIT mode
#defineVJP_ALIGN_Qtrue// dead-code eliminates bounds checks
#defineVJP_ALIGN_Ktrue
#defineVJP_DO_CAUSALfalse// dead-code eliminates causal branches
#defineVJP_HAS_BLOCK_MASKfalse// dead-code eliminates sparse mask checks

For metallib (non-JIT) builds, macros fall back to params-> runtime reads and Metal function constants, ensuring zero behavior change.

Why Fused Backward Is Not Always Faster

The fundamental trade-off

Path	Compute	Memory
Unfused	S computed once, reused from O(L^2) buffer. Uses NAX large-tile matmuls (~10.7 TFLOPS)	Materializes full L x L attention matrix per head
Fused	S recomputed in both dQ and dKV kernels (~2.5x forward FLOPs). Small BQ=16-32 STEEL tiles (~1.9 TFLOPS)	Only O(L) memory — no attention matrix

The fused path pays two penalties:

1. Recomputation: S = Q @ K^T computed twice instead of once
2. Smaller tiles: STEEL 32x32 tiles achieve lower MMA utilization than NAX 64x64+ tiles

When fused wins: causal masking

With causal attention, the attention matrix is lower-triangular. Fused kernels skip entire tile blocks in the upper triangle, eliminating ~50% of compute. Unfused still materializes the full N x N matrix then masks it:

Dense attention matrix: Causal attention matrix:
┌─────────────┐ ┌─────────────┐
│ █ █ █ █ █ █ │ │ █ · · · · · │
│ █ █ █ █ █ █ │ │ █ █ · · · · │
│ █ █ █ █ █ █ │ ──────► │ █ █ █ · · · │ ~50% tiles skipped
│ █ █ █ █ █ █ │ │ █ █ █ █ · · │
│ █ █ █ █ █ █ │ │ █ █ █ █ █ · │
│ █ █ █ █ █ █ │ │ █ █ █ █ █ █ │
└─────────────┘ └─────────────┘
All tiles computed Only lower triangle

This is why our benchmarks show fused causal at 1.18-1.36x speedup (fused faster) but fused dense at 0.39-0.70x (fused slower).

When fused wins: long sequences

At long sequence lengths, the O(L^2) attention matrix becomes a memory bandwidth bottleneck. At L=4096 with 32 heads, unfused needs 3.4 GB just for the backward intermediate. Even when fused is slower in raw compute, it avoids thrashing the memory system.

How we deal with it: auto-dispatch

Rather than always using fused (slow for short dense sequences) or always unfused (OOM risk for long sequences), we route per-configuration:

auto_dispatch(D, L, B, H, causal, gqa):
1. Hard ceiling: if attn_bytes >= 256 MB → fused (prevent OOM)
2. Per-config L thresholds:
- D=64 causal MHA: always fused (causal skip makes it faster)
- D=64 dense MHA: fused if L >= 4096
- D=128 causal MHA: fused if L >= 4096
- D=128 dense MHA: fused if L >= 8192
- GQA configs: higher thresholds (GQA adds overhead to dKV)
3. Otherwise: unfused (faster for short sequences)

Users can override with MLX_SDPA_VJP_MODE={fused|unfused}.

Benchmark Results

All benchmarks on Apple M3 Max, B=1, H=32, float16. Speedup > 1.0x means fused is faster.

Performance: Fused vs Unfused (MLX_SDPA_VJP_MODE=fused)

Causal attention — fused is competitive or faster due to block skipping:

D	L	Unfused (s)	Fused (s)	Speedup
64	512	0.013	0.011	1.18x
64	1024	0.043	0.031	1.36x
64	2048	0.162	0.125	1.29x
64	4096	0.616	0.459	1.34x
96	512	0.016	0.013	1.22x
96	1024	0.055	0.043	1.27x
96	2048	0.203	0.158	1.28x
128	512	0.016	0.020	0.80x
128	1024	0.056	0.070	0.80x
128	2048	0.213	0.279	0.76x

Dense attention — fused pays the full recomputation penalty:

D	L	Unfused (s)	Fused (s)	Speedup
64	512	0.012	0.017	0.70x
64	1024	0.040	0.057	0.70x
64	2048	0.136	0.211	0.65x
64	4096	0.664	0.996	0.67x
96	512	0.015	0.021	0.68x
96	1024	0.053	0.079	0.66x
128	512	0.015	0.036	0.43x
128	1024	0.055	0.133	0.41x
128	2048	0.196	0.504	0.39x

Pattern analysis:

• D=64/96 causal: Fused is 1.18-1.36x faster. The ~50% causal tile skip more than compensates for the 2.5x recomputation.
• D=64/96 dense: Fused is 0.63-0.70x (1.4-1.6x slower). Full recomputation penalty, but manageable.
• D=128 dense: Fused is 0.39-0.43x (2.3-2.6x slower). Larger D means more computation per tile; the recomputation penalty scales with D.
• D=128 causal: Fused is 0.76-0.80x — causal helps but cannot fully overcome the D=128 overhead. Smem aliasing improved this from 0.73x (see analysis below).

Why D=128 causal is slower despite block skipping

D=64/96 causal fused is 1.18-1.36x faster than unfused, but D=128 causal fused is only 0.76-0.80x. Both get the same ~50% causal tile-skip benefit, yet D=128 cannot overcome the recomputation penalty. This is caused by five compounding factors:

1. Occupancy collapse (dominant factor). Apple M3 Max has 32 KB threadgroup memory per GPU core. D=64 dKV uses ~7 KB smem, allowing 4 concurrent threadgroups per core — when one threadgroup stalls on a memory load, others keep the ALU busy. D=128 dKV originally used ~22.5 KB smem (Q: 4,352B + dO: 4,352B + KV: 6,144B + reduction: 8,192B), limiting occupancy to 1 threadgroup per core. After smem aliasing (red over Q+dO), this drops to 14.8 KB enabling 2 TGs/core — improving D=128 causal from 0.73x to 0.80x, but still below D=64's 4 TGs/core.

2. BQ=16 doubles iteration overhead. D=128 dKV uses BQ=16 tiles (half of D=64's BQ=32) to keep register pressure under 256 regs/thread. This means 2x more Q-tile iterations in the inner loop, each paying fixed costs: 2 global memory tile loads, 5 threadgroup barriers, and LSE/delta scalar reads. Total useful compute is the same, but the overhead-to-compute ratio doubles — from ~30% at D=64 to ~60% at D=128.

3. WM=2 threadgroup reduction. D=128 requires WM=2 (2 simdgroups, 64 threads) to split register pressure across simdgroups (~202 regs/thread vs ~364 at WM=1). This adds a threadgroup reduction with 4 extra barriers per Q-block iteration for dK/dV partial sum accumulation — overhead that D=64 (WM=1, single simdgroup) never pays.

4. Fewer MMAs per iteration. Each Q-tile iteration performs:

• D=64: TQ=4, TK=4, TD=8 → 128 MMAs per iteration (good compute density)
• D=128: TQ=1, TK=2, TD=16 → 32 MMAs per iteration (4x less)

Less compute per iteration means the fixed overhead (barriers, loads) dominates a larger fraction of wall time.

5. Unfused baseline is stronger at D=128. The unfused path uses NAX-optimized matmuls with large tiles (BM=128, BN=128). At D=128, matmul shapes like [L, L] @ [L, 128] fill NAX tiles perfectly (N=128 = BN). At D=64, N=64 wastes half the tile width. The unfused baseline is proportionally faster at D=128, widening the gap.

Combined effect: The effective fused/unfused throughput ratio is ~18% at D=64 vs ~12% at D=128 (after smem aliasing, up from ~10%). Causal tile-skipping saves ~50% of work for both: 50% of 18% = 36% effective throughput (enough to beat unfused → 1.36x), but 50% of 12% = 24% (still slower → 0.80x). Even with 2 TGs/core after aliasing, D=128's combination of BQ=16, WM=2 reduction, and stronger NAX baseline keeps it below parity.

Memory: Peak Usage During VJP

The primary motivation for fused backward is memory efficiency at scale:

Config	Unfused Peak	Fused Peak	Savings	Attn Matrix Size
D=64 L=512	69 MB	21 MB	70%	17 MB
D=64 L=1024	239 MB	42 MB	82%	67 MB
D=64 L=2048	881 MB	84 MB	90%	268 MB
D=64 L=4096	3,373 MB	169 MB	95%	1,074 MB
D=96 L=1024	258 MB	63 MB	76%	67 MB
D=96 L=2048	919 MB	126 MB	86%	268 MB
D=128 L=1024	277 MB	84 MB	70%	67 MB
D=128 L=2048	956 MB	168 MB	82%	268 MB
D=128 L=2048 GQA(32/8)	940 MB	118 MB	87%	268 MB

At L=4096, unfused requires 3.4 GB for backward alone — this is where fused becomes essential. With model weights, optimizer states, and activations competing for memory during training, the 95% reduction determines whether a training run fits in memory or not.

Auto-Dispatch Validation

The auto-dispatch correctly routes short sequences to unfused and long sequences to fused. At the L=1024 threshold boundary:

D	L	Fused/Unfused ratio	Auto routes to
64	512	1.02x (identical)	Unfused
64	1023	1.00x (identical)	Unfused
64	1024	1.00x (identical)	Unfused
64	2048	0.64x	Fused (for memory)
128	512	1.01x (identical)	Unfused
128	1024	1.00x (identical)	Unfused
128	2048	0.39x	Fused (for memory)

The 1.00x ratios for L <= 1024 confirm that auto mode is running unfused (both columns execute the same code path). The speedup drop at L=2048 confirms the switch to fused.

References

• Dao et al., "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness", NeurIPS 2022 (arXiv:2205.14135)
• Dao, "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning", ICLR 2024 (arXiv:2307.08691)
• Philip Turner, Metal Flash Attention — reference Metal GPU implementation
• Marco GVA, mlx-flashattention-steel — MFA fork with STEEL V2 forward, sparse masks, dispatch policy
• Draw Things, Metal FlashAttention 2.0 — Apple Silicon backward pass analysis

---

Note

High Risk
Adds new Metal GPU backward kernels and changes SDPA forward kernels/dispatch to produce and consume logsumexp for fused VJP, which can affect correctness/performance and memory across many attention shapes. Also introduces an environment-controlled auto-dispatch policy, making runtime behavior configuration-dependent.

Overview
Implements a fused FlashAttention backward (VJP) path on Metal by adding new STEEL kernels for dQ and dK/dV, wiring them into ScaledDotProductAttentionVJP::eval_gpu, and running the two kernels concurrently while feeding precomputed delta and forward logsumexp.

Updates SDPA forward kernels to optionally emit logsumexp (needed for fused backward) and aligns the vector path to STEEL’s log2/exp2 domain; adds JIT plumbing/caching for baked constants (e.g., scale, GQA factor, alignment/causal flags) plus new build targets/headers for the VJP kernels.

Introduces an auto-dispatch policy (with MLX_SDPA_VJP_MODE, MLX_SDPA_VJP_LONG_L_THRESHOLD, MLX_SDPA_VJP_ATTENTION_BYTES_THRESHOLD) to choose fused vs unfused backward based on causal/dense, sequence length, GQA, and an attention-matrix memory ceiling, and expands Python tests + a new benchmark script to validate correctness, unaligned shapes, and memory behavior.

^{Written by Cursor Bugbot for commit 985fdd0. This will update automatically on new commits. Configure here.}

[Copilot]

Pull request overview

This PR implements Flash Attention VJP (Vector-Jacobian Product) backward pass for scaled dot-product attention, enabling memory-efficient training with automatic differentiation. The implementation includes two kernel approaches: vector VJP for short sequences (≤8) and STEEL VJP for longer sequences, with proper handling of GQA (Grouped Query Attention), causal masking, and numeric stability.

Changes:

• Adds cached logsumexp mechanism to enable VJP access without materializing attention matrices
• Implements two-kernel STEEL VJP approach (dQ and dKV kernels) with K-row ownership model to eliminate atomic operations
• Adds vector VJP kernels with float32 accumulator support for half/bfloat16 dtypes
• Updates forward kernels to output logsumexp when training and removes the training mode fallback to enable Flash Attention during training

Reviewed changes

Copilot reviewed 14 out of 14 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
mlx/fast_primitives.h	Adds cached_logsumexp member and accessors for VJP backward pass
mlx/fast.cpp	Updates VJP dispatch logic to check mask/sinks and access cached logsumexp
mlx/backend/no_gpu/primitives.cpp	Updates use_fallback signature with mask/sinks parameters
mlx/backend/metal/scaled_dot_product_attention.cpp	Implements complete VJP eval_gpu with STEEL and vector kernel dispatch
mlx/backend/metal/kernels/steel/attn/params.h	Adds AttnVJPParams structure with gradient strides and LSE strides
mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_vjp_dq.{h,metal}	Implements dQ gradient kernel (607 lines)
mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_vjp_dkv.{h,metal}	Implements dK/dV gradient kernel (817 lines)
mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h	Adds logsumexp output support to forward kernel
mlx/backend/metal/kernels/sdpa_vector_vjp.h	Implements vector VJP kernels (507 lines)
mlx/backend/metal/kernels/scaled_dot_product_attention.metal	Adds VJP kernel instantiations
mlx/backend/metal/kernels/CMakeLists.txt	Updates build rules for new VJP kernels
mlx/backend/cuda/scaled_dot_product_attention.cpp	Updates use_fallback signature for consistency

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Brooooooklyn]

PR Review Comment Resolutions

✅ Fixed Issues

Comment	Issue	Resolution
@copilot - LSE documentation	Inconsistent shape docs	Fixed in 60d7377 - clarified 2D flattened layout
@cursor - Shared memory indexing	simdgroup data overlap	Fixed in 60d7377 - changed BD to D for dQ indexing
@cursor - Missing logsumexp output	Short sequence LSE null	Fixed in 60d7377 - added buffer null check
@cursor - Missing allocation	fill_gpu on unallocated	Fixed in 60d7377 - added set_data() before fill_gpu()
@cursor - Inverted alignment check	OOB writes	Fixed in 46fc722 - removed ! from align_Q check
@cursor - Input strides for output	Potential memory corruption	Fixed in 46fc722 & 0d4316c - added stride validation
@cursor - Unused threadgroup arrays	Memory waste	Fixed in c274852 - removed unused arrays
@cursor - D=256 memory limit	Exceeds 32KB threadgroup	Fixed in c274852 - removed D=256 from supported dims
@cursor - Wrong exponential	exp vs exp2 for log2 domain	Fixed in c274852 - changed to exp2()
@cursor - Stride assertions in debug only	Release build vulnerability	Fixed in 0d4316c - converted to runtime checks

❌ Not Bugs (False Positives)

Comment	Claim	Analysis
@copilot - Redundant scale	scale * dS * (q[j] / scale) redundant	Intentional optimization - q[j] is pre-scaled, so dS * q[j] correctly computes scale * dS * Q. Avoids division, improves numerical stability.
@cursor - VJP buffer check fails	Buffer null during graph construction	Not a bug - VJP is called AFTER forward execution completes, so outputs[1].buffer().ptr() is allocated. The check correctly validates the buffer exists.

📝 Minor/Documentation

Comment	Note
@copilot - Missing const qualifier	Valid suggestion, minor improvement for follow-up
@copilot - BD==128 barrier docs	Added context in 60d7377 explaining synchronization need

[Copilot]

@Brooooooklyn I've opened a new pull request, #2, to work on those changes. Once the pull request is ready, I'll request review from you.

[Brooooooklyn]

Additional Bug Fixes (Latest Commits)

✅ Fix: Unallocated arrays passed to kernel (c1d5375)

Issue: zeros() creates lazy arrays with null data pointer. When passed to set_output_array(), the call to a.buffer().ptr() dereferences null → SIGSEGV.

Fix: Replaced zeros() with explicit allocation using array(buffer, shape, dtype) constructor followed by fill_gpu() zero-initialization.

✅ Fix: NAX path ignores logsumexp output (4caea80)

Issue: When NAX conditions are met, the function early-returns to sdpa_full_self_attention_nax without passing output_logsumexp_ or lse_out. The logsumexp buffer is allocated but never written, causing VJP to use garbage data.

Fix: Added !output_logsumexp_ to the NAX path condition at line 184, ensuring we fall through to the STEEL kernel which properly supports logsumexp output.

Both issues affected bfloat16/half VJP correctness.

[Brooooooklyn]

@codex review

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

^{Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

^{Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

^{Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

^{Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

[cursor]

Cursor Bugbot has reviewed your changes and found 3 potential issues.

^{Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

^{Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

^{Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

^{Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

^{Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

[cursor]

LSE written after output normalization corrupts sum_score values

High Severity

The LSE output at line 505 uses sum_score[i] after row_bin_op<DivOp>(sum_score) at line 482 has already divided the output tile by sum_score. The row_bin_op template in STEEL modifies each element of sum_score as part of broadcasting the division — depending on the implementation, the sum_score array values may no longer represent the original row sums after the division operation. The LSE formula max_score + log2(sum_score) requires the original unnormalized sum of exponentials, not values modified by the division broadcast. If row_bin_op leaves sum_score unchanged this is benign, but if it modifies the array (e.g., by computing reciprocals in-place), the LSE output would be incorrect, causing the VJP backward pass to produce wrong gradients.

[cursor]

Unnecessary delta computation wastes FLOPs on CUDA path

Low Severity

The delta = sum(multiply(dO_f32, O_f32), axis=3) computation is unconditionally added to VJP inputs in the shared vjp() function. The CUDA backend's eval_gpu never reads this delta (cuDNN computes it internally), but MLX's lazy evaluation materializes all primitive inputs before eval_gpu runs. This triggers an unnecessary element-wise multiply and reduce of O(B×H×L×D) on CUDA, including two astype casts to float32.