physics.transformer

physics.transformer

Transformer FLOP accounting identities.

Functions

Name	Description
calc_attention_crossover	Find where full attention and dense block multiply-add counts match.
calc_attention_matmul_flops	Count work in both rectangular attention matrix products.
calc_attention_score_macs	Count multiply-adds in both full attention matrix products.
calc_online_softmax	Accumulate softmax normalization and weighted values in score blocks.
calc_score_matrix_intensity	Compute arithmetic intensity for materialized attention scores.
calc_transformer_block_weights	Count attention and feed-forward weights in a dense decoder block.
calc_transformer_decode_flops	Estimate parameter-dominated forward work for generated tokens.
calc_transformer_prefill_flops	Estimate projection and attention work for one prompt.
calc_transformer_training_flops	Estimate parameter-dominated Transformer training work.
calc_vla_inference_latency	Bound one vision-language-action evaluation by compute and weight traffic.
calc_weight_only_step_time	Bound forward-pass time using weight traffic and linear computation.

calc_attention_crossover

physics.transformer.calc_attention_crossover(
    width,
    ffn_multiple=4,
    attention_projections=4,
)

Find where full attention and dense block multiply-add counts match.

Parameters

Name	Type	Description	Default
width	int	Positive hidden width H in scalar features.	required
ffn_multiple	float	Positive feed-forward width expansion ratio; default four.	`4`
attention_projections	int	Nonnegative number of square attention projections; default four.	`4`

Returns

Name	Type	Description
	float	Sequence length in tokens at the algebraic crossover; not rounded to an integer token boundary.

Notes

Equation: S_cross = W_block / (2 H), obtained from W_block*S = 2 S**2 H. Assumptions: the dense block shape in calc_transformer_block_weights and full square attention. This compares operation counts, not execution time. Excludes: different kernel efficiencies, memory bandwidth, causal masking, biases, embeddings, and grouped-query or sparse-expert architectures.

References

Derivation: divide the equality of dense and attention MAC counts by the positive sequence length and solve for S. Both sides use MACs, so no FLOP conversion factor is needed.

Examples

>>> width = 8
>>> context = calc_attention_crossover(width)
>>> weights = calc_transformer_block_weights(width)[2]
>>> weights * context == calc_attention_score_macs(context, width)
True

calc_attention_matmul_flops

physics.transformer.calc_attention_matmul_flops(
    n_layers,
    hidden_dim,
    query_tokens,
    key_tokens,
    batch_size=1,
)

Count work in both rectangular attention matrix products.

Parameters

Name	Type	Description	Default
n_layers	int	Number of transformer layers evaluated.	required
hidden_dim	int	Query-head count times head width, in elements.	required
query_tokens	int or float	Query positions evaluated per request.	required
key_tokens	int or float	Key/value positions visited by each query; may be an expected context.	required
batch_size	int	Independent requests, default one.	`1`

Returns

Name	Type	Description
	Quantity	Work in flop for QK-transpose and attention-weighted V.

Notes

Equation: F = 4 L H Q S B. Each matrix product costs two FLOPs per multiply-accumulate; the two products have the same operation count. Assumptions: every query visits every supplied key position. H is the query width, including grouped-query attention’s query heads. Key and value head widths are equal. Excludes: projections, softmax, masking, data movement, and overhead. A causal prefill that skips the masked triangle applies its own reduction; this helper does not silently choose that execution convention.

References

Vaswani et al. (2017), “Attention Is All You Need,” Section 3.2.1, https://arxiv.org/abs/1706.03762 (bibliography key: vaswani2017attention). The FLOP count follows the two matrix products in the attention equation.

Examples

Partitioning queries into chunks preserves total matrix-product work.

>>> whole = calc_attention_matmul_flops(2, 16, 8, 32)
>>> chunk = calc_attention_matmul_flops(2, 16, 4, 32)
>>> bool(whole == 2 * chunk)
True

calc_attention_score_macs

physics.transformer.calc_attention_score_macs(seq_len, width)

Count multiply-adds in both full attention matrix products.

Parameters

Name	Type	Description	Default
seq_len	int	Nonnegative sequence length S in token positions.	required
width	int	Nonnegative hidden width H, summed across attention heads.	required

Returns

Name	Type	Description
	int	Multiply-add count for QK-transpose and attention-times-V in one layer.

Notes

Equation: MACs = 2 S**2 H. Assumptions: full square attention, equal query/key/value sequence lengths, and head widths summing to H. A multiply-add counts as one MAC here. Excludes: projection matrices, softmax, causal triangular savings, normalization, and conversion from MACs to FLOPs.

References

