[Performance]: unconditional fp32 logits convert + aliasing copy costs ~12% of the greedy decode step (Gemma 2B, v5e, batch 512)

[RubensZimbres]

Proposal to improve performance

Summary

In the non-spec-decode path, _sample_from_logits in tpu_runner.py converts the logits to float32 unconditionally before calling sample(). Measured consequences at batch 512 (Gemma 2B, TPU v5e-1, greedy decoding, offline batch), these are ranks 2 - 4 of the entire HLO op stats table for a steady-state generation trace on tpu-inference 0.24.0:

Rank	Op	Shape	Avg time/step	Note
2	copy.2	f32[512,256000]	1,562 µs	input-aliasing copy; identical layouts on both sides
3	convert_element_type.1	bf16→f32[512,256000]	1,232 µs	standalone dispatch, own program id (tpu_runner.py:1624 on 0.24.0)
4	jit(sample)/reduce (argmax)	f32[512,256000]	696 µs	reads f32; ~2× the bytes it needs (sampling.py:125 on 0.24.0)

Total ≈ 3.1 ms of a ~26.5 ms decode step ≈ 12%, on a workload class (greedy, high batch: structured output, extraction, agent traffic) where the sampling transforms are unused.

Confirmed present on current main as of ⟨sha⟩ (2026-07-08):

• unconditional logits = logits.astype(jnp.float32) in the non-spec branch of _sample_from_logits: ⟨permalink 1⟩
• greedy branch of sample() returns its input (ret_tokens = greedy_tokens; ret_logits = logits), which forces the input-aliasing copy: ⟨permalink 2⟩

Note that sample() already performs its own astype(jnp.float32) after the argmax, so the runner-side convert is redundant in this path.

Why fp32 is not needed for correctness here

bf16 to f32 widening is exact, so argmax over the widened tensor is mathematically identical to argmax over the original bf16, the greedy token cannot change. The f32 logits required by the logprobs consumers can be produced inside sample() (where the astype already exists), preserving the output dtype contract with no observable change downstream.

Proposed change

In the greedy branch: run the argmax on the incoming bf16 logits; perform the single bf16 to f32 convert inside sample() to produce ret_logits, so the convert itself becomes the output materialization (eliminating the aliasing copy, since the output is no longer an aliased input); rebuild the warmup/precompiled graphs together with the dtype change (see the regression evidence in the discussion section below for why this last part is essential). Expected recovery: up to the ~12% above on greedy high-batch workloads, minus whatever the pipeline needs to retain.

Report of performance regression

Not a regression. The cost is present and stable across the 0.23.0 release, the 0.24.0 release, and current main (⟨sha⟩). Filed as a proposal because it is a standing inefficiency in the greedy path, not a change in behavior between versions.

Misc discussion on performance

Why this is filed as a proposal rather than submitted as a PR

I attempted the fix externally twice, on two releases, and both attempts regressed severely with the same signature:

• 0.23.0: made the fp32 convert conditional on do_sampling at the output side to −14% throughput, batch TTFT ×5. Reverted; baseline reconfirmed.
• 0.24.0: removed the runner-side input convert, letting sample()'s internal astype do the work to −26% throughput (19,320 → 14,258 tok/s), batch TTFT ×4.4 (275 → 1,209 ms). Reverted; baseline reconfirmed (19,217 tok/s, within session variance of the pre-patch 19,320).

Diagnosis: the precompiled graph menu built during the engine's warmup pass is keyed on the dtype at this program boundary. An external edit leaves the menu built for the old dtype, so runtime calls land off the precompiled paths on every step, which both the TTFT multiplication and the throughput loss detect. The fix therefore has to rekey the warmup/precompile pass together with the dtype change, i.e., it belongs at the pipeline level, inside the project. The measurements above are offered as evidence for that constraint, not only as failed attempts: any patch for this should be checked against batch TTFT, which detected the miscompilation both times before throughput was even read.

Validation offer

I can validate a candidate fix on TPU v5e with same-session baseline discipline and the batch-TTFT regression check described above. The full measurement protocol (clean-run rules, treatment verification independent of the treatment flag, HLO shape-keyed trace analysis) is documented in DOI 10.5281/zenodo.21212010 and DOI 10.5281/zenodo.21227874, which measured this serving stack on the same hardware.

Related: #1424 (sampling-path performance).

Your current environment (if you think it is necessary)

Name: tpu_inference
Version: 0.24.0
Name: vllm-tpu
Version: 0.24.0
Name: jax
Version: 0.10.2
Name: libtpu
Version: 0.0.42.1
Name: torch
Version: 2.10.0
Python 3.12.13