From ce5eef46e0a47966eca688896f8bc223e1be2c47 Mon Sep 17 00:00:00 2001 From: "Qichao(Arlo) Wang" Date: Mon, 25 May 2026 12:35:46 +0100 Subject: [PATCH] refactor(analytic): split latency and cycle breakdown helpers --- analytic_models/__init__.py | 4 +- analytic_models/performance/latency.py | 98 ++++++++++++ analytic_models/performance/perf_model.py | 144 ++++++++++++------ analytic_models/performance/tests/conftest.py | 7 + .../tests/test_perf_model_refactor.py | 144 ++++++++++++++++++ 5 files changed, 349 insertions(+), 48 deletions(-) create mode 100644 analytic_models/performance/latency.py create mode 100644 analytic_models/performance/tests/conftest.py create mode 100644 analytic_models/performance/tests/test_perf_model_refactor.py diff --git a/analytic_models/__init__.py b/analytic_models/__init__.py index e17a5ab2..a065bb14 100644 --- a/analytic_models/__init__.py +++ b/analytic_models/__init__.py @@ -1,3 +1,3 @@ -from . import memory, performance, utilisation +from . import performance -__all__ = ["memory", "performance", "utilisation"] +__all__ = ["performance"] diff --git a/analytic_models/performance/latency.py b/analytic_models/performance/latency.py new file mode 100644 index 00000000..7bbd33c6 --- /dev/null +++ b/analytic_models/performance/latency.py @@ -0,0 +1,98 @@ +"""Instruction latency expression loading and evaluation.""" + +from __future__ import annotations + +import ast +import json +import math +import operator +from collections.abc import Mapping +from typing import Any + + +_ALLOWED_BINOPS = { + ast.Add: operator.add, + ast.Sub: operator.sub, + ast.Mult: operator.mul, + ast.Div: operator.truediv, + ast.FloorDiv: operator.floordiv, +} +_ALLOWED_UNARYOPS = { + ast.UAdd: operator.pos, + ast.USub: operator.neg, +} +_MAX_POWER_EXPONENT = 8 + + +def build_latency_context(hardware_config: Any) -> dict[str, int | float]: + """Build the variable context used by custom ISA latency expressions.""" + context = dict(hardware_config.model_dump()) + context["SA_ACC_CYCLES"] = int(math.log2(hardware_config.MLEN / hardware_config.BLEN) + 1) + return context + + +def evaluate_latency_expression(expression: str, context: Mapping[str, int | float]) -> int: + """Evaluate a restricted arithmetic latency expression.""" + expression = expression.strip() + try: + tree = ast.parse(expression, mode="eval") + value = _evaluate_node(tree, context) + except (SyntaxError, TypeError, ZeroDivisionError) as exc: + raise ValueError(f"Unsupported latency expression: {expression}") from exc + + if isinstance(value, float): + if not value.is_integer(): + raise ValueError(f"Latency expression did not evaluate to an integer: {expression}") + return int(value) + return int(value) + + +def build_pipelined_latency_map(hardware_config: Any, custom_isa_path: str) -> dict[str, int]: + """Load customISA_lib.json and evaluate pipelined latency expressions.""" + with open(custom_isa_path) as f: + custom_isa_lib = json.load(f) + + context = build_latency_context(hardware_config) + latencies = {} + for instr_name, instr_data in custom_isa_lib.items(): + if "pipelined" not in instr_data: + raise ValueError(f"Instruction '{instr_name}' missing 'pipelined' field.") + latencies[instr_name] = evaluate_latency_expression(instr_data["pipelined"], context) + return latencies + + +def _evaluate_node(node: ast.AST, context: Mapping[str, int | float]) -> int | float: + if isinstance(node, ast.Expression): + return _evaluate_node(node.body, context) + + if isinstance(node, ast.Constant) and isinstance(node.value, int | float): + return node.value + + if isinstance(node, ast.Name): + if node.id not in context: + raise ValueError(f"Unsupported latency expression variable: {node.id}") + return context[node.id] + + if isinstance(node, ast.UnaryOp): + op = _ALLOWED_UNARYOPS.get(type(node.op)) + if op is None: + raise ValueError(f"Unsupported latency expression unary operator: {type(node.op).__name__}") + return op(_evaluate_node(node.operand, context)) + + if isinstance(node, ast.BinOp): + if isinstance(node.op, ast.Pow): + left = _evaluate_node(node.left, context) + right = _evaluate_node(node.right, context) + if not isinstance(right, int | float) or int(right) != right: + raise ValueError("Power exponent must be an integer") + exponent = int(right) + if exponent < 0 or exponent > _MAX_POWER_EXPONENT: + raise ValueError(f"Power exponent must be between 0 and {_MAX_POWER_EXPONENT}") + return operator.pow(left, exponent) + + op = _ALLOWED_BINOPS.get(type(node.op)) + if op is None: + raise ValueError(f"Unsupported latency expression operator: {type(node.op).__name__}") + return op(_evaluate_node(node.left, context), _evaluate_node(node.right, context)) + + raise ValueError(f"Unsupported latency expression node: {type(node).__name__}") diff --git a/analytic_models/performance/perf_model.py b/analytic_models/performance/perf_model.py index 70a7b26a..0d4b0c61 100644 --- a/analytic_models/performance/perf_model.py +++ b/analytic_models/performance/perf_model.py @@ -5,12 +5,19 @@ This module is used by llama_model.py for LLM-level performance estimation. """ -import json import math +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType import toml from pydantic import BaseModel, Field, model_validator +try: + from .latency import build_pipelined_latency_map +except ImportError: # pragma: no cover - supports direct script execution from this directory + from latency import build_pipelined_latency_map + # ============================================================================= # Hardware Configuration Schema # ============================================================================= @@ -70,6 +77,25 @@ def items(self): return self.latencies.items() +@dataclass(frozen=True) +class CycleBreakdown: + """Named cycle components for a layer-level operation.""" + + components: Mapping[str, int] + total_cycles: int | None = None + + def __post_init__(self): + components = dict(self.components) + object.__setattr__(self, "components", MappingProxyType(components)) + if self.total_cycles is None: + object.__setattr__(self, "total_cycles", sum(components.values())) + + @property + def total(self) -> int: + assert self.total_cycles is not None + return self.total_cycles + + # ============================================================================= # Hardware Configuration Loading # ============================================================================= @@ -129,21 +155,7 @@ def build_pipelined_latency(hardware_config: HardwareConfig, custom_isa_path: st Returns: InstructionLatency: Validated instruction latencies """ - with open(custom_isa_path) as f: - custom_isa_lib = json.load(f) - - # Build config dict for eval (convert pydantic model to dict) - configs = hardware_config.model_dump() - configs["SA_ACC_CYCLES"] = int(math.log2(hardware_config.MLEN / hardware_config.BLEN) + 1) - - latencies = {} - for instr_name, instr_data in custom_isa_lib.items(): - if "pipelined" in instr_data: - latencies[instr_name] = eval(instr_data["pipelined"], {"__builtins__": {}}, configs) - else: - raise ValueError(f"Instruction '{instr_name}' missing 'pipelined' field.") - - return InstructionLatency(latencies=latencies) + return InstructionLatency(latencies=build_pipelined_latency_map(hardware_config, custom_isa_path)) # ============================================================================= @@ -489,14 +501,34 @@ def mlp_moe( intermediate_size: int, mode: str = "prefill", ) -> int: + """MoE cycle count.""" + return self.mlp_moe_breakdown( + hidden_size=hidden_size, + seq_len=seq_len, + batch_size=batch_size, + num_experts=num_experts, + expert_per_token=expert_per_token, + intermediate_size=intermediate_size, + mode=mode, + ).total + + def mlp_moe_breakdown( + self, + hidden_size: int, + seq_len: int, + batch_size: int, + num_experts: int, + expert_per_token: int, + intermediate_size: int, + mode: str = "prefill", + ) -> CycleBreakdown: """ - MoE cycle count. + MoE cycle breakdown. In MoE, tokens are routed to experts and batched per expert. Each expert processes its batch of tokens using M_MM (not per-token M_MV). - Average tokens per expert = (total_tokens * expert_per_token) / num_experts + Average tokens per expert = (total_tokens * expert_per_token) / num_experts. """ - overall_cycles = 0 if mode == "prefill": # Total tokens being processed @@ -507,21 +539,21 @@ def mlp_moe( tokens_per_expert = math.ceil((total_tokens * expert_per_token) / num_experts) # Normalize (b, s, h) -> (b, s, h) - overall_cycles += (math.ceil(hidden_size / self.vlen) * self.instr["V_BASIC"] * 4) * total_tokens + norm = (math.ceil(hidden_size / self.vlen) * self.instr["V_BASIC"] * 4) * total_tokens # Router / Gate: (b*s, h) @ (h, num_experts) -> (b*s, num_experts) # Using M_MM for batch matrix multiply - overall_cycles += ( + router = ( (4 + math.ceil(hidden_size / self.mlen) * self.instr["M_MM"] + self.instr["H_PREFETCH_M"]) * math.ceil(total_tokens / self.blen) * math.ceil(num_experts / self.blen) ) # TOP K: (b*s, num_experts) -> (b*s, expert_per_token) - overall_cycles += (4 + math.ceil(num_experts / self.vlen) * self.instr["V_TOPK"]) * total_tokens + topk = (4 + math.ceil(num_experts / self.vlen) * self.instr["V_TOPK"]) * total_tokens # Softmax over selected experts: (b*s, expert_per_token) -> (b*s, expert_per_token) - overall_cycles += ( + routing_softmax = ( total_tokens * math.ceil(expert_per_token / self.vlen) * (self.instr["V_EXP_V"] + self.instr["V_RED_MAX"] + self.instr["V_BASIC"]) @@ -531,7 +563,7 @@ def mlp_moe( # Tokens are grouped by expert and processed in batches using M_MM # Each expert: (tokens_per_expert, hidden) @ (hidden, 2*intermediate) -> (tokens_per_expert, 2*intermediate) # Run for all num_experts experts - overall_cycles += ( + expert_up_gate = ( num_experts * (4 + math.ceil(hidden_size / self.mlen) * self.instr["M_MM"] + self.instr["H_PREFETCH_M"]) * math.ceil(tokens_per_expert / self.blen) @@ -540,13 +572,13 @@ def mlp_moe( # SiLU activation + element-wise multiply (gate * up) # Total activations = total_tokens * expert_per_token (each token activates expert_per_token experts) - overall_cycles += ( + expert_activation = ( total_tokens * expert_per_token * math.ceil(intermediate_size / self.vlen) * 6 * self.instr["V_BASIC"] ) # Expert FFN Computation - MLP2 (Down projection) # Each expert: (tokens_per_expert, intermediate) @ (intermediate, hidden) -> (tokens_per_expert, hidden) - overall_cycles += ( + expert_down = ( num_experts * (4 + math.ceil(intermediate_size / self.mlen) * self.instr["M_MM"] + self.instr["H_PREFETCH_M"]) * math.ceil(tokens_per_expert / self.blen) @@ -555,7 +587,7 @@ def mlp_moe( # Weighted sum of experts # Per token: sum over expert_per_token weighted vectors of size hidden_size - overall_cycles += ( + combine = ( total_tokens * expert_per_token * math.ceil(hidden_size / self.vlen) @@ -566,21 +598,21 @@ def mlp_moe( total_tokens = batch_size # Normalize (b, h) -> (b, h) - overall_cycles += (math.ceil(hidden_size / self.vlen) * self.instr["V_BASIC"] * 4) * total_tokens + norm = (math.ceil(hidden_size / self.vlen) * self.instr["V_BASIC"] * 4) * total_tokens # Router / Gate: (b, h) @ (h, num_experts) -> (b, num_experts) # For small batch, use M_MV per token - overall_cycles += ( + router = ( total_tokens * (4 + math.ceil(hidden_size / self.mlen) * self.instr["M_MV"] + self.instr["H_PREFETCH_M"]) * math.ceil(num_experts / self.blen) ) # TOP K: (b, num_experts) -> (b, expert_per_token) - overall_cycles += (4 + math.ceil(num_experts / self.vlen) * self.instr["V_TOPK"]) * total_tokens + topk = (4 + math.ceil(num_experts / self.vlen) * self.instr["V_TOPK"]) * total_tokens # Softmax over selected experts: (b, expert_per_token) -> (b, expert_per_token) - overall_cycles += ( + routing_softmax = ( total_tokens * math.ceil(expert_per_token / self.vlen) * (self.instr["V_EXP_V"] + self.instr["V_RED_MAX"] + self.instr["V_BASIC"]) @@ -588,7 +620,7 @@ def mlp_moe( # Expert FFN Computation - MLP1 (Gate + Up projection) # In decode, few tokens so use M_MV per (token, expert) pair - overall_cycles += ( + expert_up_gate = ( total_tokens * expert_per_token * (4 + math.ceil(hidden_size / self.mlen) * self.instr["M_MV"] + self.instr["H_PREFETCH_M"]) @@ -596,12 +628,12 @@ def mlp_moe( ) # SiLU activation + element-wise multiply - overall_cycles += ( + expert_activation = ( total_tokens * expert_per_token * math.ceil(intermediate_size / self.vlen) * 6 * self.instr["V_BASIC"] ) # Expert FFN Computation - MLP2 (Down projection) - overall_cycles += ( + expert_down = ( total_tokens * expert_per_token * (4 + math.ceil(intermediate_size / self.mlen) * self.instr["M_MV"] + self.instr["H_PREFETCH_M"]) @@ -609,14 +641,25 @@ def mlp_moe( ) # Weighted sum of experts - overall_cycles += ( + combine = ( total_tokens * expert_per_token * math.ceil(hidden_size / self.vlen) * (self.instr["V_MUL_VV"] + self.instr["V_ADD_VV"]) ) - return overall_cycles + return CycleBreakdown( + { + "norm": norm, + "router": router, + "topk": topk, + "routing_softmax": routing_softmax, + "expert_up_gate": expert_up_gate, + "expert_activation": expert_activation, + "expert_down": expert_down, + "combine": combine, + } + ) def sliding_window_attention( self, @@ -788,11 +831,22 @@ def feed_forward( self, hidden_size: int, intermediate_size: int, seq_len: int, batch_size: int, mode: str = "prefill" ) -> int: """Feed-forward (MLP) layer cycle count.""" - overall_cycles = 0 + return self.feed_forward_breakdown( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + seq_len=seq_len, + batch_size=batch_size, + mode=mode, + ).total + + def feed_forward_breakdown( + self, hidden_size: int, intermediate_size: int, seq_len: int, batch_size: int, mode: str = "prefill" + ) -> CycleBreakdown: + """Feed-forward (MLP) layer cycle breakdown.""" if mode == "prefill": # Upsize Linear and Gate - overall_cycles += ( + up_gate = ( 2 * math.ceil((seq_len * batch_size) / self.blen) * math.ceil(hidden_size / self.mlen) @@ -800,11 +854,9 @@ def feed_forward( * self.instr["M_MM"] ) # SiLU - overall_cycles += ( - math.ceil(intermediate_size / self.vlen) * 6 * self.instr["V_BASIC"] * seq_len * batch_size - ) + activation = math.ceil(intermediate_size / self.vlen) * 6 * self.instr["V_BASIC"] * seq_len * batch_size # Downsize Linear - overall_cycles += ( + down = ( math.ceil((seq_len * batch_size) / self.blen) * math.ceil(intermediate_size / self.mlen) * math.ceil(hidden_size / self.blen) @@ -812,7 +864,7 @@ def feed_forward( ) else: # Upsize Linear and Gate - overall_cycles += ( + up_gate = ( 2 * math.ceil(intermediate_size / self.blen) * math.ceil(hidden_size / self.mlen) @@ -820,16 +872,16 @@ def feed_forward( * self.instr["M_MM"] ) # SiLU - overall_cycles += math.ceil(intermediate_size / self.vlen) * 6 * self.instr["V_BASIC"] * batch_size + activation = math.ceil(intermediate_size / self.vlen) * 6 * self.instr["V_BASIC"] * batch_size # Downsize Linear - overall_cycles += ( + down = ( math.ceil(batch_size / self.blen) * math.ceil(intermediate_size / self.mlen) * math.ceil(hidden_size / self.blen) * self.instr["M_MM"] ) - return overall_cycles + return CycleBreakdown({"up_gate": up_gate, "activation": activation, "down": down}) def embeddings(self, hidden_size: int, seq_len: int, batch_size: int, mode: str = "prefill") -> int: """Embedding layer cycle count.""" diff --git a/analytic_models/performance/tests/conftest.py b/analytic_models/performance/tests/conftest.py new file mode 100644 index 00000000..0a604bd8 --- /dev/null +++ b/analytic_models/performance/tests/conftest.py @@ -0,0 +1,7 @@ +import sys +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[3] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) diff --git a/analytic_models/performance/tests/test_perf_model_refactor.py b/analytic_models/performance/tests/test_perf_model_refactor.py new file mode 100644 index 00000000..f2caf3d6 --- /dev/null +++ b/analytic_models/performance/tests/test_perf_model_refactor.py @@ -0,0 +1,144 @@ +from pathlib import Path + +import pytest + +from analytic_models.performance.perf_model import CycleBreakdown, HardwareConfig, PerfModel + +_CUSTOM_ISA_PATH = Path(__file__).resolve().parents[1] / "customISA_lib.json" + + +def _hardware_config() -> HardwareConfig: + return HardwareConfig( + MLEN=64, + BLEN=4, + VLEN=64, + HLEN=16, + VECTOR_SRAM_SIZE=1024, + HBM_V_Prefetch_Amount=4, + VECTOR_BASIC_CYCLES=1, + VECTOR_ADD_CYCLES=1, + VECTOR_MUL_CYCLES=1, + VECTOR_EXP_CYCLES=3, + VECTOR_MAX_CYCLES=2, + VECTOR_SUM_CYCLES=2, + SCALAR_FP_BASIC_CYCLES=1, + SCALAR_FP_EXP_CYCLES=3, + SCALAR_FP_SQRT_CYCLES=4, + SCALAR_FP_RECI_CYCLES=3, + SCALAR_INT_BASIC_CYCLES=1, + ) + + +def _perf_model() -> PerfModel: + return PerfModel(_hardware_config(), str(_CUSTOM_ISA_PATH)) + + +def test_latency_expression_evaluator_supports_custom_isa_arithmetic(): + from analytic_models.performance.latency import build_latency_context, evaluate_latency_expression + + context = build_latency_context(_hardware_config()) + + assert context["SA_ACC_CYCLES"] == 5 + assert evaluate_latency_expression("(MLEN // BLEN)**2 * BLEN", context) == 1024 + assert evaluate_latency_expression("MLEN / BLEN", context) == 16 + assert evaluate_latency_expression("-BLEN + MLEN", context) == 60 + + +def test_latency_expression_evaluator_rejects_large_power_exponents(): + from analytic_models.performance.latency import build_latency_context, evaluate_latency_expression + + context = build_latency_context(_hardware_config()) + + with pytest.raises(ValueError, match="Power exponent"): + evaluate_latency_expression("2 ** 9", context) + + +def test_latency_expression_evaluator_rejects_function_calls(): + from analytic_models.performance.latency import build_latency_context, evaluate_latency_expression + + context = build_latency_context(_hardware_config()) + + with pytest.raises(ValueError, match="Unsupported latency expression"): + evaluate_latency_expression("__import__('os').system('echo unsafe')", context) + + +def test_cycle_breakdown_supports_explicit_total_and_immutable_components(): + breakdown = CycleBreakdown({"compute": 10, "overlapped_transfer": 7}, total_cycles=10) + + assert breakdown.total == 10 + with pytest.raises(TypeError): + breakdown.components["compute"] = 20 + + +@pytest.mark.parametrize( + ("mode", "expected_components", "expected_total"), + [ + ("prefill", {"up_gate": 8192, "activation": 768, "down": 4096}, 13056), + ("decode", {"up_gate": 1024, "activation": 48, "down": 512}, 1584), + ], +) +def test_feed_forward_breakdown_matches_fixed_cycle_baselines(mode, expected_components, expected_total): + perf = _perf_model() + + breakdown = perf.feed_forward_breakdown( + hidden_size=128, + intermediate_size=256, + seq_len=16, + batch_size=2, + mode=mode, + ) + + assert dict(breakdown.components) == expected_components + assert breakdown.total == expected_total + assert perf.feed_forward(128, 256, 16, 2, mode) == expected_total + + +@pytest.mark.parametrize( + ("mode", "expected_components", "expected_total"), + [ + ( + "prefill", + { + "norm": 256, + "router": 208, + "topk": 224, + "routing_softmax": 256, + "expert_up_gate": 26624, + "expert_activation": 1536, + "expert_down": 10752, + "combine": 256, + }, + 40112, + ), + ( + "decode", + { + "norm": 16, + "router": 60, + "topk": 14, + "routing_softmax": 16, + "expert_up_gate": 7680, + "expert_activation": 96, + "expert_down": 3200, + "combine": 16, + }, + 11098, + ), + ], +) +def test_mlp_moe_breakdown_matches_fixed_cycle_baselines(mode, expected_components, expected_total): + perf = _perf_model() + + breakdown = perf.mlp_moe_breakdown( + hidden_size=128, + seq_len=16, + batch_size=2, + num_experts=8, + expert_per_token=2, + intermediate_size=256, + mode=mode, + ) + + assert dict(breakdown.components) == expected_components + assert breakdown.total == expected_total + assert perf.mlp_moe(128, 16, 2, 8, 2, 256, mode) == expected_total