Derivation: QK-transpose computes S**2 dot products of width H across all heads. Multiplying the resulting attention matrix by V requires the same number of multiply-adds, giving two equal terms.

Examples

>>> work = calc_attention_score_macs(8, 4)
>>> calc_attention_score_macs(16, 4) == 4 * work
True

calc_online_softmax

physics.transformer.calc_online_softmax(scores, values, block_size)

Accumulate softmax normalization and weighted values in score blocks.

Parameters

Name	Type	Description	Default
scores	sequence of float	Finite, dimensionless logits; empty input returns an empty trace.	required
values	sequence of float	Scalar values paired with scores, of the same length and shared units.	required
block_size	int	Positive number of score/value pairs processed per block.	required

Returns

Name	Type	Description
	list of dict	One record per block with running max, rescale factor, exponential sum, and unnormalized weighted output. Divide output by sum to obtain the normalized weighted mean. The output retains the values’ numeric scale.

Notes

Equation: for new maximum m, rescale old sums by exp(m_old-m), then add sum(exp(x-m)) and sum(exp(x-m)*v) for the current block. Assumptions: finite scalar logits and scalar values; each block uses its maximum to avoid overflow from exponentiating large positive logits. Excludes: vector values, attention masking, normalization of the returned output, GPU execution, and memory-traffic timing. In exact arithmetic, output/sum agrees with full-row softmax; floating-point rounding can differ.

References

Milakov and Gimelshein (2018), “Online Normalizer Calculation for Softmax.” https://arxiv.org/abs/1805.02867. The running weighted numerator follows the same exponential rescaling as the normalizer.

Examples

>>> trace = calc_online_softmax([0.0, 0.0], [2.0, 4.0], 1)
>>> last = trace[-1]
>>> last["output"] / last["sum"] == sum([2.0, 4.0]) / len([2.0, 4.0])
True

calc_score_matrix_intensity

physics.transformer.calc_score_matrix_intensity(head_dim, bytes_per_value)

Compute arithmetic intensity for materialized attention scores.

Parameters

Name	Type	Description	Default
head_dim	int	Positive dot-product width H of one attention head.	required
bytes_per_value	Quantity or float	Positive storage size b of one score in bytes; bare values mean bytes.	required

Returns

Name	Type	Description
	Quantity	FLOPs per byte for score formation with one write and one subsequent read.

Notes

Equation: intensity = (2 H) / (2 b) = H / b FLOP/byte. Assumptions: two FLOPs per dot-product element and each score is written once to memory and later read once. Excludes: Q/K input traffic, softmax, value products, cache reuse, and fused attention kernels that avoid materializing the score matrix. This local intensity is not the whole layer’s arithmetic intensity.

References

Derivation: each score is one length-H dot product, costing 2 H FLOPs. Writing and reading its b-byte representation costs 2 b bytes; divide work by traffic to obtain the score-materialization intensity.

Examples

>>> intensity = calc_score_matrix_intensity(8, 2)
>>> calc_score_matrix_intensity(16, 2) == 2 * intensity
True

calc_transformer_block_weights

physics.transformer.calc_transformer_block_weights(
    width,
    ffn_multiple=4,
    attention_projections=4,
)

Count attention and feed-forward weights in a dense decoder block.

Parameters

Name	Type	Description	Default
width	int	Positive hidden width H in scalar features.	required
ffn_multiple	float	Positive expansion ratio R of the feed-forward hidden width; default four.	`4`
attention_projections	int	Nonnegative number A of square H-by-H projections; default four represents query, key, value, and output projections.	`4`

Returns

Name	Type	Description
	tuple	Attention, feed-forward, and total scalar weight counts, respectively.

Notes

Equation: W_attention = A H**2; W_ffn = 2 R H**2; W_total = W_attention + W_ffn. Assumptions: square attention projections and two dense feed-forward matrices with widths H to R*H to H. Excludes: biases, normalization parameters, embeddings, gated feed-forward branches, grouped-query dimensions, and sparse experts.

References

Derivation: a dense matrix contributes its row count times its column count. The attention matrices each have H**2 entries; each of the two feed-forward matrices has R H**2 entries. Different architectures need different shapes.

Examples

>>> attention, feed_forward, total = calc_transformer_block_weights(8)
>>> total == attention + feed_forward
True
>>> calc_transformer_block_weights(16)[2] == 4 * total
True

calc_transformer_decode_flops

physics.transformer.calc_transformer_decode_flops(n_params, n_tokens=1)

Estimate parameter-dominated forward work for generated tokens.

Parameters

Name	Type	Description	Default
n_params	Quantity or int	Nonnegative model parameter count P, in param or as a count.	required
n_tokens	Quantity or int	Nonnegative number N of token positions; default one token.	`1`

