perf(buckets): add compilation buckets below 256 tokens

[teerthsharma]

Summary

_BUCKETS starts at 256, so every input below 256 tokens is padded up to 256 before the trunk runs. This adds 64, 96, 128, 192, 384. All thirteen existing sizes are retained, so nothing is padded more than today.

No change to model code, config, outputs, or --buckets semantics.

Padding cost

tokens	bucket before → after	padded cost before → after
64	256 → 64	5.97x → 1.00x
100	256 → 128	3.52x → 1.36x
128	256 → 128	2.59x → 1.00x
300	512 → 384	2.44x → 1.48x
384	512 → 384	1.64x → 1.00x

Ratios use a whole-model operation count derived from the op shapes in modules.py, evoformer.py, diffusion_head.py, confidence_head.py:
F(N) = 976896*N^3 + 1433624448*N^2 + 585859918336*N.
Worst case over 64..5120: 5.97x → 2.15x. Mean, log-uniform: 1.45x → 1.32x.

Validation

Against the committed branch:

• Default parsed back from source is strictly increasing, as pipeline.py:53 requires; all 13 previous sizes present.
• Of 5,120 token counts in 1..5120: 0 padded more than before, 320 padded less, rest unchanged.
• python -m py_compile run_alphafold.py passes.

Limits

Operation counts, not wall-clock; at small N inference is closer to launch-bound, so the realised gain is smaller. Not exercised end to end — the model parameters are access-gated and were unavailable, so the checks above are static ones on the bucket list and padding arithmetic. Benefit is confined to inputs below ~400 tokens. Each added bucket costs one compilation the first time an input falls into it; workloads that never run small inputs are unaffected.

AI disclosure

Per CONTRIBUTING.md: AI (Claude) was used for the analysis and for this description. Topology and geometric methods are my area of specialisation, and the padding arithmetic and bucket reasoning here sit in that area. The operation counts are reproducible from the formula and bucket list above. No inference run was performed, so nothing here is a manually tested result.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.