from __future__ import annotations import json import math def causal_visibility(length: int) -> list[list[int]]: return [ [int(key <= query) for key in range(length)] for query in range(length) ] def local_causal_visibility(length: int, previous_tokens: int) -> list[list[int]]: return [ [int(max(0, query - previous_tokens) <= key <= query) for key in range(length)] for query in range(length) ] def with_global_bos(matrix: list[list[int]]) -> list[list[int]]: result = [row[:] for row in matrix] for query in range(len(result)): result[query][0] = 1 return result def edge_count(matrix: list[list[int]]) -> int: return sum(sum(row) for row in matrix) def cache_record( name: str, *, batch: int, layers: int, query_heads: int, kv_heads: int, length: int, head_width: int, ) -> dict: scalars = 2 * batch * layers * kv_heads * length * head_width return { "name": name, "query_heads": query_heads, "kv_heads": kv_heads, "shape_per_K_or_V_per_layer": [batch, kv_heads, length, head_width], "scalars": scalars, "float32_bytes": scalars * 4, } def expert_ledger(d_model: int, d_ff: int, experts: int) -> dict: one_expert_weights = 2 * d_model * d_ff one_expert_biases = d_ff + d_model one_expert_parameters = one_expert_weights + one_expert_biases router_parameters = d_model * experts return { "bias_convention": "both expert linear maps have bias; router has no bias", "dense_mlp_parameters": one_expert_parameters, "top1_moe_total_parameters": experts * one_expert_parameters + router_parameters, "top1_active_expert_parameters_per_token": one_expert_parameters, "router_parameters_read_per_token": router_parameters, "main_expert_matrix_multiply_entries_per_token": one_expert_weights, "excluded_from_equal-work_claim": [ "router computation", "dispatch and gather", "communication", "capacity padding or dropped tokens", "load imbalance", ], } def illustrative_power_law() -> dict: compute_1, loss_1 = 1.0, 2.0 compute_2, loss_2 = 4.0, 1.5 alpha = math.log(loss_1 / loss_2) / math.log(compute_2 / compute_1) coefficient = loss_1 * compute_1**alpha extrapolated_compute = 16.0 extrapolated_loss = coefficient * extrapolated_compute ** (-alpha) return { "status": "illustrative numbers, not measurements from a paper or model", "assumed_form": "loss = A * compute ** (-alpha)", "points": [ {"compute": compute_1, "loss": loss_1}, {"compute": compute_2, "loss": loss_2}, ], "fitted_alpha": alpha, "fitted_A": coefficient, "unvalidated_extrapolation": { "compute": extrapolated_compute, "loss": extrapolated_loss, }, "warning": "two points determine this log-space line but do not validate its extrapolation", } def validate_cards(cards: list[dict]) -> None: required = { "name", "changed_component", "invariant_core", "analytical_consequence", "empirical_source_boundary", "later_subject", } for card in cards: missing = required - card.keys() assert not missing, (card.get("name"), missing) assert all(card[field] for field in required) def main() -> None: length = 8 previous_tokens = 2 dense = causal_visibility(length) local = local_causal_visibility(length, previous_tokens) local_global = with_global_bos(local) assert edge_count(dense) == length * (length + 1) // 2 == 36 expected_local = sum(min(index + 1, previous_tokens + 1) for index in range(length)) assert edge_count(local) == expected_local == 21 assert all( matrix[query][key] == 0 for matrix in (dense, local, local_global) for query in range(length) for key in range(query + 1, length) ) cache_parameters = { "batch": 1, "layers": 4, "query_heads": 8, "length": 16, "head_width": 8, } caches = [ cache_record("MHA", kv_heads=8, **cache_parameters), cache_record("GQA with 2 K/V heads", kv_heads=2, **cache_parameters), cache_record("MQA", kv_heads=1, **cache_parameters), ] assert caches[0]["scalars"] > caches[1]["scalars"] > caches[2]["scalars"] experts = expert_ledger(d_model=4, d_ff=8, experts=4) assert experts["dense_mlp_parameters"] == 76 assert experts["top1_moe_total_parameters"] == 320 assert experts["top1_active_expert_parameters_per_token"] == 76 cards = [ { "name": "causal local-window attention", "changed_component": "each query sees only recent keys unless a declared global edge is added", "invariant_core": "Q/K scores, normalized value mixing, head output, residual and MLP blocks", "analytical_consequence": "allowed attention edges fall from 36 to 21 for T=8 and two preceding tokens", "empirical_source_boundary": "Longformer and Big Bird use different sparse graphs and evaluate specified models and tasks", "later_subject": "AI Engineering for kernels and serving; task subjects for empirical quality", }, { "name": "multi-query attention", "changed_component": "eight query heads share one K/V head", "invariant_core": "scaled dot-product attention and a multi-head query/output interface", "analytical_consequence": "the declared cache falls from 8192 MHA scalars to 1024 MQA scalars", "empirical_source_boundary": "Shazeer 2019 reports specified model-quality and hardware decoding comparisons", "later_subject": "AI Engineering for production latency and memory behavior", }, { "name": "top-1 mixture of experts", "changed_component": "a router sends each token to one of four MLP experts", "invariant_core": "attention and residual block interfaces remain Transformer-compatible", "analytical_consequence": "320 stored expert-plus-router parameters but 76 active expert parameters per token under this ledger", "empirical_source_boundary": "Switch Transformer results retain routing, capacity, precision, hardware, data, and task conditions", "later_subject": "AI Engineering for routing, communication, sharding, and capacity", }, { "name": "image-patch Transformer", "changed_component": "an image is converted into projected patch tokens instead of text-token embeddings", "invariant_core": "position-aware records pass through Transformer blocks", "analytical_consequence": "sequence length depends on image size and patch size", "empirical_source_boundary": "ViT results retain pretraining scale, transfer data, patch configuration, and vision tasks", "later_subject": "Multimodal Models for encoders, fusion, objectives, data, and evaluation", }, { "name": "selective state-space model", "changed_component": "input-dependent recurrent state replaces explicit pairwise self-attention", "invariant_core": "the model still maps an ordered sequence to contextual records and task outputs", "analytical_consequence": "there is no T-by-T attention matrix or Transformer KV-cache contract", "empirical_source_boundary": "Mamba results retain its selective SSM, hardware-aware implementation, scales, modalities, and benchmarks", "later_subject": "Frontier Architectures for state-space derivation and implementation", }, ] validate_cards(cards) output = { "visibility": { "length": length, "local_previous_tokens": previous_tokens, "dense_causal": {"matrix": dense, "edges": edge_count(dense)}, "local_causal": {"matrix": local, "edges": edge_count(local)}, "local_plus_global_bos": { "matrix": local_global, "edges": edge_count(local_global), }, }, "kv_cache": caches, "expert_ledger": experts, "scaling_fit": illustrative_power_law(), "architecture_cards": cards, } print(json.dumps(output, indent=2)) if __name__ == "__main__": main()