[refactor][model] Kimi-K3 DSpark MLA common base class

[gronsti-amd]

Purpose

Currently, the Kimi-K3 DSpark draft model has only one NVIDIA implementation,
and the dispatch is hard-wired to it on every platform. A ROCm module that copies that file drifts.
This PR moves the draft model into vllm/models/kimi_k3/common/dspark_mla.py and leaves NVIDIA as a thin binding.
NVIDIA behavior stays the same: the same MultiHeadLatentAttention, the same layer name, and the same kernels.
The registry still loads vllm.models.kimi_k3.nvidia.dspark_mla:
removing the hard-wired dispatch is left for a follow-up PR (#60857).

What changes

• The decoder, context-KV insert, weight mapper, Markov head, and confidence head live once, in common/dspark_mla.py. That file imports neither amd/ nor nvidia/.
• NVIDIA supplies build_self_attn and build_mlp. kv_cache_layer defaults to the attention module, so context-KV insert and get_draft_kv_cache_layer_names() still use self_attn.

Validation

• pytest tests/models/test_dspark_mla.py passes.
• Existing tests cover the registry class, the weight mapper, duplicated context-KV shards, the replicated Markov head, and the compile-free entry point. Construction monkeypatches target vllm.models.kimi_k3.common.dspark_mla.
• New tests check that the NVIDIA binding's attention class is MultiHeadLatentAttention, that kv_cache_layer returns the attention module, that a wrapper subclass's mla_attn_cls is constructed with a prefix ending in .attn, and that the default hooks return their inputs.
• The file is collected in Basic Models, including the CPU job. Main's newer DeepSeek and layer-registry entries in that job stay on the command line.

Details

AMD and NVIDIA attention and cache architecture for MLA differ: NVIDIA MultiHeadLatentAttention owns the KV cache, while the AMD stack is KimiK3MultiHeadLatentAttentionWrapper around a generic MLAAttention whose layer_name is prefix + ".attn". Geometry, norms, and context RoPE already sit on self_attn on both sides. This PR refactors so that the three reads that differ (cache, dtype and scale, layer name) go through kv_cache_layer.

• common/ does not construct the NVIDIA attention class. A default there would import nvidia.mla, which is also the NVIDIA target K3 attention.
• #58814 pointer refresh is in the shared path: _has_uniform_block_layout keys the cache on data_ptr() and drops _context_cache_ptrs when the allocation changes.
• enable_confidence_head builds DSparkConfidenceHead, keeps those checkpoint weights, and compute_confidence serves adaptive verification. The head stays None when the flag is off.
• MLAAttention stores positions as _rope_positions (assigned every forward, including None). The wrapper takes mla_attn_cls. Three hooks default to today's behavior: _prepare_kv_cache_update, _prepare_mha_inputs, and _form_decode_q. Custom-op schemas stay as they are. Positions are module state the ops read.
• CUDA-graph replay sees new positions and slot mappings only while those tensors stay the persistent buffers the runner updates in place. Replacing either tensor freezes the captured pointer. Comments on both sites say so.

A planned fusion of QK-RoPE plus cache insert (fused_qk_rope_concat_and_cache_mla) requires an epilogue to replace query formation and the KV write.
The fusion cannot live only in the draft model: Absorb runs inside MLAAttention.forward_impl, and unified_mla_kv_cache_update cannot take positions.
This PR adds three defaulted hooks as an extension point for this decode epilogue, to avoid needing a second refactor immediately after this. The defaults do not read _rope_positions.

The Kimi-K3 DSpark MLA is under active development, causing a naive duplicate-and-modify approach to quickly drift, as shown by #57942:

• [Bugfix][Kimi-K3] Refresh DSpark context KV cache pointers after the KV cache is re-bound #58814 pointer refresh (_kv_cache_ptrs_key, drop _context_cache_ptrs when the cache is rebound).'
• Grouped FP8 context-KV insert (see commit f589f93)
• Missing confidence head.

Refactoring reduces the drift, allowing both platforms to gain new features in a single change.

Out of scope

• The AMD binding and the ROCm registry dispatch ([ROCm] ROCm binding on common Kimi-K3 Dspark MLA module #60857).
• AITER fused_qk_rope_concat_and_cache_mla ([ROCm] Fuse Kimi-K3 DSpark QK RoPE with the MLA cache write #60864).

Suggested review order

Review by commit. The aggregate diff looks like a rewrite of nvidia/dspark_mla.py because that path stays as the public binding, so Git does not record a rename.

---

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

[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