vLLM adapter: index kv_caches through layer_name_to_kvcache_index

[sierraisland]

Description

This PR ensures MaxText's vLLM adapter correctly maps physical KV caches to decoder layers via layer_name_to_kvcache_index, rather than assuming physical slot order matches layer order.

Related PRs & Context:

• Supersedes / Continues: vLLM adapter: index kv_caches through layer_name_to_kvcache_index #5113 (incorporates review feedback, pyink formatting, and pylint cleanups).
• Upstream Workaround: [Fix] Preserve layer order in KV cache slots for sequential decoders (MaxText) vllm-project/tpu-inference#3512 (which worked around the issue by re-sorting slots in the cache manager).
• (Related background on hybrid cache allocation: [Perf/Fix] Configure heterogeneous KV cache sharding for GDN/Mamba in model_loader vllm-project/tpu-inference#3515)

Background & Problem

tpu-inference allocates one physical KV cache per unique slot and reports which cache each layer uses via layer_name_to_kvcache_index, which the runner passes to every JAX model as a static positional argument (kv_caches, input_ids, attention_metadata, inputs_embeds, input_positions, layer_name_to_kvcache_index, ...). Native models (e.g. Gemma4) consume it directly; the MaxText adapter previously dropped the map into *args and let the decoder index the physical list positionally (kv_caches[lyr]).
That assumption broke with the vLLM hybrid tensor layout (tpu-inference #3481): all Mamba/GDN caches are allocated before the attention caches. Consequently, on Qwen3.5, an attention layer received a 3D recurrent conv state and the RPA shard_map failed with:

ValueError: shard_map applied to '_ragged_paged_attention': in_specs[3] has length 5, but the passed args[3][0] has shape bfloat16[1028,3,8192], which has rank 3

While vllm-project/tpu-inference#3512 provided a temporary workaround by re-sorting slots in the cache manager, honoring layer_name_to_kvcache_index in MaxText is the principled solution. It also correctly handles cases that sorting cannot fix (e.g., KV-sharing layers where multiple layers redirect to the same physical cache, and lists shorter than the layer count).

Changes in this PR:

1. Model Signature: MaxTextForCausalLM.__call__ explicitly accepts inputs_embeds, _input_positions, and _layer_name_to_kvcache_index in positional alignment with the runner contract, with # pylint: disable=keyword-arg-before-vararg.
2. Hybrid Cache Helpers (hybrid_cache_utils.py):
   • resolve_layer_kv_cache_indices: Parses layer.{i} entries into per-layer physical indices with contiguity, string type, and range validation.
   • gather_layer_kv_caches: Builds the layer-ordered view the MaxText decoder expects.
   • scatter_layer_kv_caches: Validates length consistency and writes updated caches back into the runner's physical slot order (preserving runner donation and out_shardings).
3. Zero XLA Overhead: Gather and scatter are pure Python list indexings on traced values and introduce 0 XLA ops.

Tests

• Ran unit tests covering the hybrid-cache reordering helpers, scatter restoration, KV cache sharing, non-string map keys, gap/out-of-range validation, and length mismatch detection
• End-to-End TPU Run (Qwen3.5 on vLLM Adapter), verified end-to-end decode using MaxTextForCausalLM on TPU with a tpu-inference build without the [JAX] Replace jnp.clip(..., a_min=..., a_max=...) with jnp.clip(..., min=..., max=...). #3512 slot-sorting workaround (reproducing and verifying the root fix)

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request implements mapping, gathering, and scattering utilities to translate between the physical slot order of KV caches from the tpu-inference runner and the layer-ordered view expected by MaxText decoders. It updates the vLLM adapter's forward pass to use these utilities and adds comprehensive unit tests to verify the remapping logic. There are no review comments, so I have no feedback to provide.

[codecov]

Codecov Report

❌ Patch coverage is 71.15385% with 15 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...t/integration/vllm/maxtext_vllm_adapter/adapter.py	7.14%	13 Missing ⚠️
src/maxtext/integration/vllm/hybrid_cache_utils.py	94.73%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!