feat(memory): size the LoRA prefill scratch, not just adapter storage

[dittops]

Scope

Serving a LoRA-enabled model allocates two very different things, and this library reported
only the small one.

calculate_lora_adapter_memory returns the adapter A/B storage — 0.22 GiB at rank 256 on
a 0.6B model. Punica's prefill working buffers at the same rank are 16.77 GiB, 75x larger,
and were reported as nothing at all. A control plane sizing a pod off the first number
under-sizes it by that factor; the pod then OOMKills at Warming up model for the compilation,
which destroys the evidence along with the pod.

This adds ModelMemoryCalculator.calculate_lora_prefill_scratch and reports it as its own
MemoryReport field. It is deliberately not folded into the storage term: summed together it
reads as a rounding error on a number readers already believe is small, which is how a 16 GiB
allocation stayed invisible.

Calibration

Qwen3-0.6B, 28 layers, rank 256, max_num_seqs=1, vLLM CPU backend, cgroup limit 40 GiB so
nothing was clipped. LoRA-attributable is the LoRA-on peak minus a LoRA-off control pod at the
identical token budget:

batched tokens	LoRA off	LoRA on	scratch
2394	5143 MiB	22311 MiB	16.77 GiB
4273	5113 MiB	27119 MiB	21.49 GiB
8192	5096 MiB	OOM @ 50 GiB	> 45.02 GiB

Two results worth stating:

• Base memory is flat in the token budget — 47 MiB across a 3.4x range. Every bit of a LoRA
  pod's context sensitivity is this one term.
• The term is sized by max_num_batched_tokens, the tokens in one forward pass, not by the
  context length. With chunked prefill those differ, and charging the context prices an
  8192-context pod at 31 GiB it never allocates.

API

• MemoryReport gains lora_prefill_scratch_bytes and notes; total_memory_bytes includes
  the scratch.
• calculate_memory(...) gains max_num_batched_tokens, defaulting to seq_length. Callers
  that do not set it get the unchunked reading, so nothing existing changes shape.

Where it stops, and says so

Above 4273 batched tokens the term goes convex — at 8192 the linear form under-predicts a
measured floor by >13 GiB — so it returns a note marking the figure a LOWER BOUND rather
than an estimate. It also notes when a request leaves the measured 28-layer/rank-256 geometry,
since linearity in depth and in rank is assumed rather than measured.

Testing

cd llm-memory-calculator && python3 -m pytest tests/test_lora_prefill_scratch.py -q

14 passed. New file tests/test_lora_prefill_scratch.py pins both calibration points, the
batch-not-context behaviour, rank and TP scaling, both notes, and the refusal paths (unknown
layer count, unset rank).

Full suite: 8 pre-existing failures, all missing packages in the local environment
(pymoo x3; llm_memory_calculator not pip-installed, which breaks the 5 test_bud_component
tests that spawn a subprocess). None relate to this change.

Downstream

BudCluster consumes this calibration for pod sizing in BudEcosystem/bud-runtime#2422. It
restates the constants rather than importing the library — it is a control plane, and this
package pulls matplotlib/plotly/seaborn/pandas — and pins the two copies together with a test.

🤖 Generated with Claude Code

---

^{Need help on this PR? Tag @codesmith-bot with what you need. Autofix is disabled.}

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.