Returns

Name	Type	Description
	Quantity	Linear projection and feed-forward work in flop.

Notes

Equation: F = 2 P N, using the registered Chinchilla decode coefficient. Assumptions: dense parameter-dominated execution; each parameter contributes one multiply and one addition per token position. Excludes: attention products over the KV cache, softmax, normalization, communication, and runtime overhead. At long contexts, callers must add attention work separately; this is not a complete decode FLOP count.

References

Kaplan et al. (2020), “Scaling Laws for Neural Language Models.” https://arxiv.org/abs/2001.08361 (bibliography key: kaplan2020scaling).

Examples

>>> single = calc_transformer_decode_flops(1000)
>>> calc_transformer_decode_flops(1000, 2) == 2 * single
True

calc_transformer_prefill_flops

physics.transformer.calc_transformer_prefill_flops(
    n_params,
    seq_len,
    n_layers=None,
    hidden_dim=None,
    causal=True,
)

Estimate projection and attention work for one prompt.

Parameters

Name	Type	Description	Default
n_params	Quantity or int	Nonnegative model parameter count P, in param or as a count.	required
seq_len	Quantity or int	Nonnegative prompt length S in tokens, represented as a count.	required
n_layers	int	Nonnegative layer count L. Attention work is omitted if this or hidden_dim is None.	`None`
hidden_dim	int	Nonnegative model width H, summed across attention heads.	`None`
causal	bool	Use a half-square causal attention approximation when True; default True.	`True`

Returns

Name	Type	Description
	Quantity	Projection plus attention matrix-product work in flop for one prompt.

Notes

Equation: F = 2 P S + a L H S**2, where a is 2 for causal attention and 4 for bidirectional attention. If either attention dimension is omitted, only 2 P S is returned. Assumptions: two matrix products for attention, two FLOPs per multiply-add, and head widths summing to H. Causal attention uses half of the square matrix rather than the exact triangular S(S+1)/2 count. Excludes: softmax, normalization, embedding lookup, vision encoders, communication, and kernel overhead. A batch of prompts requires summing or multiplying this per-prompt work.

References

Derivation: QK-transpose and attention-times-V each require S**2 H multiply-adds per layer. Counting two FLOPs per multiply-add yields 4 L H S**2; the causal approximation halves that term. Parameter products contribute 2 P S independently.

Examples

>>> dense = calc_transformer_prefill_flops(1000, 8)
>>> causal = calc_transformer_prefill_flops(1000, 8, 2, 4)
>>> full = calc_transformer_prefill_flops(1000, 8, 2, 4, causal=False)
>>> full - dense == 2 * (causal - dense)
True

calc_transformer_training_flops

physics.transformer.calc_transformer_training_flops(
    n_params,
    n_tokens,
    activation_checkpointing=False,
)

Estimate parameter-dominated Transformer training work.

Parameters

Name	Type	Description	Default
n_params	Quantity or int	Parameter count P, as a count or a quantity in param.	required
n_tokens	Quantity or int	Total training tokens D, including repeated passes through the data.	required
activation_checkpointing	bool	Include one recomputed forward pass when True; default False.	`False`

Returns

Name	Type	Description
	Quantity	Analytical training work in flop, not an execution-time measurement.

Notes

Equation: C = 6 P D without recomputation; C = 8 P D with full recomputation. Forward work contributes 2 P per token and backward work contributes 4 P. Full recomputation adds another 2 P forward pass.

Assumptions: dense, parameter-dominated work. Excludes: attention over the context, embedding details, optimizer updates, communication, and runtime overhead. The recomputation increment is explicit accounting, not a fitted correction to the scaling-law paper. A serving solver must add context-dependent attention work separately.

References

Kaplan et al. (2020), “Scaling Laws for Neural Language Models.” https://arxiv.org/abs/2001.08361 (bibliography key: kaplan2020scaling). Hoffmann et al. (2022), “Training Compute-Optimal Large Language Models.” https://arxiv.org/abs/2203.15556 (bibliography key: hoffmann2022chinchilla).

Examples

Doubling processed tokens doubles work under identical assumptions.

>>> first = calc_transformer_training_flops(1000, 100)
>>> second = calc_transformer_training_flops(1000, 200)
>>> second == first * 2
True

calc_vla_inference_latency

physics.transformer.calc_vla_inference_latency(
    n_params,
    weight_memory,
    input_tokens,
    action_tokens,
    peak_flops,
    memory_bandwidth,
    *,
    bus_efficiency=1.0,
    compute_utilization=1.0,
    parallel_decode=False,
    n_layers=None,
    hidden_dim=None,
)

Bound one vision-language-action evaluation by compute and weight traffic.

