[Runner] Never pad a step past max_num_batched_tokens

[gutianyu-google]

[Runner] Never pad a step past max_num_batched_tokens

Problem

TPUModelRunner builds num_tokens_paddings with get_token_paddings(min, max_num_batched_tokens * dp_size, gap).
With the default VLLM_TPU_BUCKET_PADDING_GAP=0 the buckets double until they cover the budget, so a full
chunked-prefill step is padded to the next power of two whenever the budget is not one itself. Example
(Qwen3-0.6B, TP8 x DP16, max_num_batched_tokens=16640): a 16640-token prefill chunk runs at the 32768-token
shape per rank — 97% padding.

On that shape we also observe a silent correctness problem: the last-position logits of the padded prefill step
are wrong, so the first sampled token of ~90% of requests is garbage ( tắc, iquement, 人力, ...) before
the model recovers on the next decode step. Reproduced on TPU7x with page_size 128, 1 KV head/device
(tpu-inference @ be9cdf0): first token correct for 0/2048 requests at the 32768 shape, 2048/2048 when the
same chunk is padded to 18432 (VLLM_TPU_BUCKET_PADDING_GAP=32768), 16384 (max_num_batched_tokens=12544),
or 4096. The kernel-level cause at the 32768 shape is not localized here; this PR removes the trigger
and the padding waste.

Change

• runner/utils.py::build_token_paddings(min, max, gap, additional_sizes): the runner's bucket list is now the generated
  buckets capped at the budget (rounded up to a multiple of min_token_size), the budget itself, and the user's
  additional_config.compilation_sizes, merged and deduplicated in one place. The budget enters through the same merge
  that already handled compilation_sizes; no new flag on get_token_paddings, which is unchanged.
• Buckets above the budget are dropped: the scheduler never schedules more than max_num_batched_tokens per rank, so
  they are unreachable and only cost a compile each and oversize the preallocated token buffers (max_num_tokens is
  derived from the last bucket). With a 16640 budget the list ends ..., 8192, 16384, 16640 instead of
  ..., 8192, 16384, 32768. User-supplied compilation_sizes are kept as given, even above the budget.
• runner/tpu_runner.py: builds the list through the helper. A budget that is already a bucket (e.g. 8192, 16384)
  yields the stock list, so those configurations see no change.
• tests/runner/test_utils.py: capped list for the exponential and bucketed gaps, stock list for a power-of-two budget,
  rounding of the budget, and compilation_sizes preserved and deduplicated.

Validation

• Unit tests for get_token_paddings.
• End-to-end (MaxText/tunix GRPO, Qwen3-0.6B, TP8xDP16, budget 16640): with the patch and no env override,
  first token correct for 2048/2048 requests (LENDIST census), generation distribution identical to the
  VLLM_TPU_BUCKET_PADDING_GAP=32768 workaround; without it 0-256/2048. Run 51 (this patch, no env): first token correct 2048/2048 at steps 0 and 1; 29.3% / osl 4477 at step 0 (workaround run: 27.5% / 4391 — same distribution, different samples); step-1 engine_generate 84.5 s vs 88.5 s with the 18432 workaround and 82.9 s on the corrupted 32768 shape, i.e. the exact-budget bucket also pads less than the gap workaround.
• No behaviour change when the budget is already a bucket (e.g. 8192, 16384).

[wenxindongwork]

btw, you can also just use additional_config.compilation_sizes to pass in custom compilation sizes.

[gutianyu-google]

Done, thanks!