AttentionTilingScenario
Training scenarios · Implementation: mlsysim/engine/scenario_definitions/training.py.
Construct this scenario with keyword inputs. The fields below show the types and defaults declared by this class.
One head’s attention at seq_len tokens, whole or in square tiles, and what each kernel stores.
Standard attention stores the scores and the probabilities (seq_len squared each) beside Q, K, V, and the output; FlashAttention stores a running maximum and sum per row instead.
Notes
Equation: score_matrix is the worked quantity; each property below states its equation and units.
Assumptions: Dense square scores use fixed precision and tiles divide the sequence length exactly.
Excludes: Kernel workspace, inter-head buffer reuse, and measured FlashAttention memory or speed.
References
Derivation: The equations below follow the stated counts, geometric relations, or accounting definitions; scenario inputs are supplied assumptions or observations, not fitted predictions.
Examples
Execute the registered worked inputs; this is a modeled calculation.
>>> from mlsysim import Scenarios
>>> case = Scenarios.AttentionScores4096TokensFp32
>>> value = case.score_matrix
>>> print(f"{value:.3f~P}")
67.109 MBAttributes
| Name | Description |
|---|---|
| all_heads_scores | Storage in GB for the score matrices of all heads. |
| bytes_per_value | |
| head_dim | |
| heads | |
| io_bytes | Configured input/output tensor storage in MB for one head. |
| io_tensors | |
| qkv_bytes | Q, K, and V storage in MB for one head. |
| row_statistics | |
| score_matrix | Storage in MB for one head’s square attention-score matrix. |
| seq_len | |
| stored_flash | Tiled-attention storage in MB including per-row statistics instead of square matrices. |
| stored_ratio | Dimensionless standard-to-tiled storage ratio. |
| stored_square_matrices | |
| stored_standard | Standard-attention storage in MB including materialized square matrices. |
| tile | |
| tile_bytes | Storage in KB for one square score tile. |
| tile_count | Whole square-tile count, tiles per side squared. |
| tiles_per_side | Whole tiles along one sequence axis; integer division drops any partial tile. |
seq_len
Type: int.
Declared default: required.
head_dim
Type: int.
Declared default: required.
tile
Type: int.
Declared default: required.
heads
Type: int.
Declared default: required.
bytes_per_value
Type: Quantity.
Declared default: required.
io_tensors
Type: int.
Declared default: 4.
stored_square_matrices
Type: int.
Declared default: 2.
row_statistics
Type: int.
Declared default: 2.
score_matrix
engine.scenario_types.AttentionTilingScenario.score_matrix
Storage in MB for one head’s square attention-score matrix.
Equation: return (self.seq_len ** 2 * self.bytes_per_value).to(‘MB’)
all_heads_scores
engine.scenario_types.AttentionTilingScenario.all_heads_scores
Storage in GB for the score matrices of all heads.
Equation: return (self.heads * self.score_matrix).to(‘GB’)
tile_bytes
engine.scenario_types.AttentionTilingScenario.tile_bytes
Storage in KB for one square score tile.
Equation: return (self.tile ** 2 * self.bytes_per_value).to(‘KB’)
tiles_per_side
engine.scenario_types.AttentionTilingScenario.tiles_per_side
Whole tiles along one sequence axis; integer division drops any partial tile.
Equation: return self.seq_len // self.tile
tile_count
engine.scenario_types.AttentionTilingScenario.tile_count
Whole square-tile count, tiles per side squared.
Equation: return self.tiles_per_side ** 2
qkv_bytes
engine.scenario_types.AttentionTilingScenario.qkv_bytes
Q, K, and V storage in MB for one head.
Q, K, and V: three tensors of seq_len x head_dim.
Equation: return (3 * self.seq_len * self.head_dim * self.bytes_per_value).to(‘MB’)
io_bytes
engine.scenario_types.AttentionTilingScenario.io_bytes
Configured input/output tensor storage in MB for one head.
Equation: return (self.io_tensors * self.seq_len * self.head_dim * self.bytes_per_value).to(‘MB’)
stored_standard
engine.scenario_types.AttentionTilingScenario.stored_standard
Standard-attention storage in MB including materialized square matrices.
Equation: return (self.io_bytes + self.stored_square_matrices * self.score_matrix).to(‘MB’)
stored_flash
engine.scenario_types.AttentionTilingScenario.stored_flash
Tiled-attention storage in MB including per-row statistics instead of square matrices.
Equation: return (self.io_bytes + self.row_statistics * self.seq_len * self.bytes_per_value).to(‘MB’)
stored_ratio
engine.scenario_types.AttentionTilingScenario.stored_ratio
Dimensionless standard-to-tiled storage ratio.
Equation: return (self.stored_standard / self.stored_flash).to(’’).magnitude