Parameters

Name	Type	Description	Default
n_params	Quantity or int	Nonnegative decoder parameter count P, in param or as a count.	required
weight_memory	Quantity or float	Resident decoder weights in bytes; bare values are interpreted as bytes.	required
input_tokens	int	Number S of image-patch and prompt tokens, at least one.	required
action_tokens	int	Number D of generated action tokens, at least one.	required
peak_flops	Quantity or float	Positive compute capacity in flop/second at the execution precision.	required
memory_bandwidth	Quantity or float	Positive weight-memory bandwidth in byte/second.	required
bus_efficiency	float	Sustained fraction of memory bandwidth in (0, 1]; default one.	`1.0`
compute_utilization	float	Sustained fraction of peak compute in (0, 1]; default one.	`1.0`
parallel_decode	bool	Model all action placeholders in one bidirectional pass if True; otherwise prefill emits the first token and D-1 serial passes follow.	`False`
n_layers	int	Layer count for prefill attention work; omitted with hidden_dim by default.	`None`
hidden_dim	int	Model width for prefill attention work; both dimensions must be supplied.	`None`

Returns

Name	Type	Description
	dict	FLOPs in prefill_flops; millisecond quantities t_prefill_compute, t_weight_sweep, t_prefill, t_decode_step, t_decode, and t_total; frequency f_max in hertz; integer prefill_tokens and decode_steps; and prefill_bound identifying the larger prefill resource time.

Notes

Equation: T_prefill = max(F_prefill / (peak_flops * compute_utilization), weight_memory / (memory_bandwidth * bus_efficiency)). Serial decoding adds (D-1) * max(2 P / effective_compute, weight_memory / effective_bandwidth). Parallel decoding instead evaluates S+D prompt positions with bidirectional attention and no serial decode phase. Assumptions: ideal compute/memory overlap within each pass, serial passes, and one full weight sweep per pass. Efficiency factors are caller inputs, not predictions. Weight-only quantization changes weight_memory; callers must still supply compute capacity at the activation execution precision. Excludes: a separately timed vision encoder, KV-cache traffic and decode attention products, launch/framework overhead, and other bus clients. A weight-dominated pass often binds on bandwidth, but the maximum also permits compute to bind; this is not a measured robot-control latency.

References

Williams et al. (2009), “Roofline: An Insightful Visual Performance Model for Multicore Architectures.” https://doi.org/10.1145/1498765.1498785 (bibliography key: williams2009). The pass composition is a modeling choice. Kim, Finn, and Liang (2025), “Fine-Tuning Vision-Language-Action Models: Optimizing Speed and Success.” https://arxiv.org/abs/2502.19645 (bibliography key: kim2025openvlaoft). Parallel action decoding motivates the parallel_decode scenario; this function does not reproduce that paper’s measured speedup. Kim et al. (2024), “OpenVLA: An Open-Source Vision-Language-Action Model.” https://arxiv.org/abs/2406.09246 (bibliography key: kim2024openvla). Autoregressive action-token generation motivates the serial scenario.

Examples

>>> from mlsysim.core.units import Q_
>>> result = calc_vla_inference_latency(
...     1000, Q_(2000, "byte"), 8, 1,
...     Q_(1, "TFLOP/s"), Q_(1, "GB/s"))
>>> result["decode_steps"] == 0
True
>>> result["t_total"] == result["t_prefill"]
True

calc_weight_only_step_time

physics.transformer.calc_weight_only_step_time(
    n_params,
    n_tokens,
    sustained_rate,
    weight_read_time,
)

Bound forward-pass time using weight traffic and linear computation.

Parameters

Name	Type	Description	Default
n_params	Quantity or int	Model parameter count P, in param or as a count.	required
n_tokens	Quantity or int	Token positions N processed in this pass.	required
sustained_rate	Quantity	Effective computation rate R in flop/second, including any efficiency.	required
weight_read_time	Quantity	Time required to read the weights once, in time units.	required

Returns

Name	Type	Description
	Quantity	Step-time bound in the same units as weight_read_time.

Notes

Equation: T = max(T_weights, 2 P N / R). Assumptions: one weight read and ideal compute/memory overlap. Excludes: attention traffic and attention matrix products, KV-cache reads, runtime overhead, and communication. This is a weight-only bound, not a complete serving-latency prediction. R is supplied, not fitted here.

References

Examples

The step cannot finish before its weight read completes.

>>> from mlsysim.core.units import Q_
>>> read_time = Q_(1, "ms")
>>> step = calc_weight_only_step_time(
...     1000, 1, Q_(1, "TFLOP/s"), read_time)
>>> bool(step >= read_time)
True