From d34ba4ab689137ee1403ca8ddfdfc4430b2376ee Mon Sep 17 00:00:00 2001 From: Zhiwen Mo Date: Sat, 3 Oct 2026 22:41:34 +0100 Subject: [PATCH 1/2] fix(analytic): make analytic_models importable and installable analytic_models/__init__.py still imported the memory and utilisation packages that c4d09075 removed, so `import analytic_models` failed. Import only performance, and list analytic_models* in the setuptools packages (the same line #103 adds) so `uv sync` installs the package. --- analytic_models/__init__.py | 4 ++-- pyproject.toml | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/analytic_models/__init__.py b/analytic_models/__init__.py index e17a5ab2e..a065bb141 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/pyproject.toml b/pyproject.toml index 4b06e98b5..89e4d76c2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -100,6 +100,7 @@ include = [ "verification*", "compiler*", "transactional_emulator*", + "analytic_models*", ] exclude = [ "*tests*", From d7605bfb891b3f66a28cc6b33f3602a53e32edb1 Mon Sep 17 00:00:00 2001 From: Zhiwen Mo Date: Tue, 13 Oct 2026 00:07:49 +0100 Subject: [PATCH 2/2] fix(analytic): treat --device-num as data-parallel replicas llama_model.py and gpt_oss_model.py split the batch over device_num devices with floor division, then divided TTFT (and LLaDA's total time) by device_num again. Each device already prefills only its own sequences, so TTFT came out device_num times too small. A batch smaller than device_num left every device with zero sequences: llama-3.1-70b with batch 16 on 32 devices reported a TTFT of 0.00015 ms and 5,000,000 tokens/s. An uneven batch dropped sequences from the busiest device's time while TPS still counted them. Treat the devices as data-parallel replicas. The busiest one serves ceil(batch / device_num) sequences and sets TTFT and the decode time, TPS counts the whole batch, and device_num < 1 is rejected. A single device is unchanged. Communication between devices is still not modelled. --- analytic_models/performance/gpt_oss_model.py | 12 +- analytic_models/performance/llama_model.py | 14 ++- .../performance/tests/test_device_num.py | 118 ++++++++++++++++++ 3 files changed, 139 insertions(+), 5 deletions(-) create mode 100644 analytic_models/performance/tests/test_device_num.py diff --git a/analytic_models/performance/gpt_oss_model.py b/analytic_models/performance/gpt_oss_model.py index b99fe607c..19e98ee05 100644 --- a/analytic_models/performance/gpt_oss_model.py +++ b/analytic_models/performance/gpt_oss_model.py @@ -15,6 +15,7 @@ import argparse import json +import math from pathlib import Path from perf_model import PerfModel, load_hardware_config_from_toml @@ -80,7 +81,12 @@ def __init__( self.frequency = frequency_hz self.hardware_config = hardware_config self.batch_size = batch_size - self.device_batch_size = batch_size // device_num + if device_num < 1: + raise ValueError(f"device_num must be at least 1, got {device_num}") + # The devices are data-parallel replicas: the busiest one serves + # ceil(batch_size / device_num) sequences and sets the time. Communication + # between devices is not modelled. + self.device_batch_size = math.ceil(batch_size / device_num) self.device_num = device_num # Initialize PLENA hardware latency model @@ -338,7 +344,9 @@ def compute_performance(self, verbose: bool = True) -> tuple: """ prefill_time = self.compute_prefill_time(verbose) first_token_decode = self.compute_decode_time(1, verbose=False) - ttft = (prefill_time + first_token_decode) / self.device_num + # Every replica prefills its own sequences at the same time, so the time is + # already per device and is not divided by device_num again. + ttft = prefill_time + first_token_decode decode_time = self.compute_decode_time(self.output_seq_len, verbose) tps = (self.batch_size * self.output_seq_len) / decode_time diff --git a/analytic_models/performance/llama_model.py b/analytic_models/performance/llama_model.py index 5089d3f00..03922a8ee 100644 --- a/analytic_models/performance/llama_model.py +++ b/analytic_models/performance/llama_model.py @@ -11,6 +11,7 @@ import argparse import json +import math from pathlib import Path from perf_model import PerfModel, load_hardware_config_from_toml @@ -51,7 +52,12 @@ def __init__( self.frequency = frequency_hz self.hardware_config = hardware_config self.batch_size = batch_size - self.device_batch_size = batch_size // device_num + if device_num < 1: + raise ValueError(f"device_num must be at least 1, got {device_num}") + # The devices are data-parallel replicas: the busiest one serves + # ceil(batch_size / device_num) sequences and sets the time. Communication + # between devices is not modelled. + self.device_batch_size = math.ceil(batch_size / device_num) self.device_num = device_num # Initialize PLENA hardware latency model @@ -191,7 +197,9 @@ def compute_performance(self, verbose: bool = True) -> tuple: """ prefill_time = self.compute_prefill_time(verbose) first_token_decode = self.compute_decode_time(1, verbose=False) - ttft = (prefill_time + first_token_decode) / self.device_num + # Every replica prefills its own sequences at the same time, so the time is + # already per device and is not divided by device_num again. + ttft = prefill_time + first_token_decode decode_time = self.compute_decode_time(self.output_seq_len, verbose) tps = (self.batch_size * self.output_seq_len) / decode_time @@ -266,7 +274,7 @@ def compute_llada_inference(self, seq_len: int, diffusion_steps: int, verbose: b tuple: (total_time_seconds, output_tokens_per_second) """ step_time = self.compute_llada_step_time(seq_len, verbose=verbose) - total_time = (diffusion_steps * step_time) / self.device_num + total_time = diffusion_steps * step_time # Throughput: output tokens generated per second (batch * seq_len across all steps) output_tokens_per_second = (self.batch_size * seq_len) / total_time diff --git a/analytic_models/performance/tests/test_device_num.py b/analytic_models/performance/tests/test_device_num.py new file mode 100644 index 000000000..f14b7aa94 --- /dev/null +++ b/analytic_models/performance/tests/test_device_num.py @@ -0,0 +1,118 @@ +"""``--device-num``: the devices are data-parallel replicas of one PLENA chip. + +The busiest replica serves ``ceil(batch / device_num)`` sequences and sets TTFT +and the decode time; TPS counts the tokens of the whole batch. +""" + +from __future__ import annotations + +import contextlib +import importlib +import io +import json +import sys +from pathlib import Path + +import pytest + +PERF_DIR = Path(__file__).resolve().parents[1] +REPO_ROOT = PERF_DIR.parents[1] +SETTINGS = REPO_ROOT / "plena_settings.toml" +ISA_LIB = PERF_DIR / "customISA_lib.json" + +# Public HuggingFace dimensions of Llama-3.1-8B and gpt-oss-20b, as in PLENA_Compiler's Model_Lib. +LLAMA_3_1_8B = { + "hidden_size": 4096, + "num_attention_heads": 32, + "num_key_value_heads": 8, + "num_hidden_layers": 32, + "intermediate_size": 14336, + "vocab_size": 128256, +} +GPT_OSS_20B = { + "hidden_size": 2880, + "num_attention_heads": 64, + "num_key_value_heads": 8, + "head_dim": 64, + "num_hidden_layers": 24, + "intermediate_size": 2880, + "vocab_size": 201088, + "num_local_experts": 32, + "experts_per_token": 4, + "sliding_window": 128, + "layer_types": ["sliding_attention", "full_attention"] * 12, +} + + +@pytest.fixture(scope="module") +def scripts(): + sys.path.insert(0, str(PERF_DIR)) + try: + perf_model = importlib.import_module("perf_model") + llama_model = importlib.import_module("llama_model") + gpt_oss_model = importlib.import_module("gpt_oss_model") + finally: + sys.path.remove(str(PERF_DIR)) + hardware = perf_model.load_hardware_config_from_toml(str(SETTINGS)) + return { + "llama": (llama_model.LLaMAModel, LLAMA_3_1_8B), + "gpt_oss": (gpt_oss_model.GPTOssModel, GPT_OSS_20B), + }, hardware + + +def performance(scripts, tmp_path, family: str, batch: int, devices: int) -> tuple[float, float]: + models, hardware = scripts + cls, config = models[family] + config_path = tmp_path / f"{family}.json" + config_path.write_text(json.dumps(config)) + model = cls(str(config_path), hardware, str(ISA_LIB), batch, 256, 3, device_num=devices) + with contextlib.redirect_stdout(io.StringIO()): + return model.compute_performance(verbose=False) + + +@pytest.mark.parametrize("family", ["llama", "gpt_oss"]) +def test_replicas_multiply_throughput_and_keep_ttft(scripts, tmp_path, family) -> None: + ttft_one, tps_one = performance(scripts, tmp_path, family, batch=2, devices=1) + ttft, tps = performance(scripts, tmp_path, family, batch=8, devices=4) + + assert ttft == pytest.approx(ttft_one, rel=1e-12) + assert tps == pytest.approx(4 * tps_one, rel=1e-12) + + +@pytest.mark.parametrize("family", ["llama", "gpt_oss"]) +def test_more_devices_than_sequences_leaves_devices_idle(scripts, tmp_path, family) -> None: + ttft_one, tps_one = performance(scripts, tmp_path, family, batch=1, devices=1) + ttft, tps = performance(scripts, tmp_path, family, batch=3, devices=8) + + assert ttft == pytest.approx(ttft_one, rel=1e-12) + assert tps == pytest.approx(3 * tps_one, rel=1e-12) + + +@pytest.mark.parametrize("family", ["llama", "gpt_oss"]) +def test_uneven_batch_is_paced_by_the_busiest_device(scripts, tmp_path, family) -> None: + ttft_three, tps_three = performance(scripts, tmp_path, family, batch=3, devices=1) + ttft, tps = performance(scripts, tmp_path, family, batch=10, devices=4) + + assert ttft == pytest.approx(ttft_three, rel=1e-12) + assert tps == pytest.approx(10 / 3 * tps_three, rel=1e-12) + + +def test_llada_replicas_keep_the_inference_time(scripts, tmp_path) -> None: + models, hardware = scripts + cls, config = models["llama"] + config_path = tmp_path / "llama.json" + config_path.write_text(json.dumps(config)) + one = cls(str(config_path), hardware, str(ISA_LIB), 2, 256, 3, device_num=1) + four = cls(str(config_path), hardware, str(ISA_LIB), 8, 256, 3, device_num=4) + + time_one, rate_one = one.compute_llada_inference(512, 4, verbose=False) + time_four, rate_four = four.compute_llada_inference(512, 4, verbose=False) + + assert time_four == pytest.approx(time_one, rel=1e-12) + assert rate_four == pytest.approx(4 * rate_one, rel=1e-12) + + +@pytest.mark.parametrize("family", ["llama", "gpt_oss"]) +def test_zero_devices_is_rejected(scripts, tmp_path, family) -> None: + with pytest.raises(ValueError, match="device_num"): + performance(scripts, tmp_path, family, batch=4, devices=0)