[Bugfix][MLA] Keep MXFP8 prefill activations in model dtype

[GirasoleY]

Purpose

Fix MLA context attention passing FP8 activations into an MXFP8 kv_b_proj that requires FP16/BF16 inputs. With FP8 prefill and FP8 KV cache, the gathered latent can remain torch.float8_e4m3fn; selecting the projection input dtype from its stored weight causes Input dtype must be float16 or bfloat16, got torch.float8_e4m3fn.

Use kv_b_proj.params_dtype for this FP8-prefill case. Both generic and DCP context paths already share the helper. Preserve the explicit ModelOpt FP8 block weight-only input_dtype override and the existing packed-weight / non-FP8-prefill behavior. Weight and KV-cache storage formats are unchanged.

Add one CPU regression that checks the activation actually passed to kv_b_proj after DCP gather/reorganization. Tighten the existing Kimi-K3 fused-versus-generic test fixture to enforce the MXFP8 activation contract without adding another test matrix.

Duplicate check: searched open PRs/issues for MLA dtype, MXFP8 prefill, and DCP. #37245 addresses INT8 weight casts, not this MXFP8 activation contract. The ModelOpt-specific handling from #49381 is retained. No open PR addressing this fix was found. The NIXL symmetric hybrid-DCP work is already upstream in #55531 and is outside this PR.

AI assistance: Codex assisted with the upstream port, regression coverage, validation, and this description.

Test Plan

Run the existing context suites:

.venv/bin/python -m pytest \
  tests/v1/attention/test_mla_context_chunks.py \
  tests/models/kimi_k3/test_mla_prefill_context.py -q -p no:cacheprovider

Check that test_dcp_context_projects_mxfp8_latent_from_bf16 fails with the original helper from public base ba2ae9f23961ac67bc5c055da8c26fbd660989c6, and passes with the fix. The negative control restores only the original helper in memory; the test and call path are identical.

Run the repository's pre-commit hooks on the three changed files, plus pre-commit run mypy-3.12 --hook-stage manual --files <changed files>.

Test Result

• Focused regression: 1 passed. Negative control: 1 failed, with torch.float8_e4m3fn != torch.bfloat16 at the assertion on the projection input.
• Existing context suites: 14 passed, 10 skipped on macOS CPU / Python 3.12 / PyTorch 2.13. CUDA-dependent cases were skipped. The actual local invocation used the virtualenv Python with this worktree on PYTHONPATH and test-process-only overrides: vllm.utils.torch_utils.PIN_MEMORY = False and torch.accelerator.empty_host_cache = lambda: None. These avoid macOS allocator/cleanup crashes; no attention or dtype implementation was replaced in the passing run. The initial unmodified pytest invocation completed the focused assertion but crashed in PyTorch host-cache cleanup.
• All applicable pre-commit hooks passed, including mypy 3.10; the manual mypy 3.12 hook passed. git diff --check passed.
• Prior deployment evaluation of the equivalent dtype fix, not a fresh GPU evaluation of this public-main port: Kimi-K3 MXFP8 with EAGLE3, 1P1D TP8/DCP8 on 16 GB300 GPUs, FlashInfer 0.6.17. GSM8K five-shot, all 1,319 questions, concurrency 128, temperature 0, max output 2,048: 1,274/1,319 = 96.5883% strict and flexible accuracy, 0 request errors and 0 truncations. Decode batch limit was 96; NIXL transferred the prompt KV successfully with zero transfer failures. This run validated the integration-branch fix; current OSS CUDA execution is left to CI.

---

Essential Elements of an Effective PR Description Checklist

• Purpose and concrete failure described.
• Test plan and commands provided.
• Test results and validation limits provided.
• Documentation considered; no new configuration or user-facing API.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.