Complexity for Tensor Programs
A tensor expression is a resource program. Shapes determine work and storage; execution order, precision, materialization, and memory movement determine whether that work fits the system.
- Authorship
- InterviewsVector
- Published / updated
- 2026-08-14 / 2026-08-14
- Review status
- Artifact tests passing · primary sources recorded
Original InterviewsVector material. Code examples are covered by repository tests and primary references are recorded. No named human reviewer is claimed.
The decision in one pass
Estimate a tensor program from named dimensions before measuring it. For batched matmul [b…,m,k] × [b…,k,n], derive output [b…,m,n], conventional work 2·prod(b)·m·k·n per execution when a fused multiply-add counts as two FLOPs, bytes from explicit input reads and output reads/writes at declared dtypes, and peak retained activation bytes from an explicit execution-liveness policy and retained-copy count. Bind broadcast policy, contraction order, safe operand-to-accumulation compatibility, execution count, liveness, traffic boundary, residency assumptions, and owners into the plan identity. Bind peak-compute evidence to the operation class, operand dtype, accumulation dtype, and FLOP convention it measured; bind sustained-bandwidth evidence to the exact traffic boundary it measured. Arithmetic intensity is FLOPs per moved byte; a simple roofline lower bound is max(FLOPs/peak compute, bytes/sustained bandwidth). It is a bound under assumptions, not a latency forecast. Reject invalid shapes, duplicate tensor identities, implicit broadcasts, incompatible accumulation, mismatched hardware evidence, ambiguous liveness, unowned traffic, and count overflow before reporting a number.
Why this matters
FLOP counts alone routinely choose the wrong design. A lower-FLOP expression may materialize a huge intermediate or move it repeatedly; two algebraically identical contraction orders can differ by orders of magnitude in work and memory; and quadratic attention storage can dominate before compute is saturated. Without explicit shapes and traffic boundaries, performance estimates become unit-free numbers that cannot survive contact with a kernel or deployment.
You will be able to
- Read a batched tensor expression as named batch, free, and contraction dimensions.
- Derive output shape, conventional FLOPs, tensor bytes, traffic, and retained activation storage.
- Use arithmetic intensity and a bound hardware envelope without presenting a model as a benchmark.
- Expose broadcasting, contraction-order, materialization, precision, and execution-count/liveness traps.
- Create an identity-bound estimate with explicit residency, retained-copy, and hardware-evidence assumptions plus accountable owners.
Prerequisite contract
- Matrices as Small Programs →
- Numerical Stability Is Part of the Algorithm →
- — Products, units, and asymptotic notation
Your Vector Loop for this lab
- 01
Model
Name every tensor dimension, dtype, execution count, liveness policy, retained copy, materialization, and memory boundary.
- 02
Derive
Derive output elements, conventional FLOPs, moved bytes, retained bytes, and intensity from shapes.
- 03
Build
Implement a checked batched-matmul plan bound to traffic and hardware evidence.
- 04
Stress
Mismatch contraction dimensions, broadcast batches, narrow accumulation, duplicate storage, and overflow counts.
- 05
Operate
Compare estimates with kernel measurements and revise only versioned assumptions.
- 06
Defend
Explain when compute, traffic, activation capacity, or contraction order controls the design.
Model shapes as the program's resource signature
[b₁,…,bᵣ,m,k] × [b₁,…,bᵣ,k,n] → [b₁,…,bᵣ,m,n]
The batch dimensions are replicated independent products, m and n survive in the output, and k is contracted. The artifact requires exact batch equality instead of inheriting framework broadcasting implicitly.
| Dimension role | Effect when doubled | Questions to ask |
|---|---|---|
| batch b | doubles work, traffic, and output elements | is it physical batch, heads, experts, or beams? |
| row/free m | doubles work and output | does it represent tokens, queries, or channels? |
| contracted k | doubles multiply-accumulate work and input size | is projection width padded or sharded? |
| column/free n | doubles work and output | is it vocabulary, keys, features, or classes? |
| execution count | multiplies total work and traffic | are outputs not retained, serially reused, or all live together? |
Derive work, storage, and traffic as different quantities
FLOPs ≈ 2 · executions · prod(batch) · m · k · n
This is the conventional dense-matmul estimate when one fused multiply-add is counted as two operations. State the convention; hardware instruction count is a different quantity.
tensor bytes = prod(shape) · bytes(dtype)
Apply this separately to inputs, outputs, gradients, optimizer state, masks, workspaces, and any materialized intermediate. Do not multiply elements by an unnamed float size.
traffic bytes = Σ tensor bytes · full reads/writes · executions
A full-read model is a declared approximation across a named boundary such as HBM to compute. Cache residency, tiling, fusion, recomputation, and sharding change this count and therefore require a new assumption identity.
retained output bytes = output bytes · explicitly live output copies
Execution count alone does not determine peak liveness. A not-retained plan requires zero copies, serial reuse requires one, and all-executions-live requires one copy per execution; changing the count without a compatible liveness plan is rejected.
| Quantity | Capacity question | Common counting error |
|---|---|---|
| parameter bytes | does the model state fit? | forgetting optimizer and master weights |
| activation bytes | what must remain live at the peak? | summing a whole run instead of liveness |
| workspace bytes | what does the selected kernel allocate? | assuming zero because it is not a model tensor |
| traffic bytes | what crosses the named memory boundary? | counting each tensor once regardless of reuse |
| FLOPs | how much arithmetic is requested? | equating FLOPs with wall time |
Use intensity to ask which resource can dominate
arithmetic intensity = FLOPs / traffic bytes
Intensity belongs to one operation and one traffic boundary. Changing fusion or residency can change intensity without changing the mathematical result or FLOP count.
time lower bound = max(FLOPs / peak FLOP/s, traffic bytes / sustained byte/s)
This roofline-style comparison identifies the tighter modeled resource bound. Launch overhead, dependencies, occupancy, synchronization, contention, and imperfect kernels make real time larger.
For attention scores with sequence length s, materializing QKᵀ creates a batch-by-s-by-s tensor. Its storage grows quadratically in s even though the head width stays fixed. IO-aware exact attention demonstrates why reducing reads and writes through tiling can improve performance without changing the exact attention result. The lesson is broader than attention: materialization and memory hierarchy are algorithm design dimensions.
Build a checked estimator with owned assumptions
1def estimate_matmul(plan: MatmulCostPlan) -> TensorCostEstimate:2 if type(plan) is not MatmulCostPlan:3 raise TypeError("plan must be a concrete MatmulCostPlan")4 flops, traffic, retained = _derived_counts(plan)5 if traffic <= 0:6 raise ValueError("traffic model must move at least one byte")7 intensity = Decimal(flops) / Decimal(traffic)8 compute_seconds = Decimal(flops) / Decimal(9 plan.hardware.peak_flops_per_second10 )11 traffic_seconds = Decimal(traffic) / Decimal(12 plan.hardware.sustained_bytes_per_second13 )Expected output
example=illustrative_only
plan_version=tensor-cost-v1
operation=illustrative-attention-scores
lhs_shape=(8,512,64)
rhs_shape=(8,64,512)
output_shape=(8,512,512)
execution_count=1
execution_liveness=all-executions-live
retained_output_copies=1
traffic_boundary=hbm-to-compute
flops=268435456
traffic_bytes=9437184
retained_activation_bytes=8388608
arithmetic_intensity_flops_per_byte=28.444
lower_bound_us=9.437
bottleneck=TRAFFICVerify: Run python3 -m unittest discover courses/ai-engineering/reference-impl/tensor_complexity.
The excerpt is exact from the downloadable artifact. The illustrative plan multiplies query [8,512,64] by transposed key [8,64,512], uses float16 operands with float32 accumulation and output, assumes one full read of each input and one full output write across an HBM-to-compute boundary, and explicitly declares one execution under all-executions-live with one retained output copy. The invented hardware envelope is usable only for dense batched matmul with float16 operands, float32 accumulation, the two-FLOP FMA convention, and HBM-to-compute bandwidth. Concrete frozen plan, tensor, and hardware records reject mutable duck-typed aliases before accounting. The plan ID binds unique tensor identities, shapes, compatibility policy, read/write counts, execution liveness and copies, traffic boundary, hardware evidence, measured date, and owners.
Expose asymptotic and accounting traps
- 01Mismatch kReject incompatible contraction dimensions before producing output or cost counts.
- 02Broadcast a batchReject [8,…] by [1,…] under the strict policy; make expansion explicit if the real kernel uses it.
- 03Narrow accumulationUse an explicit compatibility matrix: float16 cannot accumulate bfloat16 merely because both occupy two bytes; range and precision both matter.
- 04Reuse bandwidth across boundariesReject an HBM bandwidth envelope when the plan moves bytes over PCIe; sustained bandwidth belongs to the exact measured traffic boundary.
- 05Reuse one tensor identity for two shapesReject duplicate input, output, or retained-activation names unless a future explicit alias/view contract defines ownership and accounting.
- 06Increase executions without changing livenessReject the plan under all-executions-live until retained copies also increase, or explicitly declare serial reuse so unchanged peak storage is auditable.
- 07Change contraction orderFor three or more operands, derive FLOPs and largest intermediate for each order. A locally compact notation does not imply a cheap path.
- 08Double sequence lengthCheck whether attention-score output and retained activation bytes grow by four, then test whether the implementation materializes them.
- 09Overflow a countReject shapes, executions, or read multipliers beyond the supported integer range instead of wrapping into a small estimate.
- 010Claim a cache hit without evidenceChange the traffic assumption and its owner; do not silently set reads to zero to make intensity look better.
Close the loop with measured kernels
| Model input | Measurement partner | Revisit when |
|---|---|---|
| expanded tensor shapes | runtime shape trace | batching, padding, routing, or context policy changes |
| operand and accumulation dtype | kernel and compiler metadata plus bound peak evidence | precision mode, operation class, FLOP convention, or hardware changes |
| full reads and writes | memory counters or profiler | fusion, tiling, cache, or sharding changes |
| execution liveness and retained copies | allocator peak and liveness trace | layer count, replay, checkpointing, or backward graph changes |
| hardware envelope | versioned benchmark harness | device, driver, clocks, or contention changes |
| lower-bound error | modeled versus observed time | error exceeds the decision tolerance |
Use the estimator early to eliminate impossible shapes and reveal the dominant assumption. Then profile the selected implementation with production-like sizes, warmup, synchronization, concurrency, and allocation behavior. Preserve both estimates and observations under release identity. When they diverge, update the traffic or kernel model with evidence; do not tune a hidden multiplier until the chart matches.
Operate at three altitudes
Production lens
- — Capture expanded runtime shapes, operand and accumulation dtype, kernel identity, execution count, liveness policy, retained copies, and materialization decisions by request class.
- — Measure traffic and allocator peaks at the same named boundary and workload used by the estimate.
- — Bind every hardware envelope to operation class, operand dtype, accumulation dtype, FLOP convention, exact traffic boundary, evidence source, date, and owner.
- — Alert on shape-policy drift, unexpected broadcasting, workspace growth, out-of-memory retries, and modeled-versus-observed error.
Staff lens
- — Assign shape ownership to model architecture, kernel choice to runtime, traffic assumptions to performance engineering, and capacity decisions to the serving owner.
- — Review performance changes across compute, traffic, activation, workspace, communication, latency, and engineering-complexity ledgers.
- — Standardize FLOP conventions and traffic boundaries so estimates from different teams are comparable.
- — Require execution-count changes to state whether outputs are discarded, reused serially, or simultaneously live and to price retained copies accordingly.
- — Use small shape models to challenge architecture before paying for large benchmarks, then require measurements before capacity commitments.
Interview defense
How would you estimate whether a tensor operation is compute-bound, memory-bound, or capacity-bound?
I would first name every batch, free, and contracted dimension and derive the expanded output shape. From that I would compute conventional FLOPs under a stated FMA convention, tensor bytes by dtype, peak live activations from an explicit execution-liveness and retained-copy policy, and bytes moved across a named boundary using read/write and residency assumptions. FLOPs divided by moved bytes gives arithmetic intensity; comparing against a hardware envelope bound to the same operation class, operand dtype, accumulation dtype, FLOP convention, and traffic boundary gives lower bounds, not latency predictions. I would reject duplicate tensor identities, implicit broadcasts, incompatible accumulation, mismatched hardware evidence, ambiguous liveness, invalid shapes, and count overflow, then compare with profiled kernel time, traffic, and allocation at production-like shapes.
Expect the interviewer to press on
- — Why can fewer FLOPs make an operation slower?
- — How does contraction order change the answer?
- — What changes when attention scores are tiled instead of materialized?
Misconceptions to remove
“FLOPs predict latency.”
FLOPs describe requested arithmetic under a convention. Traffic, launch cost, dependency, occupancy, synchronization, contention, and implementation quality also bound latency.
“Tensor shape tells us memory use.”
Shape and dtype tell one tensor's size. Peak memory also depends on explicitly live copies across executions, materialized intermediates, gradients, optimizer state, workspaces, fragmentation, and concurrency.
“Algebraically equivalent contraction orders have the same cost.”
They can create different intermediates and perform radically different numbers of operations. Choose and version the execution path, not only the equation.
Check your model
1. What happens to a materialized square attention-score tensor when sequence length doubles?
Its element count and storage multiply by four because both free sequence dimensions double, assuming batch and dtype stay fixed.
2. Why must arithmetic intensity name a memory boundary?
The byte count differs across register, cache, device-memory, and host/network boundaries. Without the boundary, FLOPs per byte is not a reproducible quantity.
3. Why take the maximum of compute and traffic time bounds?
The operation must satisfy both resource demands, so it cannot complete faster than either independent lower bound. Real execution can be slower for additional reasons.
Prove the mechanism
Add a second identity-bound plan for one tensor operation in your system. Derive exact expanded shapes, conventional FLOPs, input/output bytes, full-read/write traffic, retained activations from execution liveness and copies, intensity, and compute/traffic lower bounds. Bind hardware evidence to operation, precision, FLOP convention, and traffic boundary. Add adversarial tests for contraction mismatch, duplicate tensor identities, batch broadcasting, bfloat16-to-float16 accumulation, peak or bandwidth-boundary mismatch, execution-count/liveness drift, mutable record aliases, missing traffic ownership, and integer overflow, then compare with a profiler trace.
Add a production constraint
Implement a checked three-operand contraction planner. Enumerate valid pairwise orders for a bounded operand count, report FLOPs and largest live intermediate for each, reject ambiguous index or dtype contracts, and compare its selected order with an official library path plus measured execution.
Artifact: Tensor cost estimator
courses/ai-engineering/reference-impl/tensor_complexity/tensor_cost_estimator.py
Download reference implementationPrimary references and next links
References
- 1. Roofline: An Insightful Visual Performance Model for Floating-Point Programs and Multicore Architectures
Williams, Waterman, and Patterson. Primary technical report grounding arithmetic intensity and the compute-versus-memory-bandwidth bound used in this lesson.
- 2. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
Dao et al.. Primary NeurIPS paper showing why reads and writes across the memory hierarchy are a first-class algorithmic cost for exact attention.
- 3. torch.matmul
PyTorch documentation. Official framework documentation for batched matrix-multiplication dimensions and broadcasting semantics; the artifact intentionally enforces a stricter batch policy.
- 4. numpy.einsum_path
NumPy documentation. Official library documentation showing that contraction order changes FLOP count and intermediate size, and that exhaustive path search has its own complexity.
Continue through the graph
- Matrices as Small Programs →
Connect linear-map semantics to tensor execution cost.
- Numerical Stability Is Part of the Algorithm →
Relate dtype choices to numerical range and accumulation behavior.
- A First AI Workload Cost Model →
Roll kernel-level quantities into a service workload and unit-cost model.
Glossary: tensor shape · contraction dimension · FLOP · memory traffic · activation liveness · arithmetic intensity · roofline · broadcasting · materialization