[Hamerlate]

Summary

llama_model.py and gpt_oss_model.py split the batch over --device-num devices with floor division (batch_size // device_num), then divided TTFT (and LLaDA's total time) by device_num again:

• TTFT came out device_num times too small. Each device already prefills only its own sequences, so the second division counts the speed-up twice.
• A batch smaller than device_num gave every device zero sequences. For example, 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.

The devices are now data-parallel replicas:

• the busiest replica serves ceil(batch / device_num) sequences and sets TTFT and the decode time;
• TPS counts the whole batch;
• device_num < 1 is rejected.

A single device is unchanged. Communication between devices is still not modelled.

Before / after

--json, input 2048, output 1024, default plena_settings.toml:

model	batch	devices	TTFT ms before → after	tokens/s before → after
llama-3.1-8b	4	1	377.26 → 377.26	444.76 → 444.76
llama-3.1-8b	4	4	24.73 → 98.94	582.11 → 582.11
llama-3.1-8b	10	4	47.93 → 284.49	1319.44 → 1206.81
llama-3.1-70b	16	32	0.00015 → 741.89	5,000,000 → 250.63
gpt-oss-20b	4	4	17.64 → 70.57	855.36 → 855.36
gpt-oss-20b	4	8	0.00018 → 70.57	4,166,666.67 → 855.36

After the fix, N devices with batch N·b give the TTFT of one device with batch b and N times its TPS.

Tests

analytic_models/performance/tests/test_device_num.py adds 9 tests, for both scripts:

• replicas multiply TPS and keep TTFT;
• more devices than sequences;
• an uneven batch is paced by the busiest device;
• LLaDA keeps its inference time;
• device_num=0 is rejected.

All 9 fail on main and pass here. ruff format --check . and ruff check . pass with the locked ruff.

Also in this PR

d34ba4ab is the same change as #121's 093de4ea, which makes analytic_models importable. Without it, pytest cannot collect anything under analytic_models on main.

#103 and #118 carry their own version of that fix. Whichever lands first, the others resolve analytic_models/__init__.py by keeping the union of imports. The llama_model.py changes here merge cleanly with #103.