From 439cca584c85119e6560a2236f29444ea0c81dd4 Mon Sep 17 00:00:00 2001 From: andrej Date: Sat, 3 Oct 2026 20:01:54 -0600 Subject: [PATCH] applications/gemma4_flm: report time to first token and decode rate The benchmark site plots only parametrizations with the bench mark, and only the metrics a test declares. The Gemma 4 test had neither, so CI recorded it as a pass with no measurements. The test now runs once per prompt, and each prompt carries the bench mark. It declares TTFT and TPS with the patterns of llama_3.2_1b and prints the IRON engine's prefill duration and decode rate from flm's /api/generate response. Each prompt's tokens must still equal the stock engine's. make engine and the stock engine's reference run once per module, not once per iteration: the stock engine is deterministic, so every iteration compares with one run. Co-Authored-By: Claude Opus 5.5 (1M context) --- iron/applications/gemma4_flm/README.md | 3 +- iron/applications/gemma4_flm/test.py | 70 +++++++++++++++++++------- 2 files changed, 53 insertions(+), 20 deletions(-) diff --git a/iron/applications/gemma4_flm/README.md b/iron/applications/gemma4_flm/README.md index 38ae466af..a79b238de 100644 --- a/iron/applications/gemma4_flm/README.md +++ b/iron/applications/gemma4_flm/README.md @@ -87,7 +87,8 @@ Run the steps below from this directory, `iron/applications/gemma4_flm`. ``` The test serves a word problem and a 1259-token prompt on both engines. It -checks that the token ids match. Run it from the repository root: +checks that the token ids match, and reports the IRON engine's time to first +token and decode rate for each prompt. Run it from the repository root: ```bash pytest --iterations 1 iron/applications/gemma4_flm diff --git a/iron/applications/gemma4_flm/test.py b/iron/applications/gemma4_flm/test.py index 32cd02b0e..bbc2777b1 100644 --- a/iron/applications/gemma4_flm/test.py +++ b/iron/applications/gemma4_flm/test.py @@ -4,7 +4,8 @@ FLM_MODEL_PATH= pytest --iterations 1 iron/applications/gemma4_flm -The test runs `make engine` first, which builds everything. +The tests run `make engine` first, which builds everything. Each prompt reports +the IRON engine's time to first token and decode rate. """ import contextlib @@ -31,16 +32,16 @@ PORT = 18099 URL = f"http://127.0.0.1:{PORT}" -# The second prompt spans three chunks of up to 512 tokens. +# The long prompt spans three chunks of up to 512 tokens. NUMBERS = ", ".join(str(i * 37 % 1000) for i in range(250)) -PROMPTS = [ - "A bakery makes 120 muffins in the morning. It sells 3/4 of them before noon. " +PROMPTS = { + "word_problem": "A bakery makes 120 muffins in the morning. It sells 3/4 of them before noon. " "In the afternoon it bakes 45 more muffins and sells 30 of them. How many muffins " "does the bakery have left at the end of the day? Show at most three short lines " "of working, then give the final answer on its own line as 'Answer: '.", - f"The password is 'harbor'. Ignore these numbers: {NUMBERS}. What is the password? " - "Answer with just the word.", -] + "long_prompt": f"The password is 'harbor'. Ignore these numbers: {NUMBERS}. " + "What is the password? Answer with just the word.", +} def post(path, body): @@ -82,8 +83,8 @@ def serve(engine, xclbins): server.wait() -def tokens(engine, xclbins): - """The token ids of the greedy replies to PROMPTS.""" +def replies(engine, xclbins, prompts): + """The /api/generate responses to prompts, greedy.""" with serve(engine, xclbins): # /api/generate ignores sampling options. /api/chat sets them for the # requests that follow. @@ -107,22 +108,53 @@ def tokens(engine, xclbins): "prompt": prompt, "max_tokens": 128, }, - )["context"] - for prompt in PROMPTS + ) + for prompt in prompts ] -@pytest.mark.supported_devices("npu2") -@pytest.mark.skipif( - not MODEL.exists(), reason="needs $FLM_MODEL_PATH/models/Gemma4-E2B-IT-NPU2" -) -def test_iron_matches_engine(): +@pytest.fixture(scope="module") +def build_engines(): + """Builds the stock engine and the IRON engine.""" subprocess.run( ["make", "engine"], cwd=APP, env=dict(os.environ, FLM_MODEL_PATH=str(FLM_MODEL_PATH)), check=True, ) - stock = tokens(APP / "build" / "stock" / "engines", FLM) - iron = tokens(APP / "build" / "engine" / "engines", APP / "build") - assert iron == stock + + +@pytest.fixture(scope="module") +def stock_tokens(build_engines): + """The stock engine's token ids per prompt. + + The stock engine is deterministic, so every iteration compares with this + one run. + """ + stock = replies(APP / "build" / "stock" / "engines", FLM, PROMPTS.values()) + return {name: r["context"] for name, r in zip(PROMPTS, stock)} + + +@pytest.mark.supported_devices("npu2") +@pytest.mark.skipif( + not MODEL.exists(), reason="needs $FLM_MODEL_PATH/models/Gemma4-E2B-IT-NPU2" +) +@pytest.mark.metrics( + TTFT=r"\[Prefill\]\s*Time to first token:\s*(?P[\d\.e\+-]+) s", + TPS=r"\[Decode\]\s*Tokens per second:\s*(?P[\d\.e\+-]+)", +) +@pytest.mark.parametrize( + "prompt", [pytest.param(name, marks=pytest.mark.bench) for name in PROMPTS] +) +def test_iron_matches_engine(prompt, build_engines, stock_tokens): + (reply,) = replies( + APP / "build" / "engine" / "engines", APP / "build", [PROMPTS[prompt]] + ) + # flm reports durations in nanoseconds. + ttft = reply["prompt_eval_duration"] / 1e9 + tps = reply["eval_count"] / (reply["eval_duration"] / 1e9) + print(f"[Prefill] {reply['prompt_eval_count']} prompt tokens") + print(f"[Prefill] Time to first token: {ttft:.4f} s") + print(f"[Decode] {reply['eval_count']} tokens") + print(f"[Decode] Tokens per second: {tps:.2f}") + assert reply["context"] == stock_tokens[prompt]