[ROCm] Fuse Kimi-K3 DSpark QK RoPE with the MLA cache write

[gronsti-amd]

Purpose

Apply AITER fused_qk_rope_concat_and_cache_mla to the Kimi-K3 DSpark draft model.
The kernel fuses the QK RoPE, the absorbed query, and the KV cache write into one launch.
This PR fills the hooks introduced in #60054.

The shared op schemas stay unchanged. unified_mla_kv_cache_update is on the compile path of every MLA model and is a splitting op. Positions stay on _rope_positions.

What changes

• The factory returns KimiK3DSparkMLAWrapper. After super().__init__ it stores the embedding as _epilogue_rotary_emb.
• When fusion is on, the wrapper skips its own YaRN and passes positions into MLAAttention.forward. positions is None raises RuntimeError. There is no fallback that writes unrotated KV.
• Prefill context parallel is refused when fusion would otherwise be on. The gather changes the row count while positions stay local.
• Decode rows are dropped from the generic cache write; the fused kernel writes them. Prefill k_pe is cloned, rotated for the cache, and the caller's tensor stays unrotated so attention can rotate it again in _prepare_mha_inputs.
• q_out is FP8 when the KV cache is quantized and supports_quant_query_input is set (ROCM_AITER_MLA). Triton MLA clears that flag and keeps bf16. The absorbed query is made contiguous before the launch.
• The NoPE identity-RoPE path stays closed when rotary_emb is set, including inside _fused_qk_prep_supported.
• MLACommonMetadata is imported from vllm.model_executor.layers.attention.mla_attention.

Validation

• CPU hook tests: fusion refuses PCP; fusion is off without an embedding or without the AITER probe; a decode-only update skips do_kv_cache_update; prefill k_pe is re-rotated for MHA and the caller's key stays unrotated until then; q_out matches the backend query dtype; fusion requires positions.
• test_fused_qk_rope_concat_requires_fp32_cos_sin asserts the wrapper rejects a bf16 cos/sin table. The probe is false unless VLLM_ROCM_USE_AITER=1 and aiter.fused_qk_rope_concat_and_cache_mla imports. The AMD Basic Models job runs that test in its own pytest process with the flag
  set, and filters it out of the broader command.
• E2E agentic run of Kimi-K3 with this drafter, compared with the NVIDIA draft module on the same prompts.
• Perf comparison against the NVIDIA drafter: vllm bench serve / spec-bench over output length × max concurrency, with

{"method": "dspark", "model": "Inferact/Kimi-K3-DSpark", "num_speculative_tokens": 5, "attention_backend": "TRITON_MLA", "draft_sample_method": "probabilistic", "rejection_sample_method": "block"}

Out of scope

• Changing DeepseekScalingRotaryEmbedding.forward_hip to AITER Triton RoPE. That would apply to every HIP user of the embedding. It can be a follow-up or alternative to this PR.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @gronsti-amd.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork