Part 3.14 - The Irrational Unending Series
We talked about the problem, then we talked about the what. Next we are going to dive deep into the why.
Now for my next trick, I am going to make 50$ disappear! (No, really. I just bought H200 and B200 rentals to run my experiments lol)
In this follow up, I am going to present some of the more interesting results for your consideration. What happens when you run a completely reproducible test on H200 cards from different cloud providers, B200 cards, and SM120 rtx 6k’s?
Nobody’s Right When Everybody’s Wrong
The problem is, there is no objective “right” answer. Every hardware/software setup for a model no matter how precise and curated differs.
Alibaba doesn’t publish “based on this prompt, here are the exact logits your setup should produce” dataset for us to consume.
So, how do we define right in an unending series of different answers? Hopefully this video renders for you.
Note how there are no units to this picture. It is not an absolute view of how different things are. It is all relative. How different is flash attention 2 on an SM120 vs B200 vs H200? Draw a triangle. How different is flash attention 2 from flash inference from triton attention on a single H200 card? Draw another triangle. By interconnecting the triangles we get a 3d region. And now we can plot the relative difference everyone is from the center of that region. So, who is closest to the center?
This result is rather interesting for several reasons I will get into. Obviously the calculations are prompt specific so I used the widest selection of data across many context depths to calculate the divergence.
| Message |
Depth |
Tokens |
Content |
| SP01-M001 |
6,335 |
108 |
text+tool_calls |
| SP01-M002 |
13,319 |
118 |
text+tool_calls |
| SP01-M003 |
13,765 |
185 |
text+tool_calls |
| SP03-M010 |
19,873 |
331 |
text |
| SP03-M027 |
29,634 |
389 |
text plus four parallel CLI tool calls |
| SP03-M046 |
67,144 |
505 |
text plus five parallel CLI tool calls |
| SP03-M054 |
79,488 |
627 |
text plus four parallel tool calls, including a mutation retry |
| SP03-M073 |
94,146 |
225 |
text plus one exact remote CLI command |
| SP03-M097 |
122,863 |
698 |
text plus one confirmed switch-configuration apply |
| SP04-M018 |
19,548 |
145 |
text+tool_calls |
| SP04-M026 |
31,058 |
376 |
text+tool_calls |
| SP04-M031 |
39,292 |
202 |
tool_calls |
| SP04-M046 |
52,085 |
130 |
tool_calls |
| SP04-M081 |
76,659 |
342 |
text+tool_calls |
| SP04-M106 |
94,962 |
340 |
text+tool_calls |
| SP06-M008 |
18,829 |
389 |
diagnostic prose plus three parallel host/container tool calls |
| SP06-M010 |
27,525 |
936 |
diagnostic prose plus three parallel shell/Python/SQLite tool calls |
| SP06-M018 |
42,192 |
347 |
migration plan plus one mutating Podman tool call |
| SP06-M019 |
42,769 |
270 |
brief recovery prose plus one corrected Podman tool call |
| SP06-M023 |
48,626 |
1,570 |
prose plus a complete Python SQLite-to-PostgreSQL migration script in a remote-write call |
| SP06-M042 |
73,149 |
330 |
prose plus two parallel code-inspection tool calls |
| SP06-M050 |
83,233 |
497 |
architectural recommendation in prose with Mermaid and exact configuration names |
| Total |
|
9,060 |
22 ranges |
A multitude of different long prompts resulted in over 9000 positions sampled.
“All you need is attention” But… Which kind?
In our case, we are testing 3 different pluggable attention backends. FA2, FI, and TRT.
The model inputs, weights, key/value cache, and output logits used bfloat16, a 16-bit floating-point format that keeps a wide numerical range but substantially less precision than 32-bit float.
Each hardware/software coordinate was highly repeatable. H200 results were byte-identical across two providers, while B200 and local SM120 repeats were also byte-identical. The differences therefore look like stable consequences of different arithmetic programs, not random run-to-run noise.
The usual disclaimer for oversimplification applies.
The orange boxes in the figure are places where a 32-bit number is rounded to 16-bit bfloat16. Blue boxes retain 32-bit state. Purple boxes show how the input tokens are divided into parallel work groups.
!!! WARNING !!!
LONG BORING MATH AHEAD! AVERT YOUR EYES NOW!!!
This section is entirely @splifingate 's fault. He showed me how to math and now you all get to suffer. Also my math is about as accurate as your GPU’s. 
What attention is trying to calculate
For one new token, the model constructs a query vector: the information it is currently looking for. Every earlier token has a key vector, used to measure how relevant that token is, and a value vector, containing the information retrieved if that token receives weight.
For earlier token number token, the ideal real-number calculation is:
\text{attention score}_{\text{token}}
=
\text{scale}
\times
\operatorname{dot}(\text{query},\text{key}_{\text{token}}),
\text{unnormalized weight}_{\text{token}}
=
\exp(\text{attention score}_{\text{token}}),
\text{normalized weight}_{\text{token}}
=
\frac{\text{unnormalized weight}_{\text{token}}}
{\sum_{\text{earlier tokens}}\text{unnormalized weight}},
\text{attention output}
=
\sum_{\text{earlier tokens}}
\text{normalized weight}_{\text{token}}
\times
\text{value}_{\text{token}}.
Qwen uses 256 numbers per attention head, so its score scale is
\text{scale}=\frac{1}{\sqrt{256}}=\frac{1}{16}=2^{-4}.
That scale is represented exactly in binary. The differences arise from how the dot products, exponentials, sums, intermediate rounding, and parallel combination are performed.
Why the calculation is split into groups
A 27,525-token history is too large for one GPU work block. Each attention implementation divides the history into groups, calculates a partial answer for every group, and combines the partial answers afterward.
For one group, the kernel first finds the local maximum score, meaning the largest attention score in that group. Subtracting it prevents the exponential from overflowing:
\text{local exponential weight}_{\text{token}}
=
\exp(
\text{attention score}_{\text{token}}
-
\text{local maximum score}
).
In exact arithmetic, subtracting a different local maximum would not change the final result because the scale factor cancels during the final combination. In the real bfloat16 path, FlashAttention 2 uses a structure like
\text{partial weighted value sum}
=
\operatorname{float32\_sum}
\left(
\operatorname{round to bfloat16}
(\text{local exponential weight})
\times
\text{value}
\right),
while its normalization sum uses the unrounded 32-bit exponential weights:
\text{partial normalization sum}
=
\operatorname{float32 sum}
(\text{local exponential weight}).
Rounding does not commute with rescaling. In general,
\operatorname{round to bfloat16}(e^{x-c})
\ne
e^{-c}
\operatorname{round to bfloat16}(e^x).
Changing a group boundary can therefore change its local maximum, where every weight lands on the bfloat16 number grid, its weighted-value sum, and the shape of the final combination tree.
Why FlashAttention 2 selects 54, 62, or 87 groups
FlashAttention 2 operates on 64-token key/value blocks in this Qwen decode. At 27,525 tokens, the number of blocks is:
\text{key value block count}
=
\left\lceil\frac{27{,}525}{64}\right\rceil
=431.
A streaming multiprocessor is one of the GPU’s parallel processing units. H200 has 132, B200 has 148, and the tested SM120 card has 188. The pinned FlashAttention 2 heuristic estimates device capacity as:
\text{device capacity}
=
2\times\text{streaming multiprocessor count}.
Qwen has four key/value heads after its grouped-query reshape, so a candidate split count is scored using:
\text{work ratio}
=
\frac{4\times\text{candidate split count}}
{\text{device capacity}},
\text{estimated efficiency}
=
\frac{\text{work ratio}}{\lceil\text{work ratio}\rceil}.
The source chooses the smallest eligible split count whose estimated efficiency is at least 85% of the best candidate. This gives:
| Platform |
GPU processing units |
Chosen groups |
64-token blocks per group |
Maximum normal group span |
| H200 |
132 |
54 |
8 |
512 tokens |
| B200 |
148 |
62 |
7 |
448 tokens |
| SM120 |
188 |
87 |
5 |
320 tokens |
The H200 boundaries begin at tokens 0, 512, 1,024, and so on. B200 boundaries begin at 0, 448, 896. SM120 boundaries begin at 0, 320, 640. Consequently, the same token is normalized alongside different neighboring tokens on each GPU.
Plugging in actual numbers
The accompanying Python program creates a deterministic 27,525-token teaching
input. These are synthetic scores and values, not captured Qwen activations, but
the program uses the real 64-token blocks and the real 54/62/87 split counts.
The ideal real-number output for that input is
\text{exact real number output}=1.325056137248.
The scalar FlashAttention-2-like calculation produces:
| Platform |
Split count |
First group size |
First local maximum |
32-bit output before final rounding |
Final bfloat16 output |
| H200 |
54 |
512 tokens |
1.875 |
1.3241630793 |
1.3203125 |
| B200 |
62 |
448 tokens |
1.750 |
1.3242336512 |
1.3281250 |
| SM120 |
87 |
320 tokens |
1.500 |
1.3242592812 |
1.3281250 |
Only the platform-selected grouping changed in this FlashAttention 2 example. The H200 result falls on one bfloat16 value, while the B200 and SM120 results round to the next value. This example is deliberately chosen to expose the boundary and does not imply that one split count is generally more accurate.
Enjoy some vibe coded AI slop:
#!/usr/bin/env python3
"""Plain-language, dependency-free attention arithmetic walkthrough.
This is an educational scalar simulator, not a bit-for-bit CUDA emulator.
It uses the real partition geometry selected for the 27,525-token Qwen decode
example, while replacing 256-dimensional vectors and GPU reduction trees with
one scalar value per token and explicit float32 operations.
Run:
python attention_algorithms_walkthrough.py
python attention_algorithms_walkthrough.py --show-first-partitions 2
"""
from __future__ import annotations
import argparse
import math
import struct
from dataclasses import dataclass
TOKENS_PER_FLASH_ATTENTION_BLOCK = 64
TOKENS_PER_XQA_WORK_TILE = 256
TOKENS_PER_TENSOR_RT_LLM_WORK_TILE = 128
TRITON_SEGMENT_COUNT = 16
@dataclass(frozen=True)
class PlatformDetails:
platform_name: str
streaming_multiprocessor_count: int
@dataclass
class LocalAttentionSummary:
token_count: int
local_maximum_score: float
normalization_sum: float
weighted_value_sum: float
normalized_partial_output: float
log_sum_of_exponentials: float
@dataclass
class AttentionResult:
platform_name: str
algorithm_name: str
partition_count: int
float32_output_before_final_rounding: float
bfloat16_output: float
first_partition_summaries: list[LocalAttentionSummary]
qualification: str = ""
H200 = PlatformDetails("NVIDIA H200 (SM90 architecture)", 132)
B200 = PlatformDetails("NVIDIA B200 (SM100 architecture)", 148)
SM120 = PlatformDetails("NVIDIA SM120 card", 188)
def round_to_float32(number: float) -> float:
"""Round a Python number to IEEE float32."""
return struct.unpack("<f", struct.pack("<f", number))[0]
def round_to_bfloat16(number: float) -> float:
"""Round float32 to bfloat16, ties to even, then return it as a float."""
float32_bits = struct.unpack("<I", struct.pack("<f", round_to_float32(number)))[0]
tie_to_even_bias = 0x7FFF + ((float32_bits >> 16) & 1)
bfloat16_bits = (float32_bits + tie_to_even_bias) & 0xFFFF0000
return struct.unpack("<f", struct.pack("<I", bfloat16_bits))[0]
def add_as_float32(left_number: float, right_number: float) -> float:
return round_to_float32(round_to_float32(left_number) + round_to_float32(right_number))
def multiply_as_float32(left_number: float, right_number: float) -> float:
return round_to_float32(round_to_float32(left_number) * round_to_float32(right_number))
def choose_flash_attention_two_split_count(
total_token_count: int,
streaming_multiprocessor_count: int,
) -> tuple[int, int, int]:
"""Implement the pinned FlashAttention 2 split-count heuristic.
Qwen has four key/value heads after its grouped-query reshape. The FA2
source treats each streaming multiprocessor as capacity for two of these
128-thread work blocks.
"""
key_value_block_count = math.ceil(
total_token_count / TOKENS_PER_FLASH_ATTENTION_BLOCK
)
effective_device_capacity = 2 * streaming_multiprocessor_count
workers_before_splitting = 4
largest_candidate = min(128, effective_device_capacity, key_value_block_count)
eligible_candidates: list[tuple[int, float]] = []
for candidate_split_count in range(1, largest_candidate + 1):
changes_blocks_per_split = candidate_split_count == 1 or math.ceil(
key_value_block_count / candidate_split_count
) != math.ceil(key_value_block_count / (candidate_split_count - 1))
if not changes_blocks_per_split:
continue
wave_count = (
workers_before_splitting
* candidate_split_count
/ effective_device_capacity
)
estimated_efficiency = wave_count / math.ceil(wave_count)
eligible_candidates.append((candidate_split_count, estimated_efficiency))
best_efficiency = max(efficiency for _, efficiency in eligible_candidates)
selected_split_count = next(
candidate
for candidate, efficiency in eligible_candidates
if efficiency >= 0.85 * best_efficiency
)
blocks_per_split = math.ceil(key_value_block_count / selected_split_count)
return selected_split_count, key_value_block_count, blocks_per_split
def make_triton_partitions(total_token_count: int) -> list[list[int]]:
"""Triton decode: 16 contiguous segments made from 16-token tiles."""
tokens_per_tile = 16
tiles_per_segment = math.ceil(
total_token_count / (TRITON_SEGMENT_COUNT * tokens_per_tile)
)
tokens_per_segment = tiles_per_segment * tokens_per_tile
return [
list(
range(
segment_number * tokens_per_segment,
min((segment_number + 1) * tokens_per_segment, total_token_count),
)
)
for segment_number in range(TRITON_SEGMENT_COUNT)
if segment_number * tokens_per_segment < total_token_count
]
def make_flash_attention_two_partitions(
total_token_count: int,
split_count: int,
) -> list[list[int]]:
"""FA2: contiguous 64-token blocks, visited backward inside each split."""
block_count = math.ceil(total_token_count / TOKENS_PER_FLASH_ATTENTION_BLOCK)
blocks_per_split = math.ceil(block_count / split_count)
partitions: list[list[int]] = []
for split_number in range(split_count):
first_block = split_number * blocks_per_split
one_past_last_block = min(first_block + blocks_per_split, block_count)
token_indices: list[int] = []
for block_number in range(one_past_last_block - 1, first_block - 1, -1):
first_token = block_number * TOKENS_PER_FLASH_ATTENTION_BLOCK
one_past_last_token = min(
first_token + TOKENS_PER_FLASH_ATTENTION_BLOCK,
total_token_count,
)
token_indices.extend(range(first_token, one_past_last_token))
if token_indices:
partitions.append(token_indices)
return partitions
def make_xqa_partitions(
total_token_count: int,
partition_count: int,
) -> list[list[int]]:
"""FlashInfer XQA: round-robin streams of 256-token GPU work tiles."""
tile_count = math.ceil(total_token_count / TOKENS_PER_XQA_WORK_TILE)
partitions: list[list[int]] = []
for partition_number in range(partition_count):
token_indices: list[int] = []
for tile_number in range(partition_number, tile_count, partition_count):
first_token = tile_number * TOKENS_PER_XQA_WORK_TILE
one_past_last_token = min(
first_token + TOKENS_PER_XQA_WORK_TILE,
total_token_count,
)
token_indices.extend(range(first_token, one_past_last_token))
if token_indices:
partitions.append(token_indices)
return partitions
def make_tensor_rt_llm_partitions(total_token_count: int) -> list[list[int]]:
"""Visible reducer geometry: one partial for each 128-token work tile."""
return [
list(range(first_token, min(first_token + TOKENS_PER_TENSOR_RT_LLM_WORK_TILE, total_token_count)))
for first_token in range(0, total_token_count, TOKENS_PER_TENSOR_RT_LLM_WORK_TILE)
]
def calculate_local_attention_summary(
attention_scores: list[float],
scalar_values: list[float],
token_indices: list[int],
*,
normalization_uses_bfloat16_weights: bool,
partial_output_is_stored_as_bfloat16: bool,
) -> LocalAttentionSummary:
"""Calculate one partition with explicit datatype boundaries.
The input `attention_scores` are already-scaled query-key dot products.
Each `scalar_value` stands in for one component of a real 256-component
value vector.
"""
local_maximum_score = round_to_float32(
max(attention_scores[token_index] for token_index in token_indices)
)
normalization_sum = round_to_float32(0.0)
weighted_value_sum = round_to_float32(0.0)
for token_index in token_indices:
score_minus_local_maximum = round_to_float32(
attention_scores[token_index] - local_maximum_score
)
exponential_weight_float32 = round_to_float32(
math.exp(score_minus_local_maximum)
)
exponential_weight_bfloat16 = round_to_bfloat16(
exponential_weight_float32
)
weight_used_by_normalization = (
exponential_weight_bfloat16
if normalization_uses_bfloat16_weights
else exponential_weight_float32
)
normalization_sum = add_as_float32(
normalization_sum,
weight_used_by_normalization,
)
input_value_bfloat16 = round_to_bfloat16(scalar_values[token_index])
weighted_value = multiply_as_float32(
exponential_weight_bfloat16,
input_value_bfloat16,
)
weighted_value_sum = add_as_float32(weighted_value_sum, weighted_value)
normalized_partial_output = round_to_float32(
weighted_value_sum / normalization_sum
)
if partial_output_is_stored_as_bfloat16:
normalized_partial_output = round_to_bfloat16(normalized_partial_output)
log_sum_of_exponentials = round_to_float32(
local_maximum_score
+ round_to_float32(math.log(normalization_sum))
)
return LocalAttentionSummary(
token_count=len(token_indices),
local_maximum_score=local_maximum_score,
normalization_sum=normalization_sum,
weighted_value_sum=weighted_value_sum,
normalized_partial_output=normalized_partial_output,
log_sum_of_exponentials=log_sum_of_exponentials,
)
def merge_float32_numerator_partials(
partial_summaries: list[LocalAttentionSummary],
) -> float:
"""Triton-style merge of float32 numerator, maximum, and normalizer."""
global_maximum_score = round_to_float32(
max(summary.local_maximum_score for summary in partial_summaries)
)
merged_normalization_sum = round_to_float32(0.0)
merged_weighted_value_sum = round_to_float32(0.0)
for summary in partial_summaries:
correction = round_to_float32(
math.exp(
round_to_float32(
summary.local_maximum_score - global_maximum_score
)
)
)
merged_normalization_sum = add_as_float32(
merged_normalization_sum,
multiply_as_float32(summary.normalization_sum, correction),
)
merged_weighted_value_sum = add_as_float32(
merged_weighted_value_sum,
multiply_as_float32(summary.weighted_value_sum, correction),
)
return round_to_float32(merged_weighted_value_sum / merged_normalization_sum)
def merge_normalized_outputs_using_log_sums(
partial_summaries: list[LocalAttentionSummary],
) -> float:
"""FA2-style merge of normalized float32 output and float32 log-sum-exp."""
global_log_sum_maximum = round_to_float32(
max(summary.log_sum_of_exponentials for summary in partial_summaries)
)
merged_weight_sum = round_to_float32(0.0)
merged_output_sum = round_to_float32(0.0)
for summary in partial_summaries:
merge_weight = round_to_float32(
math.exp(
round_to_float32(
summary.log_sum_of_exponentials - global_log_sum_maximum
)
)
)
merged_weight_sum = add_as_float32(merged_weight_sum, merge_weight)
merged_output_sum = add_as_float32(
merged_output_sum,
multiply_as_float32(summary.normalized_partial_output, merge_weight),
)
return round_to_float32(merged_output_sum / merged_weight_sum)
def merge_xqa_normalized_bfloat16_partials(
partial_summaries: list[LocalAttentionSummary],
) -> float:
"""XQA-style merge of bfloat16 normalized output plus float32 statistics."""
global_maximum_score = round_to_float32(
max(summary.local_maximum_score for summary in partial_summaries)
)
merged_normalization_sum = round_to_float32(0.0)
merged_output_sum = round_to_float32(0.0)
for summary in partial_summaries:
maximum_correction = round_to_float32(
math.exp(
round_to_float32(
summary.local_maximum_score - global_maximum_score
)
)
)
corrected_normalization = multiply_as_float32(
summary.normalization_sum,
maximum_correction,
)
merged_normalization_sum = add_as_float32(
merged_normalization_sum,
corrected_normalization,
)
merged_output_sum = add_as_float32(
merged_output_sum,
multiply_as_float32(
summary.normalized_partial_output,
corrected_normalization,
),
)
return round_to_float32(merged_output_sum / merged_normalization_sum)
def merge_tensor_rt_llm_bfloat16_numerators(
partial_summaries: list[LocalAttentionSummary],
) -> float:
"""Visible TensorRT-LLM reducer recurrence, with already-scaled scores."""
running_maximum_score = float("-inf")
running_normalization_sum = round_to_float32(0.0)
running_weighted_value_sum = round_to_float32(0.0)
for summary in partial_summaries:
new_maximum_score = max(
running_maximum_score,
summary.local_maximum_score,
)
previous_state_correction = (
round_to_float32(0.0)
if running_maximum_score == float("-inf")
else round_to_float32(
math.exp(
round_to_float32(
running_maximum_score - new_maximum_score
)
)
)
)
new_partial_correction = round_to_float32(
math.exp(
round_to_float32(
summary.local_maximum_score - new_maximum_score
)
)
)
running_normalization_sum = round_to_float32(
multiply_as_float32(
running_normalization_sum,
previous_state_correction,
)
+ multiply_as_float32(
summary.normalization_sum,
new_partial_correction,
)
)
partial_numerator_bfloat16 = round_to_bfloat16(
summary.weighted_value_sum
)
running_weighted_value_sum = round_to_float32(
multiply_as_float32(
running_weighted_value_sum,
previous_state_correction,
)
+ multiply_as_float32(
partial_numerator_bfloat16,
new_partial_correction,
)
)
running_maximum_score = new_maximum_score
return round_to_float32(
running_weighted_value_sum / running_normalization_sum
)
def calculate_partials(
attention_scores: list[float],
scalar_values: list[float],
partitions: list[list[int]],
*,
normalization_uses_bfloat16_weights: bool,
partial_output_is_stored_as_bfloat16: bool,
) -> list[LocalAttentionSummary]:
return [
calculate_local_attention_summary(
attention_scores,
scalar_values,
token_indices,
normalization_uses_bfloat16_weights=normalization_uses_bfloat16_weights,
partial_output_is_stored_as_bfloat16=partial_output_is_stored_as_bfloat16,
)
for token_indices in partitions
]
def simulate_triton(
platform: PlatformDetails,
attention_scores: list[float],
scalar_values: list[float],
) -> AttentionResult:
partitions = make_triton_partitions(len(attention_scores))
summaries = calculate_partials(
attention_scores,
scalar_values,
partitions,
normalization_uses_bfloat16_weights=False,
partial_output_is_stored_as_bfloat16=False,
)
output = merge_float32_numerator_partials(summaries)
return AttentionResult(
platform.platform_name,
"Triton attention: 16 contiguous segments, float32 split state",
len(partitions),
output,
round_to_bfloat16(output),
summaries,
"The source-level segment formula is common; target-specific machine code still differs.",
)
def simulate_flash_attention_two(
platform: PlatformDetails,
attention_scores: list[float],
scalar_values: list[float],
) -> AttentionResult:
split_count, _, _ = choose_flash_attention_two_split_count(
len(attention_scores),
platform.streaming_multiprocessor_count,
)
partitions = make_flash_attention_two_partitions(
len(attention_scores),
split_count,
)
summaries = calculate_partials(
attention_scores,
scalar_values,
partitions,
normalization_uses_bfloat16_weights=False,
partial_output_is_stored_as_bfloat16=False,
)
output = merge_normalized_outputs_using_log_sums(summaries)
return AttentionResult(
platform.platform_name,
"FlashAttention 2: platform-selected contiguous split count",
len(partitions),
output,
round_to_bfloat16(output),
summaries,
)
def simulate_xqa(
platform: PlatformDetails,
attention_scores: list[float],
scalar_values: list[float],
) -> AttentionResult:
key_value_head_count = 4
batch_size = 1
partition_count = min(
max(
1,
platform.streaming_multiprocessor_count
// (batch_size * key_value_head_count),
),
math.ceil(len(attention_scores) / TOKENS_PER_XQA_WORK_TILE),
)
partitions = make_xqa_partitions(len(attention_scores), partition_count)
summaries = calculate_partials(
attention_scores,
scalar_values,
partitions,
normalization_uses_bfloat16_weights=True,
partial_output_is_stored_as_bfloat16=True,
)
output = merge_xqa_normalized_bfloat16_partials(summaries)
return AttentionResult(
platform.platform_name,
'FlashInfer decoder named "XQA": round-robin work tiles and bfloat16 partial outputs',
len(partitions),
output,
round_to_bfloat16(output),
summaries,
)
def simulate_tensor_rt_llm_reducer(
platform: PlatformDetails,
attention_scores: list[float],
scalar_values: list[float],
) -> AttentionResult:
partitions = make_tensor_rt_llm_partitions(len(attention_scores))
summaries = calculate_partials(
attention_scores,
scalar_values,
partitions,
normalization_uses_bfloat16_weights=False,
partial_output_is_stored_as_bfloat16=False,
)
output = merge_tensor_rt_llm_bfloat16_numerators(summaries)
return AttentionResult(
platform.platform_name,
"TensorRT-LLM generation reducer: 128-token partials and ordered merge",
len(partitions),
output,
round_to_bfloat16(output),
summaries,
"The reducer is source-backed. The packaged main attention body is represented illustratively here.",
)
def make_synthetic_27_525_token_input() -> tuple[list[float], list[float]]:
"""Deterministic teaching input chosen to expose a one-bfloat16-ULP split."""
repeating_scores = [
0.5,
1.0,
0.75,
-2.0,
-1.25,
-0.75,
-0.5,
-3.0,
-0.25,
-2.25,
0.75,
0.5,
]
repeating_values = [
1.875,
3.5,
-0.5,
6.0,
0.0,
-1.0,
5.5,
-0.5,
-3.5,
0.5,
4.0,
-2.5,
]
attention_scores: list[float] = []
scalar_values: list[float] = []
for token_index in range(27_525):
repeating_position = token_index % len(repeating_scores)
key_value_block_number = token_index // TOKENS_PER_FLASH_ATTENTION_BLOCK
score_offset = (key_value_block_number % 9) * 0.125
attention_scores.append(
repeating_scores[repeating_position] + score_offset
)
scalar_values.append(repeating_values[repeating_position])
return attention_scores, scalar_values
def calculate_exact_real_number_reference(
attention_scores: list[float],
scalar_values: list[float],
) -> float:
largest_score = max(attention_scores)
exponential_weights = [
math.exp(score - largest_score) for score in attention_scores
]
return math.fsum(
weight * value
for weight, value in zip(exponential_weights, scalar_values)
) / math.fsum(exponential_weights)
def print_result(result: AttentionResult, partition_examples_to_show: int) -> None:
print(f"\n{result.platform_name}")
print(f" Algorithm: {result.algorithm_name}")
print(f" Number of partitions: {result.partition_count}")
print(
" Float32 output before final bfloat16 rounding: "
f"{result.float32_output_before_final_rounding:.10f}"
)
print(f" Final bfloat16 output: {result.bfloat16_output:.10f}")
if result.qualification:
print(f" Important qualification: {result.qualification}")
for partition_number, summary in enumerate(
result.first_partition_summaries[:partition_examples_to_show]
):
print(f" Partition {partition_number}:")
print(f" tokens processed: {summary.token_count}")
print(f" largest attention score in this partition: {summary.local_maximum_score:.7f}")
print(f" sum used to normalize the weights: {summary.normalization_sum:.7f}")
print(f" weighted-value sum before normalization: {summary.weighted_value_sum:.7f}")
print(f" normalized partial output stored for merge: {summary.normalized_partial_output:.7f}")
print(f" log of this partition's exponential sum: {summary.log_sum_of_exponentials:.7f}")
def demonstrate_target_specific_reduction_order() -> None:
half_of_one_float32_unit = 2.0 ** -24
left_associated = add_as_float32(
add_as_float32(1.0, half_of_one_float32_unit),
half_of_one_float32_unit,
)
right_associated = add_as_float32(
1.0,
add_as_float32(
half_of_one_float32_unit,
half_of_one_float32_unit,
),
)
print("Reduction-order example using the same three real numbers:")
print(f" Add left pair first: {left_associated:.10f}")
print(f" Add right pair first: {right_associated:.10f}")
print(" The second result is one float32 unit larger because the parenthesis tree changed.\n")
def main() -> None:
argument_parser = argparse.ArgumentParser()
argument_parser.add_argument(
"--show-first-partitions",
type=int,
default=1,
help="number of initial partition summaries to print for each algorithm",
)
arguments = argument_parser.parse_args()
demonstrate_target_specific_reduction_order()
attention_scores, scalar_values = make_synthetic_27_525_token_input()
exact_reference = calculate_exact_real_number_reference(
attention_scores,
scalar_values,
)
print("Synthetic 27,525-token teaching input (not captured Qwen activations)")
print(" Each attention score is an already-scaled query-key dot product.")
print(" Each scalar value represents one component of a 256-component value vector.")
print(f" Exact real-number softmax result: {exact_reference:.12f}")
print("\n=== Triton requested backend on all three platforms ===")
for platform in (H200, B200, SM120):
print_result(
simulate_triton(platform, attention_scores, scalar_values),
arguments.show_first_partitions,
)
print("\n=== Explicit FlashAttention 2 on all three platforms ===")
for platform in (H200, B200, SM120):
print_result(
simulate_flash_attention_two(
platform,
attention_scores,
scalar_values,
),
arguments.show_first_partitions,
)
print("\n=== FlashInfer requested backend resolves differently by platform ===")
print_result(
simulate_xqa(H200, attention_scores, scalar_values),
arguments.show_first_partitions,
)
print_result(
simulate_tensor_rt_llm_reducer(B200, attention_scores, scalar_values),
arguments.show_first_partitions,
)
print_result(
simulate_xqa(SM120, attention_scores, scalar_values),
arguments.show_first_partitions,
)
if __name__ == "__main__":
main()
The slop generator above prints descriptive intermediate names such as local_maximum_score, normalization_sum, weighted_value_sum, and normalized_partial_output for each implementation.
Backend names expand into four decode algorithms
The scored SP06 workload is 99.84% one-token-at-a-time decode, so the decode algorithms dominate the comparison.
Triton attention always uses 16 contiguous token groups. It stores the weighted-value sum, local maximum, and normalization sum as 32-bit floats, then combines at most 16 group records in a 32-bit reduction tree.
FlashAttention 2 uses contiguous 64-token blocks, but its 54/62/87 group count depends on the platform’s processing-unit count. It stores each normalized partial output and its logarithmic normalization statistic as 32-bit floats.
FlashInfer’s decoder named “XQA” runs on H200 and SM120. XQA is the source’s kernel-family name. It uses 33 round-robin groups on H200 and 47 on SM120, assigning 256-token GPU work tiles cyclically among those groups. It uses bfloat16 exponential weights for both the normalization sum and the weighted-value calculation. It also rounds every normalized group output to bfloat16 before the final combination.
FlashInfer on B200 instead selects a TensorRT-LLM generation-attention kernel. Its visible reducer consumes one bfloat16 weighted-value partial plus 32-bit maximum and normalization statistics for each 128-token work tile, then combines those partials using an ordered running recurrence. The main attention body is packaged machine code, so its internal probability rounding cannot be stated without a kernel launch trace.
GPU specific lowering in Triton
Triton source contains abstract operations such as “dot product” and “sum.” Target-specific GPU lowering is the compiler stage that maps those abstract operations to the lane layout, tensor-core instructions, shuffle operations, and reduction tree supported by a particular NVIDIA architecture:
Triton program:
-> common compiler representations
-> SM90, SM100, or SM120-specific layout
-> virtual GPU instructions
-> final machine code
The pinned Triton attention source is identical through its common LLVM compiler representation on all three platforms, but the target layouts, virtual instructions, and machine-code binaries differ.
Changing only a floating-point parenthesis tree can change a result. At 1.0, 2^{-24} is half the distance to the next float32 number:
\operatorname{round to float32}
\left(
\operatorname{round to float32}(1+2^{-24})+2^{-24}
\right)
=1,
but
\operatorname{round to float32}
\left(
1+
\operatorname{round to float32}(2^{-24}+2^{-24})
\right)
=1.0000001192092896.
Both expressions contain the same three real numbers. The second performs the two small additions first, so their combined value survives rounding.
This does not prove which lowered instruction first changes a real Qwen bit. It explains why identical high-level source is not yet an identical executable floating-point specification.
OK, so where does that leave us?
The attention backend was one visible source of variability I wanted to explore in detail despite how unimportant it seems at face value. I will continue aggregating and parsing through an enormous pile of experimental results and come back with something more entertaining next time.
EDIT: The math-free version for JW. This was too funny not to post.
Okay sweetie, sit down. Put the crayons away. This is important.
When a mommy matrix and a daddy matrix love each other very much, a dot product is born. Now, in a normal family, you’d think that’s the end of the story. But no. Attention wants every single word in a sentence to hold hands with every other word in the sentence. All of them. At once. It’s like a birthday party where you invited the entire school and now you have to write a thank-you card for every kid to every other kid. That’s N-squared cards, honey. Your hand will fall off.
And here’s the really dumb part. The GPU has a tiny, super-fast desk right in front of it, and a giant, slow closet down the hall. Old attention would write every single thank-you card, carry the whole enormous pile down the hall, shove it in the closet, walk back, go “wait, I need those,” walk down the hall again, get the pile, walk back, do a little math, and then repeat this like a golden retriever who forgot where it put the ball. The GPU wasn’t slow at thinking. It was slow at walking to the closet. Groundbreaking stuff, really.
So some very tired grown-ups invented Flash Attention, which is basically what your teacher told you the first day of kindergarten: only take out a few crayons at a time. Instead of making the whole giant pile, you grab a little chunk of words, do all the math with them right there on the tiny fast desk, write down only the tiny answer, and then grab the next chunk. You never build the big pile. The big pile does not exist. The closet is barely involved. Everyone is happier.
“But how do you do the softmax if you haven’t seen everyone yet?” Great question, and also, why do you know what softmax is, you’re five. The answer is you keep a little running tally, like when Grandma keeps changing who her favorite grandchild is and you have to keep re-doing the math. Every time a new chunk shows up, you fix your earlier numbers a little bit. It’s called the online softmax, and it’s the reason this whole thing works, and nobody appreciates it enough.
The last trick is even sillier. When it’s time to learn from mistakes, the GPU needs those thank-you cards again. Does it go back to the closet? No. It just redoes the math from scratch, because it turns out doing the homework twice is faster than walking to the closet once. Let that sink in the next time an adult tells you to “work smarter, not harder.”
So that’s Flash Attention. Same~ish answer, same~ish math, no giant pile, way less walking, and everybody acts like it was obvious the whole time. Now go tell your friends and watch them not care.
.
