fp8 KV cache: pre-dispatch quantize_kv in forward_serve_vllm

[lokic233]

Problem

Attention.forward_serve_vllm passes bf16 k/v straight into the ragged-paged-attention (RPA) v3 kernel with q_scale = k_scale = v_scale = None. When the vLLM rollout allocates the KV cache in fp8 (kv_cache_dtype=fp8), the kernel raises:

ValueError: Expected kv_cache.dtype=float8_e4m3fn to be equal to k.dtype=bfloat16 and v.dtype=bfloat16.

The RPA v3 kernel's static_validate_inputs requires kv_cache.dtype == k.dtype == v.dtype, and its own comment states it "expects the kv quantization happens outside of the RPA kernel." The MaxText serving path never quantizes, so it trips the assert during capture_model/KV-init and again on any real decode step.

Why the native path doesn't hit this

The native vLLM backend (layers/vllm/backends/flash_attn.py) already quantizes before dispatch — it calls quantize_kv(...) on k/v ahead of the attention call. forward_serve_vllm is simply missing the equivalent step.

Fix

Mirror flash_attn.py on the MaxText serving path: when the KV cache is a 1-byte float (fp8 e4m3/e5m2), quantize k/v to the cache dtype via the shared quantize_kv helper and thread the scales through to the kernel. Non-fp8 caches are a no-op (scales stay None), so bf16 behavior is unchanged. Single code path, reusing the existing helper — no duplicated quantization logic and no in-kernel cast.

Validation

Validated on TPU v7x (gemma-4-2B DAPO/GRPO RL, 64 chips, 4x4x4) with kv_cache_dtype=fp8:

• Clears capture_model/KV-init — no dtype assert; real decode runs on the same RPA kernel that previously asserted (peak ~500 tok/s).
• Reward unchanged — rewards/mean = 0.1405, matching the bf16-cache baseline (~0.13); advantages centered at ~0, no NaN/collapse. Properly-scaled fp8 preserves the reward signal.
• Memory — fp8 halves per-token KV bytes (2× token capacity) at the same HBM footprint.

[gemini-code-assist]

Caution

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

[lokic233]

Added a minimal regression test (test_forward_serve_vllm_fp8_kv_quantizes) that mirrors the existing test_forward_serve_vllm — with an fp8 KV cache it asserts the k/v passed to the RPA kernel are cast to the fp8 cache dtype. Skipped like its sibling (needs vllm-tpu, not yet a MaxText dep); the behavior itself was validated on TPU (reward 0.1405, clears capture_model).