[KK] Prepare B12X Kimi KDA prefill and expose MLA partial precision

[voipmonitor]

Behavior

Enables --kda-prefill-backend b12x for Kimi-K3 through prepared recurrent-state bindings and reusable caller-owned scratch. The auto default is unchanged. Exposes VLLM_K3_DENSE_MLA_PARTIAL_DTYPE=fp32 for dense MLA split partials; BF16 remains the default and the cache dtype does not change.

The KDA path consumes state-pool indices directly instead of gathering and scattering recurrent states around the prefill kernel. FP32 MLA partials avoid a BF16 intermediate rounding step and double only the partial-output scratch allocation.

Implemented. Qualified in the Kimi-K3 EXL3 TP9/DCP9 integration with Red Hat DSpark K5, BF16 activation arithmetic, FP8 KV, and a 4096-token scheduler chunk. This is not a qualification of every KDA backend or TP topology.

Dependencies and compatibility

Targets dev/karmic-kraken directly. The FP32 selector requires B12X #441. Selecting BF16 does not require that optional precision extension. The KDA implementation uses the prepared B12X API already present on master.

This does not duplicate #692, which concerns stateless first-token classification. The implementation does not import FlashKDA.

Validation

• tests/models/kimi_k3/test_kda.py::test_b12x_kda_prefill_matches_triton_serving_batch: GPU numerical comparison passed.
• GPU dense-MLA/cache tests cover both partial dtypes, high page IDs, and CUDA Graph replay; two noncausal DFlash cache cases also pass.
• Matched 183-token coding prompt, temperature zero, seed 1, 4096 generated tokens, three runs: all output hashes and all speculative acceptance counters equal the preserved serving implementation. The composed runtime measures 100.00 median decode tok/s versus 100.38 tok/s for the control on the same nine RTX PRO 6000 Blackwell GPUs.

The performance result qualifies the composition, not an isolated speedup from this PR. Source identities, launch settings and remaining merge dependencies are tracked in #906. No broad task-quality or long-context claim is made.

AI assistance: OpenAI Codex. Maintainer review is required.

[coderabbitai]

Warning

Review limit reached

Next included review available in 58 minutes.

Check out review usage here.

View limit details

Limit details: You’ve used the included review currently available.

You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository.

Learn how review limits work.

Review configuration:

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Advanced

Run ID: 5f9d04a0-5b7c-4146-83a5-c7ae584783d2

📥 Commits

Reviewing files that changed from the base of the PR and between 502d6cb and 185c108.

📒 Files selected for processing (7)

• .lil/changes/vllm-dense-mla-partial-precision.json
• .lil/changes/vllm-kimi-b12x-kda-prefill.json
• tests/models/kimi_k3/test_kda.py
• tests/v1/attention/test_b12x.py
• vllm/envs.py
• vllm/models/kimi_k3/nvidia/kda.py
• vllm/v1/attention/backends/mla/b12x_mla.py

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[voipmonitor]

Closing: this change is already in dev/karmic-kraken, through Luke's karmic integration. The remaining diff is only formatting or test annotations.