FlashAttention Online Softmax Block Accumulatorhard

FlashAttention Online Softmax Block Accumulator

hard NumPy & Data Wrangling 20-25 min

NumPy allowed · plain Python earns 10 bonus points

Tokens attend to tokens

Q · Kᵀ / √d → softmax → weighted values, computed in the page.

Open lesson →

01tokens

02Q K V

03scores

04attention

modelslearnrelationships

models32%

learn38%

relationships30%

Q · Kᵀ / √d scores normalized with softmax

Merge two partitioned attention blocks using online softmax rescaling without materializing the full attention matrix.

Related Skills

Log-Sum-Exp

Task

Implement merge_online_softmax_blocks(block_maxes, block_sums, block_weights): given two blocks with max logits m_1, m_2, exponential sums l_1, l_2, and unnormalized output accumulators O_1, O_2, calculate new maximum m_new = max(m_1, m_2). Rescale factors alpha_1 = exp(m_1 - m_new) and alpha_2 = exp(m_2 - m_new). Compute merged sum l_new = alpha_1 * l_1 + alpha_2 * l_2 and output O_new = (alpha_1 * O_1 + alpha_2 * O_2) / l_new. Return dict with keys merged_max, merged_sum, and merged_output rounded to 4 decimals.

Function Signature

merge_online_softmax_blocks(block_maxes: list[float], block_sums: list[float], block_weights: list[list[float]]) -> dict[str, any]

Examples

Example 1: Merging Two Identical Blocks
Input: block_maxes = [2, 2]
block_sums = [1, 1]
block_weights = [[0.5, 0.5], [0.5, 0.5]]
Output: {'merged_max': 2, 'merged_sum': 2, 'merged_output': [0.5, 0.5]}
Explanation: Merge identical partial softmax blocks.

Example 2: One Dominant Block
Input: block_maxes = [10, 1]
block_sums = [2, 5]
block_weights = [[1, 0], [0, 1]]
Output: {'merged_max': 10, 'merged_sum': 2.0006, 'merged_output': [0.4998, 0.0001]}
Explanation: Block with high maximum dominates weighted average.

Constraints

block_maxes and block_sums have length 2
block_weights has 2 rows of dimension D

Python3Saved ✓

def merge_online_softmax_blocks(block_maxes: list[float], block_sums: list[float], block_weights: list[list[float]]) -> dict[str, any]:

"""Merge two tiled partial softmax blocks using online rescaling.

Return {'merged_max': float, 'merged_sum': float, 'merged_output': list[float]} rounded to 4 decimals."""

# Implement solution logic here

pass

Case 1: Merging Two Identical Blocks

Input: block_maxes = [2, 2]
block_sums = [1, 1]
block_weights = [[0.5, 0.5], [0.5, 0.5]]

Expected: {'merged_max': 2, 'merged_sum': 2, 'merged_output': [0.5, 0.5]}

Case 2: One Dominant Block

Input: block_maxes = [10, 1]
block_sums = [2, 5]
block_weights = [[1, 0], [0, 1]]

Expected: {'merged_max': 10, 'merged_sum': 2.0006, 'merged_output': [0.4998, 0.0001]}