October 2026 main drop: mirror from GitHub

2166d7e verified about 1 month ago

12.5 kB

"""
FlashAttention Golden Model - Software Reference Implementation
Implements Dao et al. FlashAttention (Algorithm 1) with Tiling and Online Softmax
This model demonstrates:
1. Block-wise (tiled) computation of attention
2. Online softmax with running max (m_i) and sum (l_i) scalars
3. Numerical validation against standard attention
Key Reference: Dao, Dan, et al. "FlashAttention: Fast and Memory-Efficient Exact Attention
with IO-Awareness" (2022)
"""

import numpy as np
from typing import Tuple, Dict, Callable
import math


class FlashAttentionGoldenModel:
    """
    FlashAttention software golden model with configurable block sizes (Br, Bc)
    and online softmax computation.
    
    Notation:
        - Q: (N, d) Query matrix
        - K: (N, d) Key matrix
        - V: (N, d) Value matrix
        - Br: Block size for Q (query block)
        - Bc: Block size for K, V (key/value block)
        - d: Embedding dimension
        - m_i: Running maximum of logits (for numerical stability)
        - l_i: Running sum of exponentials (denominator)
    """
    
    def __init__(self, scale: float = None, eps: float = 1e-6):
        """
        Args:
            scale: Scaling factor for Q*K^T (default: 1/sqrt(d))
            eps: Small constant for numerical stability
        """
        self.scale = scale
        self.eps = eps
    
    def _softmax_stable(self, logits: np.ndarray, scale: float = None) -> np.ndarray:
        """
        Numerically stable softmax computation.
        
        Args:
            logits: (batch, seq_len) raw attention scores
            scale: optional scaling factor
            
        Returns:
            Normalized attention weights (batch, seq_len)
        """
        if scale is not None:
            logits = logits * scale
        
        # Subtract max for numerical stability
        logits_shifted = logits - np.max(logits, axis=-1, keepdims=True)
        exp_logits = np.exp(logits_shifted)
        return exp_logits / (np.sum(exp_logits, axis=-1, keepdims=True) + self.eps)
    
    def attention_naive(
        self,
        Q: np.ndarray,
        K: np.ndarray,
        V: np.ndarray,
        scale: float = None
    ) -> Tuple[np.ndarray, Dict]:
        """
        Standard attention (for validation reference).
        
        Output = softmax(Q @ K^T / sqrt(d)) @ V
        
        Args:
            Q, K, V: (N, d) matrices
            scale: scaling factor for Q @ K^T
            
        Returns:
            output: (N, d) attention output
            aux: metadata dict
        """
        if scale is None:
            scale = 1.0 / math.sqrt(Q.shape[-1])
        
        N = Q.shape[0]
        
        # Compute full attention matrix: (N, N)
        scores = Q @ K.T  # (N, N)
        scores = scores * scale
        
        # Numerically stable softmax
        scores_shifted = scores - np.max(scores, axis=-1, keepdims=True)
        attn_weights = np.exp(scores_shifted)
        attn_weights = attn_weights / np.sum(attn_weights, axis=-1, keepdims=True)
        
        # Apply to values: (N, d)
        output = attn_weights @ V
        
        return output, {
            'attn_weights': attn_weights,
            'scores': scores
        }
    
    def attention_flash(
        self,
        Q: np.ndarray,
        K: np.ndarray,
        V: np.ndarray,
        Br: int = 32,
        Bc: int = 32,
        scale: float = None,
        verbose: bool = False
    ) -> Tuple[np.ndarray, Dict]:
        """
        FlashAttention Algorithm 1: Tiled computation with online softmax.
        
        This implements the block-wise attention:
        - Partition Q into blocks of size Br
        - For each Q block, iterate over all K/V blocks (size Bc)
        - Maintain running max (m_i) and sum (l_i) for online softmax
        - Update output O incrementally
        
        Args:
            Q, K, V: (N, d) matrices
            Br: Query block size (rows of Q)
            Bc: Key/Value block size (rows of K, V)
            scale: Scaling factor for scores (default: 1/sqrt(d))
            verbose: Print debug info
            
        Returns:
            output: (N, d) attention output
            aux: metadata dict with online softmax stats
        """
        if scale is None:
            scale = 1.0 / math.sqrt(Q.shape[-1])
        
        N, d = Q.shape
        num_blocks_q = math.ceil(N / Br)
        num_blocks_kv = math.ceil(N / Bc)
        
        # Initialize output matrix O (N, d)
        O = np.zeros_like(Q)
        
        # Statistics for validation and debugging
        m_history = []  # Track m_i per block
        l_history = []  # Track l_i per block
        
        if verbose:
            print(f"FlashAttention Configuration:")
            print(f"  Input shape: Q, K, V = ({N}, {d})")
            print(f"  Block sizes: Br={Br}, Bc={Bc}")
            print(f"  Num Q-blocks: {num_blocks_q}, Num KV-blocks: {num_blocks_kv}")
            print(f"  Scale factor: {scale:.6f}")
            print()
        
        # Algorithm 1: FlashAttention
        # Iterate over blocks of Q (queries)
        for i in range(num_blocks_q):
            # Extract Q block: [i*Br : min((i+1)*Br, N), :]
            q_start, q_end = i * Br, min((i + 1) * Br, N)
            Q_block = Q[q_start:q_end, :]  # (Br, d)
            
            # Initialize for this Q block:
            # m_i: running max of logits (shape: (Br,))
            # l_i: running denominator (shape: (Br,))
            # O_i: accumulator for output (shape: (Br, d))
            m_i = np.full((q_end - q_start,), -np.inf)
            l_i = np.zeros((q_end - q_start,))
            O_i = np.zeros((q_end - q_start, d))
            
            # Iterate over blocks of K, V
            for j in range(num_blocks_kv):
                kv_start, kv_end = j * Bc, min((j + 1) * Bc, N)
                K_block = K[kv_start:kv_end, :]  # (Bc, d)
                V_block = V[kv_start:kv_end, :]  # (Bc, d)
                
                # Compute block scores: Q_block @ K_block^T  [(Br, d) @ (d, Bc)] -> (Br, Bc)
                S_ij = Q_block @ K_block.T * scale
                
                # Online softmax update (Dao et al. Algorithm 1):
                # 1. Compute max of current block: m_ij = max(S_ij) per row
                m_ij_block = np.max(S_ij, axis=1)  # (Br,)
                
                # 2. New overall max: m_ij_new = max(m_i_old, m_ij_block)
                m_ij_new = np.maximum(m_i, m_ij_block)  # (Br,)
                
                # 3. Compute softmax for this block: P_ij = exp(S_ij - m_ij_new)
                P_ij = np.exp(S_ij - m_ij_new[:, np.newaxis])  # (Br, Bc)
                
                # 4. Compute sum of exponentials: l_ij_sum = sum(P_ij) per row
                l_ij_sum = np.sum(P_ij, axis=1)  # (Br,)
                
                # 5. Update running sum with rescaling for max change:
                # l_ij_new = exp(m_i - m_ij_new) * l_i + l_ij_sum
                exp_factor = np.exp(m_i - m_ij_new)  # (Br,)
                l_ij_new = exp_factor * l_i + l_ij_sum  # (Br,)
                
                # 6. Update output O_i with rescaling:
                # O_i_new = exp(m_i - m_ij_new) * O_i + P_ij @ V_block
                scale_factor = exp_factor[:, np.newaxis]  # (Br, 1)
                O_i = scale_factor * O_i + P_ij @ V_block  # (Br, d)
                
                # 7. Update running max and sum for next iteration
                m_i = m_ij_new
                l_i = l_ij_new
            
            # Final normalization: O_i = O_i / l_i
            O_i = O_i / l_i[:, np.newaxis]
            
            # Write back to output matrix
            O[q_start:q_end, :] = O_i
            
            m_history.append(m_i.copy())
            l_history.append(l_i.copy())
            
            if verbose and (i + 1) % max(1, num_blocks_q // 4) == 0:
                print(f"  Processed Q-block {i+1}/{num_blocks_q}")
        
        return O, {
            'algorithm': 'FlashAttention Algorithm 1',
            'block_sizes': {'Br': Br, 'Bc': Bc},
            'm_history': m_history,
            'l_history': l_history
        }
    
    def validate(
        self,
        Q: np.ndarray,
        K: np.ndarray,
        V: np.ndarray,
        Br: int = 32,
        Bc: int = 32,
        tolerance: float = 1e-4,
        verbose: bool = True
    ) -> Dict:
        """
        Compare FlashAttention against naive attention and verify correctness.
        
        Args:
            Q, K, V: Input matrices
            Br, Bc: Block sizes
            tolerance: Maximum allowed relative error
            verbose: Print detailed comparison
            
        Returns:
            results: Dict with error metrics and pass/fail status
        """
        scale = 1.0 / math.sqrt(Q.shape[-1])
        
        # Compute both versions
        O_naive, aux_naive = self.attention_naive(Q, K, V, scale)
        O_flash, aux_flash = self.attention_flash(Q, K, V, Br, Bc, scale, verbose=False)
        
        # Compute error metrics
        abs_error = np.abs(O_naive - O_flash)
        rel_error = np.abs(abs_error / (np.abs(O_naive) + self.eps))
        
        max_abs_error = np.max(abs_error)
        max_rel_error = np.max(rel_error)
        mean_abs_error = np.mean(abs_error)
        mean_rel_error = np.mean(rel_error)
        
        passed = max_rel_error < tolerance
        
        if verbose:
            print("\n" + "=" * 70)
            print("VALIDATION RESULTS")
            print("=" * 70)
            print(f"Input shape: {Q.shape}")
            print(f"Block config: Br={Br}, Bc={Bc}")
            print(f"Tolerance: {tolerance:.2e}")
            print()
            print("Error Metrics:")
            print(f"  Max Absolute Error:  {max_abs_error:.2e}")
            print(f"  Mean Absolute Error: {mean_abs_error:.2e}")
            print(f"  Max Relative Error:  {max_rel_error:.2e}")
            print(f"  Mean Relative Error: {mean_rel_error:.2e}")
            print()
            print(f"Status: {'PASS ✓' if passed else 'FAIL ✗'}")
            print("=" * 70 + "\n")
        
        return {
            'passed': passed,
            'max_abs_error': max_abs_error,
            'max_rel_error': max_rel_error,
            'mean_abs_error': mean_abs_error,
            'mean_rel_error': mean_rel_error,
            'O_naive': O_naive,
            'O_flash': O_flash
        }


def run_golden_model_test():
    """
    Run a comprehensive golden model test with various configurations.
    """
    print("\n" + "=" * 70)
    print("FlashAttention RTL - Software Golden Model Test Suite")
    print("=" * 70 + "\n")
    
    # Start with a very simple test case
    print("Simple Sanity Check (N=4, d=2, Br=2, Bc=2):")
    print("-" * 70)
    np.random.seed(42)
    Q_simple = np.random.randn(4, 2).astype(np.float64) * 0.1
    K_simple = np.random.randn(4, 2).astype(np.float64) * 0.1
    V_simple = np.random.randn(4, 2).astype(np.float64) * 0.1
    
    model = FlashAttentionGoldenModel()
    results_simple = model.validate(Q_simple, K_simple, V_simple, Br=2, Bc=2, tolerance=1e-10, verbose=True)
    
    # Full test suite
    test_configs = [
        {'N': 64, 'd': 32, 'Br': 16, 'Bc': 16, 'name': 'Small (N=64, d=32)'},
        {'N': 128, 'd': 64, 'Br': 32, 'Bc': 32, 'name': 'Medium (N=128, d=64)'},
        {'N': 256, 'd': 128, 'Br': 64, 'Bc': 64, 'name': 'Large (N=256, d=128)'},
    ]
    
    all_passed = True
    
    for config in test_configs:
        N = config['N']
        d = config['d']
        Br = config['Br']
        Bc = config['Bc']
        name = config['name']
        
        print(f"\nTest: {name}")
        print("-" * 70)
        
        # Generate random data (seed for reproducibility) - use float64 for golden model
        np.random.seed(42)
        Q = np.random.randn(N, d).astype(np.float64) * 0.1
        K = np.random.randn(N, d).astype(np.float64) * 0.1
        V = np.random.randn(N, d).astype(np.float64) * 0.1
        
        # Validate with higher precision
        results = model.validate(Q, K, V, Br=Br, Bc=Bc, tolerance=1e-5, verbose=True)
        all_passed = all_passed and results['passed']
    
    print("\n" + "=" * 70)
    if all_passed:
        print("ALL TESTS PASSED ✓")
    else:
        print("SOME TESTS FAILED ✗")
    print("=" * 70 + "\n")
    
    return all_passed


if __name__ == "__main__":
    run_golden_model_test()

	"""
	FlashAttention Golden Model - Software Reference Implementation
	Implements Dao et al. FlashAttention (Algorithm 1) with Tiling and Online Softmax

	This model demonstrates:
	1. Block-wise (tiled) computation of attention
	2. Online softmax with running max (m_i) and sum (l_i) scalars
	3. Numerical validation against standard attention

	Key Reference: Dao, Dan, et al. "FlashAttention: Fast and Memory-Efficient Exact Attention
	with IO-Awareness" (2022)
	"""

	import numpy as np
	from typing import Tuple, Dict, Callable
	import math


	class FlashAttentionGoldenModel:
	"""
	FlashAttention software golden model with configurable block sizes (Br, Bc)
	and online softmax computation.

	Notation:
	- Q: (N, d) Query matrix
	- K: (N, d) Key matrix
	- V: (N, d) Value matrix
	- Br: Block size for Q (query block)
	- Bc: Block size for K, V (key/value block)
	- d: Embedding dimension
	- m_i: Running maximum of logits (for numerical stability)
	- l_i: Running sum of exponentials (denominator)
	"""

	def __init__(self, scale: float = None, eps: float = 1e-6):
	"""
	Args:
	scale: Scaling factor for Q*K^T (default: 1/sqrt(d))
	eps: Small constant for numerical stability
	"""
	self.scale = scale
	self.eps = eps

	def _softmax_stable(self, logits: np.ndarray, scale: float = None) -> np.ndarray:
	"""
	Numerically stable softmax computation.

	Args:
	logits: (batch, seq_len) raw attention scores
	scale: optional scaling factor

	Returns:
	Normalized attention weights (batch, seq_len)
	"""
	if scale is not None:
	logits = logits * scale

	# Subtract max for numerical stability
	logits_shifted = logits - np.max(logits, axis=-1, keepdims=True)
	exp_logits = np.exp(logits_shifted)
	return exp_logits / (np.sum(exp_logits, axis=-1, keepdims=True) + self.eps)

	def attention_naive(
	self,
	Q: np.ndarray,
	K: np.ndarray,
	V: np.ndarray,
	scale: float = None
	) -> Tuple[np.ndarray, Dict]:
	"""
	Standard attention (for validation reference).

	Output = softmax(Q @ K^T / sqrt(d)) @ V

	Args:
	Q, K, V: (N, d) matrices
	scale: scaling factor for Q @ K^T

	Returns:
	output: (N, d) attention output
	aux: metadata dict
	"""
	if scale is None:
	scale = 1.0 / math.sqrt(Q.shape[-1])

	N = Q.shape[0]

	# Compute full attention matrix: (N, N)
	scores = Q @ K.T # (N, N)
	scores = scores * scale

	# Numerically stable softmax
	scores_shifted = scores - np.max(scores, axis=-1, keepdims=True)
	attn_weights = np.exp(scores_shifted)
	attn_weights = attn_weights / np.sum(attn_weights, axis=-1, keepdims=True)

	# Apply to values: (N, d)
	output = attn_weights @ V

	return output, {
	'attn_weights': attn_weights,
	'scores': scores
	}

	def attention_flash(
	self,
	Q: np.ndarray,
	K: np.ndarray,
	V: np.ndarray,
	Br: int = 32,
	Bc: int = 32,
	scale: float = None,
	verbose: bool = False
	) -> Tuple[np.ndarray, Dict]:
	"""
	FlashAttention Algorithm 1: Tiled computation with online softmax.

	This implements the block-wise attention:
	- Partition Q into blocks of size Br
	- For each Q block, iterate over all K/V blocks (size Bc)
	- Maintain running max (m_i) and sum (l_i) for online softmax
	- Update output O incrementally

	Args:
	Q, K, V: (N, d) matrices
	Br: Query block size (rows of Q)
	Bc: Key/Value block size (rows of K, V)
	scale: Scaling factor for scores (default: 1/sqrt(d))
	verbose: Print debug info

	Returns:
	output: (N, d) attention output
	aux: metadata dict with online softmax stats
	"""
	if scale is None:
	scale = 1.0 / math.sqrt(Q.shape[-1])

	N, d = Q.shape
	num_blocks_q = math.ceil(N / Br)
	num_blocks_kv = math.ceil(N / Bc)

	# Initialize output matrix O (N, d)
	O = np.zeros_like(Q)

	# Statistics for validation and debugging
	m_history = [] # Track m_i per block
	l_history = [] # Track l_i per block

	if verbose:
	print(f"FlashAttention Configuration:")
	print(f" Input shape: Q, K, V = ({N}, {d})")
	print(f" Block sizes: Br={Br}, Bc={Bc}")
	print(f" Num Q-blocks: {num_blocks_q}, Num KV-blocks: {num_blocks_kv}")
	print(f" Scale factor: {scale:.6f}")
	print()

	# Algorithm 1: FlashAttention
	# Iterate over blocks of Q (queries)
	for i in range(num_blocks_q):
	# Extract Q block: [iBr : min((i+1)Br, N), :]
	q_start, q_end = i * Br, min((i + 1) * Br, N)
	Q_block = Q[q_start:q_end, :] # (Br, d)

	# Initialize for this Q block:
	# m_i: running max of logits (shape: (Br,))
	# l_i: running denominator (shape: (Br,))
	# O_i: accumulator for output (shape: (Br, d))
	m_i = np.full((q_end - q_start,), -np.inf)
	l_i = np.zeros((q_end - q_start,))
	O_i = np.zeros((q_end - q_start, d))

	# Iterate over blocks of K, V
	for j in range(num_blocks_kv):
	kv_start, kv_end = j * Bc, min((j + 1) * Bc, N)
	K_block = K[kv_start:kv_end, :] # (Bc, d)
	V_block = V[kv_start:kv_end, :] # (Bc, d)

	# Compute block scores: Q_block @ K_block^T [(Br, d) @ (d, Bc)] -> (Br, Bc)
	S_ij = Q_block @ K_block.T * scale

	# Online softmax update (Dao et al. Algorithm 1):
	# 1. Compute max of current block: m_ij = max(S_ij) per row
	m_ij_block = np.max(S_ij, axis=1) # (Br,)

	# 2. New overall max: m_ij_new = max(m_i_old, m_ij_block)
	m_ij_new = np.maximum(m_i, m_ij_block) # (Br,)

	# 3. Compute softmax for this block: P_ij = exp(S_ij - m_ij_new)
	P_ij = np.exp(S_ij - m_ij_new[:, np.newaxis]) # (Br, Bc)

	# 4. Compute sum of exponentials: l_ij_sum = sum(P_ij) per row
	l_ij_sum = np.sum(P_ij, axis=1) # (Br,)

	# 5. Update running sum with rescaling for max change:
	# l_ij_new = exp(m_i - m_ij_new) * l_i + l_ij_sum
	exp_factor = np.exp(m_i - m_ij_new) # (Br,)
	l_ij_new = exp_factor * l_i + l_ij_sum # (Br,)

	# 6. Update output O_i with rescaling:
	# O_i_new = exp(m_i - m_ij_new) * O_i + P_ij @ V_block
	scale_factor = exp_factor[:, np.newaxis] # (Br, 1)
	O_i = scale_factor * O_i + P_ij @ V_block # (Br, d)

	# 7. Update running max and sum for next iteration
	m_i = m_ij_new
	l_i = l_ij_new

	# Final normalization: O_i = O_i / l_i
	O_i = O_i / l_i[:, np.newaxis]

	# Write back to output matrix
	O[q_start:q_end, :] = O_i

	m_history.append(m_i.copy())
	l_history.append(l_i.copy())

	if verbose and (i + 1) % max(1, num_blocks_q // 4) == 0:
	print(f" Processed Q-block {i+1}/{num_blocks_q}")

	return O, {
	'algorithm': 'FlashAttention Algorithm 1',
	'block_sizes': {'Br': Br, 'Bc': Bc},
	'm_history': m_history,
	'l_history': l_history
	}

	def validate(
	self,
	Q: np.ndarray,
	K: np.ndarray,
	V: np.ndarray,
	Br: int = 32,
	Bc: int = 32,
	tolerance: float = 1e-4,
	verbose: bool = True
	) -> Dict:
	"""
	Compare FlashAttention against naive attention and verify correctness.

	Args:
	Q, K, V: Input matrices
	Br, Bc: Block sizes
	tolerance: Maximum allowed relative error
	verbose: Print detailed comparison

	Returns:
	results: Dict with error metrics and pass/fail status
	"""
	scale = 1.0 / math.sqrt(Q.shape[-1])

	# Compute both versions
	O_naive, aux_naive = self.attention_naive(Q, K, V, scale)
	O_flash, aux_flash = self.attention_flash(Q, K, V, Br, Bc, scale, verbose=False)

	# Compute error metrics
	abs_error = np.abs(O_naive - O_flash)
	rel_error = np.abs(abs_error / (np.abs(O_naive) + self.eps))

	max_abs_error = np.max(abs_error)
	max_rel_error = np.max(rel_error)
	mean_abs_error = np.mean(abs_error)
	mean_rel_error = np.mean(rel_error)

	passed = max_rel_error < tolerance

	if verbose:
	print("\n" + "=" * 70)
	print("VALIDATION RESULTS")
	print("=" * 70)
	print(f"Input shape: {Q.shape}")
	print(f"Block config: Br={Br}, Bc={Bc}")
	print(f"Tolerance: {tolerance:.2e}")
	print()
	print("Error Metrics:")
	print(f" Max Absolute Error: {max_abs_error:.2e}")
	print(f" Mean Absolute Error: {mean_abs_error:.2e}")
	print(f" Max Relative Error: {max_rel_error:.2e}")
	print(f" Mean Relative Error: {mean_rel_error:.2e}")
	print()
	print(f"Status: {'PASS ✓' if passed else 'FAIL ✗'}")
	print("=" * 70 + "\n")

	return {
	'passed': passed,
	'max_abs_error': max_abs_error,
	'max_rel_error': max_rel_error,
	'mean_abs_error': mean_abs_error,
	'mean_rel_error': mean_rel_error,
	'O_naive': O_naive,
	'O_flash': O_flash
	}


	def run_golden_model_test():
	"""
	Run a comprehensive golden model test with various configurations.
	"""
	print("\n" + "=" * 70)
	print("FlashAttention RTL - Software Golden Model Test Suite")
	print("=" * 70 + "\n")

	# Start with a very simple test case
	print("Simple Sanity Check (N=4, d=2, Br=2, Bc=2):")
	print("-" * 70)
	np.random.seed(42)
	Q_simple = np.random.randn(4, 2).astype(np.float64) * 0.1
	K_simple = np.random.randn(4, 2).astype(np.float64) * 0.1
	V_simple = np.random.randn(4, 2).astype(np.float64) * 0.1

	model = FlashAttentionGoldenModel()
	results_simple = model.validate(Q_simple, K_simple, V_simple, Br=2, Bc=2, tolerance=1e-10, verbose=True)

	# Full test suite
	test_configs = [
	{'N': 64, 'd': 32, 'Br': 16, 'Bc': 16, 'name': 'Small (N=64, d=32)'},
	{'N': 128, 'd': 64, 'Br': 32, 'Bc': 32, 'name': 'Medium (N=128, d=64)'},
	{'N': 256, 'd': 128, 'Br': 64, 'Bc': 64, 'name': 'Large (N=256, d=128)'},
	]

	all_passed = True

	for config in test_configs:
	N = config['N']
	d = config['d']
	Br = config['Br']
	Bc = config['Bc']
	name = config['name']

	print(f"\nTest: {name}")
	print("-" * 70)

	# Generate random data (seed for reproducibility) - use float64 for golden model
	np.random.seed(42)
	Q = np.random.randn(N, d).astype(np.float64) * 0.1
	K = np.random.randn(N, d).astype(np.float64) * 0.1
	V = np.random.randn(N, d).astype(np.float64) * 0.1

	# Validate with higher precision
	results = model.validate(Q, K, V, Br=Br, Bc=Bc, tolerance=1e-5, verbose=True)
	all_passed = all_passed and results['passed']

	print("\n" + "=" * 70)
	if all_passed:
	print("ALL TESTS PASSED ✓")
	else:
	print("SOME TESTS FAILED ✗")
	print("=" * 70 + "\n")

	return all_passed


	if __name__ == "__main__":
	run_golden_model_test()