# examples/flashattention.py """Seeded offline CPU forward algebra for FlashAttention Algorithm 1. Grounding: https://arxiv.org/html/2205.14135v2#alg1, Section 3.1. This authored NumPy reference uses a = l * O and normalizes once. It adds standard 1/sqrt(D) scaling (absorb it into Q in the paper's unscaled QK). No dropout, backward kernel, SRAM residency, HBM traffic or timing claims. """ import argparse import hashlib import json import math from pathlib import Path import numpy as np def _validate(q, k, v, br, bc, causal, query_offset, key_offset): """Reject malformed arrays and undefined block or index conventions.""" arrays = (q, k, v) for array in arrays: if not isinstance(array, np.ndarray) or array.ndim != 2: raise ValueError("Q, K and V must be rank-two NumPy arrays") if array.dtype.kind != "f" or 0 in array.shape: raise ValueError("Arrays must be nonempty floating-point matrices") if not np.isfinite(array).all(): raise ValueError("Inputs must be finite") if q.shape[1] != k.shape[1] or k.shape[0] != v.shape[0]: raise ValueError("Q/K widths and K/V lengths must match") for size in (br, bc): if type(size) is not int or size <= 0: raise ValueError("Block sizes must be positive integers") for offset in (query_offset, key_offset): if type(offset) is not int or offset < 0: raise ValueError("Offsets must be nonnegative integers") if type(causal) is not bool: raise ValueError("causal must be boolean") # region: inputs_and_blocks def tiled_attention(q, k, v, br=3, bc=4, causal=False, key_order=None, query_offset=0, key_offset=0): """Return output, row max and denominator using float64 CPU tiles. Global causal indices allow key j only if key_offset+j <= query_offset+i. An entirely masked row returns zero output, m=-inf and l=0 by convention. key_order is a permutation of tile indices, for traversal checks only. """ _validate(q, k, v, br, bc, causal, query_offset, key_offset) q, k, v = (x.astype(np.float64, copy=False) for x in (q, k, v)) nq, d = q.shape nk, dv = v.shape query_starts = range(0, nq, br) key_starts = list(range(0, nk, bc)) order = list(range(len(key_starts))) if key_order is None else key_order if (not isinstance(order, (list, tuple)) or any(type(index) is not int for index in order) or sorted(order) != list(range(len(key_starts)))): raise ValueError("key_order must permute all key tiles exactly once") scale = 1.0 / math.sqrt(d) # endregion: inputs_and_blocks # region: row_state m = np.full(nq, -np.inf, dtype=np.float64) l = np.zeros(nq, dtype=np.float64) a = np.zeros((nq, dv), dtype=np.float64) # Original Algorithm 1: outer K/V loop, inner Q loop. for tile in order: j = key_starts[tile] kj, vj = k[j:j + bc], v[j:j + bc] for i in query_starts: rows = slice(i, min(i + br, nq)) # endregion: row_state # region: score_tile scores = (q[rows] @ kj.T) * scale if not np.isfinite(scores).all(): raise ValueError("Scaled score computation overflowed") if causal: qi = query_offset + np.arange(i, rows.stop) ki = key_offset + np.arange(j, j + len(kj)) scores = np.where(ki[None, :] <= qi[:, None], scores, -np.inf) # endregion: score_tile # region: online_max old_m = m[rows] new_m = np.maximum(old_m, scores.max(axis=1)) # Never evaluate -inf - (-inf), even for an initially masked tile. finite_old = np.isfinite(old_m) alpha = np.zeros_like(old_m) alpha[finite_old] = np.exp(old_m[finite_old] - new_m[finite_old]) safe_m = np.where(np.isfinite(new_m), new_m, 0.0) # endregion: online_max # region: probabilities_and_sum p = np.exp(scores - safe_m[:, None]) l[rows] = alpha * l[rows] + p.sum(axis=1) # endregion: probabilities_and_sum # region: value_accumulator a[rows] = alpha[:, None] * a[rows] + p @ vj m[rows] = new_m # endregion: value_accumulator # region: normalize_and_compare output = np.divide(a, l[:, None], out=np.zeros_like(a), where=l[:, None] > 0) return output, m, l def dense_attention(q, k, v, causal=False, query_offset=0, key_offset=0): """Independent dense stable softmax oracle; not used by tiled forward.""" _validate(q, k, v, 1, 1, causal, query_offset, key_offset) scores = (q.astype(np.float64) @ k.astype(np.float64).T) scores /= math.sqrt(q.shape[1]) if causal: mask = (key_offset + np.arange(len(k))[None, :] <= query_offset + np.arange(len(q))[:, None]) scores = np.where(mask, scores, -np.inf) maximum = scores.max(axis=1) shifted = scores - np.where(np.isfinite(maximum), maximum, 0)[:, None] weights = np.exp(shifted) sums = weights.sum(axis=1, keepdims=True) probabilities = np.divide(weights, sums, out=np.zeros_like(weights), where=sums > 0) return probabilities @ v.astype(np.float64) # endregion: normalize_and_compare def smoke(seed=7): """Run real reduced fixtures without changing NumPy's global RNG.""" rng = np.random.default_rng(seed) q = rng.normal(size=(7, 4)) k = rng.normal(size=(9, 4)) v = rng.normal(size=(9, 3)) errors = [] observations = [] for causal in (False, True): dense = dense_attention(q, k, v, causal) for br, bc in ((3, 4), (8, 2), (1, 12)): forward, m, l = tiled_attention(q, k, v, br, bc, causal) order = list(reversed(range(math.ceil(len(k) / bc)))) reverse, _, _ = tiled_attention(q, k, v, br, bc, causal, order) error = float(np.max(np.abs(forward - dense))) errors.append(error) observations.append({ "causal": causal, "blocks": [br, bc], "dense_max_abs_error": error, "reverse_max_abs_error": float(np.max( np.abs(reverse - forward))), "row_max_range": [float(m.min()), float(m.max())], "denominator_range": [float(l.min()), float(l.max())], "output_norm": float(np.linalg.norm(forward)), }) extreme, _, _ = tiled_attention(q * 1000, k * 1000, v, 2, 2) extreme_dense = dense_attention(q * 1000, k * 1000, v) masked, _, masked_l = tiled_attention( q, k, v, 2, 2, True, query_offset=0, key_offset=20) return { "source_sha256": hashlib.sha256( Path(__file__).read_bytes()).hexdigest(), "seed": seed, "dtype": "float64", "shapes": { "q": list(q.shape), "k": list(k.shape), "v": list(v.shape), "output": [len(q), v.shape[1]], }, "reductions": {"softmax": "sum over keys per query row", "comparison": "maximum absolute element error"}, "observations": observations, "max_dense_error": max(errors), "extreme_max_abs_error": float(np.max( np.abs(extreme - extreme_dense))), "fully_masked_output_norm": float(np.linalg.norm(masked)), "fully_masked_denominator_sum": float(masked_l.sum()), "finite_outputs": bool(np.isfinite(extreme).all() and np.isfinite(masked).all()), "boundary": "offline CPU forward algebra, not a GPU kernel benchmark", } def main(): """Print actual smoke JSON, with nonfinite serialization prohibited.""" parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--smoke", action="store_true", required=True) parser.add_argument("--seed", type=int, default=7) args = parser.parse_args() print(json.dumps(smoke(args.seed), indent=2, allow_nan=False)) if __name__ == "__main__": main()